Skip to content

Commit 871b713

Browse files
committed
feat: Array Literals nested support
1 parent ef92f52 commit 871b713

File tree

1 file changed

+32
-29
lines changed

1 file changed

+32
-29
lines changed

native/core/src/execution/planner.rs

Lines changed: 32 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -85,11 +85,7 @@ use datafusion::physical_expr::window::WindowExpr;
8585
use datafusion::physical_expr::LexOrdering;
8686

8787
use crate::parquet::parquet_exec::init_datasource_exec;
88-
use arrow::array::{
89-
Array, ArrayRef, BinaryBuilder, BooleanArray, Date32Array, Decimal128Array, Float32Array,
90-
Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray, NullArray,
91-
StringBuilder, TimestampMicrosecondArray,
92-
};
88+
use arrow::array::{new_empty_array, Array, ArrayRef, BinaryBuilder, BooleanArray, Date32Array, Decimal128Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray, NullArray, StringBuilder, TimestampMicrosecondArray};
9389
use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer};
9490
use datafusion::common::utils::SingleRowListArrayBuilder;
9591
use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
@@ -2844,7 +2840,7 @@ fn literal_to_array_ref(
28442840
DataType::List(f) if !matches!(f.data_type(), DataType::List(_)) => {
28452841
literal_to_array_ref(f.data_type().clone(), list_literal)
28462842
}
2847-
DataType::List(f) => {
2843+
DataType::List(ref f) => {
28482844
let dt = f.data_type().clone();
28492845
let child_arrays: Vec<ArrayRef> = list_literal
28502846
.list_values
@@ -2867,14 +2863,21 @@ fn literal_to_array_ref(
28672863

28682864
// Concatenate all child arrays' *values* into one array
28692865
// Example: [[1,2,3], [4,5,6]] → values = [1,2,3,4,5,6]
2870-
let concat = arrow::compute::concat(&child_refs)?;
2866+
let output_array = if !child_refs.is_empty() {
2867+
arrow::compute::concat(&child_refs)?
2868+
} else {
2869+
let x = new_empty_array(&dt.clone());
2870+
offsets = vec![x.offset() as i32];
2871+
dbg!(&offsets);
2872+
x
2873+
};
28712874

28722875
// Create and return the parent ListArray
28732876
Ok(Arc::new(ListArray::new(
28742877
// Field: item type matches the concatenated child's type
2875-
FieldRef::from(Field::new("item", concat.data_type().clone(), true)),
2878+
FieldRef::from(Field::new("item", output_array.data_type().clone(), true)),
28762879
OffsetBuffer::new(offsets.into()),
2877-
concat,
2880+
output_array,
28782881
Some(NullBuffer::from(list_literal.null_mask)),
28792882
)))
28802883
}
@@ -3775,26 +3778,26 @@ mod tests {
37753778
true, // outer list nullable
37763779
)));
37773780

3778-
let data = ListLiteral {
3779-
list_values: vec![
3780-
ListLiteral {
3781-
int_values: vec![1, 2],
3782-
null_mask: vec![true, true],
3783-
..Default::default()
3784-
},
3785-
ListLiteral {
3786-
..Default::default()
3787-
}
3788-
],
3789-
null_mask: vec![true, false],
3790-
..Default::default()
3791-
};
3792-
3793-
let nested_type = DataType::List(FieldRef::from(Field::new(
3794-
"item",
3795-
DataType::List(
3796-
Field::new("item", DataType::Int32, true).into(),
3797-
), true)));
3781+
// let data = ListLiteral {
3782+
// list_values: vec![
3783+
// ListLiteral {
3784+
// int_values: vec![1, 2],
3785+
// null_mask: vec![true, true],
3786+
// ..Default::default()
3787+
// },
3788+
// ListLiteral {
3789+
// ..Default::default()
3790+
// }
3791+
// ],
3792+
// null_mask: vec![true, false],
3793+
// ..Default::default()
3794+
// };
3795+
//
3796+
// let nested_type = DataType::List(FieldRef::from(Field::new(
3797+
// "item",
3798+
// DataType::List(
3799+
// Field::new("item", DataType::Int32, true).into(),
3800+
// ), true)));
37983801

37993802
let array = literal_to_array_ref(nested_type, data)?;
38003803

0 commit comments

Comments
 (0)