From b18bd6f04d27bd8f9847c534b950313b7d9d2183 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 16:51:29 -0600 Subject: [PATCH 1/7] 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 cf8baac011d..b0ad79e6014 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 e96df28b126..faef35a512a 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -125,12 +125,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 d86f8076bde..0ec14e2d747 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -314,28 +314,28 @@ own Rust plan over its child's output. `CometTakeOrderedAndProject` is another: Rust plan, and when its input has several partitions, it first gathers the top rows of each partition through a single-partition Comet shuffle. -| 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 Rust Parquet writer) | -| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental Rust Parquet writer) | -| `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 Rust Parquet writer) | +| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental Rust Parquet writer) | +| `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 603bba872bc..97b528358df 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -5317,6 +5317,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 28d94e78bf3..36458686076 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1121,6 +1121,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 dd7c59f6556..8896af6eb9c 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 @@ -31,7 +31,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} @@ -1938,6 +1938,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")) } @@ -1961,10 +1968,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 } @@ -1986,11 +2002,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 117eee57e9cc245f271aeb5b83610f0f6d9dee92 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 17:06:49 -0600 Subject: [PATCH 2/7] 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 8896af6eb9c..a25b5a3bd1b 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.annotation.tailrec import scala.collection.mutable import scala.collection.mutable.ArrayBuffer @@ -31,7 +29,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} @@ -65,7 +63,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 @@ -1939,28 +1937,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}")) } } @@ -1968,19 +1961,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 } @@ -1995,8 +1987,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 ba3f5d4e928e5d5ffad722b279aa800b11c13c50 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 27 Oct 2026 06:57:47 -0600 Subject: [PATCH 3/7] address review: restore sliced-input flatten test, drop liftFallbackReasons for stack --- native/operators/src/explode.rs | 39 ++++++++++++++----- .../apache/spark/sql/comet/operators.scala | 5 +-- 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 33abb13836d..70d11933da3 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -1880,18 +1880,39 @@ mod tests { #[test] fn flatten_struct_cols_nulls_fields_under_a_null_struct() { - // Row 1's struct is null but its field holds a value, which Arrow allows. Spark's - // `inline` turns a null struct into a row of nulls, so the value must not surface. - let field = Arc::new(Field::new("a", DataType::Int32, true)); + // Row 1's struct is null but its fields hold values, which Arrow allows. Spark's `inline` + // turns a null struct into a row of nulls, so none of those values may surface. + let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30, 40])); + let fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Int32, false)), + ]; let structs: ArrayRef = Arc::new(StructArray::new( - vec![Arc::clone(&field)].into(), - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - Some(NullBuffer::from(vec![true, false, true])), + fields.clone().into(), + vec![a, b], + Some(NullBuffer::from(vec![true, false, true, true])), )); - let schema = Arc::new(Schema::new(vec![field])); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let indices = HashSet::from([0]); + + let batch = flatten_struct_cols(&[Arc::clone(&structs)], &schema, &indices).unwrap(); + let a = batch.column(0).as_primitive::(); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [Some(1), None, None, Some(4)]); + assert_eq!( + b.iter().collect::>(), + [Some(10), None, Some(30), Some(40)] + ); - let batch = flatten_struct_cols(&[structs], &schema, &HashSet::from([0])).unwrap(); + // A slice keeps the struct's mask aligned with its fields. + let batch = flatten_struct_cols(&[structs.slice(1, 2)], &schema, &indices).unwrap(); let a = batch.column(0).as_primitive::(); - assert_eq!(a.iter().collect::>(), [Some(1), None, Some(3)]); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [None, None]); + assert_eq!(b.iter().collect::>(), [None, Some(30)]); } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index a25b5a3bd1b..b5225c5512c 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 @@ -63,7 +63,7 @@ import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde -import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, liftFallbackReasons, supportedSortType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.MergeRowsMetricsShim @@ -1970,9 +1970,6 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { val childExprProto = exprToProto(childExpr, op.child.output) if (childExprProto.isEmpty) { - // `stack` builds its array of structs here rather than taking it from the plan, so carry any - // reason recorded on the built nodes to the generator, where the operator collects it. - liftFallbackReasons(childExpr, op.generator) return None } From a7a754e024e8d65edb3d5bcab71eb7f3c06b473f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 3 Nov 2026 07:19:35 -0700 Subject: [PATCH 4/7] fix: keep large stack row counts and zero-field inline on Spark Return Unsupported for stack above Spark's 50-row codegen limit, since the native rewrite builds one struct expression per output row, and for inline over an array of structs with no fields, which cannot build a zero-column batch. --- .../apache/spark/sql/comet/operators.scala | 18 ++++++++++-- .../comet/exec/CometGenerateExecSuite.scala | 28 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 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 b5225c5512c..87cccfc64b7 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 @@ -1939,8 +1939,22 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { op.generator match { // Spark only accepts an array of structs for `inline`, and `stack` is planned as `inline`; // see `stackAsStructArray`. - case _: Inline | _: Stack => - Compatible() + case inline: Inline => + inline.child.dataType match { + case ArrayType(struct: StructType, _) if struct.isEmpty => + // An empty struct carries no columns, so the unnested rows have no data to hold their + // count and the batch cannot be built. + Unsupported(Some("inline over an array of structs with no fields is not supported")) + case _ => Compatible() + } + case stack: Stack => + // The rewrite builds one struct expression per output row, so keep large row counts on + // Spark, which stops generating code for them too. + if (stack.supportCodegen) { + Compatible() + } else { + Unsupported(Some("stack with more than 50 rows is not supported")) + } case explode: ExplodeBase => explode.child.dataType match { case _: ArrayType => 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..63c5042a5df 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -251,6 +251,34 @@ class CometGenerateExecSuite extends CometTestBase { } } + test("stack with more than 50 rows falls back") { + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val values = (1 to 51).map(i => s"v + $i").mkString(", ") + val df = Seq((1, 10), (2, 20)).toDF("id", "v").selectExpr("id", s"stack(51, $values)") + checkSparkAnswerAndFallbackReason(df, "stack with more than 50 rows is not supported") + } + } + + test("stack with 50 rows runs natively") { + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val values = (1 to 50).map(i => s"v + $i").mkString(", ") + val df = Seq((1, 10), (2, 20)).toDF("id", "v").selectExpr("id", s"stack(50, $values)") + checkSparkAnswerAndOperator(df) + } + } + + test("inline over an array of structs with no fields falls back") { + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + Seq("inline", "inline_outer").foreach { fn => + val df = + Seq((1, 10), (2, 20)).toDF("id", "v").selectExpr(s"$fn(array(struct(), struct()))") + checkSparkAnswerAndFallbackReason( + df, + "inline over an array of structs with no fields is not supported") + } + } + } + test("explode with nullable projected column") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", From a2899256b3ece95b7b76384c27645c9e57f412f4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 17:44:54 -0600 Subject: [PATCH 5/7] 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 | 45 +++- .../comet/exec/CometGenerateExecSuite.scala | 79 ++++++ .../spark/sql/hive/CometHiveUDTFs.scala | 63 +++++ 15 files changed, 708 insertions(+), 65 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 b0ad79e6014..22c35e6ef4c 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 faef35a512a..13bb28fc4e9 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -125,12 +125,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 0ec14e2d747..2b7bcda3b01 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -314,28 +314,28 @@ own Rust plan over its child's output. `CometTakeOrderedAndProject` is another: Rust plan, and when its input has several partitions, it first gathers the top rows of each partition through a single-partition Comet shuffle. -| 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 Rust Parquet writer) | -| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental Rust Parquet writer) | -| `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 Rust Parquet writer) | +| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental Rust Parquet writer) | +| `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 97b528358df..89d21ebd7cd 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -5318,6 +5318,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 70d11933da3..97d0ec45794 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 fields hold values, which Arrow allows. Spark's `inline` diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 36458686076..dc177beff91 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1124,6 +1124,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 824118f3efe..754fd806110 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -812,6 +812,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 87cccfc64b7..349cf6c7586 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 @@ -1967,7 +1967,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() + } } } @@ -1975,13 +1978,35 @@ 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) { return None @@ -1998,12 +2023,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 63c5042a5df..25be15ec657 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 @@ -764,4 +765,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 d7f3e42b9f657122097639aa6f3e3694ba50502e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 26 Oct 2026 20:33:14 -0600 Subject: [PATCH 6/7] 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 349cf6c7586..bee5b51d1a9 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 @@ -2012,10 +2012,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 25be15ec657..2949ef7cddd 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -843,4 +843,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)) + } + } } From 4c335a8a5b471fc50ca5b64d5385e99635769100 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 3 Nov 2026 07:49:36 -0700 Subject: [PATCH 7/7] fix: address review of JVM generators Skip the placeholder row when terminate emits nothing and mark the carried columns nullable, so an empty partition with a non-nullable column no longer fails. Snapshot generator rows into UnsafeRow so a reused buffer cannot alias earlier rows, and drop the unused partition index. Lower ReplicateRows to inline over array_repeat so EXCEPT ALL and INTERSECT ALL skip the JVM. Document that the JVM generator works a batch at a time. Add regression tests, including a state check across batch boundaries, and SQL file tests for json_tuple and EXCEPT ALL. Adapt the stack case to the updated inline and stack PR. --- .../user-guide/latest/scala_java_udfs.md | 13 ++- native/core/src/execution/planner/explode.rs | 10 +- native/operators/src/explode.rs | 72 +++++++++++- .../spark/sql/comet/CometGeneratorUDF.scala | 23 ++-- .../apache/spark/sql/comet/operators.scala | 62 +++++++--- .../sql-tests/expressions/json/json_tuple.sql | 59 ++++++++++ .../sql-tests/operators/except_all.sql | 108 ++++++++++++++++++ .../comet/exec/CometGenerateExecSuite.scala | 76 +++++++++--- .../spark/sql/comet/CometTestGenerators.scala | 68 +++++++++++ 9 files changed, 439 insertions(+), 52 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/json/json_tuple.sql create mode 100644 spark/src/test/resources/sql-tests/operators/except_all.sql create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/CometTestGenerators.scala diff --git a/docs/source/user-guide/latest/scala_java_udfs.md b/docs/source/user-guide/latest/scala_java_udfs.md index 246c610450d..6ebf2e094bb 100644 --- a/docs/source/user-guide/latest/scala_java_udfs.md +++ b/docs/source/user-guide/latest/scala_java_udfs.md @@ -44,18 +44,21 @@ This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabl - Aggregate UDFs (`ScalaAggregator`, `TypedImperativeAggregate`, the legacy `UserDefinedAggregateFunction`). - Python UDTFs (`@udtf`). - Python `@udf` and Pandas `@pandas_udf`. +- 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. + +When a UDF is rejected, the reason surfaces through Comet's standard fallback diagnostics; the query still runs on Spark. ## 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. +The operator works a batch at a time rather than a row at a time. It evaluates the generator for every row of a batch, and holds all the rows one input row generates, before it passes any of them on. A `LIMIT` above the generator does not stop evaluation partway through a batch, so a generator that raises an error on a later row can fail a query Spark would have satisfied from earlier rows, and a generator that emits a very large number of rows for one input row needs memory for all of them. -- 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. +Spark's `EXCEPT ALL` and `INTERSECT ALL` use an internal generator, `ReplicateRows`. Comet runs it natively as `inline` over `array_repeat`, without the JVM. -When a UDF is rejected, the reason surfaces through Comet's standard fallback diagnostics; the query still runs on Spark. +Set `spark.comet.exec.jvmGenerator.enabled` to `false` to run such a generator in Spark instead. ## Behavior diff --git a/native/core/src/execution/planner/explode.rs b/native/core/src/execution/planner/explode.rs index 70bb69e9e12..953e9e534f9 100644 --- a/native/core/src/execution/planner/explode.rs +++ b/native/core/src/execution/planner/explode.rs @@ -139,8 +139,16 @@ impl OperatorBuilder for ExplodeBuilder { let mut output_fields: Vec = Vec::new(); // Add all projection columns (non-array columns) + // The rows `terminate` yields carry the last input row's columns, which are all null when + // the input is empty, so those columns are nullable whatever the input says. + let carried_nullable = explode.terminate.is_some(); for i in 0..projections.len() { - output_fields.push(project_schema.field(i).clone()); + let field = project_schema.field(i).clone(); + output_fields.push(if carried_nullable { + field.with_nullable(true) + } else { + field + }); } let array_input_index = if explode.position { diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 97d0ec45794..dcddb37fd3e 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -524,7 +524,23 @@ impl ExplodeStream { 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(); + // Nothing to emit, so there is no need to build the placeholder row below. That row's + // columns are null when no input row arrived, which a non-nullable column rejects. + let list = rows.as_list::(); + if list.is_null(0) || list.value_length(0) == 0 { + return Ok(None); + } + + // The output schema marks every column terminate may null as nullable (see the planner), + // so the placeholder batch has to as well. + let input_schema = Arc::new(Schema::new( + self.input + .schema() + .fields() + .iter() + .map(|f| f.as_ref().clone().with_nullable(true)) + .collect::>(), + )); let mut columns = match self.last_row.take() { Some(row) => row.columns().to_vec(), None => input_schema @@ -2065,6 +2081,60 @@ mod tests { ); } + /// `terminate` over an empty input whose carried column is non-nullable, with the nullable + /// output column the planner declares for it. + async fn terminate_over_empty_non_nullable_input(terminate: Vec) -> usize { + let input_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("l", DataType::new_list(DataType::Int32, true), true), + ])); + let source = MemorySourceConfig::try_new_exec( + &[vec![RecordBatch::new_empty(Arc::clone(&input_schema))]], + 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(NullHandling::Drop), + ) + .unwrap() + .with_terminate(terminate) + .unwrap(); + datafusion::physical_plan::common::collect(explode.execute(0, Arc::default()).unwrap()) + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum() + } + + #[tokio::test] + async fn empty_terminate_over_empty_non_nullable_input_adds_nothing() { + assert_eq!(terminate_over_empty_non_nullable_input(vec![]).await, 0); + } + + #[tokio::test] + async fn terminate_over_empty_non_nullable_input_emits_null_carried_column() { + assert_eq!(terminate_over_empty_non_nullable_input(vec![7, 8]).await, 2); + } + #[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` 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 index f1b6822b028..8057c58ad2a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala @@ -31,7 +31,7 @@ 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.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Generator, GenericInternalRow, Literal, Unevaluable, UnsafeProjection} import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.comet.execution.arrow.ArrowWriter import org.apache.spark.sql.execution.ExecSubqueryExpression @@ -84,15 +84,17 @@ class CometGeneratorUDF extends CometUDF { val digest = ByteBuffer.wrap(binaryScalar(inputs(0))) val bound = generators.getOrElseUpdate( (digest, planId), - BoundGenerator(deserialize(binaryScalar(inputs(1))), partitionIndex)) + new BoundGenerator(deserialize(binaryScalar(inputs(1))))) val output = CometBatchKernelCodegen.allocateOutput(bound.outputField, numRows, 0) try { val writer = ArrowWriter.createFieldWriter(output) val outputRow = new GenericInternalRow(1) def writeRows(rows: Any): Unit = { + // A generator may reuse a row's buffers from one `next` to the next, which + // `GenericInternalRow.copy` shares rather than copies. val copies = new ArrayBuffer[Any] - iteratorOf(rows).foreach(row => copies += row.copy()) + iteratorOf(rows).foreach(row => copies += bound.snapshot(row)) outputRow.update(0, new GenericArrayData(copies.toArray)) writer.write(outputRow, 0) } @@ -217,16 +219,13 @@ object CometGeneratorUDF { /** * 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 class BoundGenerator(val generator: Generator) { + val outputField: Field = CometGeneratorUDF.outputField(generator) - private object BoundGenerator { - def apply(generator: Generator, partitionIndex: Int): BoundGenerator = { - generator.foreach { - case n: Nondeterministic => n.initialize(partitionIndex) - case _ => - } - BoundGenerator(generator, outputField(generator)) - } + private val projection = UnsafeProjection.create(generator.elementSchema) + + /** A copy of `row` that owns its bytes, so it survives the generator reusing its buffers. */ + def snapshot(row: InternalRow): InternalRow = projection(row).copy() } /** The type of each element of the output: the rows one input row generates. */ 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 bee5b51d1a9..4ed16af9a13 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 @@ -29,7 +29,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, ExplodeBase, Expression, ExpressionSet, Generator, Inline, LeafExpression, Literal, NamedExpression, SortOrder, Stack, XXH64} +import org.apache.spark.sql.catalyst.expressions.{ArrayRepeat, Ascending, Attribute, AttributeSeq, AttributeSet, Cast, CodegenObjectFactoryMode, CreateArray, CreateNamedStruct, ExplodeBase, Expression, ExpressionSet, Generator, GreaterThan, If, Inline, LeafExpression, Literal, NamedExpression, ReplicateRows, 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} @@ -60,7 +60,7 @@ import org.apache.comet.{CometConf, CometExecIterator, CometRuntimeException, Co import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled, isSpark35Plus, withFallbackReason} import org.apache.comet.parquet.{CometParquetUtils, HadoopFileOutput} import org.apache.comet.rules.CometExecRule -import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} +import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, ExprOuterClass, 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} @@ -1985,27 +1985,18 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { 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 - } + exprToProto(stackAsStructArray(stack), op.child.output) 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. + case rows: ReplicateRows => 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 - } + // The rewrite is native end to end, so it is preferred to the JVM generator. + exprToProto(replicateRowsAsStructArray(rows), op.child.output) + .orElse(jvmGeneratorCall(op, explodeBuilder)) + case _ => + explodeBuilder.setInline(true) + jvmGeneratorCall(op, explodeBuilder) } if (childExprProto.isEmpty) { @@ -2029,6 +2020,39 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { Some(builder.setExplode(explodeBuilder).build()) } + /** + * The expression that runs `op.generator` on the JVM, which yields for each input row an array + * of the structs it emits, and sets the one that runs its `terminate` on `explodeBuilder`. + */ + private def jvmGeneratorCall( + op: GenerateExec, + explodeBuilder: OperatorOuterClass.Explode.Builder): Option[ExprOuterClass.Expr] = + CometGeneratorUDF.calls(op.generator, op.child.output) match { + case Right((process, terminate)) => + explodeBuilder.setTerminate(terminate) + Some(process) + case Left(reason) => + withFallbackReason(op, reason) + None + } + + /** + * `ReplicateRows(n, v1, ..., vk)` emits `n` copies of the row `(v1, ..., vk)`, which is + * `inline` over `array_repeat` of that row as a struct. A NULL or negative `n` gives no rows in + * both. The count is a `Long` and `array_repeat` takes an `Int`, so a count beyond + * `Int.MaxValue` is clamped rather than wrapped; no array that long fits in memory. + */ + private def replicateRowsAsStructArray(rows: ReplicateRows): Expression = { + val count = rows.children.head + val row = CreateNamedStruct(rows.elementSchema.fields.toSeq.zip(rows.children.tail).flatMap { + case (field, value) => Seq(Literal(field.name), value) + }) + val maxInt = Literal(Int.MaxValue.toLong) + ArrayRepeat( + row, + If(GreaterThan(count, maxInt), Literal(Int.MaxValue), Cast(count, IntegerType))) + } + /** * `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. diff --git a/spark/src/test/resources/sql-tests/expressions/json/json_tuple.sql b/spark/src/test/resources/sql-tests/expressions/json/json_tuple.sql new file mode 100644 index 00000000000..e78c3cc4353 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/json/json_tuple.sql @@ -0,0 +1,59 @@ +-- 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. + +-- json_tuple has no native implementation. It runs on the JVM inside the native generate operator, +-- one row per input row, with NULL for a field the document lacks and a row of NULLs for a NULL or +-- invalid document. + +statement +CREATE TABLE test_json_tuple(id int, js string) USING parquet + +statement +INSERT INTO test_json_tuple VALUES + (1, '{"a": 1, "b": "x", "c": true}'), + (2, '{"a": null, "b": "y"}'), + (3, '{"b": "multi word", "d": {"e": 5}}'), + (4, '{"a": [1, 2, 3], "b": 1.5}'), + (5, 'not json'), + (6, ''), + (7, NULL), + (8, '{"a": "é日本", "b": ""}'), + (NULL, '{"a": 9}') + +query +SELECT id, json_tuple(js, 'a', 'b') FROM test_json_tuple + +query +SELECT json_tuple(js, 'a', 'b', 'c', 'd') FROM test_json_tuple + +-- a field the documents lack, and a repeated field +query +SELECT id, json_tuple(js, 'missing', 'a', 'a') FROM test_json_tuple + +query +SELECT id, a, b FROM test_json_tuple LATERAL VIEW json_tuple(js, 'a', 'b') j AS a, b + +-- the outer form keeps the rows that yield no values +query +SELECT id, a FROM test_json_tuple LATERAL VIEW OUTER json_tuple(js, 'a') j AS a + +-- a filter and an aggregate above the generator +query +SELECT id, a FROM test_json_tuple LATERAL VIEW json_tuple(js, 'a') j AS a WHERE a IS NOT NULL + +query +SELECT b, count(*) FROM test_json_tuple LATERAL VIEW json_tuple(js, 'b') j AS b GROUP BY b diff --git a/spark/src/test/resources/sql-tests/operators/except_all.sql b/spark/src/test/resources/sql-tests/operators/except_all.sql new file mode 100644 index 00000000000..4425b8c780e --- /dev/null +++ b/spark/src/test/resources/sql-tests/operators/except_all.sql @@ -0,0 +1,108 @@ +-- 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. + +-- EXCEPT ALL and INTERSECT ALL are planned with the internal ReplicateRows generator, which +-- repeats each group of rows as many times as the counts left after the set operation. It has no +-- native implementation, so it runs on the JVM inside the native generate operator and carries +-- every column of the row through. + +statement +CREATE TABLE test_set_left( + id int, + d decimal(20, 2), + bin binary, + f float, + dbl double, + dt date, + ts timestamp, + ts_ntz timestamp_ntz, + arr array, + st struct, + s string +) USING parquet + +statement +CREATE TABLE test_set_right LIKE test_set_left + +statement +INSERT INTO test_set_left VALUES + (1, 1.50, X'01', 1.5, 1.5, DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', array(1, 2), named_struct('a', 1, 'b', 'x'), 'a'), + (1, 1.50, X'01', 1.5, 1.5, DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', array(1, 2), named_struct('a', 1, 'b', 'x'), 'a'), + (1, 1.50, X'01', 1.5, 1.5, DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', array(1, 2), named_struct('a', 1, 'b', 'x'), 'a'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (3, 0.00, X'', float('-0.0'), double('-0.0'), DATE '1970-01-01', TIMESTAMP '1970-01-01 00:00:00', TIMESTAMP_NTZ '1970-01-01 00:00:00', array(), named_struct('a', 3, 'b', NULL), ''), + (4, 99999999999999999.99, X'AB', CAST('Infinity' AS FLOAT), CAST('-Infinity' AS DOUBLE), DATE '2000-02-29', TIMESTAMP '2000-02-29 12:34:56', TIMESTAMP_NTZ '2000-02-29 12:34:56', array(7), named_struct('a', 4, 'b', 'z'), 'w'), + (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL) + +statement +INSERT INTO test_set_right VALUES + (1, 1.50, X'01', 1.5, 1.5, DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', array(1, 2), named_struct('a', 1, 'b', 'x'), 'a'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (3, 0.00, X'', float('0.0'), double('0.0'), DATE '1970-01-01', TIMESTAMP '1970-01-01 00:00:00', TIMESTAMP_NTZ '1970-01-01 00:00:00', array(), named_struct('a', 3, 'b', NULL), ''), + (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL) + +-- every column at once +query +SELECT * FROM test_set_left EXCEPT ALL SELECT * FROM test_set_right + +query +SELECT * FROM test_set_left INTERSECT ALL SELECT * FROM test_set_right + +-- one type at a time, so a type the carried-column path mishandles shows up on its own +query +SELECT d FROM test_set_left EXCEPT ALL SELECT d FROM test_set_right + +query +SELECT bin FROM test_set_left EXCEPT ALL SELECT bin FROM test_set_right + +query +SELECT f FROM test_set_left EXCEPT ALL SELECT f FROM test_set_right + +query +SELECT dbl FROM test_set_left EXCEPT ALL SELECT dbl FROM test_set_right + +query +SELECT dt FROM test_set_left EXCEPT ALL SELECT dt FROM test_set_right + +query +SELECT ts FROM test_set_left EXCEPT ALL SELECT ts FROM test_set_right + +query +SELECT ts_ntz FROM test_set_left EXCEPT ALL SELECT ts_ntz FROM test_set_right + +query +SELECT arr FROM test_set_left EXCEPT ALL SELECT arr FROM test_set_right + +query +SELECT st FROM test_set_left EXCEPT ALL SELECT st FROM test_set_right + +query +SELECT s FROM test_set_left EXCEPT ALL SELECT s FROM test_set_right + +query +SELECT s FROM test_set_left INTERSECT ALL SELECT s FROM test_set_right + +-- an empty result and an empty right side +query +SELECT s FROM test_set_left EXCEPT ALL SELECT s FROM test_set_left + +query +SELECT s FROM test_set_left EXCEPT ALL SELECT s FROM test_set_right WHERE id > 100 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 2949ef7cddd..3e886c07f18 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -20,7 +20,7 @@ package org.apache.comet.exec import org.apache.spark.sql.{CometTestBase, DataFrame, Row} -import org.apache.spark.sql.comet.CometExplodeExec +import org.apache.spark.sql.comet.{CometExplodeExec, CometTestGenerators} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.functions.col @@ -788,18 +788,21 @@ class CometGenerateExecSuite extends CometTestBase { 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(""" + // Two rows per batch, so the running count and the `close` total cross a batch boundary. + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "2") { + 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)) + } } } @@ -807,7 +810,8 @@ class CometGenerateExecSuite extends CometTestBase { CometHiveUDTFs.register(spark, "counting_udtf", classOf[CountingUDTF]) withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.SHUFFLE_PARTITIONS.key -> "3") { + SQLConf.SHUFFLE_PARTITIONS.key -> "3", + CometConf.COMET_BATCH_SIZE.key -> "1") { 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. @@ -819,6 +823,30 @@ class CometGenerateExecSuite extends CometTestBase { } } + test("generator rows that share a reused buffer are kept apart") { + CometTestGenerators.registerReusedBuffer(spark, "reused_buffer") + withGeneratorTable { + checkSparkAnswerWithNativeGenerate( + sql("SELECT id, reused_buffer(encode(s, 'UTF-8')) FROM t")) + checkSparkAnswerWithNativeGenerate( + sql("SELECT id, b FROM t LATERAL VIEW OUTER reused_buffer(encode(s, 'UTF-8')) r AS b")) + } + } + + test("terminate with no rows over a non-nullable carried column adds nothing") { + // `k` is never null, so the column that stands in for a last input row cannot hold the nulls + // a partition with no input would give it, and `json_tuple` has no `terminate` rows to emit. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withGeneratorTable { + val df = sql(""" + |SELECT k, json_tuple(js, 'a') FROM ( + | SELECT 1 AS k, js FROM t WHERE id > 10)""".stripMargin) + checkSparkAnswerWithNativeGenerate(df) + assert(df.collect().isEmpty) + } + } + } + test("json_tuple runs natively") { withGeneratorTable { checkSparkAnswerAndOperator(sql("SELECT id, json_tuple(js, 'a', 'b') FROM t")) @@ -827,10 +855,30 @@ class CometGenerateExecSuite extends CometTestBase { } } - test("EXCEPT ALL replicates rows through a JVM generator") { + test("EXCEPT ALL replicates rows with array_repeat, not a JVM generator") { withGeneratorTable { - checkSparkAnswerWithNativeGenerate( - sql("SELECT s FROM t UNION ALL SELECT s FROM t EXCEPT ALL SELECT 'a'")) + val (_, cometPlan) = + checkSparkAnswer(sql("SELECT s FROM t UNION ALL SELECT s FROM t EXCEPT ALL SELECT 'a'")) + val generate = collectFirst(cometPlan) { case generate: CometExplodeExec => generate } + assert(generate.isDefined, s"Expected a native generator in\n$cometPlan") + assert(!generate.get.nativeOp.getExplode.getChild.hasJvmScalarUdf) + } + } + + test("EXCEPT ALL repeats rows by a count of zero, one and many") { + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withTempView("l", "r") { + Seq(("a", 1), ("b", 1), ("b", 1), ("b", 1), ("c", 1), (null, 1), (null, 1)) + .toDF("k", "v") + .createOrReplaceTempView("l") + Seq(("a", 1), ("b", 1), ("c", 1), ("c", 1), (null, 1)) + .toDF("k", "v") + .createOrReplaceTempView("r") + checkSparkAnswerWithNativeGenerate( + sql("SELECT k, v FROM l EXCEPT ALL SELECT k, v FROM r")) + checkSparkAnswerWithNativeGenerate( + sql("SELECT k, v FROM l INTERSECT ALL SELECT k, v FROM r")) + } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTestGenerators.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTestGenerators.scala new file mode 100644 index 00000000000..81eb52909a4 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTestGenerators.scala @@ -0,0 +1,68 @@ +/* + * 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 org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo, Generator, GenericInternalRow} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.types.{BinaryType, StructField, StructType} + +/** Catalyst generators for tests, registered as SQL functions. */ +object CometTestGenerators { + def register(spark: SparkSession, name: String, build: Seq[Expression] => Generator): Unit = + spark.sessionState.functionRegistry.registerFunction( + FunctionIdentifier(name), + new ExpressionInfo(name, name), + (children: Seq[Expression]) => build(children)) + + def registerReusedBuffer(spark: SparkSession, name: String): Unit = + register(spark, name, ReusedBufferGenerator(_)) +} + +/** + * Emits the bytes of its argument three times, `x1`, `x2` and `x3`, through one row and one byte + * array it overwrites between rows, as a generator that avoids allocating a row per output does. + * Only a copy of the bytes made before the iterator advances keeps the three rows apart. + */ +case class ReusedBufferGenerator(children: Seq[Expression]) + extends Generator + with CodegenFallback { + + override def elementSchema: StructType = StructType(Seq(StructField("b", BinaryType))) + + override def eval(input: InternalRow): Iterator[InternalRow] = { + val prefix = children.head.eval(input) match { + case null => Array.empty[Byte] + case bytes: Array[Byte] => bytes + } + val buffer = new Array[Byte](prefix.length + 1) + prefix.copyToArray(buffer) + val row = new GenericInternalRow(Array[Any](buffer)) + Iterator.tabulate(3) { i => + buffer(buffer.length - 1) = ('1' + i).toByte + row + } + } + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[Expression]): ReusedBufferGenerator = copy(children = newChildren) +}