Repository navigation
feat: run inline, inline_outer, and stack generators natively #7251
+390
−64
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -51,7 +51,9 @@ | |
| //! * dropping upstream's `ListUnnest` declaration in favor of importing the public one; | ||
| //! * `find_longest_length` applies the empty-list bump once after the row-wise maximum rather | ||
| //! than once per array, which is equivalent because `max` is associative; | ||
| //! * `list_output_lens` and the contiguous-run fast path described above; | ||
| //! * `flatten_struct_cols` folds a struct's null mask into each field it flattens, since the | ||
| //! field values under a null struct are unspecified and Spark's `inline` emits nulls there. | ||
| //! | ||
| //! `ExplodeExec` and `ExplodeStream` were always Comet's own. | ||
|
|
||
|
|
@@ -82,6 +84,7 @@ use datafusion::physical_plan::metrics::{ | |
| SplitMetrics, | ||
| }; | ||
| use datafusion::physical_plan::stream::BatchSplitStream; | ||
| use datafusion_comet_common::children_with_parent_nulls; | ||
| // `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the | ||
| // kernels, upstream exports it publicly. | ||
| use datafusion::physical_plan::unnest::ListUnnest; | ||
|
|
@@ -573,7 +576,9 @@ fn flatten_struct_cols( | |
| Some(_) => match column_data.data_type() { | ||
| DataType::Struct(_) => { | ||
| let struct_arr = column_data.as_any().downcast_ref::<StructArray>().unwrap(); | ||
| // Arrow leaves the field values under a null struct unspecified, so they | ||
| // must not leak out once the struct is flattened. | ||
| columns_expanded.extend(children_with_parent_nulls(struct_arr)?); | ||
| } | ||
| data_type => { | ||
| return internal_err!( | ||
|
|
@@ -1872,4 +1877,42 @@ mod tests { | |
| ); | ||
| assert_output_lens(with_nulls.slice(1, 3), [&[0, 0, 3], &[1, 0, 3], &[1, 1, 3]]); | ||
| } | ||
|
|
||
| #[test] | ||
| fn flatten_struct_cols_nulls_fields_under_a_null_struct() { | ||
|
Member
There was a problem hiding this comment. The second commit dropped the sliced-input case and the non-nullable field from this test, but the PR description still says it covers a sliced input. I think the slice case is worth keeping. The contiguous fast path in
Member
Author
There was a problem hiding this comment. Restored the sliced-input case and the non-nullable field in f5cfc51. |
||
| // Row 1's struct is null but its fields hold values, which Arrow allows. Spark's `inline` | ||
| // turns a null struct into a row of nulls, so none of those values may surface. | ||
| let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])); | ||
| let b: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30, 40])); | ||
| let fields = vec![ | ||
| Arc::new(Field::new("a", DataType::Int32, true)), | ||
| Arc::new(Field::new("b", DataType::Int32, false)), | ||
| ]; | ||
| let structs: ArrayRef = Arc::new(StructArray::new( | ||
| fields.clone().into(), | ||
| vec![a, b], | ||
| Some(NullBuffer::from(vec![true, false, true, true])), | ||
| )); | ||
| let schema = Arc::new(Schema::new(vec![ | ||
| Field::new("a", DataType::Int32, true), | ||
| Field::new("b", DataType::Int32, true), | ||
| ])); | ||
| let indices = HashSet::from([0]); | ||
|
|
||
| let batch = flatten_struct_cols(&[Arc::clone(&structs)], &schema, &indices).unwrap(); | ||
| let a = batch.column(0).as_primitive::<Int32Type>(); | ||
| let b = batch.column(1).as_primitive::<Int32Type>(); | ||
| assert_eq!(a.iter().collect::<Vec<_>>(), [Some(1), None, None, Some(4)]); | ||
| assert_eq!( | ||
| b.iter().collect::<Vec<_>>(), | ||
| [Some(10), None, Some(30), Some(40)] | ||
| ); | ||
|
|
||
| // A slice keeps the struct's mask aligned with its fields. | ||
| let batch = flatten_struct_cols(&[structs.slice(1, 2)], &schema, &indices).unwrap(); | ||
| let a = batch.column(0).as_primitive::<Int32Type>(); | ||
| let b = batch.column(1).as_primitive::<Int32Type>(); | ||
| assert_eq!(a.iter().collect::<Vec<_>>(), [None, None]); | ||
| assert_eq!(b.iter().collect::<Vec<_>>(), [None, Some(30)]); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pass row counts to
RecordBatchOptions;inline(array(_partition,_partition))fails on unpartitioned Iceberg tables when no columns are carried.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Confirmed the native failure at
f5cfc51e29with the extractedflatten_struct_colsfunction and Arrow 59.3.1 from this PR's lockfile. AStructArray::new_empty_fields(2, None)with no carried columns fails withInvalid argument error: must either specify a row count or at least one column. The control usingRecordBatchOptions::with_row_count(Some(2))returns the expected two rows.The Iceberg case is reachable:
CometScanRuleexplicitly exempts metadata columns such as empty_partitionstructs from schema validation, and the attribute/array serdes accept that type. Could we preserve the unnested row count when constructing the batch and add a zero-fieldinline/inline_outerregression? I verified the Arrow failure and traced the serde path, but did not run an Iceberg integration test.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'd lean toward falling back here rather than threading a row count through
flatten_struct_cols. Even with a row count, a zero-column batch still has to leave the native plan and be exported to the JVM, and I'm not sure that path handles a batch with no columns. ReturningUnsupportedfromgetSupportLevelwhen theinlineelement struct has no fields keeps the previous Spark behavior and costs nothing, sincearray<struct<>>carries no data.stackcannot hit this because it always has at least one column. Would a plainSELECT inline(array(struct())) FROM talso reach it? Spark acceptsstruct(), so if Comet converts the zero-fieldCreateNamedStructnatively, that would make a simpler regression test than the Iceberg case.