Skip to content

Commit 6933034

Browse files
zhang2014zhang2014
authored
fix(query): fix panic if miss match empty data type (#18894)
1 parent cd9cb4e commit 6933034

File tree

1 file changed

+167
-80
lines changed

1 file changed

+167
-80
lines changed

‎src/query/expression/src/kernels/stream_partition.rs‎

Lines changed: 167 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -250,91 +250,178 @@ impl BlockPartitionStream {
250250
}
251251

252252
pub fn copy_column<I: Index>(indices: &[I], from: &Column, to: &mut ColumnBuilder) {
253-
match (to, from) {
254-
(ColumnBuilder::Nullable(builder), Column::Null { .. }) => {
255-
builder.push_repeat_null(indices.len());
256-
}
257-
(ColumnBuilder::Array(builder), Column::EmptyArray { .. }) => {
258-
for _ in 0..indices.len() {
259-
builder.commit_row();
253+
match to {
254+
ColumnBuilder::EmptyArray { len } => match from {
255+
Column::EmptyArray { .. } => *len += indices.len(),
256+
Column::Array(column) => {
257+
let capacity = *len + column.len();
258+
match ColumnBuilder::with_capacity(&from.data_type(), capacity) {
259+
ColumnBuilder::Array(mut builder) => {
260+
builder.offsets.extend(vec![0; *len]);
261+
copy_array(&mut builder, column, indices);
262+
*to = ColumnBuilder::Array(builder);
263+
}
264+
_ => unreachable!(
265+
"ColumnBuilder::with_capacity for Array type should return ColumnBuilder::Array, \
266+
but got different variant. data_type: {}, capacity: {}",
267+
from.data_type(), capacity
268+
),
269+
}
260270
}
261-
}
262-
(ColumnBuilder::Map(builder), Column::EmptyMap { .. }) => {
263-
for _ in 0..indices.len() {
264-
builder.commit_row();
271+
_ => unreachable!(
272+
"EmptyArray builder can only copy from EmptyArray or Array, but got from type: {}",
273+
from.data_type()
274+
),
275+
},
276+
ColumnBuilder::Array(builder) => match from {
277+
Column::EmptyArray { .. } => {
278+
for _ in 0..indices.len() {
279+
builder.commit_row();
280+
}
265281
}
266-
}
267-
(ColumnBuilder::Number(builder), Column::Number(number_column)) => {
268-
with_number_mapped_type!(|NUM_TYPE| match (builder, number_column) {
269-
(NumberColumnBuilder::NUM_TYPE(b), NumberColumn::NUM_TYPE(c)) => {
270-
copy_primitive_type(b, c, indices);
282+
Column::Array(column) => {
283+
copy_array(builder, column, indices);
284+
}
285+
_ => unreachable!(
286+
"Array builder can only copy from EmptyArray or Array, but got from type: {}",
287+
from.data_type()
288+
),
289+
},
290+
ColumnBuilder::Null { len } => match from {
291+
Column::Null { .. } => *len += indices.len(),
292+
Column::Nullable(column) => {
293+
let capacity = *len + column.len();
294+
295+
match ColumnBuilder::with_capacity(&from.data_type(), capacity) {
296+
ColumnBuilder::Nullable(mut builder) => {
297+
builder.push_repeat_null(*len);
298+
copy_nullable(&mut builder, column, indices);
299+
*to = ColumnBuilder::Nullable(builder);
300+
}
301+
_ => unreachable!("ColumnBuilder::with_capacity for Nullable type should return ColumnBuilder::Nullable, \
302+
but got different variant. data_type: {}, capacity: {}",
303+
from.data_type(), capacity),
271304
}
272-
_ => unreachable!(),
273-
})
274-
}
275-
(ColumnBuilder::Decimal(builder), Column::Decimal(column)) => {
276-
with_decimal_type!(|DECIMAL_TYPE| match (builder, column) {
277-
(
278-
DecimalColumnBuilder::DECIMAL_TYPE(builder, _),
279-
DecimalColumn::DECIMAL_TYPE(column, _),
280-
) => {
281-
copy_primitive_type(builder, column, indices);
305+
}
306+
_ => unreachable!(
307+
"Null builder can only copy from Null or Nullable, but got from type: {}",
308+
from.data_type()
309+
),
310+
},
311+
ColumnBuilder::Nullable(builder) => match from {
312+
Column::Null { .. } => {
313+
builder.push_repeat_null(indices.len());
314+
}
315+
Column::Nullable(column) => {
316+
copy_nullable(builder, column, indices);
317+
}
318+
_ => unreachable!(
319+
"Nullable builder can only copy from Null or Nullable, but got from type: {}",
320+
from.data_type()
321+
),
322+
},
323+
ColumnBuilder::EmptyMap { len } => match from {
324+
Column::EmptyMap { .. } => *len += indices.len(),
325+
Column::Map(column) => {
326+
let capacity = *len + indices.len();
327+
match ColumnBuilder::with_capacity(&from.data_type(), capacity) {
328+
ColumnBuilder::Map(mut builder) => {
329+
builder.offsets.extend(vec![0; *len]);
330+
copy_array(&mut builder, column, indices);
331+
*to = ColumnBuilder::Map(builder);
332+
}
333+
_ => unreachable!("ColumnBuilder::with_capacity for Map type should return ColumnBuilder::Map, \
334+
but got different variant. data_type: {}, capacity: {}",
335+
from.data_type(), capacity),
282336
}
283-
_ => unreachable!(),
284-
});
285-
}
286-
(ColumnBuilder::Boolean(builder), Column::Boolean(column)) => {
287-
copy_boolean(builder, column, indices)
288-
}
289-
(ColumnBuilder::Date(builder), Column::Date(column)) => {
290-
copy_primitive_type(builder, column, indices);
291-
}
292-
(ColumnBuilder::Interval(builder), Column::Interval(column)) => {
293-
copy_primitive_type(builder, column, indices);
294-
}
295-
(ColumnBuilder::Timestamp(builder), Column::Timestamp(column)) => {
296-
copy_primitive_type(builder, column, indices);
297-
}
298-
(ColumnBuilder::Bitmap(builder), Column::Bitmap(column)) => {
299-
copy_binary(builder, column, indices);
300-
}
301-
(ColumnBuilder::Binary(builder), Column::Binary(column)) => {
302-
copy_binary(builder, column, indices);
303-
}
304-
(ColumnBuilder::Variant(builder), Column::Variant(column)) => {
305-
copy_binary(builder, column, indices);
306-
}
307-
(ColumnBuilder::Geometry(builder), Column::Geometry(column)) => {
308-
copy_binary(builder, column, indices);
309-
}
310-
(ColumnBuilder::Geography(builder), Column::Geography(column)) => {
311-
copy_binary(builder, &column.0, indices);
312-
}
313-
(ColumnBuilder::String(builder), Column::String(column)) => {
314-
copy_string(builder, column, indices);
315-
}
316-
(ColumnBuilder::Array(builder), Column::Array(column)) => {
317-
copy_array(builder, column, indices);
318-
}
319-
(ColumnBuilder::Map(builder), Column::Map(column)) => {
320-
copy_array(builder, column, indices);
321-
}
322-
(ColumnBuilder::Nullable(builder), Column::Nullable(column)) => {
323-
copy_nullable(builder, column, indices);
324-
}
325-
(ColumnBuilder::Vector(builder), Column::Vector(column)) => {
326-
copy_vector(indices, builder, column);
327-
}
328-
(ColumnBuilder::Opaque(builder), Column::Opaque(column)) => {
329-
copy_opaque(indices, builder, column);
330-
}
331-
(ColumnBuilder::Tuple(builders), Column::Tuple(columns)) => {
332-
for (builder, column) in builders.iter_mut().zip(columns.iter()) {
333-
copy_column(indices, column, builder)
334337
}
335-
}
336-
_ => unreachable!(),
337-
}
338+
_ => unreachable!(
339+
"EmptyMap builder can only copy from EmptyMap or Map, but got from type: {}",
340+
from.data_type()
341+
),
342+
},
343+
ColumnBuilder::Map(builder) => match from {
344+
Column::Map(column) => {
345+
copy_array(builder, column, indices);
346+
}
347+
Column::EmptyMap { .. } => {
348+
for _ in 0..indices.len() {
349+
builder.commit_row();
350+
}
351+
}
352+
_ => unreachable!(
353+
"Map builder can only copy from EmptyMap or Map, but got from type: {}",
354+
from.data_type()
355+
),
356+
},
357+
_ => match (to, from) {
358+
(ColumnBuilder::Number(builder), Column::Number(number_column)) => {
359+
with_number_mapped_type!(|NUM_TYPE| match (builder, number_column) {
360+
(NumberColumnBuilder::NUM_TYPE(b), NumberColumn::NUM_TYPE(c)) => {
361+
copy_primitive_type(b, c, indices);
362+
}
363+
_ => unreachable!(),
364+
})
365+
}
366+
(ColumnBuilder::Decimal(builder), Column::Decimal(column)) => {
367+
with_decimal_type!(|DECIMAL_TYPE| match (builder, column) {
368+
(
369+
DecimalColumnBuilder::DECIMAL_TYPE(builder, _),
370+
DecimalColumn::DECIMAL_TYPE(column, _),
371+
) => {
372+
copy_primitive_type(builder, column, indices);
373+
}
374+
_ => unreachable!(),
375+
});
376+
}
377+
(ColumnBuilder::Boolean(builder), Column::Boolean(column)) => {
378+
copy_boolean(builder, column, indices)
379+
}
380+
(ColumnBuilder::Date(builder), Column::Date(column)) => {
381+
copy_primitive_type(builder, column, indices);
382+
}
383+
(ColumnBuilder::Interval(builder), Column::Interval(column)) => {
384+
copy_primitive_type(builder, column, indices);
385+
}
386+
(ColumnBuilder::Timestamp(builder), Column::Timestamp(column)) => {
387+
copy_primitive_type(builder, column, indices);
388+
}
389+
(ColumnBuilder::Bitmap(builder), Column::Bitmap(column)) => {
390+
copy_binary(builder, column, indices);
391+
}
392+
(ColumnBuilder::Binary(builder), Column::Binary(column)) => {
393+
copy_binary(builder, column, indices);
394+
}
395+
(ColumnBuilder::Variant(builder), Column::Variant(column)) => {
396+
copy_binary(builder, column, indices);
397+
}
398+
(ColumnBuilder::Geometry(builder), Column::Geometry(column)) => {
399+
copy_binary(builder, column, indices);
400+
}
401+
(ColumnBuilder::Geography(builder), Column::Geography(column)) => {
402+
copy_binary(builder, &column.0, indices);
403+
}
404+
(ColumnBuilder::String(builder), Column::String(column)) => {
405+
copy_string(builder, column, indices);
406+
}
407+
(ColumnBuilder::Vector(builder), Column::Vector(column)) => {
408+
copy_vector(indices, builder, column);
409+
}
410+
(ColumnBuilder::Opaque(builder), Column::Opaque(column)) => {
411+
copy_opaque(indices, builder, column);
412+
}
413+
(ColumnBuilder::Tuple(builders), Column::Tuple(columns)) => {
414+
for (builder, column) in builders.iter_mut().zip(columns.iter()) {
415+
copy_column(indices, column, builder)
416+
}
417+
}
418+
(to, from) => unreachable!(
419+
"Unsupported column builder type for copy_column. to type: {:?}, from type: {}",
420+
to.data_type(),
421+
from.data_type()
422+
),
423+
},
424+
};
338425
}
339426

340427
fn copy_boolean<I: Index>(to: &mut MutableBitmap, from: &Bitmap, indices: &[I]) {

0 commit comments

Comments
 (0)