diff --git a/docs/source/contributor-guide/expression-audits/generator_funcs.md b/docs/source/contributor-guide/expression-audits/generator_funcs.md index 3023cfaa774..140830c346d 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. 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/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index f9292b27343..33162f5ec4a 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -336,20 +336,21 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci ## generator_funcs -`explode`, `explode_outer`, `posexplode`, and `posexplode_outer` are supported via -`CometExplodeExec` (operator-level, not expression-level) for array input; map input falls back -to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via +`explode`, `explode_outer`, `posexplode`, `posexplode_outer`, `inline`, `inline_outer`, and +`stack` are supported via `CometExplodeExec` (operator-level, not expression-level). `explode` and +`posexplode` support array input; map input falls back to Spark +([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via `spark.comet.exec.explode.enabled`. | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `explode` | ✅ | — | via `CometExplodeExec` | | `explode_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#explode_outer)) | -| `inline` | 🔜 | — | Operator-level generator (like `explode`) | -| `inline_outer` | 🔜 | — | Operator-level generator (like `explode`) | +| `inline` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#inline)) | +| `inline_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#inline_outer)) | | `posexplode` | ✅ | — | via `CometExplodeExec` | | `posexplode_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#posexplode_outer)) | -| `stack` | 🔜 | — | Operator-level generator | +| `stack` | ✅ | — | via `CometExplodeExec`, planned as `inline` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#stack)) | --- diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 334d34673f2..dc5ecdf743d 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -124,12 +124,12 @@ combined with `OR`, which Spark plans as an existence `BroadcastNestedLoopJoinEx ## Generators and set operations -| Operator | Status | Notes | -| -------------- | ------ | ---------------------------------------------------------------------------------------------------------------- | -| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays. `inline` / `stack` fall back. | -| `ExpandExec` | ✅ | | -| `UnionExec` | ✅ | | -| `CoalesceExec` | ✅ | | +| Operator | Status | Notes | +| -------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------- | +| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays, `inline`, `inline_outer`, and `stack`. Other generators fall back. | +| `ExpandExec` | ✅ | | +| `UnionExec` | ✅ | | +| `CoalesceExec` | ✅ | | ## Writes diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 7789e8eacc2..f7840966970 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -310,28 +310,28 @@ consecutively in a plan, they execute as a single fused block. The write operato (`CometWriteFiles`, `CometNativeWrite` and `CometIcebergWrite`) are the exception: each runs as its own native plan over its child's output. -| Node | Spark equivalent | -| ------------------------------ | --------------------------------------------------------------------------------- | -| `CometProject` | `ProjectExec` | -| `CometFilter` | `FilterExec` | -| `CometSort` | `SortExec` | -| `CometLocalLimit` | `LocalLimitExec` | -| `CometGlobalLimit` | `GlobalLimitExec` | -| `CometExpand` | `ExpandExec` | -| `CometExplode` | `GenerateExec` (for `explode`, `explode_outer`, `posexplode`, `posexplode_outer`) | -| `CometHashAggregate` | `HashAggregateExec`, `ObjectHashAggregateExec` | -| `CometHashJoin` | `ShuffledHashJoinExec` | -| `CometBroadcastHashJoin` | `BroadcastHashJoinExec` | -| `CometBroadcastNestedLoopJoin` | `BroadcastNestedLoopJoinExec` | -| `CometSortMergeJoin` | `SortMergeJoinExec` | -| `CometMergeRows` | `MergeRowsExec` (Spark 3.5+; opt-in) | -| `CometWindowExec` | `WindowExec` | -| `CometWindowGroupLimitExec` | `WindowGroupLimitExec` (Spark 3.5 and later) | -| `CometSample` | `SampleExec` (sampling without replacement) | -| `CometTakeOrderedAndProject` | `TakeOrderedAndProjectExec` | -| `CometWriteFiles` | `WriteFilesExec` (Spark 4.0 and later, experimental native Parquet writes) | -| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental native Parquet writes) | -| `CometIcebergWrite` | `IcebergWrite` (native Iceberg data-file writes, opt-in) | +| Node | Spark equivalent | +| ------------------------------ | -------------------------------------------------------------------------------------------------------------------- | +| `CometProject` | `ProjectExec` | +| `CometFilter` | `FilterExec` | +| `CometSort` | `SortExec` | +| `CometLocalLimit` | `LocalLimitExec` | +| `CometGlobalLimit` | `GlobalLimitExec` | +| `CometExpand` | `ExpandExec` | +| `CometExplode` | `GenerateExec` (for `explode`, `explode_outer`, `posexplode`, `posexplode_outer`, `inline`, `inline_outer`, `stack`) | +| `CometHashAggregate` | `HashAggregateExec`, `ObjectHashAggregateExec` | +| `CometHashJoin` | `ShuffledHashJoinExec` | +| `CometBroadcastHashJoin` | `BroadcastHashJoinExec` | +| `CometBroadcastNestedLoopJoin` | `BroadcastNestedLoopJoinExec` | +| `CometSortMergeJoin` | `SortMergeJoinExec` | +| `CometMergeRows` | `MergeRowsExec` (Spark 3.5+; opt-in) | +| `CometWindowExec` | `WindowExec` | +| `CometWindowGroupLimitExec` | `WindowGroupLimitExec` (Spark 3.5 and later) | +| `CometSample` | `SampleExec` (sampling without replacement) | +| `CometTakeOrderedAndProject` | `TakeOrderedAndProjectExec` | +| `CometWriteFiles` | `WriteFilesExec` (Spark 4.0 and later, experimental native Parquet writes) | +| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental native Parquet writes) | +| `CometIcebergWrite` | `IcebergWrite` (native Iceberg data-file writes, opt-in) | ### JVM-Side Operators diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 0241e1f1ce8..7cdb0e26e6e 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -5088,6 +5088,7 @@ mod tests { outer, position, project_list: vec![bound], + inline: false, })), ..Default::default() }; diff --git a/native/core/src/execution/planner/explode.rs b/native/core/src/execution/planner/explode.rs index 1159f242a53..35b9e0dca9e 100644 --- a/native/core/src/execution/planner/explode.rs +++ b/native/core/src/execution/planner/explode.rs @@ -15,8 +15,9 @@ // specific language governing permissions and limitations // under the License. -//! Builder for the native `Explode` operator, which plans Spark's `explode`, `posexplode`, and -//! their `_outer` variants as a projection followed by Comet's `ExplodeExec`. +//! Builder for the native `Explode` operator, which plans Spark's `explode`, `posexplode`, +//! `inline`, and their `_outer` variants as a projection followed by Comet's `ExplodeExec`. +//! Spark's `stack` arrives here as `inline` over an array of structs built by the JVM serde. use std::sync::Arc; @@ -162,11 +163,30 @@ impl OperatorBuilder for ExplodeBuilder { } }; - output_fields.push(Field::new( - array_field.name(), - element_type, - true, // Element is nullable after unnesting - )); + // `inline` flattens each struct element into one column per field. Every field is + // nullable after unnesting: a null struct element, or the row `inline_outer` emits for a + // NULL or empty array, yields a row of nulls. + let struct_column_indices = if explode.inline { + let DataType::Struct(struct_fields) = &element_type else { + return Err(ExecutionError::GeneralError(format!( + "Expected List of Struct for inline, got {:?}", + array_field.data_type() + ))); + }; + output_fields.extend( + struct_fields + .iter() + .map(|f| f.as_ref().clone().with_nullable(true)), + ); + vec![array_input_index] + } else { + output_fields.push(Field::new( + array_field.name(), + element_type, + true, // Element is nullable after unnesting + )); + vec![] + }; let output_schema = Arc::new(Schema::new(output_fields)); @@ -196,7 +216,7 @@ impl OperatorBuilder for ExplodeBuilder { let unnest_exec = Arc::new(ExplodeExec::new( project_exec, list_unnests, - vec![], // No struct columns to unnest + struct_column_indices, output_schema, unnest_options, )?); diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 6aae0e2d349..70d11933da3 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -51,7 +51,9 @@ //! * dropping upstream's `ListUnnest` declaration in favor of importing the public one; //! * `find_longest_length` applies the empty-list bump once after the row-wise maximum rather //! than once per array, which is equivalent because `max` is associative; -//! * `list_output_lens` and the contiguous-run fast path described above. +//! * `list_output_lens` and the contiguous-run fast path described above; +//! * `flatten_struct_cols` folds a struct's null mask into each field it flattens, since the +//! field values under a null struct are unspecified and Spark's `inline` emits nulls there. //! //! `ExplodeExec` and `ExplodeStream` were always Comet's own. @@ -82,6 +84,7 @@ use datafusion::physical_plan::metrics::{ SplitMetrics, }; use datafusion::physical_plan::stream::BatchSplitStream; +use datafusion_comet_common::children_with_parent_nulls; // `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the // kernels, upstream exports it publicly. use datafusion::physical_plan::unnest::ListUnnest; @@ -573,7 +576,9 @@ fn flatten_struct_cols( Some(_) => match column_data.data_type() { DataType::Struct(_) => { let struct_arr = column_data.as_any().downcast_ref::().unwrap(); - columns_expanded.extend_from_slice(struct_arr.columns()); + // Arrow leaves the field values under a null struct unspecified, so they + // must not leak out once the struct is flattened. + columns_expanded.extend(children_with_parent_nulls(struct_arr)?); } data_type => { return internal_err!( @@ -1872,4 +1877,42 @@ mod tests { ); assert_output_lens(with_nulls.slice(1, 3), [&[0, 0, 3], &[1, 0, 3], &[1, 1, 3]]); } + + #[test] + fn flatten_struct_cols_nulls_fields_under_a_null_struct() { + // Row 1's struct is null but its fields hold values, which Arrow allows. Spark's `inline` + // turns a null struct into a row of nulls, so none of those values may surface. + let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30, 40])); + let fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Int32, false)), + ]; + let structs: ArrayRef = Arc::new(StructArray::new( + fields.clone().into(), + vec![a, b], + Some(NullBuffer::from(vec![true, false, true, true])), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let indices = HashSet::from([0]); + + let batch = flatten_struct_cols(&[Arc::clone(&structs)], &schema, &indices).unwrap(); + let a = batch.column(0).as_primitive::(); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [Some(1), None, None, Some(4)]); + assert_eq!( + b.iter().collect::>(), + [Some(10), None, Some(30), Some(40)] + ); + + // A slice keeps the struct's mask aligned with its fields. + let batch = flatten_struct_cols(&[structs.slice(1, 2)], &schema, &indices).unwrap(); + let a = batch.column(0).as_primitive::(); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [None, None]); + assert_eq!(b.iter().collect::>(), [None, Some(30)]); + } } diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 406592cbce8..b050a8eead4 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1114,6 +1114,9 @@ message Explode { repeated spark.spark_expression.Expr project_list = 3; // Whether to emit a position column alongside the exploded values (posexplode) bool position = 4; + // Whether the array holds structs whose fields become separate output columns (inline). A + // null struct element yields a row of nulls. + bool inline = 5; } message HashJoin { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index da26c91f9c7..23a83f10261 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -19,8 +19,6 @@ package org.apache.spark.sql.comet -import java.util.Locale - import scala.collection.mutable import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ @@ -30,7 +28,7 @@ import org.apache.spark.broadcast.Broadcast import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, CodegenObjectFactoryMode, Expression, ExpressionSet, Generator, LeafExpression, Literal, NamedExpression, SortOrder, 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} @@ -1920,22 +1918,24 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { if (!op.generator.deterministic) { return Unsupported(Some("Only deterministic generators are supported")) } - 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 => + 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 _: 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}")) } } @@ -1943,7 +1943,12 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { op: GenerateExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - val childExpr = op.generator.children.head + 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) { @@ -1961,18 +1966,34 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { return None } - val isPosExplode = op.generator.nodeName.toLowerCase(Locale.ROOT) == "posexplode" - val explodeBuilder = OperatorOuterClass.Explode .newBuilder() .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..95d7c7b0951 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/inline.sql @@ -0,0 +1,120 @@ +-- 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 + +-- 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 new file mode 100644 index 00000000000..98a5b792212 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/stack.sql @@ -0,0 +1,91 @@ +-- 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 + +-- 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 + +-- 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..773affe0298 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,20 @@ 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, + s"${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false") + } + } + test("explode with nullable projected column") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true",