From efdba13ec96e971605324f9eef181aa6fe418619 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 16:51:29 -0600 Subject: [PATCH 1/3] feat: run inline, inline_outer, and stack generators natively Plan `inline` through the native explode operator by flattening each unnested struct into one column per field, and plan `stack` as `inline` over an array holding one struct per output row. Flattening now folds a struct's null mask into its fields, so a null struct element yields a row of nulls as Spark's `inline` does. Closes #5125. --- .../expression-audits/generator_funcs.md | 12 ++ docs/source/user-guide/latest/expressions.md | 13 +- docs/source/user-guide/latest/operators.md | 12 +- .../latest/understanding-comet-plans.md | 44 +++---- native/core/src/execution/planner.rs | 1 + native/core/src/execution/planner/explode.rs | 36 ++++-- native/operators/src/explode.rs | 68 +++++++++- native/proto/src/proto/operator.proto | 3 + .../apache/spark/sql/comet/operators.scala | 38 +++++- .../sql-tests/expressions/array/inline.sql | 116 ++++++++++++++++++ .../sql-tests/expressions/misc/stack.sql | 87 +++++++++++++ .../comet/exec/CometGenerateExecSuite.scala | 12 ++ 12 files changed, 393 insertions(+), 49 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/inline.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/misc/stack.sql diff --git a/docs/source/contributor-guide/expression-audits/generator_funcs.md b/docs/source/contributor-guide/expression-audits/generator_funcs.md index 3023cfaa774..d4ab82675c2 100644 --- a/docs/source/contributor-guide/expression-audits/generator_funcs.md +++ b/docs/source/contributor-guide/expression-audits/generator_funcs.md @@ -37,4 +37,16 @@ - Same `CometExplodeExec` path as `posexplode`. Compatible for array inputs; empty and NULL arrays both emit one row with null `pos` and null `value` per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. +## inline + +- Same `CometExplodeExec` path as `explode`, with the unnested struct flattened into one column per field. Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1 share one implementation: a NULL struct element yields a row of NULLs, which the native flattening reproduces by folding the struct's validity into each field. Every generated column is reported nullable. + +## inline_outer + +- Same path as `inline`. Empty and NULL arrays both emit one row of NULLs, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. + +## stack + +- Planned as `inline` over an array holding one struct per output row, so `stack(n, v1, ..., vk)` evaluates each argument once per input row, as Spark does, and pads the last row with NULLs. Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1 share one implementation. Arguments the native `array` or `named_struct` cannot handle fall back with the reason `stack arguments are not supported natively`. + [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index f9292b27343..33162f5ec4a 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -336,20 +336,21 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci ## generator_funcs -`explode`, `explode_outer`, `posexplode`, and `posexplode_outer` are supported via -`CometExplodeExec` (operator-level, not expression-level) for array input; map input falls back -to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via +`explode`, `explode_outer`, `posexplode`, `posexplode_outer`, `inline`, `inline_outer`, and +`stack` are supported via `CometExplodeExec` (operator-level, not expression-level). `explode` and +`posexplode` support array input; map input falls back to Spark +([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via `spark.comet.exec.explode.enabled`. | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `explode` | ✅ | — | via `CometExplodeExec` | | `explode_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#explode_outer)) | -| `inline` | 🔜 | — | Operator-level generator (like `explode`) | -| `inline_outer` | 🔜 | — | Operator-level generator (like `explode`) | +| `inline` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#inline)) | +| `inline_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#inline_outer)) | | `posexplode` | ✅ | — | via `CometExplodeExec` | | `posexplode_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#posexplode_outer)) | -| `stack` | 🔜 | — | Operator-level generator | +| `stack` | ✅ | — | via `CometExplodeExec`, planned as `inline` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#stack)) | --- diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 334d34673f2..dc5ecdf743d 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -124,12 +124,12 @@ combined with `OR`, which Spark plans as an existence `BroadcastNestedLoopJoinEx ## Generators and set operations -| Operator | Status | Notes | -| -------------- | ------ | ---------------------------------------------------------------------------------------------------------------- | -| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays. `inline` / `stack` fall back. | -| `ExpandExec` | ✅ | | -| `UnionExec` | ✅ | | -| `CoalesceExec` | ✅ | | +| Operator | Status | Notes | +| -------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------- | +| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays, `inline`, `inline_outer`, and `stack`. Other generators fall back. | +| `ExpandExec` | ✅ | | +| `UnionExec` | ✅ | | +| `CoalesceExec` | ✅ | | ## Writes diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 7789e8eacc2..f7840966970 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -310,28 +310,28 @@ consecutively in a plan, they execute as a single fused block. The write operato (`CometWriteFiles`, `CometNativeWrite` and `CometIcebergWrite`) are the exception: each runs as its own native plan over its child's output. -| Node | Spark equivalent | -| ------------------------------ | --------------------------------------------------------------------------------- | -| `CometProject` | `ProjectExec` | -| `CometFilter` | `FilterExec` | -| `CometSort` | `SortExec` | -| `CometLocalLimit` | `LocalLimitExec` | -| `CometGlobalLimit` | `GlobalLimitExec` | -| `CometExpand` | `ExpandExec` | -| `CometExplode` | `GenerateExec` (for `explode`, `explode_outer`, `posexplode`, `posexplode_outer`) | -| `CometHashAggregate` | `HashAggregateExec`, `ObjectHashAggregateExec` | -| `CometHashJoin` | `ShuffledHashJoinExec` | -| `CometBroadcastHashJoin` | `BroadcastHashJoinExec` | -| `CometBroadcastNestedLoopJoin` | `BroadcastNestedLoopJoinExec` | -| `CometSortMergeJoin` | `SortMergeJoinExec` | -| `CometMergeRows` | `MergeRowsExec` (Spark 3.5+; opt-in) | -| `CometWindowExec` | `WindowExec` | -| `CometWindowGroupLimitExec` | `WindowGroupLimitExec` (Spark 3.5 and later) | -| `CometSample` | `SampleExec` (sampling without replacement) | -| `CometTakeOrderedAndProject` | `TakeOrderedAndProjectExec` | -| `CometWriteFiles` | `WriteFilesExec` (Spark 4.0 and later, experimental native Parquet writes) | -| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental native Parquet writes) | -| `CometIcebergWrite` | `IcebergWrite` (native Iceberg data-file writes, opt-in) | +| Node | Spark equivalent | +| ------------------------------ | -------------------------------------------------------------------------------------------------------------------- | +| `CometProject` | `ProjectExec` | +| `CometFilter` | `FilterExec` | +| `CometSort` | `SortExec` | +| `CometLocalLimit` | `LocalLimitExec` | +| `CometGlobalLimit` | `GlobalLimitExec` | +| `CometExpand` | `ExpandExec` | +| `CometExplode` | `GenerateExec` (for `explode`, `explode_outer`, `posexplode`, `posexplode_outer`, `inline`, `inline_outer`, `stack`) | +| `CometHashAggregate` | `HashAggregateExec`, `ObjectHashAggregateExec` | +| `CometHashJoin` | `ShuffledHashJoinExec` | +| `CometBroadcastHashJoin` | `BroadcastHashJoinExec` | +| `CometBroadcastNestedLoopJoin` | `BroadcastNestedLoopJoinExec` | +| `CometSortMergeJoin` | `SortMergeJoinExec` | +| `CometMergeRows` | `MergeRowsExec` (Spark 3.5+; opt-in) | +| `CometWindowExec` | `WindowExec` | +| `CometWindowGroupLimitExec` | `WindowGroupLimitExec` (Spark 3.5 and later) | +| `CometSample` | `SampleExec` (sampling without replacement) | +| `CometTakeOrderedAndProject` | `TakeOrderedAndProjectExec` | +| `CometWriteFiles` | `WriteFilesExec` (Spark 4.0 and later, experimental native Parquet writes) | +| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental native Parquet writes) | +| `CometIcebergWrite` | `IcebergWrite` (native Iceberg data-file writes, opt-in) | ### JVM-Side Operators diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 0241e1f1ce8..7cdb0e26e6e 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -5088,6 +5088,7 @@ mod tests { outer, position, project_list: vec![bound], + inline: false, })), ..Default::default() }; diff --git a/native/core/src/execution/planner/explode.rs b/native/core/src/execution/planner/explode.rs index 1159f242a53..35b9e0dca9e 100644 --- a/native/core/src/execution/planner/explode.rs +++ b/native/core/src/execution/planner/explode.rs @@ -15,8 +15,9 @@ // specific language governing permissions and limitations // under the License. -//! Builder for the native `Explode` operator, which plans Spark's `explode`, `posexplode`, and -//! their `_outer` variants as a projection followed by Comet's `ExplodeExec`. +//! Builder for the native `Explode` operator, which plans Spark's `explode`, `posexplode`, +//! `inline`, and their `_outer` variants as a projection followed by Comet's `ExplodeExec`. +//! Spark's `stack` arrives here as `inline` over an array of structs built by the JVM serde. use std::sync::Arc; @@ -162,11 +163,30 @@ impl OperatorBuilder for ExplodeBuilder { } }; - output_fields.push(Field::new( - array_field.name(), - element_type, - true, // Element is nullable after unnesting - )); + // `inline` flattens each struct element into one column per field. Every field is + // nullable after unnesting: a null struct element, or the row `inline_outer` emits for a + // NULL or empty array, yields a row of nulls. + let struct_column_indices = if explode.inline { + let DataType::Struct(struct_fields) = &element_type else { + return Err(ExecutionError::GeneralError(format!( + "Expected List of Struct for inline, got {:?}", + array_field.data_type() + ))); + }; + output_fields.extend( + struct_fields + .iter() + .map(|f| f.as_ref().clone().with_nullable(true)), + ); + vec![array_input_index] + } else { + output_fields.push(Field::new( + array_field.name(), + element_type, + true, // Element is nullable after unnesting + )); + vec![] + }; let output_schema = Arc::new(Schema::new(output_fields)); @@ -196,7 +216,7 @@ impl OperatorBuilder for ExplodeBuilder { let unnest_exec = Arc::new(ExplodeExec::new( project_exec, list_unnests, - vec![], // No struct columns to unnest + struct_column_indices, output_schema, unnest_options, )?); diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 6aae0e2d349..ba7f02dbb5a 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -51,15 +51,18 @@ //! * 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. +//! * `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. use arrow::array::{ - new_null_array, Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, - LargeListArray, LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, - StructArray, + make_array, new_null_array, Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, + Int64Array, LargeListArray, LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, + Scalar, StructArray, }; +use arrow::buffer::NullBuffer; use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; @@ -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::().unwrap(); - columns_expanded.extend_from_slice(struct_arr.columns()); + for field in struct_arr.columns() { + columns_expanded.push(with_parent_nulls(field, struct_arr.nulls())?); + } } data_type => { return internal_err!( @@ -587,6 +592,21 @@ fn flatten_struct_cols( Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?) } +/// Returns `field` with every slot null where its parent struct is null. Arrow leaves the field +/// values under a null struct unspecified, so they must not leak out once the struct is flattened. +fn with_parent_nulls(field: &ArrayRef, parent_nulls: Option<&NullBuffer>) -> Result { + let Some(parent_nulls) = parent_nulls.filter(|n| n.null_count() > 0) else { + return Ok(Arc::clone(field)); + }; + // A `Null` field is already all null and carries no validity buffer. + if field.data_type() == &DataType::Null { + return Ok(Arc::clone(field)); + } + let nulls = NullBuffer::union(Some(parent_nulls), field.nulls()); + let data = field.to_data().into_builder().nulls(nulls).build()?; + Ok(make_array(data)) +} + /// This function is used to execute the unnesting on multiple columns all at once, but /// one level at a time, and is called n times, where n is the highest recursion level among /// the unnest exprs in the query. @@ -1872,4 +1892,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() { + // 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::(); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [Some(1), None, None, Some(4)]); + assert_eq!( + b.iter().collect::>(), + [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::(); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [None, None]); + assert_eq!(b.iter().collect::>(), [None, Some(30)]); + } } diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 406592cbce8..b050a8eead4 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1114,6 +1114,9 @@ message Explode { repeated spark.spark_expression.Expr project_list = 3; // Whether to emit a position column alongside the exploded values (posexplode) bool position = 4; + // Whether the array holds structs whose fields become separate output columns (inline). A + // null struct element yields a row of nulls. + bool inline = 5; } message HashJoin { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index da26c91f9c7..7e12f9c8156 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -30,7 +30,7 @@ import org.apache.spark.broadcast.Broadcast import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, CodegenObjectFactoryMode, Expression, ExpressionSet, Generator, LeafExpression, Literal, NamedExpression, SortOrder, XXH64} +import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, CodegenObjectFactoryMode, CreateArray, CreateNamedStruct, Expression, ExpressionSet, Generator, Inline, LeafExpression, Literal, NamedExpression, SortOrder, Stack, XXH64} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, Average, CollectList, CollectSet, Final, ImperativeAggregate, Mode, Partial, PartialMerge, Percentile, Sum} import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} @@ -1920,6 +1920,13 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { if (!op.generator.deterministic) { return Unsupported(Some("Only deterministic generators are supported")) } + op.generator match { + // Spark only accepts an array of structs for `inline`. + case _: Inline => return Compatible() + // Planned as `inline` over an array of structs; see `stackAsStructArray`. + case _: Stack => return Compatible() + case _ => + } if (op.generator.children.length != 1) { return Unsupported(Some("generators with multiple inputs are not supported")) } @@ -1943,10 +1950,19 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { op: GenerateExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - val childExpr = op.generator.children.head + val (childExpr, isInline) = op.generator match { + case inline: Inline => (inline.child, true) + case stack: Stack => (stackAsStructArray(stack), true) + case generator => (generator.children.head, false) + } val childExprProto = exprToProto(childExpr, op.child.output) if (childExprProto.isEmpty) { + if (op.generator.isInstanceOf[Stack]) { + // The array of structs is built here, not in the plan, so its fallback reason would not + // surface on its own. + withFallbackReason(op, "stack arguments are not supported natively") + } return None } @@ -1968,11 +1984,29 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { .setChild(childExprProto.get) .setOuter(op.outer) .setPosition(isPosExplode) + .setInline(isInline) .addAllProjectList(projectExprs.map(_.get).asJava) Some(builder.setExplode(explodeBuilder).build()) } + /** + * `stack(n, v1, ..., vk)` emits `n` rows of `ceil(k / n)` columns, filled row by row and padded + * with nulls. That is `inline` over an array holding one struct per output row. + */ + private def stackAsStructArray(stack: Stack): Expression = { + val numRows = stack.children.head.eval().asInstanceOf[Int] + val values = stack.children.tail + val fields = stack.elementSchema.fields.toSeq + CreateArray(Seq.tabulate(numRows) { row => + CreateNamedStruct(fields.zipWithIndex.flatMap { case (field, col) => + val index = row * fields.length + col + val value = if (index < values.length) values(index) else Literal(null, field.dataType) + Seq(Literal(field.name), value) + }) + }) + } + override def createExec(nativeOp: Operator, op: GenerateExec): CometNativeExec = { CometExplodeExec( nativeOp, diff --git a/spark/src/test/resources/sql-tests/expressions/array/inline.sql b/spark/src/test/resources/sql-tests/expressions/array/inline.sql new file mode 100644 index 00000000000..8d155e00f07 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/inline.sql @@ -0,0 +1,116 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- inline and inline_outer run natively: the planner unnests the array of +-- structs and flattens each struct into one column per field. A NULL struct +-- element yields a row of NULLs. inline_outer emits one row of NULLs for a +-- NULL or empty array, and plain inline drops both. + +statement +CREATE TABLE test_inline(id int, arr array>) USING parquet + +statement +INSERT INTO test_inline VALUES + (1, array(named_struct('a', 1, 'b', 'x'), named_struct('a', 2, 'b', 'y'))), + (2, array(named_struct('a', 3, 'b', NULL))), + (3, array(named_struct('a', NULL, 'b', NULL))), + (4, array(NULL, named_struct('a', 4, 'b', 'z'), NULL)), + (5, array()), + (6, NULL), + (NULL, array(named_struct('a', 5, 'b', 'w'))) + +query +SELECT id, inline(arr) FROM test_inline + +query +SELECT id, inline_outer(arr) FROM test_inline + +query +SELECT id, a, b FROM test_inline LATERAL VIEW inline(arr) t AS a, b + +query +SELECT id, a, b FROM test_inline LATERAL VIEW OUTER inline(arr) t AS a, b + +-- only the generated columns +query +SELECT inline(arr) FROM test_inline + +-- a filter and an aggregate above the generator +query +SELECT id, a FROM test_inline LATERAL VIEW inline(arr) t AS a, b WHERE a > 1 + +query +SELECT b, count(*), sum(a) FROM test_inline LATERAL VIEW OUTER inline(arr) t AS a, b GROUP BY b + +-- a struct field named like a carried column +statement +CREATE TABLE test_inline_dup(id int, arr array>) USING parquet + +statement +INSERT INTO test_inline_dup VALUES + (1, array(named_struct('id', 10, 'v', 'a'), named_struct('id', 11, 'v', 'b'))), + (2, array(named_struct('id', 20, 'v', NULL))) + +query +SELECT * FROM test_inline_dup LATERAL VIEW inline(arr) t AS id2, v + +-- single-field structs +statement +CREATE TABLE test_inline_single(id int, arr array>) USING parquet + +statement +INSERT INTO test_inline_single VALUES + (1, array(named_struct('v', 1L), NULL, named_struct('v', NULL))), + (2, array()) + +query +SELECT id, inline_outer(arr) FROM test_inline_single + +-- nested and wide field types +statement +CREATE TABLE test_inline_nested( + id int, + arr array, + m: map, + s: struct>>) USING parquet + +statement +INSERT INTO test_inline_nested VALUES + (1, array( + named_struct('d', 1.5D, 'dec', 12345.678BD, 'dt', date'2024-01-01', + 'ts', timestamp'2024-01-01 12:34:56', 'bin', X'CAFE', 'ints', array(1, NULL, 3), + 'm', map('k', 1), 's', named_struct('x', 1, 'y', 'one')), + NULL, + named_struct('d', double('NaN'), 'dec', NULL, 'dt', NULL, 'ts', NULL, 'bin', NULL, + 'ints', array(), 'm', NULL, 's', NULL))), + (2, NULL) + +query +SELECT id, inline_outer(arr) FROM test_inline_nested + +-- an array built in the query rather than read from the table +query +SELECT id, inline(array(named_struct('k', id, 'v', 'p'), named_struct('k', id * 10, 'v', 'q'))) FROM test_inline WHERE id IS NOT NULL + +query +SELECT id, inline(arrays_zip(array(1, 2, 3), array('a', 'b'))) FROM test_inline WHERE id = 1 diff --git a/spark/src/test/resources/sql-tests/expressions/misc/stack.sql b/spark/src/test/resources/sql-tests/expressions/misc/stack.sql new file mode 100644 index 00000000000..e4ed7f201b6 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/stack.sql @@ -0,0 +1,87 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- stack(n, v1, ..., vk) emits n rows of ceil(k / n) columns for every input +-- row, filling them row by row and padding the last row with NULLs. Comet plans +-- it as inline over an array holding one struct per output row. + +statement +CREATE TABLE test_stack(id int, x int, y int, s1 string, s2 string) USING parquet + +statement +INSERT INTO test_stack VALUES + (1, 10, 20, 'a', 'b'), + (2, NULL, 21, NULL, 'c'), + (3, 12, NULL, 'd', NULL), + (NULL, NULL, NULL, NULL, NULL) + +-- unpivot two columns into (name, value) pairs +query +SELECT id, stack(2, 'x', x, 'y', y) AS (name, value) FROM test_stack + +query +SELECT id, name, value FROM test_stack LATERAL VIEW stack(2, 'x', x, 'y', y) t AS name, value + +-- uneven argument count: the last row is padded with NULLs +query +SELECT id, stack(3, x, y, id, x, y) FROM test_stack + +-- more rows than values +query +SELECT id, stack(3, x) FROM test_stack + +-- a single row +query +SELECT id, stack(1, x, s1, y, s2) FROM test_stack + +-- one column per row +query +SELECT id, stack(4, s1, s2, s1, s2) FROM test_stack + +-- a NULL literal takes the column's type +query +SELECT id, stack(2, x, NULL) FROM test_stack + +query +SELECT id, stack(2, NULL, s1, 'k', NULL) FROM test_stack + +-- expressions as arguments +query +SELECT id, stack(2, x + 1, upper(s1), y * 2, concat(s1, s2)) FROM test_stack + +-- a filter and an aggregate above the generator +query +SELECT name, count(value), sum(value) FROM test_stack LATERAL VIEW stack(2, 'x', x, 'y', y) t AS name, value WHERE id IS NOT NULL GROUP BY name + +-- only the generated columns +query +SELECT stack(2, x, y) FROM test_stack + +-- complex values +statement +CREATE TABLE test_stack_nested(id int, a1 array, a2 array, st1 struct, st2 struct) USING parquet + +statement +INSERT INTO test_stack_nested VALUES + (1, array(1, 2), array(3), named_struct('p', 1, 'q', 'one'), named_struct('p', 2, 'q', NULL)), + (2, NULL, array(), NULL, named_struct('p', NULL, 'q', 'two')) + +query +SELECT id, stack(2, a1, st1, a2, st2) FROM test_stack_nested + +query +SELECT id, stack(3, a1, a2) FROM test_stack_nested diff --git a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala index 1ccfacdffe5..e4975805f8a 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -237,6 +237,18 @@ class CometGenerateExecSuite extends CometTestBase { } } + test("stack with an argument Comet cannot run natively falls back") { + withSQLConf( + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + spark.udf.register("plus_one", (i: Int) => i + 1) + val df = Seq((1, 10), (2, 20)) + .toDF("id", "v") + .selectExpr("id", "stack(2, v, plus_one(v))") + checkSparkAnswerAndFallbackReason(df, "stack arguments are not supported natively") + } + } + test("explode with nullable projected column") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", From 164e51cd7b98c95748decba62c87004b27a9b20d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 17:06:49 -0600 Subject: [PATCH 2/3] refactor: reuse shared struct-null helper and simplify generator dispatch - Flatten structs with `children_with_parent_nulls` from comet-common, which skips the full revalidation that rebuilding each field did. - Pick the generator in one match on `Inline`, `Stack` and `ExplodeBase` instead of mixing type matches with node-name checks. - Lift the fallback reasons of the array `stack` builds onto the generator, so the operator reports the failing argument's own reason. - Test a foldable `stack` row count and NULL structs built by `IF`. --- .../expression-audits/generator_funcs.md | 2 +- native/operators/src/explode.rs | 68 +++++-------------- .../apache/spark/sql/comet/operators.scala | 60 +++++++--------- .../sql-tests/expressions/array/inline.sql | 4 ++ .../sql-tests/expressions/misc/stack.sql | 4 ++ .../comet/exec/CometGenerateExecSuite.scala | 4 +- 6 files changed, 53 insertions(+), 89 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/generator_funcs.md b/docs/source/contributor-guide/expression-audits/generator_funcs.md index d4ab82675c2..140830c346d 100644 --- a/docs/source/contributor-guide/expression-audits/generator_funcs.md +++ b/docs/source/contributor-guide/expression-audits/generator_funcs.md @@ -47,6 +47,6 @@ ## stack -- Planned as `inline` over an array holding one struct per output row, so `stack(n, v1, ..., vk)` evaluates each argument once per input row, as Spark does, and pads the last row with NULLs. Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1 share one implementation. Arguments the native `array` or `named_struct` cannot handle fall back with the reason `stack arguments are not supported natively`. +- Planned as `inline` over an array holding one struct per output row, so `stack(n, v1, ..., vk)` evaluates each argument once per input row, as Spark does, and pads the last row with NULLs. Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1 share one implementation. An argument that cannot run natively falls back with that argument's own reason. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index ba7f02dbb5a..33abb13836d 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -58,11 +58,10 @@ //! `ExplodeExec` and `ExplodeStream` were always Comet's own. use arrow::array::{ - make_array, new_null_array, Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, - Int64Array, LargeListArray, LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, - Scalar, StructArray, + new_null_array, Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, + LargeListArray, LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, + StructArray, }; -use arrow::buffer::NullBuffer; use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; @@ -85,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; @@ -576,9 +576,9 @@ fn flatten_struct_cols( Some(_) => match column_data.data_type() { DataType::Struct(_) => { let struct_arr = column_data.as_any().downcast_ref::().unwrap(); - for field in struct_arr.columns() { - columns_expanded.push(with_parent_nulls(field, struct_arr.nulls())?); - } + // 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!( @@ -592,21 +592,6 @@ fn flatten_struct_cols( Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?) } -/// Returns `field` with every slot null where its parent struct is null. Arrow leaves the field -/// values under a null struct unspecified, so they must not leak out once the struct is flattened. -fn with_parent_nulls(field: &ArrayRef, parent_nulls: Option<&NullBuffer>) -> Result { - let Some(parent_nulls) = parent_nulls.filter(|n| n.null_count() > 0) else { - return Ok(Arc::clone(field)); - }; - // A `Null` field is already all null and carries no validity buffer. - if field.data_type() == &DataType::Null { - return Ok(Arc::clone(field)); - } - let nulls = NullBuffer::union(Some(parent_nulls), field.nulls()); - let data = field.to_data().into_builder().nulls(nulls).build()?; - Ok(make_array(data)) -} - /// This function is used to execute the unnesting on multiple columns all at once, but /// one level at a time, and is called n times, where n is the highest recursion level among /// the unnest exprs in the query. @@ -1895,39 +1880,18 @@ mod tests { #[test] fn flatten_struct_cols_nulls_fields_under_a_null_struct() { - // 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)), - ]; + // Row 1's struct is null but its field holds a value, which Arrow allows. Spark's + // `inline` turns a null struct into a row of nulls, so the value must not surface. + let field = Arc::new(Field::new("a", DataType::Int32, true)); let structs: ArrayRef = Arc::new(StructArray::new( - fields.clone().into(), - vec![a, b], - Some(NullBuffer::from(vec![true, false, true, true])), + vec![Arc::clone(&field)].into(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + Some(NullBuffer::from(vec![true, false, 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::(); - let b = batch.column(1).as_primitive::(); - assert_eq!(a.iter().collect::>(), [Some(1), None, None, Some(4)]); - assert_eq!( - b.iter().collect::>(), - [Some(10), None, Some(30), Some(40)] - ); + let schema = Arc::new(Schema::new(vec![field])); - // A slice keeps the struct's mask aligned with its fields. - let batch = flatten_struct_cols(&[structs.slice(1, 2)], &schema, &indices).unwrap(); + let batch = flatten_struct_cols(&[structs], &schema, &HashSet::from([0])).unwrap(); let a = batch.column(0).as_primitive::(); - let b = batch.column(1).as_primitive::(); - assert_eq!(a.iter().collect::>(), [None, None]); - assert_eq!(b.iter().collect::>(), [None, Some(30)]); + assert_eq!(a.iter().collect::>(), [Some(1), None, Some(3)]); } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 7e12f9c8156..a43be0e6c3a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -19,8 +19,6 @@ package org.apache.spark.sql.comet -import java.util.Locale - import scala.collection.mutable import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ @@ -30,7 +28,7 @@ import org.apache.spark.broadcast.Broadcast import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, CodegenObjectFactoryMode, CreateArray, CreateNamedStruct, Expression, ExpressionSet, Generator, Inline, LeafExpression, Literal, NamedExpression, SortOrder, Stack, XXH64} +import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, CodegenObjectFactoryMode, CreateArray, CreateNamedStruct, ExplodeBase, Expression, ExpressionSet, Generator, Inline, LeafExpression, Literal, NamedExpression, SortOrder, Stack, XXH64} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, Average, CollectList, CollectSet, Final, ImperativeAggregate, Mode, Partial, PartialMerge, Percentile, Sum} import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} @@ -64,7 +62,7 @@ import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde -import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, liftFallbackReasons, supportedSortType} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.MergeRowsMetricsShim @@ -1921,28 +1919,23 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { return Unsupported(Some("Only deterministic generators are supported")) } op.generator match { - // Spark only accepts an array of structs for `inline`. - case _: Inline => return Compatible() - // Planned as `inline` over an array of structs; see `stackAsStructArray`. - case _: Stack => return Compatible() - case _ => - } - if (op.generator.children.length != 1) { - return Unsupported(Some("generators with multiple inputs are not supported")) - } - val nodeName = op.generator.nodeName.toLowerCase(Locale.ROOT) - if (nodeName != "explode" && nodeName != "posexplode") { - return Unsupported(Some(s"Unsupported generator: ${op.generator.nodeName}")) - } - op.generator.children.head.dataType match { - case _: ArrayType => + // Spark only accepts an array of structs for `inline`, and `stack` is planned as `inline`; + // see `stackAsStructArray`. + case _: Inline | _: Stack => Compatible() - case _: MapType => - // TODO add support for map types - // https://github.com/apache/datafusion-comet/issues/2837 - Unsupported(Some("Comet only supports explode/explode_outer for arrays, not maps")) + case explode: ExplodeBase => + explode.child.dataType match { + case _: ArrayType => + Compatible() + case _: MapType => + // TODO add support for map types + // https://github.com/apache/datafusion-comet/issues/2837 + Unsupported(Some("Comet only supports explode/explode_outer for arrays, not maps")) + case other => + Unsupported(Some(s"Unsupported data type: $other")) + } case other => - Unsupported(Some(s"Unsupported data type: $other")) + Unsupported(Some(s"Unsupported generator: ${other.nodeName}")) } } @@ -1950,19 +1943,18 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { op: GenerateExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - val (childExpr, isInline) = op.generator match { - case inline: Inline => (inline.child, true) - case stack: Stack => (stackAsStructArray(stack), true) - case generator => (generator.children.head, false) + val (childExpr, isPosExplode, isInline) = op.generator match { + case inline: Inline => (inline.child, false, true) + case stack: Stack => (stackAsStructArray(stack), false, true) + case explode: ExplodeBase => (explode.child, explode.position, false) + case other => throw new IllegalStateException(s"Unsupported generator: ${other.nodeName}") } val childExprProto = exprToProto(childExpr, op.child.output) if (childExprProto.isEmpty) { - if (op.generator.isInstanceOf[Stack]) { - // The array of structs is built here, not in the plan, so its fallback reason would not - // surface on its own. - withFallbackReason(op, "stack arguments are not supported natively") - } + // `stack` builds its array of structs here rather than taking it from the plan, so carry any + // reason recorded on the built nodes to the generator, where the operator collects it. + liftFallbackReasons(childExpr, op.generator) return None } @@ -1977,8 +1969,6 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { return None } - val isPosExplode = op.generator.nodeName.toLowerCase(Locale.ROOT) == "posexplode" - val explodeBuilder = OperatorOuterClass.Explode .newBuilder() .setChild(childExprProto.get) diff --git a/spark/src/test/resources/sql-tests/expressions/array/inline.sql b/spark/src/test/resources/sql-tests/expressions/array/inline.sql index 8d155e00f07..95d7c7b0951 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/inline.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/inline.sql @@ -114,3 +114,7 @@ SELECT id, inline(array(named_struct('k', id, 'v', 'p'), named_struct('k', id * query SELECT id, inline(arrays_zip(array(1, 2, 3), array('a', 'b'))) FROM test_inline WHERE id = 1 + +-- NULL struct elements produced by an expression rather than read from Parquet +query +SELECT id, inline_outer(array(IF(id % 2 = 0, NULL, named_struct('a', id, 'b', concat('v', id))), named_struct('a', -1, 'b', NULL))) FROM test_inline diff --git a/spark/src/test/resources/sql-tests/expressions/misc/stack.sql b/spark/src/test/resources/sql-tests/expressions/misc/stack.sql index e4ed7f201b6..98a5b792212 100644 --- a/spark/src/test/resources/sql-tests/expressions/misc/stack.sql +++ b/spark/src/test/resources/sql-tests/expressions/misc/stack.sql @@ -59,6 +59,10 @@ SELECT id, stack(2, x, NULL) FROM test_stack query SELECT id, stack(2, NULL, s1, 'k', NULL) FROM test_stack +-- a row count given as a foldable expression rather than a literal +query +SELECT id, stack(1 + 1, x, y) FROM test_stack + -- expressions as arguments query SELECT id, stack(2, x + 1, upper(s1), y * 2, concat(s1, s2)) FROM test_stack diff --git a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala index e4975805f8a..773affe0298 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -245,7 +245,9 @@ class CometGenerateExecSuite extends CometTestBase { val df = Seq((1, 10), (2, 20)) .toDF("id", "v") .selectExpr("id", "stack(2, v, plus_one(v))") - checkSparkAnswerAndFallbackReason(df, "stack arguments are not supported natively") + checkSparkAnswerAndFallbackReason( + df, + s"${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false") } } From f5cfc51e29f6d252f7b4c314e6f470ccefe5b52e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 27 Oct 2026 06:57:47 -0600 Subject: [PATCH 3/3] address review: restore sliced-input flatten test, drop liftFallbackReasons for stack --- native/operators/src/explode.rs | 39 ++++++++++++++----- .../apache/spark/sql/comet/operators.scala | 5 +-- 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 33abb13836d..70d11933da3 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -1880,18 +1880,39 @@ mod tests { #[test] fn flatten_struct_cols_nulls_fields_under_a_null_struct() { - // Row 1's struct is null but its field holds a value, which Arrow allows. Spark's - // `inline` turns a null struct into a row of nulls, so the value must not surface. - let field = Arc::new(Field::new("a", DataType::Int32, true)); + // 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( - vec![Arc::clone(&field)].into(), - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - Some(NullBuffer::from(vec![true, false, true])), + fields.clone().into(), + vec![a, b], + Some(NullBuffer::from(vec![true, false, true, true])), )); - let schema = Arc::new(Schema::new(vec![field])); + 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::(); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [Some(1), None, None, Some(4)]); + assert_eq!( + b.iter().collect::>(), + [Some(10), None, Some(30), Some(40)] + ); - let batch = flatten_struct_cols(&[structs], &schema, &HashSet::from([0])).unwrap(); + // 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::(); - assert_eq!(a.iter().collect::>(), [Some(1), None, Some(3)]); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [None, None]); + assert_eq!(b.iter().collect::>(), [None, Some(30)]); } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index a43be0e6c3a..23a83f10261 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -62,7 +62,7 @@ import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde -import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, liftFallbackReasons, supportedSortType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.MergeRowsMetricsShim @@ -1952,9 +1952,6 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { val childExprProto = exprToProto(childExpr, op.child.output) if (childExprProto.isEmpty) { - // `stack` builds its array of structs here rather than taking it from the plan, so carry any - // reason recorded on the built nodes to the generator, where the operator collects it. - liftFallbackReasons(childExpr, op.generator) return None }