From efdba13ec96e971605324f9eef181aa6fe418619 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 16:51:29 -0600 Subject: [PATCH 1/4] 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/4] 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 5e80bef5a29f42f59af7ad08116a39c6f6e70037 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 17:44:54 -0600 Subject: [PATCH 3/4] feat: run user-defined generators inside the native generate operator A generator Comet has no native implementation of, such as a Hive UDTF, `json_tuple`, or the `ReplicateRows` that `EXCEPT ALL` plans, now runs on the JVM inside `CometExplodeExec` instead of falling back to Spark. `CometGeneratorUDF` is a `CometUDF` that returns, for each input row, the rows the generator emits as one `array` element, which the operator unnests and flattens as it does for `inline`. Each native plan gets its own generator instance, so state carries across batches. Once the input ends, `ExplodeExec` evaluates a new `terminate` expression and emits the rows of `Generator.terminate` with the last input row's other columns, as Spark's `GenerateExec` does. Gated by `spark.comet.exec.jvmGenerator.enabled`, on by default. Closes #7252. --- docs/source/user-guide/latest/expressions.md | 5 +- docs/source/user-guide/latest/operators.md | 12 +- .../user-guide/latest/scala_java_udfs.md | 11 +- .../latest/understanding-comet-plans.md | 44 +-- native/core/src/execution/planner.rs | 1 + native/core/src/execution/planner/explode.rs | 12 +- native/operators/src/explode.rs | 209 +++++++++++++- native/proto/src/proto/operator.proto | 4 + .../scala/org/apache/comet/CometConf.scala | 11 + .../apache/comet/serde/CometScalaUDF.scala | 18 +- .../spark/sql/comet/CometGeneratorUDF.scala | 255 ++++++++++++++++++ .../comet/execution/arrow/ArrowWriters.scala | 4 +- .../apache/spark/sql/comet/operators.scala | 48 ++-- .../comet/exec/CometGenerateExecSuite.scala | 79 ++++++ .../spark/sql/hive/CometHiveUDTFs.scala | 63 +++++ 15 files changed, 708 insertions(+), 68 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/hive/CometHiveUDTFs.scala diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 33162f5ec4a..2e9c85bb692 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -340,7 +340,8 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci `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`. +`spark.comet.exec.explode.enabled`. Other deterministic generators run on the JVM inside +`CometExplodeExec`; see [Generators](scala_java_udfs.md#generators). | Function | Status | Implementation | Notes | | --- | --- | --- | --- | @@ -376,7 +377,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `get_json_object` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the native path is opt-in via allowIncompatible ([audit](../../contributor-guide/expression-audits/json_funcs.md#get_json_object)) | | `json_array_length` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the native path, which differs for single-quoted JSON, unescaped control characters, and trailing content, is opt-in via allowIncompatible ([audit](../../contributor-guide/expression-audits/json_funcs.md#json_array_length)) | | `json_object_keys` | ✅ | Codegen dispatch | | -| `json_tuple` | 🔜 | — | [#3160](https://github.com/apache/datafusion-comet/issues/3160) | +| `json_tuple` | ✅ | — | Runs on the JVM inside `CometExplodeExec` ([Generators](scala_java_udfs.md#generators)); [#3160](https://github.com/apache/datafusion-comet/issues/3160) tracks a native implementation | | `schema_of_json` | ✅ | Codegen dispatch | | | `to_json` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default, including options and map/array inputs; the native path is opt-in via allowIncompatible ([audit](../../contributor-guide/expression-audits/json_funcs.md#to_json)) | diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index dc5ecdf743d..d61e70671c4 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`, `inline_outer`, and `stack`. Other generators fall back. | -| `ExpandExec` | ✅ | | -| `UnionExec` | ✅ | | -| `CoalesceExec` | ✅ | | +| Operator | Status | Notes | +| -------------- | ------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays, `inline`, `inline_outer`, and `stack`. Other deterministic generators, including Hive UDTFs, run on the JVM inside the operator; see [Generators](scala_java_udfs.md#generators). | +| `ExpandExec` | ✅ | | +| `UnionExec` | ✅ | | +| `CoalesceExec` | ✅ | | ## Writes diff --git a/docs/source/user-guide/latest/scala_java_udfs.md b/docs/source/user-guide/latest/scala_java_udfs.md index 43a30374578..246c610450d 100644 --- a/docs/source/user-guide/latest/scala_java_udfs.md +++ b/docs/source/user-guide/latest/scala_java_udfs.md @@ -21,7 +21,7 @@ Comet executes Spark's Scala and Java [scalar user-defined functions (UDFs)](https://spark.apache.org/docs/latest/sql-ref-functions-udf-scalar.html) within the Comet pipeline. The presence of a UDF does not force the enclosing operator out of the Comet pipeline; surrounding Rust-implemented operators stay in the pipeline. -This page covers Spark's `ScalaUDF` (Scala `udf(...)`, `spark.udf.register(...)` over Scala or Java functional interfaces, and SQL `CREATE FUNCTION ... AS 'com.example.MyUDF'`). Other UDF kinds (Python / Pandas, Hive, aggregate) are out of scope and continue to fall back to Spark. +This page covers Spark's `ScalaUDF` (Scala `udf(...)`, `spark.udf.register(...)` over Scala or Java functional interfaces, and SQL `CREATE FUNCTION ... AS 'com.example.MyUDF'`). Other UDF kinds (Python / Pandas, Hive scalar UDFs, aggregate) are out of scope and continue to fall back to Spark. Hive UDTFs and other generators are covered in [Generators](#generators). This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabled` to `false` to route plans containing a `ScalaUDF` back to Spark for the enclosing operator. @@ -42,8 +42,15 @@ This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabl ## Not supported - Aggregate UDFs (`ScalaAggregator`, `TypedImperativeAggregate`, the legacy `UserDefinedAggregateFunction`). -- Table UDFs and generators. +- Python UDTFs (`@udtf`). - Python `@udf` and Pandas `@pandas_udf`. + +## Generators + +A deterministic generator that Comet has no native implementation of, such as a Hive `GenericUDTF` registered with `CREATE FUNCTION` or Spark's `json_tuple`, runs inside Comet's native generate operator. The operator calls the generator on the JVM for each batch and keeps one generator instance per partition, so state the generator holds carries across batches. Once a partition's input ends, the rows of the generator's `terminate()` (a Hive UDTF's `close()`) follow, carrying the partition's last input row, as in Spark. + +Set `spark.comet.exec.jvmGenerator.enabled` to `false` to run such a generator in Spark instead. + - Hive `GenericUDF` and `SimpleUDF`. - `NullType` and `UserDefinedType` arguments and return types. UDT-typed columns fall back to Spark; to keep execution in the Comet pipeline, store and read the underlying representation directly (e.g. write MLlib `Vector` outputs as `Struct, values: Array>` rather than `VectorUDT`). - Trees whose total nested-field count (output plus all input columns the UDF tree references) exceeds `spark.sql.codegen.maxFields` (default 100). Comet refuses these at plan time and the operator falls back to Spark. diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index f7840966970..bbe4102e65d 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`, `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) | +| Node | Spark equivalent | +| ------------------------------ | ----------------------------------------------------------------------------------- | +| `CometProject` | `ProjectExec` | +| `CometFilter` | `FilterExec` | +| `CometSort` | `SortExec` | +| `CometLocalLimit` | `LocalLimitExec` | +| `CometGlobalLimit` | `GlobalLimitExec` | +| `CometExpand` | `ExpandExec` | +| `CometExplode` | `GenerateExec`; generators without a native implementation run on the JVM inside it | +| `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 7cdb0e26e6e..4861b7be6bd 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -5089,6 +5089,7 @@ mod tests { position, project_list: vec![bound], inline: false, + terminate: None, })), ..Default::default() }; diff --git a/native/core/src/execution/planner/explode.rs b/native/core/src/execution/planner/explode.rs index 35b9e0dca9e..70bb69e9e12 100644 --- a/native/core/src/execution/planner/explode.rs +++ b/native/core/src/execution/planner/explode.rs @@ -17,7 +17,8 @@ //! 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. +//! Spark's `stack` arrives here as `inline` over an array of structs built by the JVM serde, and +//! any other generator as `inline` over the rows a JVM `CometUDF` produces for each input row. use std::sync::Arc; @@ -213,13 +214,18 @@ impl OperatorBuilder for ExplodeBuilder { let unnest_options = UnnestOptions::new().with_null_handling(null_handling); // Comet's specialized fork of `UnnestExec`; see `operators::explode`. - let unnest_exec = Arc::new(ExplodeExec::new( + let mut unnest_exec = ExplodeExec::new( project_exec, list_unnests, struct_column_indices, output_schema, unnest_options, - )?); + )?; + if let Some(terminate) = &explode.terminate { + let terminate = planner.create_expr(terminate, Arc::new(Schema::empty()))?; + unnest_exec = unnest_exec.with_terminate(terminate)?; + } + let unnest_exec = Arc::new(unnest_exec); Ok(( scans, diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 33abb13836d..b4860f2fef8 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -64,16 +64,16 @@ use arrow::array::{ }; use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; -use arrow::compute::{cast, is_not_null, kernels, sum}; -use arrow::datatypes::{DataType, Int64Type, SchemaRef}; -use arrow::record_batch::RecordBatch; +use arrow::compute::{cast, is_not_null, kernels, sum, take_record_batch}; +use arrow::datatypes::{DataType, Int64Type, Schema, SchemaRef, UInt32Type}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; // Upstream imports these as `arrow_ord::cmp::{eq, lt}`; Comet reaches them through `arrow`, // which does not have `arrow_ord` as a direct dependency. use arrow::compute::kernels::cmp::{eq, lt}; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{ - exec_datafusion_err, exec_err, internal_err, Constraints, HashMap, HashSet, Result, - UnnestOptions, + exec_datafusion_err, exec_err, internal_err, Constraints, HashMap, HashSet, NullHandling, + Result, UnnestOptions, }; use datafusion::execution::TaskContext; use datafusion::physical_expr::equivalence::ProjectionMapping; @@ -107,6 +107,8 @@ pub struct ExplodeExec { list_column_indices: Vec, struct_column_indices: Vec, options: UnnestOptions, + /// See [`ExplodeExec::with_terminate`]. + terminate: Option>, metrics: ExecutionPlanMetricsSet, cache: Arc, } @@ -132,11 +134,24 @@ impl ExplodeExec { list_column_indices, struct_column_indices, options, + terminate: None, metrics: ExecutionPlanMetricsSet::new(), cache: Arc::new(cache), }) } + /// Unnest one more row once the input is exhausted: `terminate`, evaluated against a one-row + /// batch with no columns, takes the place of the list column in the last input row. This is + /// how a JVM generator's `terminate()` output joins the last row's other columns, as in + /// Spark's `GenerateExec`. Requires exactly one list column to unnest. + pub fn with_terminate(mut self, terminate: Arc) -> Result { + if self.list_column_indices.len() != 1 { + return internal_err!("ExplodeExec terminate requires exactly one list column"); + } + self.terminate = Some(terminate); + Ok(self) + } + /// Compute the plan properties, keeping whatever the child guarantees about the columns /// that unnesting passes through untouched. /// @@ -235,11 +250,13 @@ impl ExecutionPlan for ExplodeExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // ExplodeExec holds no physical expressions of its own; the columns to unnest are - // identified by index, so there is nothing to visit here. - Ok(TreeNodeRecursion::Continue) + // The columns to unnest are identified by index, so `terminate` is the only expression. + match &self.terminate { + Some(terminate) => f(terminate), + None => Ok(TreeNodeRecursion::Continue), + } } fn with_new_children( @@ -249,13 +266,17 @@ impl ExecutionPlan for ExplodeExec { if children.len() != 1 { return internal_err!("ExplodeExec expects exactly one child"); } - Ok(Arc::new(ExplodeExec::new( + let mut exec = ExplodeExec::new( Arc::clone(&children[0]), self.list_column_indices.clone(), self.struct_column_indices.clone(), Arc::clone(&self.schema), self.options.clone(), - )?)) + )?; + if let Some(terminate) = &self.terminate { + exec = exec.with_terminate(Arc::clone(terminate))?; + } + Ok(Arc::new(exec)) } fn execute( @@ -277,6 +298,8 @@ impl ExecutionPlan for ExplodeExec { input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition), batch_size, pending_input: None, + terminate: self.terminate.clone(), + last_row: None, }); // Chunking the input bounds each build to roughly `batch_size` rows, but two cases can @@ -372,6 +395,10 @@ struct ExplodeStream { /// longer than `batch_size`, and recursive unnesting (where the expansion cannot be /// predicted up front), both still materialize their full expansion in one build. pending_input: Option, + /// See [`ExplodeExec::with_terminate`]. Taken once it has run. + terminate: Option>, + /// The last input row, copied out of its batch, kept while `terminate` is pending. + last_row: Option, } impl RecordBatchStream for ExplodeStream { @@ -440,6 +467,14 @@ impl ExplodeStream { self.input_rows.add(batch.num_rows()); if batch.num_rows() > 0 { let timer = self.baseline_metrics.elapsed_compute().timer(); + if self.terminate.is_some() { + // A copy rather than a slice, so the rest of the batch is not held + // until the input ends. + let last = PrimitiveArray::::from(vec![ + batch.num_rows() as u32 - 1, + ]); + self.last_row = Some(take_record_batch(&batch, &last)?); + } let output_lens = self.predict_output_lens(&batch); timer.done(); self.pending_input = Some(PendingInput { @@ -458,6 +493,18 @@ impl ExplodeStream { if other.is_none() { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + if let Some(terminate) = self.terminate.take() { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = self.terminate_batch(&terminate); + timer.done(); + if let Some(batch) = result.transpose() { + if let Ok(batch) = &batch { + batch.record_output(&self.baseline_metrics); + } + return Poll::Ready(Some(batch)); + } + } } return Poll::Ready(other); } @@ -465,6 +512,40 @@ impl ExplodeStream { } } + /// Unnest the rows `terminate` yields in place of the last input row's list. With no input + /// rows, the other columns are null. Spark's `outer` handling does not apply to these rows, so + /// an empty result adds nothing. + fn terminate_batch( + &mut self, + terminate: &Arc, + ) -> Result> { + let options = RecordBatchOptions::new().with_row_count(Some(1)); + let no_columns = + RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options)?; + let rows = terminate.evaluate(&no_columns)?.into_array(1)?; + + let input_schema = self.input.schema(); + let mut columns = match self.last_row.take() { + Some(row) => row.columns().to_vec(), + None => input_schema + .fields() + .iter() + .map(|f| new_null_array(f.data_type(), 1)) + .collect(), + }; + columns[self.list_type_columns[0].index_in_input_schema] = rows; + let batch = RecordBatch::try_new(input_schema, columns)?; + + build_batch( + &batch, + &self.schema, + &self.list_type_columns, + &self.struct_column_indices, + &UnnestOptions::new().with_null_handling(NullHandling::Drop), + None, + ) + } + /// Compute how many output rows each input row of `batch` will expand into, so the input /// can be chunked to keep each build bounded. /// @@ -1574,6 +1655,8 @@ mod tests { input_rows: MetricBuilder::new(&metrics).counter("input_rows", 0), batch_size: 4, pending_input: None, + terminate: None, + last_row: None, }; let mut rows = 0; @@ -1878,6 +1961,110 @@ mod tests { assert_output_lens(with_nulls.slice(1, 3), [&[0, 0, 3], &[1, 0, 3], &[1, 1, 3]]); } + /// Explode column "l" of `(id, l)` batches, ending with the rows `terminate` yields. + async fn explode_with_terminate( + input: Vec, + terminate: Vec, + null_handling: NullHandling, + ) -> Vec<(Option, Option)> { + let input_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("l", DataType::new_list(DataType::Int32, true), true), + ])); + let source = + MemorySourceConfig::try_new_exec(&[input], Arc::clone(&input_schema), None).unwrap(); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("l", DataType::Int32, true), + ])); + let rows = ListArray::from_iter_primitive::(vec![Some( + terminate.into_iter().map(Some).collect::>(), + )]); + let terminate = Arc::new(datafusion::physical_expr::expressions::Literal::new( + datafusion::common::ScalarValue::List(Arc::new(rows)), + )); + let explode = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 1, + depth: 1, + }], + vec![], + output_schema, + UnnestOptions::new().with_null_handling(null_handling), + ) + .unwrap() + .with_terminate(terminate) + .unwrap(); + let batches = + datafusion::physical_plan::common::collect(explode.execute(0, Arc::default()).unwrap()) + .await + .unwrap(); + batches + .iter() + .flat_map(|b| { + let id = b.column(0).as_primitive::(); + let l = b.column(1).as_primitive::(); + id.iter().zip(l.iter()).collect::>() + }) + .collect() + } + + fn id_list_batch(rows: Vec<(i32, Vec)>) -> RecordBatch { + let (ids, lists): (Vec<_>, Vec<_>) = rows.into_iter().unzip(); + let lists = ListArray::from_iter_primitive::( + lists + .into_iter() + .map(|l| Some(l.into_iter().map(Some).collect::>())), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("l", lists.data_type().clone(), true), + ])); + RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(ids)), Arc::new(lists)], + ) + .unwrap() + } + + #[tokio::test] + async fn terminate_rows_carry_the_last_input_row() { + let input = vec![ + id_list_batch(vec![(1, vec![10]), (2, vec![])]), + id_list_batch(vec![(3, vec![30, 31])]), + ]; + assert_eq!( + explode_with_terminate(input, vec![100, 101], NullHandling::Drop).await, + [ + (Some(1), Some(10)), + (Some(3), Some(30)), + (Some(3), Some(31)), + (Some(3), Some(100)), + (Some(3), Some(101)), + ] + ); + } + + #[tokio::test] + async fn terminate_without_input_rows_has_null_columns() { + let input = vec![id_list_batch(vec![])]; + assert_eq!( + explode_with_terminate(input, vec![100], NullHandling::Drop).await, + [(None, Some(100))] + ); + } + + #[tokio::test] + async fn empty_terminate_adds_no_outer_row() { + // `outer` keeps an input row whose list is empty, but not an empty `terminate()`. + let input = vec![id_list_batch(vec![(1, vec![])])]; + assert_eq!( + explode_with_terminate(input, vec![], NullHandling::PreserveAndExpandEmpty).await, + [(Some(1), None)] + ); + } + #[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 diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index b050a8eead4..ce75904a884 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1117,6 +1117,10 @@ message Explode { // Whether the array holds structs whose fields become separate output columns (inline). A // null struct element yields a row of nulls. bool inline = 5; + // Evaluated once after the input is exhausted, against a one-row batch with no columns. Yields + // an array of the same type as `child` holding the rows a JVM generator's `terminate()` emits, + // which carry the last input row's projected columns. + spark.spark_expression.Expr terminate = 6; } message HashJoin { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6003703a630..75eba338879 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -780,6 +780,17 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_JVM_GENERATOR_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.jvmGenerator.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether to run a generator that Comet has no native implementation of, such as a " + + "Hive UDTF or `json_tuple`, inside Comet's native generate operator, which calls the " + + "generator on the JVM for each batch. When disabled, such a generator falls back to " + + "Spark.") + .booleanConf + .createWithDefault(true) + val COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.native.partitioning.hash.enabled") .withAlternative("spark.comet.native.shuffle.partitioning.hash.enabled") diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 7244875e2cb..98fb495427a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -116,6 +116,18 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { case _ => false } + /** + * `expr` serialized with Spark's closure serializer, which respects the task context + * ClassLoader so classes from user jars resolve, and is how the JVM side of a `JvmScalarUdf` + * reads it back. Throws when the tree holds something the serializer refuses. + */ + def closureSerialize(expr: Expression): Array[Byte] = { + val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(expr) + val serialized = new Array[Byte](buffer.remaining()) + buffer.get(serialized) + serialized + } + /** * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through * [[CometScalaUDFCodegen]] so that native execution evaluates the expression inside the @@ -202,11 +214,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // kernels and state on the executor (see `DispatchOccurrence`). val bytes = try { - val serializer = SparkEnv.get.closureSerializer.newInstance() - val buffer = serializer.serialize(DispatchOccurrence.tag(boundExpr)) - val serialized = new Array[Byte](buffer.remaining()) - buffer.get(serialized) - serialized + closureSerialize(DispatchOccurrence.tag(boundExpr)) } catch { // `NonFatal` rather than `NotSerializableException`: Java serialization reports an // unserializable object graph as one of several exception types depending on where in diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala new file mode 100644 index 00000000000..f1b6822b028 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala @@ -0,0 +1,255 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import java.nio.ByteBuffer +import java.util.concurrent.atomic.AtomicLong + +import scala.collection.mutable +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.apache.arrow.vector.{BitVector, ValueVector, VarBinaryVector} +import org.apache.arrow.vector.types.pojo.Field +import org.apache.spark.SparkEnv +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Generator, GenericInternalRow, Literal, Nondeterministic, Unevaluable} +import org.apache.spark.sql.catalyst.util.GenericArrayData +import org.apache.spark.sql.comet.execution.arrow.ArrowWriter +import org.apache.spark.sql.execution.ExecSubqueryExpression +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.comet.CometConf +import org.apache.comet.codegen.CometBatchKernelCodegen +import org.apache.comet.serde.{CometInvokeTargets, ExprOuterClass} +import org.apache.comet.serde.CometScalaUDF.closureSerialize +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} +import org.apache.comet.udf.{CometUDF, NativeUdfCall} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen +import org.apache.comet.vector.CometVector + +/** + * Runs a Spark [[Generator]] for Comet's native generate operator. For each input row it returns + * the rows the generator emits as one `array` element, which the native operator unnests + * and flattens the way it does for `inline`. An input row that emits nothing gets an empty array, + * so `outer` keeps it as one row of nulls. + * + * Arg 0 is the digest of the closure-serialized bound generator, arg 1 the serialized bytes, and + * arg 2 a boolean scalar that is true for the call the native operator makes once its input is + * exhausted. That call returns one element holding the rows of `Generator.terminate`. Args 3..N + * are the columns the bound generator reads, in ordinal order. + * + * A generator can keep state across rows and emit rows from `terminate`, as a Hive UDTF does, so + * every native plan gets its own instance: the key is the digest and the plan, and `releasePlan` + * drops the plan's instances. The serde ships the bound generator rather than evaluating it on + * the driver, so this works in cluster mode without driver-side state. `evaluate` is synchronized + * for the reason `CometScalaUDFCodegen.evaluate` is. + */ +class CometGeneratorUDF extends CometUDF { + import CometGeneratorUDF._ + + private val generators = mutable.HashMap.empty[(ByteBuffer, Long), BoundGenerator] + + override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = + evaluate(inputs, numRows, partitionIndex = 0, planId = NoPlan) + + override def evaluate( + inputs: Array[ValueVector], + numRows: Int, + partitionIndex: Int, + planId: Long): ValueVector = this.synchronized { + require( + inputs.length >= 3, + "CometGeneratorUDF requires the generator digest, its serialized bytes and the terminate " + + s"flag, got ${inputs.length} inputs") + val digest = ByteBuffer.wrap(binaryScalar(inputs(0))) + val bound = generators.getOrElseUpdate( + (digest, planId), + BoundGenerator(deserialize(binaryScalar(inputs(1))), partitionIndex)) + + val output = CometBatchKernelCodegen.allocateOutput(bound.outputField, numRows, 0) + try { + val writer = ArrowWriter.createFieldWriter(output) + val outputRow = new GenericInternalRow(1) + def writeRows(rows: Any): Unit = { + val copies = new ArrayBuffer[Any] + iteratorOf(rows).foreach(row => copies += row.copy()) + outputRow.update(0, new GenericArrayData(copies.toArray)) + writer.write(outputRow, 0) + } + + if (inputs(2).asInstanceOf[BitVector].get(0) == 1) { + writeRows(bound.generator.terminate()) + } else { + val columns: Array[ColumnVector] = + inputs.drop(3).map(v => CometVector.getVector(v, null): ColumnVector) + val batch = new ColumnarBatch(columns, numRows) + var i = 0 + while (i < numRows) { + writeRows(bound.generator.eval(batch.getRow(i))) + i += 1 + } + } + writer.finish() + output + } catch { + case t: Throwable => + output.close() + throw t + } + } + + override def releasePlan(planId: Long): Unit = this.synchronized { + generators.keys.filter(_._2 == planId).toList.foreach(generators.remove) + } + + private def deserialize(bytes: Array[Byte]): Generator = { + val loader = Option(Thread.currentThread().getContextClassLoader) + .getOrElse(classOf[Generator].getClassLoader) + SparkEnv.get.closureSerializer + .newInstance() + .deserialize[Generator](ByteBuffer.wrap(bytes), loader) + } +} + +object CometGeneratorUDF { + + /** The plan id of calls from outside a native plan, such as unit tests. */ + private val NoPlan = -1L + + /** Numbers each generator the serde ships; see [[calls]]. */ + private val nextOccurrence = new AtomicLong() + + /** Why `generator` cannot run through [[CometGeneratorUDF]], or `None` when it can. */ + def declineReason(generator: Generator): Option[String] = { + val name = generator.nodeName + if (!CometConf.COMET_EXEC_JVM_GENERATOR_ENABLED.get()) { + return Some( + s"Unsupported generator: $name. Set ${CometConf.COMET_EXEC_JVM_GENERATOR_ENABLED.key}" + + "=true to run it on the JVM inside Comet") + } + // The generator is serialized at plan time, before any subquery in it has run, and an + // unevaluable or native-only node has nothing to run on the JVM. Attributes are unevaluable + // only until they are bound, which `calls` does. + generator + .find { + case _: Attribute => false + case _: ExecSubqueryExpression | _: Unevaluable | _: NativeUdfCall => true + case _ => false + } + .map(node => s"$name: ${node.nodeName} cannot run inside a JVM generator") + .orElse(CometInvokeTargets.declineReason(generator)) + .orElse(Some(s"$name: unsupported output type ${outputType(generator).catalogString}") + .filter(_ => serializeDataType(outputType(generator)).isEmpty)) + } + + /** + * The `JvmScalarUdf` calls that run `generator` through [[CometGeneratorUDF]]: one evaluated + * for each input batch, and the one the native operator makes once its input is exhausted. + * Fails with the reason when `generator` cannot be shipped. + */ + def calls( + generator: Generator, + inputs: Seq[Attribute]): Either[String, (ExprOuterClass.Expr, ExprOuterClass.Expr)] = { + val name = generator.nodeName + // Bind against only the attributes the generator reads, which become the data args. + val attrs = generator.collect { case a: AttributeReference => a }.distinct + val bound = BindReferences.bindReference(generator, AttributeSeq(attrs)) + val bytes = + try { + closureSerialize(bound) + } catch { + case NonFatal(e) => + return Left( + s"$name could not be closure-serialized (${e.getClass.getSimpleName}: " + + s"${e.getMessage})") + } + // The digest also covers a number unique to this generator, so two identical generators in + // one plan get their own instances and state. + val occurrence = ByteBuffer.allocate(8).putLong(nextOccurrence.getAndIncrement()).array() + val digest = CometScalaUDFCodegen.digest(bytes ++ occurrence) + val returnType = serializeDataType(outputType(generator)).getOrElse { + return Left(s"$name: unsupported output type ${outputType(generator).catalogString}") + } + + def call( + terminate: Boolean, + dataArgs: Seq[Expression]): Either[String, ExprOuterClass.Expr] = { + val args = (Seq(Literal(digest), Literal(bytes), Literal(terminate)) ++ dataArgs) + .map(exprToProto(_, inputs)) + if (args.exists(_.isEmpty)) { + Left(s"$name: an argument could not be passed to the JVM generator") + } else { + val udf = ExprOuterClass.JvmScalarUdf + .newBuilder() + .setClassName(classOf[CometGeneratorUDF].getName) + .addAllArgs(args.flatten.asJava) + .setReturnType(returnType) + .setReturnNullable(false) + Right(ExprOuterClass.Expr.newBuilder().setJvmScalarUdf(udf).build()) + } + } + for { + process <- call(terminate = false, attrs) + terminate <- call(terminate = true, Nil) + } yield (process, terminate) + } + + /** + * A generator ready to run in one native plan, and the Arrow field its output is exported as. + */ + private case class BoundGenerator(generator: Generator, outputField: Field) + + private object BoundGenerator { + def apply(generator: Generator, partitionIndex: Int): BoundGenerator = { + generator.foreach { + case n: Nondeterministic => n.initialize(partitionIndex) + case _ => + } + BoundGenerator(generator, outputField(generator)) + } + } + + /** The type of each element of the output: the rows one input row generates. */ + def outputType(generator: Generator): ArrayType = + ArrayType(generator.elementSchema, containsNull = false) + + private def outputField(generator: Generator): Field = + CometBatchKernelCodegen.toFfiArrowField("generated", outputType(generator), nullable = false) + + private def binaryScalar(v: ValueVector): Array[Byte] = { + val vec = v.asInstanceOf[VarBinaryVector] + require( + vec.getValueCount >= 1 && !vec.isNull(0), + "CometGeneratorUDF requires a binary scalar") + vec.get(0) + } + + /** `Generator.eval` returns a `TraversableOnce` on Scala 2.12 and an `IterableOnce` on 2.13. */ + private def iteratorOf(rows: Any): Iterator[InternalRow] = rows match { + case it: Iterator[_] => it.asInstanceOf[Iterator[InternalRow]] + case it: Iterable[_] => it.iterator.asInstanceOf[Iterator[InternalRow]] + case opt: Option[_] => opt.iterator.asInstanceOf[Iterator[InternalRow]] + case other => + throw new IllegalStateException(s"Unexpected generator output ${other.getClass.getName}") + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala index 53b4ae7888d..c0e865302f8 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala @@ -43,7 +43,7 @@ import org.apache.spark.unsafe.Platform * potential conflicts with Spark's Arrow dependencies, hence we cannot reuse Spark's ArrowWriter * directly. */ -private[arrow] object ArrowWriter { +private[comet] object ArrowWriter { def create(root: VectorSchemaRoot, fixedWidthCapacity: Int): ArrowWriter = { require(fixedWidthCapacity >= 0, "Fixed-width capacity must be non-negative") val children = root.getFieldVectors().asScala.map { vector => @@ -668,7 +668,7 @@ private[arrow] final class DictionaryCache { } } -private[arrow] abstract class ArrowFieldWriter { +private[comet] abstract class ArrowFieldWriter { def valueVector: ValueVector 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..554fa5f65a2 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 @@ -1935,7 +1935,10 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { Unsupported(Some(s"Unsupported data type: $other")) } case other => - Unsupported(Some(s"Unsupported generator: ${other.nodeName}")) + CometGeneratorUDF.declineReason(other) match { + case Some(reason) => Unsupported(Some(reason)) + case None => Compatible() + } } } @@ -1943,18 +1946,37 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { op: GenerateExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - 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 explodeBuilder = OperatorOuterClass.Explode.newBuilder().setOuter(op.outer) + val childExprProto = op.generator match { + case inline: Inline => + explodeBuilder.setInline(true) + exprToProto(inline.child, op.child.output) + case stack: Stack => + explodeBuilder.setInline(true) + val array = stackAsStructArray(stack) + exprToProto(array, op.child.output).orElse { + // The array of structs is built here rather than taken from the plan, so carry any + // reason recorded on the built nodes to the generator, where the operator collects it. + liftFallbackReasons(array, op.generator) + None + } + case explode: ExplodeBase => + explodeBuilder.setPosition(explode.position) + exprToProto(explode.child, op.child.output) + case generator => + // The JVM generator yields, for each input row, an array of the structs it emits. + explodeBuilder.setInline(true) + CometGeneratorUDF.calls(generator, op.child.output) match { + case Right((process, terminate)) => + explodeBuilder.setTerminate(terminate) + Some(process) + case Left(reason) => + withFallbackReason(op, reason) + None + } } - 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 } @@ -1969,12 +1991,8 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { return None } - val explodeBuilder = OperatorOuterClass.Explode - .newBuilder() + explodeBuilder .setChild(childExprProto.get) - .setOuter(op.outer) - .setPosition(isPosExplode) - .setInline(isInline) .addAllProjectList(projectExprs.map(_.get).asJava) Some(builder.setExplode(explodeBuilder).build()) 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 773affe0298..154e53930aa 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -24,6 +24,7 @@ import org.apache.spark.sql.comet.CometExplodeExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.functions.col +import org.apache.spark.sql.hive.{CometHiveUDTFs, CountingUDTF} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -736,4 +737,82 @@ class CometGenerateExecSuite extends CometTestBase { } } + /** Run `f` over a single-file Parquet table `t` of `(id, s, js)` rows. */ + private def withGeneratorTable(f: => Unit): Unit = { + val rows = Seq( + (1, "a", """{"a": 1, "b": "x"}"""), + (2, null, """{"a": null}"""), + (3, "c", "not json"), + (4, "d", null)) + withTempPath { path => + rows.toDF("id", "s", "js").repartition(1).write.parquet(path.getAbsolutePath) + withParquetTable(path.getAbsolutePath, "t")(f) + } + } + + /** Check `df` against Spark, and that its generator ran in Comet's native operator. */ + private def checkSparkAnswerWithNativeGenerate(df: DataFrame): Unit = { + val (_, cometPlan) = checkSparkAnswer(df) + assert( + collectFirst(cometPlan) { case generate: CometExplodeExec => generate }.isDefined, + s"Expected a native generator in\n$cometPlan") + } + + test("Hive UDTF runs natively and emits its terminate rows") { + CometHiveUDTFs.register(spark, "counting_udtf", classOf[CountingUDTF]) + withGeneratorTable { + // The row `close` emits carries the last input row's `id`, as in Spark. + checkSparkAnswerAndOperator(sql("SELECT id, counting_udtf(s) FROM t")) + checkSparkAnswerAndOperator( + sql("SELECT id, value, n FROM t LATERAL VIEW OUTER counting_udtf(s) u AS value, n")) + // With no input rows only the `close` row remains. + checkSparkAnswerAndOperator(sql("SELECT counting_udtf(s) FROM t WHERE id > 10")) + // Two identical UDTFs in one plan each keep their own count. + checkSparkAnswerAndOperator(sql(""" + |SELECT id, a.value, a.n, b.value, b.n FROM t + |LATERAL VIEW counting_udtf(s) a AS value, n + |LATERAL VIEW counting_udtf(s) b AS value, n""".stripMargin)) + } + } + + test("Hive UDTF keeps separate state per partition") { + CometHiveUDTFs.register(spark, "counting_udtf", classOf[CountingUDTF]) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "3") { + withGeneratorTable { + // Each partition ends with one `close` row holding its own count, so the counts add up + // to the input rows whichever partition each row lands in. + checkSparkAnswerWithNativeGenerate(sql(""" + |SELECT count(*), sum(n) FROM ( + | SELECT counting_udtf(s) FROM (SELECT * FROM t DISTRIBUTE BY id)) + |WHERE value = 'total'""".stripMargin)) + } + } + } + + test("json_tuple runs natively") { + withGeneratorTable { + checkSparkAnswerAndOperator(sql("SELECT id, json_tuple(js, 'a', 'b') FROM t")) + checkSparkAnswerAndOperator( + sql("SELECT id, a FROM t LATERAL VIEW OUTER json_tuple(js, 'a') j AS a WHERE a IS NULL")) + } + } + + test("EXCEPT ALL replicates rows through a JVM generator") { + withGeneratorTable { + checkSparkAnswerWithNativeGenerate( + sql("SELECT s FROM t UNION ALL SELECT s FROM t EXCEPT ALL SELECT 'a'")) + } + } + + test("a generator without a native implementation falls back when the JVM path is off") { + withSQLConf(CometConf.COMET_EXEC_JVM_GENERATOR_ENABLED.key -> "false") { + withGeneratorTable { + checkSparkAnswerAndFallbackReason( + sql("SELECT id, json_tuple(js, 'a') FROM t"), + "Unsupported generator: JsonTuple") + } + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/hive/CometHiveUDTFs.scala b/spark/src/test/scala/org/apache/spark/sql/hive/CometHiveUDTFs.scala new file mode 100644 index 00000000000..8c9b7684230 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/hive/CometHiveUDTFs.scala @@ -0,0 +1,63 @@ +/* + * 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. + */ + +package org.apache.spark.sql.hive + +import java.util.Arrays + +import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF +import org.apache.hadoop.hive.serde2.objectinspector.{ObjectInspector, ObjectInspectorFactory, StructObjectInspector} +import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo} +import org.apache.spark.sql.hive.HiveShim.HiveFunctionWrapper + +/** Registers Hive UDTFs in a session that has no Hive support, as `HiveSessionCatalog` would. */ +object CometHiveUDTFs { + def register(spark: SparkSession, name: String, udtfClass: Class[_ <: GenericUDTF]): Unit = + spark.sessionState.functionRegistry.registerFunction( + FunctionIdentifier(name), + new ExpressionInfo(udtfClass.getName, name), + (children: Seq[Expression]) => + HiveGenericUDTF(name, new HiveFunctionWrapper(udtfClass.getName), children)) +} + +/** + * Emits each input value with a running count of the rows it has processed, and from `close` one + * more row holding the total. The count is per instance, so it shows whether the rows of a + * partition reach one instance, and the `close` row whether `terminate` runs. + */ +class CountingUDTF extends GenericUDTF { + private var count = 0L + + override def initialize(argOIs: StructObjectInspector): StructObjectInspector = + ObjectInspectorFactory.getStandardStructObjectInspector( + Arrays.asList("value", "n"), + Arrays.asList[ObjectInspector]( + PrimitiveObjectInspectorFactory.javaStringObjectInspector, + PrimitiveObjectInspectorFactory.javaLongObjectInspector)) + + override def process(args: Array[AnyRef]): Unit = { + count += 1 + forward(Array[AnyRef](Option(args(0)).map(_.toString).orNull, Long.box(count))) + } + + override def close(): Unit = forward(Array[AnyRef]("total", Long.box(count))) +} From edb56380fb95f1c57ca45fffedd4186442b46772 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 20:33:14 -0600 Subject: [PATCH 4/4] fix: carry the generate's required child columns, not its output minus the generator's --- .../org/apache/spark/sql/comet/operators.scala | 7 +++---- .../comet/exec/CometGenerateExecSuite.scala | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 4 deletions(-) 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 554fa5f65a2..167f86359b3 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 @@ -1980,10 +1980,9 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { return None } - // Convert the projection expressions (columns to carry forward) - // These are the non-generator output columns - val generatorOutput = op.generatorOutput.toSet - val projectExprs = op.output.filterNot(generatorOutput.contains).map { attr => + // The child columns carried forward. Not derived from `op.output` minus the generator output: + // a generator may reuse its child's attribute ids, as `ReplicateRows` does for EXCEPT ALL. + val projectExprs = op.requiredChildOutput.map { attr => exprToProto(attr, op.child.output) } 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 154e53930aa..1ce827d3452 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -815,4 +815,20 @@ class CometGenerateExecSuite extends CometTestBase { } } } + + test("EXCEPT ALL in a correlated subquery keeps the carried columns") { + // The decorrelated plan carries t0a and t0b past `ReplicateRows`, whose output reuses + // those attribute ids. + withTempView("t0", "t1", "t2") { + sql("CREATE TEMP VIEW t0(t0a, t0b) AS VALUES (1, 1), (2, 0)") + sql("CREATE TEMP VIEW t1(t1a, t1c) AS VALUES (1, 3)") + sql("CREATE TEMP VIEW t2(t2b, t2c) AS VALUES (1, 5), (2, 7)") + checkSparkAnswerWithNativeGenerate(sql(""" + |SELECT t0a, (SELECT count(t1c) FROM + | (SELECT t1c FROM t1 WHERE t1a = t0a + | EXCEPT ALL + | SELECT t2c FROM t2 WHERE t2b = t0b)) + |FROM t0""".stripMargin)) + } + } }