From 3357bdb882833d14997ff05595904e4f0eb5a9e8 Mon Sep 17 00:00:00 2001 From: grorge Date: Fri, 28 Aug 2026 17:27:06 +0800 Subject: [PATCH 1/6] feat: support NullType output types in codegen dispatch Untyped constructors such as map(), map('a', NULL) and array() leave NullType children in their output type, and the codegen dispatch gate rejected any output type containing NullType, so the whole operator fell back to Spark. Make the type gate asymmetric: canHandle now accepts NullType (top-level or nested in array/struct/map) for the output type while still rejecting it for BoundReference inputs, since CometScalaUDFCodegen.specFor cannot build an ArrowColumnSpec for a NullVector. The output emitter maps NullType to NullVector and writes it with setNull only. Also update the Scala/Java UDF guide: NullType arguments remain unsupported, NullType return types are now supported; CalendarIntervalType is removed from the unsupported list since it has been supported since Closes #5525 Assisted-by: Claude Code (claude-fable-5) --- .../user-guide/latest/scala_java_udfs.md | 2 +- .../codegen/CometBatchKernelCodegen.scala | 19 ++- .../CometBatchKernelCodegenOutput.scala | 7 + .../array/array_sort_comparator.sql | 4 + .../sql-tests/expressions/array/transform.sql | 4 + .../sql-tests/expressions/array/zip_with.sql | 4 + .../sql-tests/expressions/map/create_map.sql | 34 +++++ .../sql-tests/expressions/map/map_concat.sql | 4 + .../sql-tests/expressions/map/map_filter.sql | 4 + .../expressions/map/map_zip_with.sql | 4 + .../expressions/map/transform_keys.sql | 4 + .../expressions/map/transform_values.sql | 4 + .../comet/CometCodegenSourceSuite.scala | 121 ++++++++++++++++++ 13 files changed, 209 insertions(+), 6 deletions(-) diff --git a/docs/source/user-guide/latest/scala_java_udfs.md b/docs/source/user-guide/latest/scala_java_udfs.md index 3a55982bbcf..5b1b2098c41 100644 --- a/docs/source/user-guide/latest/scala_java_udfs.md +++ b/docs/source/user-guide/latest/scala_java_udfs.md @@ -45,7 +45,7 @@ This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabl - Table UDFs and generators. - Python `@udf` and Pandas `@pandas_udf`. - Hive `GenericUDF` and `SimpleUDF`. -- `CalendarIntervalType`, `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`). +- `UserDefinedType` arguments and return types, and `NullType` arguments. 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`). A `NullType` *return* type is supported: Comet writes an all-null Arrow vector for it. - 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. diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 83fbca6b635..6b02a2a498d 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -87,8 +87,15 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come * single child and the generated writer NPEs on the missing ordinal-1 vector. * `CometCreateNamedStruct` declines them on the native path for the same reason, but a struct * nested inside a dispatcher-built value (a `CreateMap` value) never reaches that check. + * + * `NullType` is output-only: [[CometBatchKernelCodegenOutput]] can write an all-null Arrow + * `NullVector`, but `CometScalaUDFCodegen.specFor` cannot build an [[ArrowColumnSpec]] for one, + * so a `NullType` input (nested or not) has to keep falling back to Spark. */ - def isSupportedDataType(dt: DataType): Boolean = dt match { + def isSupportedDataType(dt: DataType): Boolean = isSupportedDataType(dt, allowNullType = false) + + private def isSupportedDataType(dt: DataType, allowNullType: Boolean): Boolean = dt match { + case NullType => allowNullType case BooleanType | ByteType | ShortType | IntegerType | LongType => true case FloatType | DoubleType => true case _: DecimalType => true @@ -96,13 +103,15 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come case DateType | TimestampType | TimestampNTZType => true case dt if isTimeType(dt) => true case _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType => true - case ArrayType(inner, _) => isSupportedDataType(inner) + case ArrayType(inner, _) => isSupportedDataType(inner, allowNullType) case st: StructType => // `fieldNames` rebuilds an array on each call, so read it once. val names = st.fieldNames names.distinct.length == names.length && - st.fields.forall(f => isSupportedDataType(f.dataType)) - case mt: MapType => isSupportedDataType(mt.keyType) && isSupportedDataType(mt.valueType) + st.fields.forall(f => isSupportedDataType(f.dataType, allowNullType)) + case mt: MapType => + isSupportedDataType(mt.keyType, allowNullType) && + isSupportedDataType(mt.valueType, allowNullType) case _ => false } @@ -127,7 +136,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come * nested-field count on `spark.sql.codegen.maxFields`. */ def canHandle(boundExpr: Expression): Option[String] = { - if (!isSupportedDataType(boundExpr.dataType)) { + if (!isSupportedDataType(boundExpr.dataType, allowNullType = true)) { return Some(s"codegen dispatch: unsupported output type ${boundExpr.dataType}") } // Mirror WSCG's `spark.sql.codegen.maxFields` gate. Wide schemas blow the generated class's diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index 33e6c0c0355..3b124859abe 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -179,6 +179,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { case _: ArrayType => classOf[ListVector].getName case _: StructType => classOf[StructVector].getName case _: MapType => classOf[MapVector].getName + case NullType => classOf[NullVector].getName case other => throw new UnsupportedOperationException( s"CometBatchKernelCodegen.outputVectorClass: unsupported output type $other") @@ -209,6 +210,11 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { dataType: DataType, ctx: CodegenContext, nested: Boolean = false): OutputEmit = dataType match { + case NullType => + // A NullType value is null by definition: nothing to read from `source`, and `NullVector` + // has no data buffer. `setNull` is a no-op; the all-null semantics come from + // `CometScalaUDFCodegen.evaluate`'s post-`process` `setValueCount`. + OutputEmit("", s"$targetVec.setNull($idx);") case BooleanType => val set = if (nested) "setSafe" else "set" OutputEmit("", s"$targetVec.$set($idx, $source ? 1 : 0);") @@ -407,6 +413,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { */ private def emitSpecializedGetterExpr(target: String, idx: String, elemType: DataType): String = elemType match { + case NullType => "null" case BooleanType => s"$target.getBoolean($idx)" case ByteType => s"$target.getByte($idx)" case ShortType => s"$target.getShort($idx)" diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql b/spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql index 67905bb74ef..863bcf643fa 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql @@ -42,3 +42,7 @@ FROM test_array_sort -- all literals query SELECT array_sort(array(3, 1, 2), (l, r) -> CASE WHEN l < r THEN 1 WHEN l > r THEN -1 ELSE 0 END) + +-- an untyped empty array leaves the element type as NullType +query +SELECT array_sort(array(), (l, r) -> 0) diff --git a/spark/src/test/resources/sql-tests/expressions/array/transform.sql b/spark/src/test/resources/sql-tests/expressions/array/transform.sql index 46490917619..b23e827535e 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/transform.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/transform.sql @@ -80,3 +80,7 @@ SELECT a FROM test_transform WHERE transform(a, x -> array(x)) = nested -- all literals (constant folding is disabled by the test harness) query SELECT transform(array(1, 2, 3), x -> x * x) + +-- an untyped empty array leaves the element type as NullType +query +SELECT transform(array(), x -> x) diff --git a/spark/src/test/resources/sql-tests/expressions/array/zip_with.sql b/spark/src/test/resources/sql-tests/expressions/array/zip_with.sql index 9345dbdfcde..6a8943a034e 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/zip_with.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/zip_with.sql @@ -43,3 +43,7 @@ SELECT zip_with(a, b, (x, y) -> struct(x AS l, y AS r)) FROM test_zip_with -- all literals query SELECT zip_with(array(1, 2), array(3, 4), (x, y) -> x * y) + +-- untyped empty arrays leave the element type as NullType +query +SELECT zip_with(array(), array(), (x, y) -> x) diff --git a/spark/src/test/resources/sql-tests/expressions/map/create_map.sql b/spark/src/test/resources/sql-tests/expressions/map/create_map.sql index 8e479412453..07df37963dd 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/create_map.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/create_map.sql @@ -39,3 +39,37 @@ SELECT map(1, named_struct('x', 10, 'x', 20)) -- Distinct field names are unaffected. query SELECT map(1, named_struct('x', 10, 'y', 20)) + +-- ===== untyped constructors leave NullType children ===== +-- `map()` is MapType(NullType, NullType) and `map(k, NULL)` is MapType(_, NullType). + +query +SELECT map() + +query +SELECT map('a', NULL) + +query +SELECT map(k, NULL) FROM test_create_map + +query +SELECT id, map() FROM (SELECT explode(sequence(1, 3)) AS id) + +query +SELECT size(map()), size(map('a', NULL)) + +query +SELECT map_keys(map()), map_values(map('a', NULL)) + +query +SELECT array(map()), struct(map(), map('a', NULL)) + +query +SELECT map('a', array(NULL)), map('a', map()) + +query +SELECT map_from_arrays(array(), array()) + +-- carry the NullType children through a sort so the vector survives copy/spill paths +query +SELECT k, map(), map('a', NULL) FROM test_create_map ORDER BY k diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_concat.sql b/spark/src/test/resources/sql-tests/expressions/map/map_concat.sql index 453c76aa94e..130a3fd3594 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_concat.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_concat.sql @@ -49,3 +49,7 @@ SELECT map_concat(map(1, 'x', 2, 'y'), map(3, 'z')) -- a NULL literal map makes the whole result NULL query SELECT map_concat(map('a', 1), CAST(NULL AS map)) + +-- untyped operands leave the result as MapType(NullType, NullType) +query +SELECT map_concat(map(), map()) diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_filter.sql b/spark/src/test/resources/sql-tests/expressions/map/map_filter.sql index edf04847107..1aaf3e4dbdf 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_filter.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_filter.sql @@ -42,3 +42,7 @@ SELECT map_filter(m, (k, v) -> v > threshold) FROM test_map_filter -- all literals query SELECT map_filter(map('a', 1, 'b', 2), (k, v) -> v > 1) + +-- an untyped value leaves the value type as NullType +query +SELECT map_filter(map('a', NULL), (k, v) -> true) diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql b/spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql index 181e6a7c159..74dc324ec02 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql @@ -39,3 +39,7 @@ SELECT map_zip_with(m, n, (k, v1, v2) -> struct(v1 AS left, v2 AS right)) FROM t -- all literals query SELECT map_zip_with(map('a', 1), map('a', 2, 'b', 3), (k, v1, v2) -> coalesce(v1, 0) + coalesce(v2, 0)) + +-- untyped operands leave both key and value types as NullType +query +SELECT map_zip_with(map(), map(), (k, v1, v2) -> v1) diff --git a/spark/src/test/resources/sql-tests/expressions/map/transform_keys.sql b/spark/src/test/resources/sql-tests/expressions/map/transform_keys.sql index cc91075a720..526de0e17e0 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/transform_keys.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/transform_keys.sql @@ -38,3 +38,7 @@ SELECT transform_keys(m, (k, v) -> concat(k, suffix)) FROM test_transform_keys -- all literals query SELECT transform_keys(map('a', 1, 'b', 2), (k, v) -> upper(k)) + +-- an untyped value leaves the value type as NullType +query +SELECT transform_keys(map('a', NULL), (k, v) -> k) diff --git a/spark/src/test/resources/sql-tests/expressions/map/transform_values.sql b/spark/src/test/resources/sql-tests/expressions/map/transform_values.sql index 50e49d8959a..23efb5ed0e5 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/transform_values.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/transform_values.sql @@ -42,3 +42,7 @@ SELECT transform_values(m, (k, v) -> v + delta) FROM test_transform_values -- all literals query SELECT transform_values(map('a', 1, 'b', 2), (k, v) -> v * 100) + +-- a lambda returning NULL leaves the value type as NullType +query +SELECT transform_values(map('a', NULL), (k, v) -> v) diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index 400ec6fcd2f..c4936cd2639 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala @@ -287,6 +287,127 @@ class CometCodegenSourceSuite extends AnyFunSuite { s"expected reason to name the rejected expression class; got: ${reason.get}") } + test("canHandle accepts NullType outputs, including nested in complex types") { + // Spark leaves the children of untyped constructors as NullType: `map()` is + // MapType(NullType, NullType), `map('a', NULL)` is MapType(StringType, NullType) and + // `array()` is ArrayType(NullType). Rejecting those forced the whole projection back to + // Spark even though a NullType output only ever has to write nulls. + Seq( + Literal(null, NullType), + Literal.create(Map.empty[Any, Any], MapType(NullType, NullType, valueContainsNull = false)), + Literal.create(Map("a" -> null), MapType(StringType, NullType, valueContainsNull = true)), + Literal.create(Seq(null), ArrayType(NullType)), + Literal.create( + Seq(Map.empty[Any, Any]), + ArrayType(MapType(NullType, NullType, valueContainsNull = false)))).foreach { expr => + assert( + CometBatchKernelCodegen.canHandle(expr).isEmpty, + s"expected canHandle to accept output type ${expr.dataType}") + } + } + + test("canHandle rejects NullType inputs, including nested in complex types") { + // `CometScalaUDFCodegen.specFor` has no ArrowColumnSpec for a NullVector, so a NullType input + // must keep falling back at plan time rather than throwing once the plan has committed to a + // kernel. + Seq( + NullType, + ArrayType(NullType), + MapType(StringType, NullType), + StructType(Seq(StructField("f", NullType)))).foreach { dt => + val reason = CometBatchKernelCodegen.canHandle(BoundReference(0, dt, nullable = true)) + assert(reason.isDefined, s"expected canHandle to reject input type $dt") + assert( + reason.get.contains("unsupported"), + s"expected an unsupported-type reason for $dt; got: ${reason.get}") + } + } + + test("NullType output writes setNull without reading a source value") { + // A NullType leaf has no Arrow data buffer, so the emitted write must be `setNull` only. + val src = CometBatchKernelCodegen + .generateSource(Literal(null, NullType), IndexedSeq(nullableString)) + .body + assert( + src.contains("org.apache.arrow.vector.NullVector"), + s"expected the output vector to be a NullVector; got:\n$src") + assert( + !src.contains("getDataVector"), + s"expected no child-vector access for a scalar NullType output; got:\n$src") + } + + test("nested NullType output casts the child vector and writes setNull into it") { + // The scalar case above cannot distinguish `emitWrite`'s NullType branch from `defaultBody`'s + // own `ev.isNull -> output.setNull(i)` short-circuit, which emits the same text. A NullType + // *value* child inside a map is only reachable through `emitWrite`, so assert on that: the + // child vector must be cast to NullVector and written through `setNull`. + val src = CometBatchKernelCodegen + .generateSource( + Literal.create(Map("a" -> null), MapType(StringType, NullType, valueContainsNull = true)), + IndexedSeq(nullableString)) + .body + val childCast = + """org\.apache\.arrow\.vector\.NullVector\s+(\w+)\s*=\s*\(org\.apache\.arrow\.vector\.NullVector\)""".r + val childVar = childCast + .findFirstMatchIn(src) + .map(_.group(1)) + .getOrElse(fail(s"expected a NullVector child-vector cast; got:\n$src")) + assert( + src.contains(s"$childVar.setNull("), + s"expected a setNull write into the NullVector child `$childVar`; got:\n$src") + } + + test("gate and output emitters agree across the whole accepted type surface") { + // `CometBatchKernelCodegen.isSupportedDataType`, `outputVectorClass`, `emitWrite` and + // `emitSpecializedGetterExpr` each carry a doc comment saying they must stay in step, but + // nothing enforced it. Assert the implication directly: if `canHandle` greenlights an + // output type, generating the kernel for it must not throw. + val leaves: Seq[DataType] = Seq( + NullType, + BooleanType, + ByteType, + ShortType, + IntegerType, + LongType, + FloatType, + DoubleType, + DecimalType(10, 2), + DecimalType(38, 18), + StringType, + BinaryType, + DateType, + TimestampType, + TimestampNTZType, + YearMonthIntervalType(), + DayTimeIntervalType(), + CalendarIntervalType) + + val candidates: Seq[DataType] = leaves.flatMap { dt => + Seq( + dt, + ArrayType(dt), + ArrayType(ArrayType(dt)), + MapType(StringType, dt), + StructType(Seq(StructField("f", dt))), + ArrayType(StructType(Seq(StructField("f", dt)))), + MapType(StringType, ArrayType(dt))) + } + + val accepted = + candidates.filter(dt => CometBatchKernelCodegen.canHandle(Literal.create(null, dt)).isEmpty) + assert( + accepted.size > leaves.size, + s"expected the gate to accept nested shapes too; only got ${accepted.size}") + + accepted.foreach { dt => + withClue(s"canHandle accepted output type $dt, so the emitters must handle it: ") { + CometBatchKernelCodegen.generateSource( + Literal.create(null, dt), + IndexedSeq(nullableString)) + } + } + } + test("CSE collapses a repeated subtree to one evaluation in the generated body") { // `Add(Length(Upper(c0)), Length(Upper(c0)))` has `Length(Upper(c0))` as a common subtree. // Length.doGenCode emits `$value.numChars()` on every Spark version the project targets, From 0d0e7bf13071ee0c793817430ea9ada5e55347eb Mon Sep 17 00:00:00 2001 From: grorge Date: Sat, 29 Aug 2026 20:15:12 +0800 Subject: [PATCH 2/6] fix: carry NullType map columns through the JVM IPC paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the codegen gate change: admitting NullType outputs exposes four problems on the JVM/FFI paths, fixed here. * Arrow Java's MinorType.NULL factory drops the field it is handed, so a NullType map key Comet declared non-nullable comes back nullable and the map schema fails on read. `Utils.withNonNullableMapKeys` repairs the key flag, and `Utils.newArrowStreamWriter` — now the only way to build an `ArrowStreamWriter` (scalastyle-enforced) — applies it on every IPC writer: broadcast, getByteArrayRdd, the PyArrow UDF runner. `CometArrowStream.actualFieldOf` repairs the schema handed to native. * `VectorSchemaRootAppender` loops forever on a NullVector that is a direct child of a struct (a struct's capacity is the minimum over its direct children, and `NullVector.reAlloc()` is a no-op). `Utils.coalesceBroadcastBatches` ships such schemas uncoalesced. A list insulates whatever sits below it, so `array` and `map(k, array(NULL))` keep coalescing. * `CometBatchKernelCodegen.canHandle` rejects duplicate struct field names, recursively, in the output type and in BoundReference inputs: Arrow structs key children by name, so `named_struct('a', x, 'a', NULL)` collapses to one child and the generated ordinal casts fail. Whole-expression dispatch skipped `CometCreateNamedStruct`'s rule. * A NullType child (array element, map value, struct field) is always declared nullable on both sides of the FFI boundary (`Utils.declaredChildNullability`): Spark leaves `containsNull` false on `filter(array(), ...)`, and native kernels that rebuild a list around the input's actual child fail on the nullability mismatch — `map_entries(map_filter(map(), (k, v) -> true))` panicked. Tests: UtilsSuite (key repair, IPC round trip, coalesce bypass rule checked exhaustively over every NullType shape), CometCodegenSourceSuite (duplicate names rejected, NullType children nullable), CometJoinSuite, CometColumnarShuffleSuite, test_pyarrow_udf.py, and expect_fallback / NullType-input queries in create_named_struct.sql, map_entries.sql, slice.sql, array_repeat.sql, array_union.sql, transform.sql. Assisted-by: Claude Code (claude-fable-5) --- dev/scalastyle-config.xml | 14 + .../user-guide/latest/scala_java_udfs.md | 2 +- .../codegen/CometBatchKernelCodegen.scala | 53 +++- .../CometBatchKernelCodegenOutput.scala | 2 + .../apache/comet/serde/QueryPlanSerde.scala | 9 +- .../arrow/CometNativeArrowSource.scala | 6 +- .../apache/spark/sql/comet/util/Utils.scala | 209 +++++++++++-- .../python/CometArrowPythonRunnerBase.scala | 19 +- .../resources/pyspark/test_pyarrow_udf.py | 85 ++++++ .../expressions/array/array_repeat.sql | 5 + .../expressions/array/array_union.sql | 5 + .../sql-tests/expressions/array/slice.sql | 5 + .../sql-tests/expressions/array/transform.sql | 11 + .../sql-tests/expressions/map/create_map.sql | 3 - .../sql-tests/expressions/map/map_entries.sql | 8 + .../expressions/map/map_from_arrays.sql | 6 +- .../struct/create_named_struct.sql | 11 + .../comet/CometCodegenSourceSuite.scala | 76 ++++- .../exec/CometColumnarShuffleSuite.scala | 16 + .../apache/comet/exec/CometJoinSuite.scala | 38 ++- .../spark/sql/comet/util/UtilsSuite.scala | 283 +++++++++++++++++- 21 files changed, 813 insertions(+), 53 deletions(-) diff --git a/dev/scalastyle-config.xml b/dev/scalastyle-config.xml index 9de6df51ef7..e58e6b73581 100644 --- a/dev/scalastyle-config.xml +++ b/dev/scalastyle-config.xml @@ -193,6 +193,20 @@ This file is divided into 3 sections: ]]> + + new ArrowStreamWriter + + + (\.toUpperCase|\.toLowerCase)(?!(\(|\(Locale.ROOT\))) , values: Array>` rather than `VectorUDT`). A `NullType` *return* type is supported: Comet writes an all-null Arrow vector for it. +- `UserDefinedType` arguments and return types, and `NullType` arguments. 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`). A `NullType` _return_ type is supported: Comet writes an all-null Arrow vector for it. - 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. diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 6b02a2a498d..2966d5ce855 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -115,6 +115,26 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come case _ => false } + /** + * Names that repeat within one struct, searched recursively through `dataType`. Spark keeps + * duplicate struct field names as distinct positional fields, but Arrow's `StructVector` keys + * its children by name (`ConflictPolicy.CONFLICT_REPLACE` by default), so + * `initializeChildrenFromFields` collapses the duplicates and the generated ordinal-based child + * casts hit a missing or differently typed vector. `CometCreateNamedStruct` refuses the same + * shape at the serde level, but whole-expression dispatch never consults that rule for a + * `named_struct` nested inside e.g. a `transform` lambda, so [[canHandle]] re-checks here. + */ + private def duplicateStructFieldNames(dataType: DataType): Seq[String] = dataType match { + case st: StructType => + val names = st.fieldNames.toSeq + val dups = names.diff(names.distinct).distinct + if (dups.nonEmpty) dups else st.fields.flatMap(f => duplicateStructFieldNames(f.dataType)) + case ArrayType(inner, _) => duplicateStructFieldNames(inner) + case MapType(keyType, valueType, _) => + duplicateStructFieldNames(keyType) ++ duplicateStructFieldNames(valueType) + case _ => Nil + } + /** * Mirrors `WholeStageCodegenExec.numOfNestedFields` so [[canHandle]] can reuse * `spark.sql.codegen.maxFields`. @@ -132,13 +152,19 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come * back cleanly rather than crashing the Janino compile at execute time. * * Checks every `BoundReference`'s data type and the root `expr.dataType` against - * [[isSupportedDataType]], rejects aggregates / generators / `Unevaluable`, and gates total - * nested-field count on `spark.sql.codegen.maxFields`. + * [[isSupportedDataType]] and [[duplicateStructFieldNames]], rejects aggregates / generators / + * `Unevaluable`, and gates total nested-field count on `spark.sql.codegen.maxFields`. */ def canHandle(boundExpr: Expression): Option[String] = { if (!isSupportedDataType(boundExpr.dataType, allowNullType = true)) { return Some(s"codegen dispatch: unsupported output type ${boundExpr.dataType}") } + val outputDups = duplicateStructFieldNames(boundExpr.dataType) + if (outputDups.nonEmpty) { + return Some( + s"codegen dispatch: duplicate struct field name ${outputDups.mkString(", ")} " + + s"in output type ${boundExpr.dataType}") + } // Mirror WSCG's `spark.sql.codegen.maxFields` gate. Wide schemas blow the generated class's // typed input field count, the typed-getter switch, and the constant pool. Refuse here so the // operator falls back to Spark cleanly rather than tripping a Janino compile failure @@ -187,12 +213,23 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come "(aggregate, generator, or unevaluable)") case None => } - val badRef = boundExpr.collectFirst { - case b: BoundReference if !isSupportedDataType(b.dataType) => - b - } - badRef.map(b => - s"codegen dispatch: unsupported input type ${b.dataType} at ordinal ${b.ordinal}") + boundExpr.collectFirst(Function.unlift(inputRejection)) + } + + /** Why `expr` cannot be read as a codegen input, if it cannot. */ + private def inputRejection(expr: Expression): Option[String] = expr match { + case b: BoundReference if !isSupportedDataType(b.dataType) => + Some(s"codegen dispatch: unsupported input type ${b.dataType} at ordinal ${b.ordinal}") + case b: BoundReference => + val dups = duplicateStructFieldNames(b.dataType) + if (dups.isEmpty) { + None + } else { + Some( + s"codegen dispatch: duplicate struct field name ${dups.mkString(", ")} " + + s"in input type ${b.dataType} at ordinal ${b.ordinal}") + } + case _ => None } /** diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index 3b124859abe..6be7371a4c8 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -413,6 +413,8 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { */ private def emitSpecializedGetterExpr(target: String, idx: String, elemType: DataType): String = elemType match { + // Placeholder: [[emitWrite]]'s NullType branch only emits `setNull` and ignores its + // source, so this never reaches the generated Java. It just keeps the match total. case NullType => "null" case BooleanType => s"$target.getBoolean($idx)" case ByteType => s"$target.getByte($idx)" diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index be4bc9c3412..2ab308b540c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate._ import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.expressions.xml.{XPathBoolean, XPathDouble, XPathFloat, XPathInt, XPathList, XPathLong, XPathShort, XPathString} import org.apache.spark.sql.comet.DecimalPrecision +import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.{ScalarSubquery, SparkPlan} import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils import org.apache.spark.sql.internal.SQLConf @@ -645,7 +646,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val info = DataTypeInfo.newBuilder() val list = ListInfo.newBuilder() list.setElementType(elementType.get) - list.setContainsNull(a.containsNull) + // NullType children are always nullable; see Utils.declaredChildNullability. + list.setContainsNull(Utils.declaredChildNullability(a.elementType, a.containsNull)) nestedParquetFieldId(parentField, elementPath, includeFieldIds) .foreach(list.setElementFieldId) @@ -669,7 +671,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val map = MapInfo.newBuilder() map.setKeyType(keyType.get) map.setValueType(valueType.get) - map.setValueContainsNull(m.valueContainsNull) + map.setValueContainsNull(Utils.declaredChildNullability(m.valueType, m.valueContainsNull)) nestedParquetFieldId(parentField, keyPath, includeFieldIds).foreach(map.setKeyFieldId) nestedParquetFieldId(parentField, valuePath, includeFieldIds).foreach(map.setValueFieldId) @@ -686,7 +688,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val nestedFieldPath = if (nestedParentField.isDefined) Seq(field.name) else Seq.empty serializeDataType(field.dataType, nestedParentField, nestedFieldPath, includeFieldIds) } - val fieldNullable = s.map(f => Boolean.box(f.nullable)).asJava + val fieldNullable = + s.map(f => Boolean.box(Utils.declaredChildNullability(f.dataType, f.nullable))).asJava if (fieldDatatypes.exists(_.isEmpty)) { return None diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala index e2454f51322..18bff5bf856 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala @@ -199,6 +199,10 @@ object CometArrowStream extends Logging { * whose actual buffer carries validity bits must stay nullable even if Spark thought otherwise. * Taking only `raw.isNullable` here would advertise non-nullable when the next batch does carry * a null and crash native validation. + * + * Children come from the vector too, so a `NullType` map key arrives nullable (see + * `Utils.withNonNullableMapKeys`); `ArrowReader.getVectorSchemaRoot` would reject it when it + * rebuilds the `MapVector` from this schema. */ private def actualFieldOf(col: CometVector, expected: Field): Field = { val raw = col match { @@ -211,7 +215,7 @@ object CometArrowStream extends Logging { val nullable = expected.isNullable || raw.isNullable val fieldType = new FieldType(nullable, raw.getType, raw.getDictionary, expected.getMetadata) - new Field(expected.getName, fieldType, raw.getChildren) + Utils.withNonNullableMapKeys(new Field(expected.getName, fieldType, raw.getChildren)) } /** diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index d70fdab35e6..bf5bd48673c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -21,7 +21,7 @@ package org.apache.spark.sql.comet.util import java.io.{DataInputStream, DataOutputStream, File} import java.nio.ByteBuffer -import java.nio.channels.Channels +import java.nio.channels.{Channels, WritableByteChannel} import scala.jdk.CollectionConverters._ @@ -181,6 +181,20 @@ object Utils extends CometTypeShim with Logging { s"Unsupported data type: [${dt.getClass.getName}] ${dt.catalogString}") } + /** + * Nullability to declare for a nested child (array element, struct field, map value) of type + * `dataType`. A `NullType` child is always declared nullable, whatever Spark's `containsNull` / + * `valueContainsNull` / `StructField.nullable` says: every value is null, so a non-nullable + * flag is a contradiction, and native kernels that rebuild a list around the input's actual + * child (DataFusion's `map_entries` and `array_repeat`, Comet's `spark_array_slice`, ...) + * compare that child's nullability with the one they assume and fail on the mismatch + * (`map_entries(map_filter(map(), ...))`, `slice(filter(array(), ...), 1, 1)`). Applied here + * and in `QueryPlanSerde.serializeDataType` so the JVM-exported field and the type declared to + * native agree. Map keys are not children in this sense: Arrow requires them non-nullable. + */ + def declaredChildNullability(dataType: DataType, nullable: Boolean): Boolean = + nullable || dataType == NullType + /** Maps field from Spark to Arrow. NOTE: timeZoneId required for TimestampType */ def toArrowField(name: String, dt: DataType, nullable: Boolean, timeZoneId: String): Field = { dt match { @@ -189,7 +203,12 @@ object Utils extends CometTypeShim with Logging { new Field( name, fieldType, - Seq(toArrowField("element", elementType, containsNull, timeZoneId)).asJava) + Seq( + toArrowField( + "element", + elementType, + declaredChildNullability(elementType, containsNull), + timeZoneId)).asJava) case StructType(fields) => val fieldType = new FieldType(nullable, ArrowType.Struct.INSTANCE, null) new Field( @@ -197,24 +216,30 @@ object Utils extends CometTypeShim with Logging { fieldType, fields .map { field => - toArrowField(field.name, field.dataType, field.nullable, timeZoneId) + toArrowField( + field.name, + field.dataType, + declaredChildNullability(field.dataType, field.nullable), + timeZoneId) } .toSeq .asJava) case MapType(keyType, valueType, valueContainsNull) => val mapType = new FieldType(nullable, new ArrowType.Map(false), null) - // Note: Map Type struct can not be null, Struct Type key field can not be null - new Field( - name, - mapType, + // Note: Map Type struct can not be null, Struct Type key field can not be null (so the + // key is built here rather than through the StructType case and + // `declaredChildNullability`) + val entries = new Field( + MapVector.DATA_VECTOR_NAME, + new FieldType(false, ArrowType.Struct.INSTANCE, null), Seq( + toArrowField(MapVector.KEY_NAME, keyType, nullable = false, timeZoneId), toArrowField( - MapVector.DATA_VECTOR_NAME, - new StructType() - .add(MapVector.KEY_NAME, keyType, nullable = false) - .add(MapVector.VALUE_NAME, valueType, nullable = valueContainsNull), - nullable = false, + MapVector.VALUE_NAME, + valueType, + declaredChildNullability(valueType, valueContainsNull), timeZoneId)).asJava) + new Field(name, mapType, Seq(entries).asJava) case dataType => val fieldType = new FieldType(nullable, toArrowType(dataType, timeZoneId), null) new Field(name, fieldType, Seq.empty[Field].asJava) @@ -230,6 +255,125 @@ object Utils extends CometTypeShim with Logging { }.asJava) } + /** + * Returns `field` with every map key field, at any nesting depth, marked non-nullable. + * + * Arrow requires map keys to be non-nullable and `MapVector.initializeChildrenFromFields` + * enforces it, so an IPC reader rejects a schema whose key field is nullable. Comet always + * declares keys non-nullable (`toArrowField`), but Arrow Java's `MinorType.NULL` factory builds + * a `NullVector` from the field name alone, so a `NullType` key (e.g. `map()`) reports a + * nullable field once a vector for it has been created by `Field.createVector`, i.e. after a + * native import, codegen dispatch output, or row-to-Arrow conversion. Apply this to the fields + * written to IPC. + */ + def withNonNullableMapKeys(field: Field): Field = { + val children = field.getChildren.asScala.toSeq + if (children.isEmpty) { + return field + } + val newChildren = field.getType match { + case _: ArrowType.Map => + children.map { entries => + entries.getChildren.asScala.toSeq match { + case Seq(key, value) => + val nonNullKey = withNonNullableMapKeys(key) + val repairedKey = if (nonNullKey.isNullable) { + new Field( + nonNullKey.getName, + new FieldType( + false, + nonNullKey.getType, + nonNullKey.getDictionary, + nonNullKey.getMetadata), + nonNullKey.getChildren) + } else { + nonNullKey + } + new Field( + entries.getName, + entries.getFieldType, + Seq(repairedKey, withNonNullableMapKeys(value)).asJava) + case _ => withNonNullableMapKeys(entries) + } + } + case _ => children.map(withNonNullableMapKeys) + } + if (newChildren == children) { + field + } else { + new Field(field.getName, field.getFieldType, newChildren.asJava) + } + } + + /** + * Returns `root` unchanged when its declared schema already satisfies Arrow's non-nullable + * map-key invariant, otherwise a new root that shares the same vectors but advertises repaired + * fields. + * + * The check is against `root.getSchema`, not the live vectors' fields: a caller that built the + * root with already-repaired fields keeps its own root object. That matters when the row count + * is set after the writer is created (see `CometArrowPythonRunnerBase.startWriter`), because a + * replacement root copies the row count at construction and does not track the original. + */ + private def withNonNullableMapKeys(root: VectorSchemaRoot): VectorSchemaRoot = { + val declared = root.getSchema.getFields.asScala.toSeq + val repaired = declared.map(withNonNullableMapKeys) + if (repaired == declared) { + root + } else { + new VectorSchemaRoot(repaired.asJava, root.getFieldVectors, root.getRowCount) + } + } + + /** + * The only supported way to build an `ArrowStreamWriter` in Comet; enforced by the scalastyle + * `arrowstreamwriter` rule. + * + * Arrow requires map keys to be non-nullable and rejects a stream whose schema violates that + * ("Map data key type should be a non-nullable"). Comet always declares keys non-nullable in + * `toArrowField`, but Arrow's `MinorType.NULL` factory discards the field it is handed and + * rebuilds a nullable one (`Types.java` returns `new NullVector(field.getName())` even though + * `NullVector(Field)` exists), so any `NullType` map key silently turns nullable once a vector + * exists for it. Repairing here, rather than at each call site, means a new IPC writer cannot + * reintroduce the bug by forgetting to ask. + * + * Returns the writer together with the root it is bound to, which is `root` itself unless the + * declared schema needed repairing. Callers must use the returned root: a writer serializes the + * root it was constructed with, so a row count set on a superseded root is not seen (it + * surfaces in the Python worker as "Array length did not match record batch length"). + */ + def newArrowStreamWriter( + root: VectorSchemaRoot, + provider: DictionaryProvider, + channel: WritableByteChannel): (VectorSchemaRoot, ArrowStreamWriter) = { + val bound = withNonNullableMapKeys(root) + // scalastyle:off arrowstreamwriter + (bound, new ArrowStreamWriter(bound, provider, channel)) + // scalastyle:on arrowstreamwriter + } + + /** + * Whether an Arrow `Null` field is a direct child of a struct (map entries included) anywhere + * in `schema`. Arrow's `VectorAppender` cannot grow such a column: a struct's capacity is the + * minimum over its *direct* children, a `NullVector`'s capacity equals its value count and its + * `reAlloc()` is a no-op, so the struct's capacity loop never terminates. + * + * A list breaks that chain from both sides, because `ListVector` overrides + * `BaseRepeatedValueVector.getValueCapacity` with one that only looks at its own offset and + * validity buffers. So a `Null` under a list appends fine (`array(NULL)`), and so does a list + * under a struct even when the list holds nulls (`map(k, array(NULL))`) - the struct only sees + * the list's own, growable capacity. Top-level `NullVector`s are fine too. + */ + private def hasNullDirectlyUnderStruct(schema: Schema): Boolean = { + def check(field: Field): Boolean = { + val isStruct = field.getType.isInstanceOf[ArrowType.Struct] + field.getChildren.asScala.exists { child => + (isStruct && child.getType.isInstanceOf[ArrowType.Null]) || check(child) + } + } + schema.getFields.asScala.exists(check) + } + /** * Build a `StructType` from a sequence of Spark `Attribute`s. Avoids * `StructType.fromAttributes` (removed in Spark 4) and `DataTypeUtils.fromAttributes` (only on @@ -264,10 +408,10 @@ object Utils extends CometTypeShim with Logging { } val provider = batchProviderOpt.getOrElse(dictionaryProvider) - val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) + val (writeRoot, writer) = newArrowStreamWriter(root, provider, Channels.newChannel(out)) writer.start() writer.writeBatch() - root.clear() + writeRoot.clear() writer.close() if (out.size() > 0) { @@ -300,10 +444,10 @@ object Utils extends CometTypeShim with Logging { val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) val root = new VectorSchemaRoot(Seq(fieldVector).asJava) - val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) + val (writeRoot, writer) = newArrowStreamWriter(root, provider, Channels.newChannel(out)) writer.start() writer.writeBatch() - root.clear() + writeRoot.clear() writer.close() cbbos.toChunkedByteBuffer @@ -388,19 +532,24 @@ object Utils extends CometTypeShim with Logging { val reader = new ArrowStreamReader(Channels.newChannel(compressedInputStream), allocator) try { - // Comet decodes dictionaries during execution, so this shouldn't happen. - // If it does, fall back to the original uncoalesced buffers because each - // partition can have a different dictionary, and appending index vectors - // would silently mix indices from incompatible dictionaries. - if (!reader.getDictionaryVectors.isEmpty) { - logWarning( - "Unexpected dictionary-encoded column during BroadcastExchange coalescing; " + - "skipping coalesce") - reader.close() - if (targetRoot != null) { - targetRoot.close() - targetRoot = null + // Schemas that cannot be appended fall back to the original uncoalesced buffers: + // - Comet decodes dictionaries during execution, so a dictionary-encoded column + // shouldn't happen. If it does, each partition can have a different dictionary, + // and appending index vectors would silently mix incompatible dictionaries. + // - `VectorSchemaRootAppender` loops forever on a `NullVector` that is a direct + // child of a struct or of a map entry, e.g. `map(k, NULL)` or `map()`; see + // `hasNullDirectlyUnderStruct` for the Arrow mechanics. + val skipReason = + if (!reader.getDictionaryVectors.isEmpty) { + Some("unexpected dictionary-encoded column") + } else if (hasNullDirectlyUnderStruct(reader.getVectorSchemaRoot.getSchema)) { + Some("NullType directly under a struct or map entry") + } else { + None } + if (skipReason.isDefined) { + logWarning( + s"${skipReason.get} during BroadcastExchange coalescing; skipping coalesce") return (buffers, 0L, 0L) } while (reader.loadNextBatch()) { @@ -447,8 +596,8 @@ object Utils extends CometTypeShim with Logging { val outputStream = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) val compressedOutputStream = new DataOutputStream(codec.compressedOutputStream(outputStream)) - val writer = - new ArrowStreamWriter(targetRoot, null, Channels.newChannel(compressedOutputStream)) + val (_, writer) = + newArrowStreamWriter(targetRoot, null, Channels.newChannel(compressedOutputStream)) try { writer.start() writer.writeBatch() diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala index 651a9f25bc7..45f4bb22e0d 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala @@ -151,8 +151,17 @@ private[python] trait CometArrowPythonRunnerBase new FieldType(false, ArrowType.Struct.INSTANCE, null), childFields.asJava) val structVec = structField.createVector(allocator).asInstanceOf[StructVector] - writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava) - arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut)) + // Declare the root's schema from `structField`, not from `structVec.getField`: Arrow's + // `MinorType.NULL` factory builds a NullType map key as a nullable `NullVector`, so the + // live vector reports an invalid key field even when `structField` is valid. + // `Utils.newArrowStreamWriter` repairs the declared schema if needed and returns the root + // the writer is bound to, which is the one to close. + val declaredRoot = + new VectorSchemaRoot(Seq(structField).asJava, Seq[FieldVector](structVec).asJava, 0) + val (boundRoot, writer) = + Utils.newArrowStreamWriter(declaredRoot, null, Channels.newChannel(dataOut)) + writeRoot = boundRoot + arrowWriter = writer arrowWriter.start() } @@ -200,7 +209,11 @@ private[python] trait CometArrowPythonRunnerBase // identical. val childNames = inputStructType.fieldNames streamFields = batchFields.zipWithIndex.map { case (field, i) => - renamed(field, childNames(i), forceNullable = true) + // A NullType map key comes back from Arrow as a nullable `NullVector`, which + // `MapVector.initializeChildrenFromFields` rejects when `createVector` rebuilds the + // struct in `startWriter`. Repair the key nullability the same way + // `Utils.serializeBatches` does. + Utils.withNonNullableMapKeys(renamed(field, childNames(i), forceNullable = true)) } startWriter(streamFields, dataOut) } diff --git a/spark/src/test/resources/pyspark/test_pyarrow_udf.py b/spark/src/test/resources/pyspark/test_pyarrow_udf.py index 9a54300eed7..d4f0b505549 100644 --- a/spark/src/test/resources/pyspark/test_pyarrow_udf.py +++ b/spark/src/test/resources/pyspark/test_pyarrow_udf.py @@ -1091,6 +1091,91 @@ def _normalize(row): assert out == expected +def test_map_in_arrow_null_typed_map_children(spark, tmp_path, accelerated): + """ + `transform_values(map(), ...)` yields MapType(NullType, LongType) and `map(id, NULL)` + yields MapType(LongType, NullType). Arrow's MinorType.NULL factory rebuilds a NullType map + key as *nullable*, which MapVector.initializeChildrenFromFields rejects ("Map data key type + should be a non-nullable"), so the accelerated runner has to repair the key field before it + builds its destination struct from the live Comet vectors. + + pyarrow refuses to *declare* a non-nullable null-typed field (`A null type field may not + be non-nullable`), so the NullType-key map can only ever be a UDF input, never part of the + return schema; the UDF drops it and returns the NullType-value map unchanged. + """ + src = str(tmp_path / "src.parquet") + spark.range(0, 20, 1, 2).write.parquet(src) + df = spark.read.parquet(src).selectExpr( + "id", + "transform_values(map(), (k, v) -> id) AS null_key", + "map(id, NULL) AS null_value", + ) + schema_out = T.StructType( + [ + T.StructField("id", T.LongType()), + T.StructField("null_value", T.MapType(T.LongType(), T.NullType())), + ] + ) + + def drop_null_key(iterator): + for batch in iterator: + null_key_type = batch.schema.field("null_key").type + assert pa.types.is_map(null_key_type) + assert pa.types.is_null(null_key_type.key_type) + yield batch.select(["id", "null_value"]) + + result_df = df.mapInArrow(drop_null_key, schema_out) + _assert_plan_matches_mode(_executed_plan(result_df), accelerated) + + out = sorted((r["id"], r["null_value"]) for r in result_df.collect()) + assert out == [(i, {i: None}) for i in range(20)] + + +def test_map_in_arrow_null_typed_list_and_struct_children(spark, tmp_path, accelerated): + """ + `transform(array(id), x -> NULL)` yields ArrayType(NullType) and `named_struct('a', id, 'b', + NULL)` yields a struct with a NullType field. Both reach the runner as Comet vectors with a + NullVector child (no buffers, every slot null), so the destination struct root and the IPC + schema must describe a null-typed child pyarrow can read. Unlike the NullType map key these + can also be returned: pyarrow accepts a *nullable* null-typed field, and Comet declares a + NullType child nullable whatever Spark's containsNull / field nullability says. + """ + src = str(tmp_path / "src.parquet") + spark.range(0, 20, 1, 2).write.parquet(src) + df = spark.read.parquet(src).selectExpr( + "id", + "transform(array(id), x -> NULL) AS null_list", + "named_struct('a', id, 'b', NULL) AS null_struct", + ) + schema_out = T.StructType( + [ + T.StructField("id", T.LongType()), + T.StructField("null_list", T.ArrayType(T.NullType())), + T.StructField( + "null_struct", + T.StructType( + [T.StructField("a", T.LongType()), T.StructField("b", T.NullType())] + ), + ), + ] + ) + + def check_and_pass(iterator): + for batch in iterator: + assert pa.types.is_null(batch.schema.field("null_list").type.value_type) + assert pa.types.is_null(batch.schema.field("null_struct").type.field("b").type) + yield batch + + result_df = df.mapInArrow(check_and_pass, schema_out) + _assert_plan_matches_mode(_executed_plan(result_df), accelerated) + + out = sorted( + (r["id"], r["null_list"], (r["null_struct"]["a"], r["null_struct"]["b"])) + for r in result_df.collect() + ) + assert out == [(i, [None], (i, None)) for i in range(20)] + + def test_map_in_arrow_deeply_nested(spark, tmp_path, accelerated): """ Exercises direct source-vector serialization at depth > 1, in every nesting combination: diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql b/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql index 58272d164f1..6d58a62ac8f 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql @@ -118,3 +118,8 @@ SELECT array_repeat(CAST(NULL AS INT), cnt) FROM test_array_repeat query SELECT array_repeat(CAST(NULL AS STRING), cnt) FROM test_array_repeat + +-- A NullType element built by the JVM codegen dispatcher (containsNull=false in Spark) reaches +-- native declared nullable; the kernel's result must still match the planned type +query +SELECT array_repeat(filter(array(), x -> true), 2) FROM test_array_repeat diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_union.sql b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql index 1f3a9dc9dde..622606fc262 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_union.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql @@ -246,3 +246,8 @@ SELECT array_union(array(NULL, 99), b) FROM test_array_union -- conditional (CASE WHEN) arrays query SELECT array_union(CASE WHEN a IS NOT NULL THEN a ELSE array(0) END, b) FROM test_array_union + +-- A NullType element built by the JVM codegen dispatcher (containsNull=false in Spark) reaches +-- native declared nullable; the kernel's result must still match the planned type +query +SELECT array_union(filter(array(), x -> true), filter(array(), x -> true)) FROM test_array_union diff --git a/spark/src/test/resources/sql-tests/expressions/array/slice.sql b/spark/src/test/resources/sql-tests/expressions/array/slice.sql index 03a4bec0480..d3234ba391d 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/slice.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/slice.sql @@ -277,3 +277,8 @@ INSERT INTO test_slice_map VALUES (1, 10), (2, NULL), (3, 30) query SELECT slice(array(map(k, v), map(k + 1, v)), 1, 1) FROM test_slice_map + +-- A NullType element built by the JVM codegen dispatcher (containsNull=false in Spark) reaches +-- native declared nullable; the kernel's result must still match the planned type +query +SELECT slice(filter(array(), x -> true), 1, 1) FROM test_slice diff --git a/spark/src/test/resources/sql-tests/expressions/array/transform.sql b/spark/src/test/resources/sql-tests/expressions/array/transform.sql index b23e827535e..f7184a616ac 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/transform.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/transform.sql @@ -84,3 +84,14 @@ SELECT transform(array(1, 2, 3), x -> x * x) -- an untyped empty array leaves the element type as NullType query SELECT transform(array(), x -> x) + +-- Non-empty NullType results nested in list / struct outputs (codegen writes an Arrow NullVector +-- child; `array()` above only covers the empty case) +query +SELECT transform(a, x -> NULL) FROM test_transform + +query +SELECT transform(a, x -> named_struct('v', x, 'n', NULL)) FROM test_transform + +query +SELECT transform(array(NULL), x -> x) diff --git a/spark/src/test/resources/sql-tests/expressions/map/create_map.sql b/spark/src/test/resources/sql-tests/expressions/map/create_map.sql index 07df37963dd..4236c54d435 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/create_map.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/create_map.sql @@ -67,9 +67,6 @@ SELECT array(map()), struct(map(), map('a', NULL)) query SELECT map('a', array(NULL)), map('a', map()) -query -SELECT map_from_arrays(array(), array()) - -- carry the NullType children through a sort so the vector survives copy/spill paths query SELECT k, map(), map('a', NULL) FROM test_create_map ORDER BY k diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql b/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql index 2d172f436f8..11403fdfe65 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql @@ -63,3 +63,11 @@ SELECT array( map_entries(element_at(map(1, map(1, true)), id))[0], named_struct('key', 1, 'value', id IS NOT NULL)) AS a FROM test_map_entries_nested + +-- A NullType map value built by the JVM codegen dispatcher (valueContainsNull=false in Spark) +-- must reach DataFusion's map_entries declared nullable, or its ListArray build panics +query +SELECT map_entries(map_filter(map(), (k, v) -> true)) FROM test_map_entries + +query +SELECT map_entries(map('a', CAST(NULL AS int))) FROM test_map_entries diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql index 178c07f432a..9ccfe34d573 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql @@ -58,4 +58,8 @@ query SELECT map_from_arrays(array('a'), NULL) query -SELECT map_from_arrays(NULL, NULL) \ No newline at end of file +SELECT map_from_arrays(NULL, NULL) + +-- empty arrays produce MapType(NullType, NullType) +query +SELECT map_from_arrays(array(), array()) diff --git a/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql b/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql index a3ad834c3ff..d6353efa28c 100644 --- a/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql +++ b/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql @@ -33,3 +33,14 @@ SELECT named_struct('x', 1, 'y', 'hello', 'z', 3.14) query SELECT named_struct('x', a, 'y', 'fixed_val', 'z', c) FROM test_named_struct + +-- Spark keeps duplicate field names as distinct positional fields, but Arrow's StructVector +-- keys children by name, so both the native serde and the codegen dispatcher must fall back. +query expect_fallback(duplicate field names) +SELECT named_struct('x', a, 'x', b) FROM test_named_struct + +query expect_fallback(duplicate struct field name) +SELECT transform(array(a), v -> named_struct('x', v, 'x', NULL)) FROM test_named_struct + +query expect_fallback(duplicate struct field name) +SELECT transform(array(a), v -> named_struct('x', v, 'x', v + 1)) FROM test_named_struct diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index c4936cd2639..088ebeb01e7 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala @@ -19,18 +19,21 @@ package org.apache.comet +import scala.jdk.CollectionConverters._ + import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.SparkConf import org.apache.spark.serializer.JavaSerializer import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Add, AddMonths, BoundReference, Cast, Coalesce, Concat, CreateArray, CreateMap, DateFormatClass, ElementAt, Expression, GetStructField, IntegralDivide, LeafExpression, Length, Literal, MakeTimestamp, MicrosToTimestamp, MillisToTimestamp, MonthsBetween, Nondeterministic, Rand, Size, Substring, ToUnixTimestamp, Unevaluable, UnixMicros, UnixMillis, UnixSeconds, Upper} +import org.apache.spark.sql.catalyst.expressions.{Add, AddMonths, BoundReference, Cast, Coalesce, Concat, CreateArray, CreateMap, CreateNamedStruct, DateFormatClass, ElementAt, Expression, GetStructField, IntegralDivide, IsNull, LeafExpression, Length, Literal, MakeTimestamp, MicrosToTimestamp, MillisToTimestamp, MonthsBetween, Nondeterministic, Rand, Size, Substring, ToUnixTimestamp, Unevaluable, UnixMicros, UnixMillis, UnixSeconds, Upper} import org.apache.spark.sql.catalyst.expressions.codegen.{CodeFormatter, CodegenContext, CodegenFallback, ExprCode} import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.{ArrayColumnSpec, ArrowColumnSpec, MapColumnSpec, ScalarColumnSpec, StructColumnSpec, StructFieldSpec} +import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.udf.codegen.CometScalaUDFCodegen // Resolve Arrow vector classes through the codegen object so tests see the same `Class` objects @@ -323,6 +326,77 @@ class CometCodegenSourceSuite extends AnyFunSuite { } } + test("canHandle rejects duplicate struct field names, in outputs and inputs") { + // Spark keeps `named_struct('a', x, 'a', y)` as two positional fields, but Arrow's + // StructVector keys children by name and collapses them on allocation, so the generated + // ordinal-based child casts hit a missing or differently typed vector. `CometCreateNamedStruct` + // rejects this at serde level; the gate must do the same for a struct buried inside a larger + // expression that is dispatched as a whole (e.g. a `transform` lambda body). + val x = BoundReference(0, IntegerType, nullable = false) + val dupNull = CreateNamedStruct(Seq(Literal("a"), x, Literal("a"), Literal(null, NullType))) + val dupInt = CreateNamedStruct(Seq(Literal("a"), x, Literal("a"), Add(x, Literal(1)))) + val dupInput = IsNull( + BoundReference( + 0, + StructType(Seq(StructField("a", IntegerType), StructField("a", StringType))), + nullable = true)) + Seq(dupNull, dupInt, CreateArray(Seq(dupInt)), dupInput).foreach { expr => + val reason = CometBatchKernelCodegen.canHandle(expr) + assert( + reason.exists(_.contains("duplicate struct field name a")), + s"expected canHandle to reject $expr; got: $reason") + } + val distinct = + CreateNamedStruct(Seq(Literal("a"), x, Literal("b"), Literal(null, NullType))) + assert(CometBatchKernelCodegen.canHandle(distinct).isEmpty) + } + + test("NullType children are declared nullable whatever the Spark flags say") { + // Spark leaves containsNull / valueContainsNull / field nullability false on some untyped + // shapes (`filter(array(), ...)`, `map_filter(map(), ...)`), but a non-nullable Null child is a + // contradiction, and native kernels that rebuild a list around the input's actual child + // (map_entries, array_repeat, spark_array_slice) fail on the nullability mismatch. Non-Null + // children keep their flags; map keys stay non-nullable. + val exact = ArrayType( + StructType( + Seq( + StructField("n", NullType, nullable = false), + StructField("i", IntegerType, nullable = false), + StructField( + "m", + MapType(NullType, NullType, valueContainsNull = false), + nullable = false))), + containsNull = false) + val field = CometBatchKernelCodegen.toFfiArrowField("out", exact, nullable = false) + assert(!field.isNullable, "top-level nullability is the caller's") + val item = field.getChildren.get(0) + assert(!item.isNullable, "a non-Null element keeps containsNull = false") + val Seq(n, i, m) = item.getChildren.asScala.toSeq + assert(n.isNullable && !i.isNullable && !m.isNullable) + val Seq(key, value) = m.getChildren.get(0).getChildren.asScala.toSeq + assert(!key.isNullable, "map keys stay non-nullable") + assert(value.isNullable) + + val nullElements = CometBatchKernelCodegen + .toFfiArrowField("out", ArrayType(NullType, containsNull = false), nullable = true) + assert(nullElements.getChildren.get(0).isNullable) + + // The serde declares the same nullability to native. + val proto = QueryPlanSerde.serializeDataType(exact).get.getTypeInfo + assert(!proto.getList.getContainsNull) + val struct = proto.getList.getElementType.getTypeInfo.getStruct + assert( + struct.getFieldNullable(0) && !struct.getFieldNullable(1) && !struct.getFieldNullable(2)) + assert(struct.getFieldDatatypes(2).getTypeInfo.getMap.getValueContainsNull) + assert( + QueryPlanSerde + .serializeDataType(ArrayType(NullType, containsNull = false)) + .get + .getTypeInfo + .getList + .getContainsNull) + } + test("NullType output writes setNull without reading a source value") { // A NullType leaf has no Arrow data buffer, so the emitted write must be `setNull` only. val src = CometBatchKernelCodegen diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala index bf354356386..1acb28650bc 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -102,6 +102,22 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar checkShuffleAnswer(shuffled, 1) } + test("columnar shuffle with Map[NullType, _] column") { + // map() leaves a NullType key, which Arrow requires to be non-nullable in the IPC schema; + // transform_values keeps a second copy non-foldable so it is built by the child rather than + // constant-folded into a literal. `map()` is MapType(NullType, NullType) only while the + // legacy flag is off. + withSQLConf(SQLConf.LEGACY_CREATE_EMPTY_COLLECTION_USING_STRING_TYPE.key -> "false") { + val df = sql( + "SELECT id, map() AS m1, transform_values(map(), (k, v) -> id) AS m2 " + + "FROM VALUES (1), (2), (3) AS t(id)") + assert(df.schema("m1").dataType.asInstanceOf[MapType].keyType === NullType) + assert(df.schema("m2").dataType.asInstanceOf[MapType].keyType === NullType) + val shuffled = df.repartition(2, $"id") + checkShuffleAnswer(shuffled, 1) + } + } + test("columnar shuffle on nested struct including nulls") { Seq(10, 201).foreach { numPartitions => Seq("1.0", "10.0").foreach { ratio => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 5af9a37c267..50dd9b97f4f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHas import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, MetadataBuilder, NullType, StructField, StructType} import org.apache.comet.CometConf @@ -1382,4 +1382,40 @@ class CometJoinSuite extends CometTestBase { } } } + + test("Broadcast HashJoin with NullType map columns on the build side") { + // map(k, NULL) and transform_values(map(), ...) leave NullType children in the map type + // (a bare map() would be constant-folded into a literal the optimizer hoists above the + // join). The build side goes through CometBroadcastExchangeExec's batch coalescing and an + // Arrow IPC round trip on the JVM, where a nested NullVector used to hang the appender and + // the IPC reader rejected the NullType map key (see #5525). + Seq(true, false).foreach { aqe => + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "100", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + SQLConf.PREFER_SORTMERGEJOIN.key -> "false", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + // `map()` is MapType(NullType, NullType) only while this legacy flag is off. + SQLConf.LEGACY_CREATE_EMPTY_COLLECTION_USING_STRING_TYPE.key -> "false") { + withParquetTable((0 until 1000).map(i => (i, i % 5)), "tbl_a") { + withParquetTable((0 until 300).map(i => (i % 10, i + 2)), "tbl_b") { + val df = sql(""" + |SELECT /*+ BROADCAST(b) */ tbl_a._1, b.m1, b.m2 + |FROM tbl_a JOIN ( + | SELECT _1, map(_1, NULL) AS m1, transform_values(map(), (k, v) -> _1) AS m2 + | FROM tbl_b) b + |ON tbl_a._2 = b._1""".stripMargin) + // Guard the premise: without a NullType child this test still passes while + // covering nothing. + assert(df.schema("m1").dataType.asInstanceOf[MapType].valueType === NullType) + assert(df.schema("m2").dataType.asInstanceOf[MapType].keyType === NullType) + checkSparkAnswerAndOperator( + df, + Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec])) + } + } + } + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index 4510f9d0ac1..89c38aa0552 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -19,11 +19,26 @@ package org.apache.spark.sql.comet.util +import java.io.{ByteArrayInputStream, ByteArrayOutputStream} +import java.nio.channels.Channels + +import scala.concurrent.{Await, ExecutionContext, Future} +import scala.concurrent.duration.DurationInt +import scala.jdk.CollectionConverters._ + import org.apache.arrow.c.CDataDictionaryProvider +import org.apache.arrow.vector.{FieldVector, IntVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} +import org.apache.arrow.vector.ipc.ArrowStreamReader +import org.apache.arrow.vector.types.pojo.{ArrowType, Field} import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} +import org.apache.spark.sql.comet.execution.arrow.CometArrowConverters import org.apache.spark.sql.execution.vectorized.ConstantColumnVector -import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType, TimestampType} +import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, MapType, NullType, StringType, StructField, StructType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.util.io.ChunkedByteBuffer import org.apache.comet.CometArrowAllocator import org.apache.comet.vector.CometVector @@ -197,4 +212,270 @@ class UtilsSuite extends CometTestBase { } } } + + /** + * One map column of `numRows` rows. With an `IntegerType` key every row is a single entry `i -> + * NULL` (a `NullVector` map value); with a `NullType` key every row is an empty map, as `map()` + * produces (a `NullVector` map key). Both nest a `NullVector` inside the entries struct. + */ + private def nullTypeMapBatch(numRows: Int, keyType: DataType): ColumnarBatch = { + val field = Utils.toArrowField("m", MapType(keyType, NullType), nullable = true, "UTC") + val vector = field.createVector(CometArrowAllocator).asInstanceOf[MapVector] + vector.allocateNew() + val entries = vector.getDataVector.asInstanceOf[StructVector] + (0 until numRows).foreach { i => + vector.startNewValue(i) + keyType match { + case NullType => + vector.endValue(i, 0) + case _ => + entries.setIndexDefined(i) + entries.getChild(MapVector.KEY_NAME).asInstanceOf[IntVector].setSafe(i, i) + vector.endValue(i, 1) + } + } + entries.setValueCount(if (keyType == NullType) 0 else numRows) + vector.setValueCount(numRows) + new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector, null)), numRows) + } + + private def mapKeyField(field: Field): Field = field.getChildren.get(0).getChildren.get(0) + + /** One `array` column; row `i` holds `i` nulls. */ + private def nullListBatch(numRows: Int): ColumnarBatch = { + val field = Utils.toArrowField("l", ArrayType(NullType), nullable = true, "UTC") + val vector = field.createVector(CometArrowAllocator).asInstanceOf[ListVector] + vector.allocateNew() + (0 until numRows).foreach { i => + vector.startNewValue(i) + vector.endValue(i, i) + } + vector.setValueCount(numRows) + new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector, null)), numRows) + } + + /** One `array>` column; every row holds one struct. */ + private def nullStructListBatch(numRows: Int): ColumnarBatch = { + val elementType = StructType(Seq(StructField("a", NullType))) + val field = Utils.toArrowField("l", ArrayType(elementType), nullable = true, "UTC") + val vector = field.createVector(CometArrowAllocator).asInstanceOf[ListVector] + vector.allocateNew() + val elements = vector.getDataVector.asInstanceOf[StructVector] + (0 until numRows).foreach { i => + vector.startNewValue(i) + elements.setIndexDefined(i) + vector.endValue(i, 1) + } + elements.setValueCount(numRows) + vector.setValueCount(numRows) + new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector, null)), numRows) + } + + test("withNonNullableMapKeys restores the non-nullable key flag a NullVector drops") { + val batch = nullTypeMapBatch(2, NullType) + val field = batch.column(0).asInstanceOf[CometVector].getValueVector.getField + // `toArrowField` declares the key non-nullable, but Arrow's `MinorType.NULL` factory builds the + // key `NullVector` from the name alone, so the vector reports a nullable key. If this assertion + // starts failing, Arrow fixed that and `withNonNullableMapKeys` can go. + assert(mapKeyField(field).isNullable) + + val repaired = Utils.withNonNullableMapKeys(field) + assert(!mapKeyField(repaired).isNullable) + assert(mapKeyField(repaired).getType.isInstanceOf[ArrowType.Null]) + assert(repaired.getName == field.getName) + assert(repaired.getFieldType == field.getFieldType) + assert(repaired.getChildren.get(0).getFieldType == field.getChildren.get(0).getFieldType) + // Idempotent, and a no-op on fields that already satisfy the invariant. + assert(Utils.withNonNullableMapKeys(repaired) eq repaired) + batch.close() + } + + test("newArrowStreamWriter keeps a root whose declared schema is already valid") { + val batch = nullTypeMapBatch(2, NullType) + val vector = + batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector] + val declared = Utils.withNonNullableMapKeys(vector.getField) + // The live vector still reports a nullable key, so a root declared from it would be swapped + // for a repaired copy. One declared from `declared` must be kept as-is: the row count is set + // only after the writer exists, and a swapped root would not see it. + val root = new VectorSchemaRoot(Seq(declared).asJava, Seq(vector).asJava, 0) + val out = new ByteArrayOutputStream() + val (bound, writer) = Utils.newArrowStreamWriter(root, null, Channels.newChannel(out)) + assert(bound eq root) + root.setRowCount(2) + writer.start() + writer.writeBatch() + writer.end() + + val reader = + new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray), CometArrowAllocator) + assert(reader.loadNextBatch()) + assert(reader.getVectorSchemaRoot.getRowCount == 2) + assert(!mapKeyField(reader.getVectorSchemaRoot.getSchema.getFields.get(0)).isNullable) + reader.close() + batch.close() + } + + test("newArrowStreamWriter returns the root a later row count must be set on") { + val batch = nullTypeMapBatch(2, NullType) + val vector = + batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector] + // Declared from the live vector, so the key is nullable and the root must be swapped. Setting + // the row count on the returned root has to reach the writer; setting it on the original one + // would ship an empty batch ("Array length did not match record batch length" downstream). + val root = new VectorSchemaRoot(Seq(vector).asJava) + val out = new ByteArrayOutputStream() + val (bound, writer) = Utils.newArrowStreamWriter(root, null, Channels.newChannel(out)) + assert(bound ne root) + bound.setRowCount(2) + writer.start() + writer.writeBatch() + writer.end() + + val reader = + new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray), CometArrowAllocator) + assert(reader.loadNextBatch()) + assert(reader.getVectorSchemaRoot.getRowCount == 2) + assert(!mapKeyField(reader.getVectorSchemaRoot.getSchema.getFields.get(0)).isNullable) + reader.close() + batch.close() + } + + test("serializeBatches round-trips a NullType map key through Arrow IPC") { + // The IPC reader rebuilds a MapVector from the stream's schema and rejects a nullable key + // ("Map data key type should be a non-nullable"), which is exactly what a NullVector key + // reports unless the written schema is repaired. + val numRows = 3 + val batch = nullTypeMapBatch(numRows, NullType) + val (rowCount, buf) = Utils.serializeBatches(Iterator(batch)).next() + assert(rowCount == numRows) + + val decoded = Utils.decodeBatches(buf, "test").toSeq + assert(decoded.map(_.numRows()).sum == numRows) + decoded.foreach(_.close()) + } + + test("coalesceBroadcastBatches ships struct-nested NullType uncoalesced") { + // VectorSchemaRootAppender cannot grow a NullVector nested in a struct, including the map + // entries struct (NullVector.reAlloc is a no-op), so such buffers must be passed through, + // not appended. The list case pins that a struct below a list is still a struct. + val cases: Seq[(String, Int => ColumnarBatch)] = Seq( + "map" -> (nullTypeMapBatch(_, IntegerType)), + "map" -> (nullTypeMapBatch(_, NullType)), + "array>" -> nullStructListBatch) + cases.foreach { case (name, batch) => + val numRows = 4 + val numBatches = 3 + val batches = (0 until numBatches).map(_ => batch(numRows)) + val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq + + val (result, batchCount, totalRows) = Utils.coalesceBroadcastBatches(bufs.iterator) + // The pass-through signature: original buffers, nothing coalesced. + assert(batchCount == 0 && totalRows == 0, name) + assert(result.length == numBatches, name) + + val decoded = result.iterator.flatMap(b => Utils.decodeBatches(b, "test")).toSeq + assert(decoded.map(_.numRows()).sum == numRows.toLong * numBatches, name) + decoded.foreach(_.close()) + } + } + + test("coalesceBroadcastBatches bypasses exactly the schemas with a NullType under a struct") { + // Exhaustive over the shape space of the bypass rule: VectorAppender hangs only when a + // NullVector is a *direct* child of a struct (see `Utils.hasNullDirectlyUnderStruct` for + // the Arrow mechanics). Each shape runs the real appender under a timeout, so a rule that + // is too narrow shows up as a timeout on the hanging shapes instead of a hung build, and + // one that is too wide shows up as a needless bypass. + val nullStruct = StructType(Seq(StructField("a", NullType))) + val shapes: Seq[(DataType, Any)] = Seq( + NullType -> null, + ArrayType(NullType) -> new GenericArrayData(Array[Any](null)), + ArrayType(ArrayType(NullType)) -> + new GenericArrayData(Array[Any](new GenericArrayData(Array[Any](null)))), + nullStruct -> InternalRow(null), + ArrayType(nullStruct) -> new GenericArrayData(Array[Any](InternalRow(null))), + StructType(Seq(StructField("l", ArrayType(NullType)))) -> + InternalRow(new GenericArrayData(Array[Any](null))), + MapType(IntegerType, NullType) -> ArrayBasedMapData(Array[Any](1), Array[Any](null)), + // `map(k, array(NULL))`: the entry struct's direct child is a list, not a NullVector, so + // this still coalesces. + MapType(IntegerType, ArrayType(NullType)) -> + ArrayBasedMapData(Array[Any](1), Array[Any](new GenericArrayData(Array[Any](null)))), + MapType(NullType, NullType) -> ArrayBasedMapData(Array.empty[Any], Array.empty[Any])) + // A list insulates whatever is below it, so `inStruct` resets when descending into one. + def nullUnderStruct(dt: DataType, inStruct: Boolean): Boolean = dt match { + case NullType => inStruct + case ArrayType(element, _) => nullUnderStruct(element, inStruct = false) + case StructType(fields) => fields.exists(f => nullUnderStruct(f.dataType, inStruct = true)) + case MapType(k, v, _) => + nullUnderStruct(k, inStruct = true) || nullUnderStruct(v, inStruct = true) + case _ => false + } + + val numRows = 4 + val numBatches = 3 + shapes.foreach { case (dataType, value) => + val name = dataType.simpleString + val schema = StructType(Seq(StructField("c", dataType))) + val batches = (0 until numBatches).map { _ => + CometArrowConverters + .rowToArrowBatchIter( + Iterator.fill(numRows)(InternalRow(value)), + schema, + numRows, + "UTC", + CometArrowAllocator) + .next() + } + val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq + batches.foreach(_.close()) + + val (result, batchCount, totalRows) = Await.result( + Future(Utils.coalesceBroadcastBatches(bufs.iterator))(ExecutionContext.global), + 10.seconds) + + val expectBypass = nullUnderStruct(dataType, inStruct = false) + assert( + (batchCount == 0) == expectBypass, + s"$name: batchCount=$batchCount but bypass expected=$expectBypass") + assert(result.length == (if (expectBypass) numBatches else 1), name) + if (!expectBypass) assert(totalRows == numRows.toLong * numBatches, name) + + val decoded = result.iterator.flatMap(b => Utils.decodeBatches(b, "test")).toSeq + assert(decoded.map(_.numRows()).sum == numRows.toLong * numBatches, name) + decoded.foreach(_.close()) + } + } + + test("coalesceBroadcastBatches keeps coalescing plain null lists") { + // A NullVector directly under a list is safe to append: the list's capacity loop only looks + // at its own offset and validity buffers, and ListVector.setValueCount sets the child's + // count from the last offset. Bypassing coalescing here would cost every consuming task one + // IPC stream per original buffer instead of one. + val numRows = 4 + val numBatches = 3 + val batches = (0 until numBatches).map(_ => nullListBatch(numRows)) + val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq + + val (result, batchCount, totalRows) = Utils.coalesceBroadcastBatches(bufs.iterator) + assert(batchCount == numBatches) + assert(totalRows == numRows.toLong * numBatches) + assert(result.length == 1) + + // Every list keeps its length across the append. Read each batch before the stream moves + // on, since the reader reclaims a batch's buffers when it loads the next one. + def listLengths(bufs: Iterator[ChunkedByteBuffer]): Seq[Int] = + bufs + .flatMap(b => Utils.decodeBatches(b, "test")) + .flatMap { batch => + val lengths = + (0 until batch.numRows()).map(r => batch.column(0).getArray(r).numElements()) + batch.close() + lengths + } + .toSeq + val expected = Seq.fill(numBatches)(0 until numRows).flatten + assert(listLengths(bufs.iterator) == expected, "uncoalesced input") + assert(listLengths(result.iterator) == expected, "coalesced output") + } } From 4a037a880257175249212e074b08bbc8ac541c6a Mon Sep 17 00:00:00 2001 From: grorge Date: Fri, 4 Sep 2026 14:33:24 +0800 Subject: [PATCH 3/6] fix: gate the NullType compositions that reach native kernel gaps Admitting NullType outputs from the codegen dispatcher lets values reach native kernels that assume NullType never arrives. The serde now refuses: * make_array (builds a single row), array_union (drops entries) and array_intersect (returns the other side's entries; reported Unsupported so the codegen dispatcher runs it at every setting, since the Incompatible branch has no dispatcher fallback under allowIncompatible), array_except's unsupported element types likewise, array_repeat and slice (non-nullable item promised nullable), collect_list/collect_set (nested nullability mismatch) and hash/xxhash64 (no Null arm) over NullType-bearing inputs, and map_from_arrays with a literal array beside a per-row one (the native map kernel reads a scalar list through its first row; pre-existing, found while probing the NullType flavour). * Any non-deterministic child under the null guards of CometElementAt (ANSI), CometArrayAppend, CometMapFromArrays, CometArraysZip, CometCoalesce and CometSize, whichever argument is the nullable one: native CASE evaluates the THEN copy on the rows the predicate selected, and a codegen-dispatched child (any lambda) is one cached kernel shared by both copies, so the kernel never sees the values Spark's single evaluation produces. CometSize builds no guard for a non-nullable child or in legacy mode, where native already answers -1, and needs no gate there. CometIf, CometCaseWhen and CometCoalesce refuse a NullType result: native CASE merges its branches' rows through Arrow's merge_n, which cannot build a NullArray with a validity bitmap. Native GetStructField returns a scalar for a scalar struct input instead of a one-row array, which a CASE result builder would slice past ("range end index 2 out of range for slice of length 1"), and rebuilds a Null-typed field as a fresh NullArray, since a kernel that grew the struct through MutableArrayData (element_at on an out-of-range index) hands over a Null child carrying a validity bitmap that fails validation once projected. The first two shapes surface once the sweep stops the optimizer from folding them away. array_union/intersect/except cast both sides to a deeply-nullable element type, since the native set-op kernel asserts identical nested nullability and a lambda variable arrives nullable where a literal field does not. The native row shuffle writer's field-major paths gain the Null struct field case their row-major path already had, so a struct with a NullType field now shuffles through the JVM columnar shuffle instead of panicking, and the writer recreates every Null-bearing builder after each batch: NullBuilder::finish keeps its length, so a second batch used to panic on a Null struct field longer than its parent and miscount a top-level Null column, a Null map value or a Null list element. Native to_csv yields NULL for a row that renders to an empty string and reports itself nullable, as Spark does: Spark hands the row to univocity's writeRowToString with skipEmptyLines on, so a struct with a lone null field (under the default empty nullValue) is NULL, not "". Pre-existing and independent of NullType; the allowIncompatible sweep profile found it on Spark 3.4, whose interpreted StructsToCsv shows the NULL, while Spark 3.5+ crashes in its own generated code on that NULL and the sweep counts the case as invalid there. CometNullTypeCompositionSuite sweeps the NullType producers under consumers, operators and nesting containers, across ANSI, nullable, non-deterministic and cross-input (stateful argument beside a nullable one) settings, with the optimizer's null and comparison simplifications excluded so no consumer folds to a literal, and under physical profiles that vary how rows are batched and which exchange path they take (native batches of two rows, the JVM shuffle's bypass and sort-based writers with and without forced spills, native shuffle, AQE, native columnar-to-row) plus one that opts every registered serde into its native kernel through allowIncompatible, so Incompatible serdes do not hide behind the codegen dispatcher; the operators include hash partitioning, a null-safe join key and a scalar subquery over the value itself. It runs with no tolerated failures and floors on the compared and natively executed counts, and checks that its templates reach every registered array, map, struct and any-type aggregate serde. CometInMemoryCacheSuite round-trips NullType columns and children through Comet's Arrow cache serializer, and a remote-shuffle decode unit test covers Null columns and children. UtilsSuite forces serialization eagerly for Scala 2.12. Verified on Spark 4.1 / Scala 2.13, Spark 3.5 / Scala 2.12 and Spark 3.4 / Scala 2.12. Assisted-by: Claude Code (claude-fable-5) --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + dev/scalastyle-config.xml | 5 +- native/common/src/struct_nulls.rs | 27 +- native/shuffle/src/remote_schema_tests.rs | 55 +- native/shuffle/src/spark_unsafe/row.rs | 141 +++- native/spark-expr/src/array_funcs/size.rs | 9 +- native/spark-expr/src/csv_funcs/to_csv.rs | 87 ++- .../src/struct_funcs/get_struct_field.rs | 60 +- .../codegen/CometBatchKernelCodegen.scala | 79 +- .../CometBatchKernelCodegenOutput.scala | 5 +- .../apache/comet/serde/QueryPlanSerde.scala | 3 +- .../org/apache/comet/serde/SupportLevel.scala | 33 + .../org/apache/comet/serde/aggregates.scala | 43 +- .../scala/org/apache/comet/serde/arrays.scala | 243 ++++-- .../org/apache/comet/serde/conditional.scala | 33 + .../scala/org/apache/comet/serde/hash.scala | 11 +- .../scala/org/apache/comet/serde/maps.scala | 10 +- .../apache/spark/sql/comet/util/Utils.scala | 30 +- .../python/CometArrowPythonRunnerBase.scala | 14 +- .../resources/pyspark/test_pyarrow_udf.py | 6 +- .../expressions/aggregate/collect_list.sql | 6 + .../expressions/aggregate/collect_set.sql | 4 + .../expressions/array/array_except.sql | 8 + .../expressions/array/array_intersect.sql | 20 + .../expressions/array/array_repeat.sql | 10 + .../expressions/array/array_union.sql | 18 +- .../expressions/array/arrays_zip.sql | 6 + .../expressions/array/create_array.sql | 5 + .../expressions/array/element_at_ansi.sql | 8 + .../sql-tests/expressions/array/slice.sql | 6 + .../sql-tests/expressions/array/transform.sql | 8 - .../expressions/conditional/case_when.sql | 5 + .../expressions/conditional/coalesce.sql | 11 + .../expressions/conditional/if_expr.sql | 6 + .../sql-tests/expressions/hash/hash.sql | 4 + .../sql-tests/expressions/map/create_map.sql | 3 - .../sql-tests/expressions/map/map_entries.sql | 3 - .../expressions/map/map_from_arrays.sql | 11 + .../struct/create_named_struct.sql | 3 - .../expressions/struct/get_struct_field.sql | 6 + .../comet/CometArrayExpressionSuite.scala | 27 + .../comet/CometCodegenSourceSuite.scala | 35 +- .../comet/CometCsvExpressionSuite.scala | 26 +- .../comet/CometNullTypeCompositionSuite.scala | 728 ++++++++++++++++++ .../exec/CometColumnarShuffleSuite.scala | 29 + .../comet/exec/CometInMemoryCacheSuite.scala | 33 + .../apache/comet/exec/CometJoinSuite.scala | 3 +- .../spark/sql/comet/util/UtilsSuite.scala | 164 ++-- 49 files changed, 1760 insertions(+), 332 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index a019e9a721f..0930d93356b 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -499,6 +499,7 @@ jobs: org.apache.comet.serde.CometScalarFunctionSuite org.apache.comet.serde.CometLiteralSuite org.apache.comet.CometFallbackInvarianceSuite + org.apache.comet.CometNullTypeCompositionSuite fail-fast: false name: ${{ matrix.profile.name }} [${{ matrix.suite.name }}] runs-on: ubuntu-24.04 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 66a1b62ba3d..63cb7023523 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -253,6 +253,7 @@ jobs: org.apache.comet.serde.CometScalarFunctionSuite org.apache.comet.serde.CometLiteralSuite org.apache.comet.CometFallbackInvarianceSuite + org.apache.comet.CometNullTypeCompositionSuite fail-fast: false name: ${{ matrix.os }}/${{ matrix.profile.name }} [${{ matrix.suite.name }}] diff --git a/dev/scalastyle-config.xml b/dev/scalastyle-config.xml index e58e6b73581..edd3de30730 100644 --- a/dev/scalastyle-config.xml +++ b/dev/scalastyle-config.xml @@ -197,9 +197,8 @@ This file is divided into 3 sections: new ArrowStreamWriter ). -use arrow::array::{make_array, Array, ArrayRef, StructArray}; +use arrow::array::{make_array, Array, ArrayRef, NullArray, StructArray}; use arrow::buffer::NullBuffer; +use arrow::datatypes::DataType; use arrow::error::ArrowError; use std::sync::Arc; @@ -48,6 +49,14 @@ pub fn child_with_parent_nulls( ordinal: usize, ) -> Result { let child = struct_array.column(ordinal); + // A Null-typed child is all-null whatever the parent's mask says, and a `NullArray` may not + // carry a validity bitmap ("Arrays of type Null cannot contain a null bitmask"). Rebuild it + // rather than reuse the child: a kernel that grew the struct through `MutableArrayData` + // (`element_at` on an out-of-range index) can hand over a child that already carries such a + // bitmap, which only fails validation once it is projected out. + if child.data_type() == &DataType::Null { + return Ok(Arc::new(NullArray::new(child.len()))); + } match struct_array.nulls() { Some(parent) if parent.null_count() > 0 => { let combined = NullBuffer::union(Some(parent), child.nulls()); @@ -85,6 +94,22 @@ mod tests { StructArray::new(fields, vec![child], nulls) } + /// A Null-typed child cannot carry a validity bitmap, so it is rebuilt all-null instead of + /// receiving the parent's mask; a `MutableArrayData`-grown struct can even hand over a child + /// that already carries one, which only fails validation once projected out. + #[test] + fn null_typed_child_is_rebuilt_without_a_bitmap() { + let fields: Fields = vec![Arc::new(Field::new("n", DataType::Null, true))].into(); + let parent = NullBuffer::from(vec![true, false, true]); + let struct_array = + StructArray::new(fields, vec![Arc::new(NullArray::new(3))], Some(parent)); + let out = child_with_parent_nulls(&struct_array, 0).unwrap(); + assert_eq!(out.data_type(), &DataType::Null); + assert_eq!(out.len(), 3); + assert!(out.nulls().is_none(), "a NullArray must not carry a validity bitmap"); + out.to_data().validate_full().unwrap(); + } + /// The case both bugs came from: the parent is null where the child still holds a value. #[test] fn hidden_child_value_under_a_null_parent_becomes_null() { diff --git a/native/shuffle/src/remote_schema_tests.rs b/native/shuffle/src/remote_schema_tests.rs index 52fafa0cf6c..ae2089e2379 100644 --- a/native/shuffle/src/remote_schema_tests.rs +++ b/native/shuffle/src/remote_schema_tests.rs @@ -18,9 +18,9 @@ use crate::{decode_remote_shuffle_batch, CompressionCodec, ShuffleBlockWriter}; use arrow::array::{ Array, ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, FixedSizeListArray, - Int16Array, Int32Array, LargeListArray, ListArray, MapArray, PrimitiveDictionaryBuilder, - RecordBatch, RecordBatchOptions, StringArray, StringDictionaryBuilder, StructArray, - UInt16Array, + Int16Array, Int32Array, Int64Array, LargeListArray, ListArray, MapArray, NullArray, + PrimitiveDictionaryBuilder, RecordBatch, RecordBatchOptions, StringArray, + StringDictionaryBuilder, StructArray, UInt16Array, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{ @@ -570,3 +570,52 @@ fn remote_shuffle_preserves_row_count_without_columns() { assert_eq!(decoded.num_columns(), 0); assert_eq!(decoded.num_rows(), 3); } + +// A `NullType` column, and one nested under a list, a map value and a struct field, decode +// unchanged: a `NullArray` owns no buffers, so the encoding, the dictionary decoding and the +// nested-nullability reconciliation all have to pass it through by length alone. +#[test] +fn null_type_columns_and_children_survive_remote_shuffle() { + let rows = 3; + let null_field = |name: &str| Arc::new(Field::new(name, DataType::Null, true)); + let top_level: ArrayRef = Arc::new(NullArray::new(rows)); + let list: ArrayRef = Arc::new(ListArray::new( + null_field("element"), + OffsetBuffer::from_lengths([1, 0, 2]), + Arc::new(NullArray::new(3)), + None, + )); + let struct_fields = Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("n", DataType::Null, true), + ]); + let structs: ArrayRef = Arc::new(StructArray::new( + struct_fields, + vec![ + Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])), + Arc::new(NullArray::new(rows)), + ], + Some(NullBuffer::from(vec![true, false, true])), + )); + let entry_fields = Fields::from(vec![ + Field::new("key", DataType::Int64, false), + Field::new("value", DataType::Null, true), + ]); + let entries = StructArray::new( + entry_fields.clone(), + vec![ + Arc::new(Int64Array::from(vec![10, 20, 30])), + Arc::new(NullArray::new(3)), + ], + None, + ); + let map: ArrayRef = Arc::new(MapArray::new( + Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)), + OffsetBuffer::from_lengths([2, 0, 1]), + entries, + None, + false, + )); + let columns = vec![top_level, list, structs, map]; + assert_roundtrip(columns.clone(), columns); +} diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index a3f3b3d36fc..263b7a33662 100644 --- a/native/shuffle/src/spark_unsafe/row.rs +++ b/native/shuffle/src/spark_unsafe/row.rs @@ -572,6 +572,11 @@ fn append_nested_struct_fields_field_major( } } } + // A Null field carries no data: every row is null, whether or not the struct is. + DataType::Null => { + let field_builder = get_field_builder!(struct_builder, NullBuilder, field_idx); + field_builder.append_nulls(num_rows); + } _ => { unreachable!( "Unsupported data type of struct field: {:?}", @@ -1020,6 +1025,11 @@ fn append_struct_fields_field_major( } } } + // A Null field carries no data: every row is null, whether or not the struct is. + DataType::Null => { + let field_builder = get_field_builder!(struct_builder, NullBuilder, field_idx); + field_builder.append_nulls(row_end - row_start); + } _ => { unreachable!( "Unsupported data type of struct field: {:?}", @@ -1179,9 +1189,7 @@ fn append_columns( } DataType::Null => { let null_builder = downcast_builder_ref!(NullBuilder, builder); - for _ in row_start..row_end { - null_builder.append_null(); - } + null_builder.append_nulls(row_end - row_start); } DataType::Timestamp(TimeUnit::Microsecond, _) => { append_column_to_builder!( @@ -1418,6 +1426,14 @@ pub fn process_sorted_row_partition( .collect(); let batch = make_batch(array_refs?, n)?; + // A finished `NullBuilder` keeps its length; see `recreate_null_type_builders`. + recreate_null_type_builders( + &mut data_builders, + schema, + batch_size, + prefer_dictionary_ratio, + )?; + frozen.clear(); let mut cursor = Cursor::new(&mut frozen); @@ -1440,6 +1456,37 @@ pub fn process_sorted_row_partition( )) } +/// Whether `dt` is `Null` or nests a `Null` anywhere below a list, struct or map. +fn contains_null_type(dt: &DataType) -> bool { + match dt { + DataType::Null => true, + DataType::List(field) | DataType::LargeList(field) | DataType::Map(field, _) => { + contains_null_type(field.data_type()) + } + DataType::Struct(fields) => fields.iter().any(|f| contains_null_type(f.data_type())), + _ => false, + } +} + +/// Replaces every builder whose type holds a `Null` somewhere, after its batch has been +/// finished. `NullBuilder::finish` keeps its length (a `NullArray` owns no buffers to hand +/// over), so such a builder would carry this batch's rows into the next: a top-level Null column +/// comes out longer than the batch, and a Null struct field longer than its parent panics in +/// `StructBuilder::finish`. Every other builder resets on finish and is kept. +fn recreate_null_type_builders( + builders: &mut [Box], + schema: &[DataType], + batch_size: usize, + prefer_dictionary_ratio: f64, +) -> Result<(), CometError> { + for (builder, datatype) in builders.iter_mut().zip(schema.iter()) { + if contains_null_type(datatype) { + *builder = make_builders(datatype, batch_size, prefer_dictionary_ratio)?; + } + } + Ok(()) +} + fn builder_to_array( builder: &mut Box, datatype: &DataType, @@ -1503,10 +1550,98 @@ fn make_batch(arrays: Vec, row_count: usize) -> Result> = schema + .iter() + .map(|dt| make_builders(dt, batch_size, 1.0).unwrap()) + .collect(); + + for batch in 0..2 { + let struct_builder = builders[0] + .as_any_mut() + .downcast_mut::() + .unwrap(); + for row in 0..batch_size { + struct_builder + .field_builder::(0) + .unwrap() + .append_value(row as i64); + struct_builder + .field_builder::(1) + .unwrap() + .append_null(); + struct_builder.append(true); + } + let nested_builder = builders[1] + .as_any_mut() + .downcast_mut::() + .unwrap(); + for _ in 0..batch_size { + let inner = nested_builder.field_builder::(0).unwrap(); + inner + .field_builder::(0) + .unwrap() + .append_null(); + inner.field_builder::(1).unwrap().append_null(); + inner.append(true); + nested_builder.append(true); + } + builders[2] + .as_any_mut() + .downcast_mut::() + .unwrap() + .append_nulls(batch_size); + + let arrays: Vec = builders + .iter_mut() + .zip(schema.iter()) + .map(|(builder, dt)| builder_to_array(builder, dt, 1.0).unwrap()) + .collect(); + for (array, dt) in arrays.iter().zip(schema.iter()) { + assert_eq!(array.len(), batch_size, "batch {batch} of {dt}"); + } + let outer = arrays[0].as_any().downcast_ref::().unwrap(); + assert_eq!( + outer.column(1).len(), + batch_size, + "batch {batch} Null field" + ); + + recreate_null_type_builders(&mut builders, &schema, batch_size, 1.0).unwrap(); + for (builder, dt) in builders.iter().zip(schema.iter()) { + assert_eq!(builder.len(), 0, "builder for {dt} after batch {batch}"); + } + } + } + #[test] fn test_append_null_row_to_struct_builder() { let data_type = DataType::Struct(Fields::from(vec![ diff --git a/native/spark-expr/src/array_funcs/size.rs b/native/spark-expr/src/array_funcs/size.rs index 2311da5fa9e..a1776d0f308 100644 --- a/native/spark-expr/src/array_funcs/size.rs +++ b/native/spark-expr/src/array_funcs/size.rs @@ -130,11 +130,10 @@ fn spark_size_list_like(array: &ArrayRef) -> Result { } }; - // Fast path for the production shape: `CometSize.convert` wraps size() in a - // `CASE WHEN isnotnull(child)` that filters null rows out before the THEN - // branch runs, so this function only ever sees a null-free array in a real - // Comet plan. Return the length kernel output as-is; skip the downcast and - // `Int32Array::clone` that `rewrite_nulls_to_minus_one` would otherwise pay. + // Fast path for a null-free array: return the length kernel output as-is and + // skip the downcast and `Int32Array::clone` that `rewrite_nulls_to_minus_one` + // would otherwise pay. `CometSize.convert` sends nullable input here only in + // legacy mode, where -1 is the answer for a null collection. if array.null_count() == 0 { return Ok(lengths); } diff --git a/native/spark-expr/src/csv_funcs/to_csv.rs b/native/spark-expr/src/csv_funcs/to_csv.rs index 01fdc901cb7..7e74594aea4 100644 --- a/native/spark-expr/src/csv_funcs/to_csv.rs +++ b/native/spark-expr/src/csv_funcs/to_csv.rs @@ -80,8 +80,10 @@ impl PhysicalExpr for ToCsv { Ok(DataType::Utf8) } - fn nullable(&self, input_schema: &Schema) -> Result { - self.expr.nullable(input_schema) + fn nullable(&self, _: &Schema) -> Result { + // Spark's `StructsToCsv.nullable` is unconditionally true: a row that renders to an empty + // string (a lone null field) yields NULL, whatever the input struct's nullability. + Ok(true) } fn evaluate(&self, batch: &RecordBatch) -> Result { @@ -201,7 +203,15 @@ pub fn to_csv_inner( } } } - builder.append_value(&csv_string); + // Spark renders a row through univocity's `writeRowToString`, whose `skipEmptyLines` + // setting (always on for Spark's CSV writer) turns an empty rendered row into `null`: + // a struct with a single null field (with the default empty `nullValue`) is NULL, + // not "". A row with two null fields renders as the delimiter and is kept. + if csv_string.is_empty() { + builder.append_null(); + } else { + builder.append_value(&csv_string); + } } } Ok(Arc::new(builder.finish())) @@ -216,3 +226,74 @@ fn escape_value(value: &str, quote_char: char, escape_char: char, output: &mut S output.push(ch); } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::datatypes::{Field, Fields}; + + fn default_write_options() -> CsvWriteOptions { + CsvWriteOptions::new( + ",".to_string(), + "\"".to_string(), + "\\".to_string(), + "".to_string(), + false, + true, + true, + ) + } + + fn to_csv_strings(array: &StructArray, write_options: &CsvWriteOptions) -> Vec> { + let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + cast_options.null_string = write_options.null_value.clone(); + let out = to_csv_inner(array, &cast_options, write_options).unwrap(); + as_string_array(&out) + .iter() + .map(|v| v.map(str::to_string)) + .collect() + } + + // Spark's `to_csv` hands the row to univocity's `writeRowToString` with `skipEmptyLines`, so a + // row whose rendering is empty comes back NULL rather than "". Only a lone null field (or a + // lone empty `nullValue`) renders empty; two null fields render as the delimiter. + #[test] + fn empty_rendered_row_is_null() { + let one_field: Fields = Fields::from(vec![Field::new("a", DataType::Int64, true)]); + let single = StructArray::new( + one_field, + vec![Arc::new(Int64Array::from(vec![Some(7_i64), None]))], + None, + ); + assert_eq!( + to_csv_strings(&single, &default_write_options()), + vec![Some("7".to_string()), None] + ); + + let two_fields: Fields = Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("n", DataType::Null, true), + ]); + let pair = StructArray::new( + two_fields, + vec![ + Arc::new(Int64Array::from(vec![Some(7_i64), None])), + Arc::new(arrow::array::NullArray::new(2)), + ], + None, + ); + assert_eq!( + to_csv_strings(&pair, &default_write_options()), + vec![Some("7,".to_string()), Some(",".to_string())] + ); + + // A non-empty `nullValue` renders, so the row is kept. + let mut named_null = default_write_options(); + named_null.null_value = "N".to_string(); + assert_eq!( + to_csv_strings(&single, &named_null), + vec![Some("7".to_string()), Some("N".to_string())] + ); + } +} diff --git a/native/spark-expr/src/struct_funcs/get_struct_field.rs b/native/spark-expr/src/struct_funcs/get_struct_field.rs index b815b4ed9a0..a5eee75b914 100644 --- a/native/spark-expr/src/struct_funcs/get_struct_field.rs +++ b/native/spark-expr/src/struct_funcs/get_struct_field.rs @@ -99,9 +99,17 @@ impl PhysicalExpr for GetStructField { self.ordinal, )?)) } - ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => Ok(ColumnarValue::Array( - child_with_parent_nulls(&struct_array, self.ordinal)?, - )), + // A scalar struct (a NULL branch of CASE, a scalar subquery) holds one row, so its + // field is a scalar too. Returning the projected one-row array instead makes the + // consumer believe it has a full column: a CASE result builder then slices row + // ranges out of it and panics ("range end index 2 out of range for slice of length + // 1"). + ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => { + let projected = child_with_parent_nulls(&struct_array, self.ordinal)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &projected, 0, + )?)) + } value => Err(DataFusionError::Execution(format!( "Expected a struct array, got {value:?}" ))), @@ -136,16 +144,58 @@ impl Display for GetStructField { #[cfg(test)] mod tests { use super::*; - use arrow::array::{ArrayRef, Int64Array}; + use arrow::array::{ArrayRef, Int64Array, NullArray}; use arrow::buffer::NullBuffer; use arrow::datatypes::Fields; - use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::expressions::{Column, Literal}; // A field of a NULL struct must be NULL (Spark semantics) even when the child buffer holds a // non-null value at that row -- Arrow stores child validity independently of the parent // struct's null mask, so a logically-null struct column read from parquet can still carry a // populated child buffer. Without propagating the parent null mask, `isnotnull(struct.field)` // wrongly evaluates TRUE for a null struct. + // A scalar struct input yields a scalar field, so a consumer that broadcasts it over the batch + // (a CASE branch, a projection beside a column) sees one value per row rather than a one-row + // array. A null struct scalar yields a null scalar, for a typed field and a Null-typed one. + #[test] + fn field_of_scalar_struct_is_scalar() { + let fields: Fields = Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("n", DataType::Null, true), + ]); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![0_i64, 1, 2, 3]))], + ) + .unwrap(); + + let present = ScalarValue::Struct(Arc::new(StructArray::new( + fields.clone(), + vec![ + Arc::new(Int64Array::from(vec![7_i64])), + Arc::new(NullArray::new(1)), + ], + None, + ))); + let absent = ScalarValue::Struct(Arc::new(StructArray::new_null(fields, 1))); + + for (scalar, ordinal, expected) in [ + (present.clone(), 0, ScalarValue::Int64(Some(7))), + (present, 1, ScalarValue::Null), + (absent.clone(), 0, ScalarValue::Int64(None)), + (absent, 1, ScalarValue::Null), + ] { + let expr = GetStructField::new(Arc::new(Literal::new(scalar)), ordinal); + match expr.evaluate(&batch).unwrap() { + ColumnarValue::Scalar(value) => assert_eq!(value, expected), + ColumnarValue::Array(array) => { + panic!("expected a scalar for ordinal {ordinal}, got array {array:?}") + } + } + } + } + #[test] fn field_of_null_struct_is_null() { // Child is non-null at every row; the struct itself is null at rows 1 and 3. diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 2966d5ce855..57de9105d47 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -82,11 +82,8 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come * Type surface the kernel covers on both input and output sides. Recursive: complex types are * supported when their children are. * - * Duplicate struct field names are excluded, an output-side rule: Arrow addresses a - * `StructVector`'s children by name, so `named_struct('x', 10, 'x', 20)` collapses into a - * single child and the generated writer NPEs on the missing ordinal-1 vector. - * `CometCreateNamedStruct` declines them on the native path for the same reason, but a struct - * nested inside a dispatcher-built value (a `CreateMap` value) never reaches that check. + * Duplicate struct field names are a separate rule, [[hasDuplicateStructFieldNames]], so that + * [[canHandle]] can name them in its reason. * * `NullType` is output-only: [[CometBatchKernelCodegenOutput]] can write an all-null Arrow * `NullVector`, but `CometScalaUDFCodegen.specFor` cannot build an [[ArrowColumnSpec]] for one, @@ -104,11 +101,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come case dt if isTimeType(dt) => true case _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType => true case ArrayType(inner, _) => isSupportedDataType(inner, allowNullType) - case st: StructType => - // `fieldNames` rebuilds an array on each call, so read it once. - val names = st.fieldNames - names.distinct.length == names.length && - st.fields.forall(f => isSupportedDataType(f.dataType, allowNullType)) + case st: StructType => st.fields.forall(f => isSupportedDataType(f.dataType, allowNullType)) case mt: MapType => isSupportedDataType(mt.keyType, allowNullType) && isSupportedDataType(mt.valueType, allowNullType) @@ -116,25 +109,34 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come } /** - * Names that repeat within one struct, searched recursively through `dataType`. Spark keeps - * duplicate struct field names as distinct positional fields, but Arrow's `StructVector` keys - * its children by name (`ConflictPolicy.CONFLICT_REPLACE` by default), so - * `initializeChildrenFromFields` collapses the duplicates and the generated ordinal-based child - * casts hit a missing or differently typed vector. `CometCreateNamedStruct` refuses the same - * shape at the serde level, but whole-expression dispatch never consults that rule for a - * `named_struct` nested inside e.g. a `transform` lambda, so [[canHandle]] re-checks here. + * Spark keeps duplicate struct field names as distinct positional fields, but Arrow's + * `StructVector` keys its children by name (`ConflictPolicy.CONFLICT_REPLACE`), so + * `named_struct('x', 10, 'x', 20)` collapses into a single child and the generated + * ordinal-based child casts hit a missing or differently typed vector. `CometCreateNamedStruct` + * refuses the same shape on the native path, but whole-expression dispatch never consults that + * rule for a `named_struct` nested inside e.g. a `transform` lambda or a `CreateMap` value. */ - private def duplicateStructFieldNames(dataType: DataType): Seq[String] = dataType match { + private def hasDuplicateStructFieldNames(dt: DataType): Boolean = dt match { case st: StructType => - val names = st.fieldNames.toSeq - val dups = names.diff(names.distinct).distinct - if (dups.nonEmpty) dups else st.fields.flatMap(f => duplicateStructFieldNames(f.dataType)) - case ArrayType(inner, _) => duplicateStructFieldNames(inner) - case MapType(keyType, valueType, _) => - duplicateStructFieldNames(keyType) ++ duplicateStructFieldNames(valueType) - case _ => Nil + // `fieldNames` rebuilds an array on each call, so read it once. + val names = st.fieldNames + names.distinct.length != names.length || + st.fields.exists(f => hasDuplicateStructFieldNames(f.dataType)) + case ArrayType(inner, _) => hasDuplicateStructFieldNames(inner) + case MapType(k, v, _) => hasDuplicateStructFieldNames(k) || hasDuplicateStructFieldNames(v) + case _ => false } + /** Why `dt` cannot cross the kernel boundary as `side` ("output" or "input"), if it cannot. */ + private def typeRejection(dt: DataType, allowNullType: Boolean, side: String): Option[String] = + if (!isSupportedDataType(dt, allowNullType)) { + Some(s"codegen dispatch: unsupported $side type $dt") + } else if (hasDuplicateStructFieldNames(dt)) { + Some(s"codegen dispatch: unsupported $side type $dt (duplicate struct field name)") + } else { + None + } + /** * Mirrors `WholeStageCodegenExec.numOfNestedFields` so [[canHandle]] can reuse * `spark.sql.codegen.maxFields`. @@ -152,18 +154,13 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come * back cleanly rather than crashing the Janino compile at execute time. * * Checks every `BoundReference`'s data type and the root `expr.dataType` against - * [[isSupportedDataType]] and [[duplicateStructFieldNames]], rejects aggregates / generators / - * `Unevaluable`, and gates total nested-field count on `spark.sql.codegen.maxFields`. + * [[isSupportedDataType]] and [[hasDuplicateStructFieldNames]], rejects aggregates, generators + * and `Unevaluable`, and gates total nested-field count on `spark.sql.codegen.maxFields`. */ def canHandle(boundExpr: Expression): Option[String] = { - if (!isSupportedDataType(boundExpr.dataType, allowNullType = true)) { - return Some(s"codegen dispatch: unsupported output type ${boundExpr.dataType}") - } - val outputDups = duplicateStructFieldNames(boundExpr.dataType) - if (outputDups.nonEmpty) { - return Some( - s"codegen dispatch: duplicate struct field name ${outputDups.mkString(", ")} " + - s"in output type ${boundExpr.dataType}") + typeRejection(boundExpr.dataType, allowNullType = true, "output") match { + case Some(reason) => return Some(reason) + case None => } // Mirror WSCG's `spark.sql.codegen.maxFields` gate. Wide schemas blow the generated class's // typed input field count, the typed-getter switch, and the constant pool. Refuse here so the @@ -218,17 +215,9 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come /** Why `expr` cannot be read as a codegen input, if it cannot. */ private def inputRejection(expr: Expression): Option[String] = expr match { - case b: BoundReference if !isSupportedDataType(b.dataType) => - Some(s"codegen dispatch: unsupported input type ${b.dataType} at ordinal ${b.ordinal}") case b: BoundReference => - val dups = duplicateStructFieldNames(b.dataType) - if (dups.isEmpty) { - None - } else { - Some( - s"codegen dispatch: duplicate struct field name ${dups.mkString(", ")} " + - s"in input type ${b.dataType} at ordinal ${b.ordinal}") - } + typeRejection(b.dataType, allowNullType = false, "input") + .map(reason => s"$reason at ordinal ${b.ordinal}") case _ => None } diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index 6be7371a4c8..68bdcd46ab5 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -413,8 +413,9 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { */ private def emitSpecializedGetterExpr(target: String, idx: String, elemType: DataType): String = elemType match { - // Placeholder: [[emitWrite]]'s NullType branch only emits `setNull` and ignores its - // source, so this never reaches the generated Java. It just keeps the match total. + // Computed eagerly for every child, then dropped by [[emitWrite]]'s NullType branch, which + // emits `setNull` instead. Dead in the generated Java but not here: without this case an + // `array` / `struct<.., null>` output throws instead of dispatching. case NullType => "null" case BooleanType => s"$target.getBoolean($idx)" case ByteType => s"$target.getByte($idx)" diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 2ab308b540c..3c03d15318c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -73,7 +73,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { classOf[ArrayMin] -> CometArrayMin, classOf[ArrayPosition] -> CometArrayPosition, classOf[ArrayRemove] -> CometArrayRemove, - classOf[ArrayRepeat] -> CometScalarFunction("array_repeat"), + classOf[ArrayRepeat] -> CometArrayRepeat, classOf[Slice] -> CometSlice, classOf[SortArray] -> CometSortArray, classOf[ArraysOverlap] -> CometArraysOverlap, @@ -646,7 +646,6 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val info = DataTypeInfo.newBuilder() val list = ListInfo.newBuilder() list.setElementType(elementType.get) - // NullType children are always nullable; see Utils.declaredChildNullability. list.setContainsNull(Utils.declaredChildNullability(a.elementType, a.containsNull)) nestedParquetFieldId(parentField, elementPath, includeFieldIds) .foreach(list.setElementFieldId) diff --git a/spark/src/main/scala/org/apache/comet/serde/SupportLevel.scala b/spark/src/main/scala/org/apache/comet/serde/SupportLevel.scala index 60f2ff79345..1c64896227a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/SupportLevel.scala +++ b/spark/src/main/scala/org/apache/comet/serde/SupportLevel.scala @@ -19,6 +19,7 @@ package org.apache.comet.serde +import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.types._ import org.apache.comet.CometConf @@ -103,3 +104,35 @@ object SupportLevel { } } } + +/** + * Serdes that wrap a kernel in `CASE WHEN guarded IS NOT NULL THEN kernel(children) ELSE NULL` + * serialize the guarded children twice, and a stateful child (one built over + * monotonically_increasing_id(), say) then produces values Spark's single evaluation never would, + * for either of two reasons: + * + * - A child evaluated natively gets an independent instance per copy, but native CASE evaluates + * the THEN branch on the rows the predicate selected, so whenever the guard filters, the THEN + * copy's counter runs over a different row sequence than Spark's. + * - A child evaluated by the JVM codegen dispatcher (any lambda function) is keyed in the + * kernel cache by its serialized bytes, so both copies run one kernel instance and share its + * state: the predicate copy consumes the counter for the whole batch and the THEN copy + * continues from there, even when nothing filters. + * + * Neither the stateful child nor the nullable one needs to be the same child, and the stateful + * child need not be nullable at all; a serde that builds no guard for a non-nullable child + * (`CometSize`) needs no gate for it either. + */ +object NullGuard { + val reason = + "non-deterministic child under a null guard is evaluated on different rows than Spark's" + + def doubleEvaluationReason(evaluated: Seq[Expression]): Option[String] = + if (evaluated.exists(!_.deterministic)) Some(reason) else None + + /** [[Unsupported]] when a child inside the guard is non-deterministic, else [[Compatible]]. */ + def supportLevel(evaluated: Expression*): SupportLevel = + doubleEvaluationReason(evaluated) + .map(reason => Unsupported(Some(reason))) + .getOrElse(Compatible()) +} diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index aef649debc7..487dac1d30d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression, L import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ApproximatePercentile, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectList, CollectSet, Corr, Count, Covariance, CovPopulation, CovSample, First, HyperLogLogPlusPlus, Last, Max, MaxBy, MaxMinBy, Min, MinBy, Mode, Percentile, RegrIntercept, RegrR2, RegrReplacement, RegrSlope, RegrSXY, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp} import org.apache.spark.sql.catalyst.util.ArrayData import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, NumericType, ShortType, StringType, TimestampNTZType, TimestampType} +import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, NullType, NumericType, ShortType, StringType, TimestampNTZType, TimestampType} import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus, isSpark42Plus, withFallbackReason} @@ -1053,13 +1053,15 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { if (!CometCollectShim.ignoreNulls(expr)) { Unsupported(Some("collect_set with RESPECT NULLS (ignoreNulls = false) is not supported")) } else { - SupportLevel - .strictFloatingPointReason( - expr.children.head.dataType, - "collect_set on floating-point types " + - "(Comet deduplicates NaN values while Spark treats each NaN as distinct)") - .map(reason => Incompatible(Some(reason))) - .getOrElse(Compatible()) + CometCollectAggregate.nestedNullLevel(expr.children.head.dataType).getOrElse { + SupportLevel + .strictFloatingPointReason( + expr.children.head.dataType, + "collect_set on floating-point types " + + "(Comet deduplicates NaN values while Spark treats each NaN as distinct)") + .map(reason => Incompatible(Some(reason))) + .getOrElse(Compatible()) + } } } @@ -1092,6 +1094,29 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { } } +/** + * Shared gate for collect_list / collect_set. Native collect accumulates the input's own child + * array, so every nested element keeps the non-nullability its producer gave it (`map_entries` + * yields a non-null entry struct with a non-null key), while the declared output type is built + * from Spark's `dataType` widened by `Utils.declaredChildNullability`. Where the two disagree + * DataFusion rejects the batch at runtime ("column types must match schema types"). Gated on any + * NullType-bearing input rather than the exact failing shapes; the same mismatch exists for + * non-Null nesting (`collect_list(array(array(id)))` fails on main) and is a native-aggregate + * bug. + */ +object CometCollectAggregate { + + def nestedNullLevel(dt: DataType): Option[SupportLevel] = + if (SupportLevel.containsType(dt, classOf[NullType])) { + Some( + Unsupported( + Some("native collect_list/collect_set rebuilds a NullType-bearing element with a " + + "nullability that does not match the declared output type"))) + } else { + None + } +} + object CometCollectList extends CometAggregateExpressionSerde[CollectList] { override def getSupportLevel(expr: CollectList): SupportLevel = { @@ -1104,7 +1129,7 @@ object CometCollectList extends CometAggregateExpressionSerde[CollectList] { if (!CometCollectShim.ignoreNulls(expr)) { Unsupported(Some("collect_list with RESPECT NULLS (ignoreNulls = false) is not supported")) } else { - Compatible() + CometCollectAggregate.nestedNullLevel(expr.children.head.dataType).getOrElse(Compatible()) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index cca9f63f8bf..45a076caf6e 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -22,7 +22,7 @@ package org.apache.comet.serde import scala.annotation.tailrec import scala.jdk.CollectionConverters._ -import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, BoundReference, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, IsNull, LambdaFunction, Literal, NamedLambdaVariable, Reverse, Sequence, Size, Slice, SortArray, ZipWith} +import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArrayRepeat, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, BoundReference, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, IsNull, LambdaFunction, Literal, NamedLambdaVariable, Reverse, Sequence, Size, Slice, SortArray, ZipWith} import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -53,6 +53,9 @@ object CometArrayRemove object CometArrayAppend extends CometExpressionSerde[ArrayAppend] with ArraysBase { + override def getSupportLevel(expr: ArrayAppend): SupportLevel = + NullGuard.supportLevel(expr.children: _*) + override def convert( expr: ArrayAppend, inputs: Seq[Attribute], @@ -214,13 +217,26 @@ object CometArrayIntersect override def getIncompatibleReasons(): Seq[String] = Seq(incompatReason, collationReason) + private val nullElementReason: String = + "native array_intersect returns the other side's entries for a NullType-element array" + override def getSupportLevel(expr: ArrayIntersect): SupportLevel = { - // The native array_intersect dedups by raw bytes, which is wrong under non-default collations, - // so report Incompatible rather than Unsupported: the JVM codegen dispatcher (Spark's own - // doGenCode) performs collation-aware set membership and keeps execution native, matching - // Spark. Only the output elements' collation metadata is dropped, consistent with CometReverse - // and CometArrayJoin. - if (hasNonDefaultStringCollation(expr.dataType)) { + // The native array_intersect dedups by raw bytes, which is wrong under non-default collations. + // That is Incompatible rather than Unsupported because there is something real to opt into: a + // user who does not need collation-aware set membership can take the native kernel. Under the + // default config the JVM codegen dispatcher (Spark's own doGenCode) runs it instead, keeping + // execution native and matching Spark; only the output elements' collation metadata is + // dropped, consistent with CometReverse and CometArrayJoin. + // + // A NullType-element side is Unsupported rather than Incompatible: the kernel's short-circuit + // is wrong for an intersection at every setting, so there is nothing to opt into. Unsupported + // reaches the same dispatcher, and unlike Incompatible it does so whatever `allowIncompatible` + // says. Reporting Incompatible would make EXPLAIN advertise the opt-in and then punish it: + // that branch calls `convert` with no dispatcher fallback, so declining there drops the whole + // projection back to Spark exactly when the user asked for more native execution. + if (NullElementSetOp.hasNullElementSide(expr)) { + Unsupported(Some(nullElementReason)) + } else if (hasNonDefaultStringCollation(expr.dataType)) { Incompatible(Some(collationReason)) } else { Incompatible(Some(incompatReason)) @@ -231,8 +247,16 @@ object CometArrayIntersect expr: ArrayIntersect, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val leftArrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) - val rightArrayExprProto = exprToProtoInternal(expr.children(1), inputs, binding) + // Defensive: `getSupportLevel` reports Unsupported for this shape, which never calls + // `convert`. Kept so a future support-level change cannot silently reach the kernel's + // short-circuit. + if (NullElementSetOp.hasNullElementSide(expr)) { + withFallbackReason(expr, nullElementReason) + return None + } + val Seq(left, right) = SetOpArguments.unify(expr.children) + val leftArrayExprProto = exprToProtoInternal(left, inputs, binding) + val rightArrayExprProto = exprToProtoInternal(right, inputs, binding) val arraysIntersectScalarExpr = scalarFunctionExprToProto("array_intersect", leftArrayExprProto, rightArrayExprProto) @@ -307,15 +331,17 @@ object CometArrayExcept override def getIncompatibleReasons(): Seq[String] = Seq(incompatReason) override def getSupportLevel(expr: ArrayExcept): SupportLevel = { - // Surface the native element-type restriction in EXPLAIN. We report Incompatible (not - // Unsupported) for these types so the JVM codegen dispatcher still evaluates them natively - // under the default config; the convert-time guard below is only reached under - // allowIncompatible=true, where the native array_except cannot handle them. - val reason = expr.children.map(_.dataType).find(dt => !isTypeSupported(dt)) match { - case Some(dt) => s"native array_except does not support element type $dt: $incompatReason" - case None => incompatReason + // Surface the native element-type restriction in EXPLAIN. Unsupported rather than + // Incompatible for these types: the JVM codegen dispatcher evaluates them natively and does + // so whatever `allowIncompatible` says, whereas the Incompatible branch calls `convert` with + // no dispatcher fallback, so the guard below would drop the projection back to Spark for a + // user who opted in. Only the element-type restriction is Unsupported; the ordering and null + // handling differences below remain a genuine opt-in. + expr.children.map(_.dataType).find(dt => !isTypeSupported(dt)) match { + case Some(dt) => + Unsupported(Some(s"native array_except does not support element type $dt")) + case None => Incompatible(Some(incompatReason)) } - Incompatible(Some(reason)) } @tailrec @@ -337,17 +363,18 @@ object CometArrayExcept expr: ArrayExcept, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - // Defensive: only reached under allowIncompatible=true (the default-config Incompatible path - // routes through the codegen dispatcher before convert). Native array_except cannot handle - // these element types, so decline and let Spark evaluate. + // Defensive: `getSupportLevel` reports Unsupported for these element types, which never + // calls `convert`. Kept so a future support-level change cannot reach a kernel that would + // raise on them. expr.children.map(_.dataType).find(dt => !isTypeSupported(dt)) match { case Some(dt) => withFallbackReason(expr, s"data type not supported: $dt") return None case None => } - val leftArrayExprProto = exprToProtoInternal(expr.left, inputs, binding) - val rightArrayExprProto = exprToProtoInternal(expr.right, inputs, binding) + val Seq(left, right) = SetOpArguments.unify(expr.children) + val leftArrayExprProto = exprToProtoInternal(left, inputs, binding) + val rightArrayExprProto = exprToProtoInternal(right, inputs, binding) val arrayExceptScalarExpr = scalarFunctionExprToProto("array_except", leftArrayExprProto, rightArrayExprProto) @@ -475,6 +502,25 @@ object CometArrayInsert extends CometExpressionSerde[ArrayInsert] with ArraysBas } object CometSlice extends CometExpressionSerde[Slice] { + + override def getSupportLevel(expr: Slice): SupportLevel = { + expr.x.dataType match { + // Native spark_array_slice rebuilds the sliced list around the input's actual child, + // whose non-NullType item keeps Spark's containsNull, while `convert` promises a + // nullable item; for containsNull = false the two disagree and the native plan is + // rejected (e.g. slice(map_entries(map(k, NULL)), 1, 1)). A bare NullType item is + // declared nullable at the FFI boundary (Utils.declaredChildNullability) and slices fine. + case ArrayType(elementType, false) + if elementType != NullType && + SupportLevel.containsType(elementType, classOf[NullType]) => + Unsupported( + Some( + "native spark_array_slice keeps a non-nullable list item where a nullable " + + "one is promised")) + case _ => Compatible() + } + } + override def convert( expr: Slice, inputs: Seq[Attribute], @@ -494,13 +540,60 @@ object CometSlice extends CometExpressionSerde[Slice] { } } +/** + * DataFusion's set-op kernel (`array_union` / `array_intersect` / `array_except`) asserts that + * both sides carry the identical element type, nested field nullability included, and native + * kernels widen nested nullability differently from what Spark planned (a lambda variable over a + * list arrives nullable, a `monotonically_increasing_id()` struct field does not). Casting both + * sides to a deeply-nullable element type only widens metadata and never changes values, so the + * kernel always sees one type; a primitive element type needs nothing, since only nested fields + * take part in the comparison. + */ +private[serde] object SetOpArguments { + def unify(children: Seq[Expression]): Seq[Expression] = + children.map { c => + c.dataType match { + case ArrayType(et, _) if isComplexType(et) => + val unified = ArrayType(deepNullable(et), containsNull = true) + if (c.dataType == unified) c else Cast(c, unified) + case _ => c + } + } +} + +/** + * DataFusion's set-op kernel (`general_set_op`, shared by array_union and array_intersect) + * short-circuits when either side's element type is Null: it returns distinct(other side). For a + * union that drops the NULL entries the Null-typed list actually holds (Spark keeps one NULL); + * for an intersection it returns the other side's entries although nothing can be common to both + * (Spark returns an empty list). Only a bare `array` takes the branch; a nested Null + * (`array>`) goes through the row converter, which handles it. + */ +private[serde] object NullElementSetOp { + def hasNullElementSide(expr: Expression): Boolean = + expr.children.exists(_.dataType match { + case ArrayType(NullType, _) => true + case _ => false + }) +} + object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { + + override def getSupportLevel(expr: ArrayUnion): SupportLevel = { + if (NullElementSetOp.hasNullElementSide(expr)) { + Unsupported(Some("native array_union drops the entries of a NullType-element array")) + } else { + Compatible() + } + } + override def convert( expr: ArrayUnion, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val leftArrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) - val rightArrayExprProto = exprToProtoInternal(expr.children(1), inputs, binding) + val Seq(left, right) = SetOpArguments.unify(expr.children) + val leftArrayExprProto = exprToProtoInternal(left, inputs, binding) + val rightArrayExprProto = exprToProtoInternal(right, inputs, binding) val arraysUnionScalarExpr = scalarFunctionExprToProto("array_union", leftArrayExprProto, rightArrayExprProto) @@ -509,6 +602,23 @@ object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { } object CometCreateArray extends CometExpressionSerde[CreateArray] with ArraysBase { + + override def getSupportLevel(expr: CreateArray): SupportLevel = { + // DataFusion's make_array funnels an argument list that is entirely Null-typed into + // SingleRowListArrayBuilder, producing ONE list row regardless of the input row count. Spark + // coerces CreateArray's children to a common type, so a NullType child means every child is + // NullType and that branch is the one taken. A non-scalar NullType argument (e.g. + // `aggregate(arr, NULL, (acc, x) -> NULL)` admitted by the JVM codegen dispatcher) therefore + // fails the scalar-function row-count check for batches with more than one row. All-literal + // NULL arguments arrive as scalars and broadcast correctly, and empty `array()` never reaches + // make_array (`convert` emits a literal). + if (expr.children.exists(c => c.dataType == NullType && !c.foldable)) { + Unsupported(Some("native make_array builds a single row from a NullType batch")) + } else { + Compatible() + } + } + override def convert( expr: CreateArray, inputs: Seq[Attribute], @@ -548,6 +658,30 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] with ArraysBas } } +object CometArrayRepeat extends CometExpressionSerde[ArrayRepeat] { + + override def getSupportLevel(expr: ArrayRepeat): SupportLevel = { + expr.left.dataType match { + // DataFusion's list repeat rebuilds the repeated list's item field as nullable. Comet + // declares a non-NullType item with Spark's containsNull, so for containsNull = false + // the planned and produced types disagree and the native plan is rejected. A NullType + // item is declared nullable on the FFI boundary (Utils.declaredChildNullability) and + // matches the rebuild, so plain array stays native. + case ArrayType(elementType, false) if elementType != NullType => + Unsupported(Some("native array_repeat rebuilds a non-nullable list item as nullable")) + case _ => Compatible() + } + } + + override def convert( + expr: ArrayRepeat, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + val childExprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) + scalarFunctionExprToProto("array_repeat", childExprs: _*) + } +} + object CometGetArrayItem extends CometExpressionSerde[GetArrayItem] { override def convert( @@ -768,8 +902,15 @@ object CometSize extends CometExpressionSerde[Size] { override def getSupportLevel(expr: Size): SupportLevel = { expr.child.dataType match { - case _: ArrayType => Compatible() - case _: MapType => Compatible() + case _: ArrayType | _: MapType => + // The null guard is only built for a nullable child in non-legacy mode; see `convert`. + // A non-nullable stateful child (`shuffle(arr)`, `filter(arr, x -> x < rand())`) is + // serialized once and evaluated once, so it needs no gate. + if (expr.legacySizeOfNull || !expr.child.nullable) { + Compatible() + } else { + NullGuard.supportLevel(expr.child) + } case other => Unsupported(Some(s"Unsupported child data type: $other")) } @@ -780,21 +921,30 @@ object CometSize extends CometExpressionSerde[Size] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val arrayExprProto = exprToProtoInternal(expr.child, inputs, binding) - for { - isNotNullExprProto <- createIsNotNullExprProto(expr, inputs, binding) - sizeScalarExprProto <- scalarFunctionExprToProto("size", arrayExprProto) - emptyLiteralExprProto <- createLiteralExprProto(expr.legacySizeOfNull) - } yield { - val caseWhenExpr = ExprOuterClass.CaseWhen - .newBuilder() - .addWhen(isNotNullExprProto) - .addThen(sizeScalarExprProto) - .setElseExpr(emptyLiteralExprProto) - .build() - ExprOuterClass.Expr - .newBuilder() - .setCaseWhen(caseWhenExpr) - .build() + // Native size already returns -1 for a null collection, which is the legacy answer; only + // the non-legacy NULL answer needs the null guard, and only when the child can be null. + if (expr.legacySizeOfNull || !expr.child.nullable) { + scalarFunctionExprToProto("size", arrayExprProto) + } else { + for { + isNotNullExprProto <- createIsNotNullExprProto(expr, inputs, binding) + sizeScalarExprProto <- scalarFunctionExprToProto("size", arrayExprProto) + nullLiteralExprProto <- exprToProtoInternal( + Literal(null, IntegerType), + Seq.empty, + binding = true) + } yield { + val caseWhenExpr = ExprOuterClass.CaseWhen + .newBuilder() + .addWhen(isNotNullExprProto) + .addThen(sizeScalarExprProto) + .setElseExpr(nullLiteralExprProto) + .build() + ExprOuterClass.Expr + .newBuilder() + .setCaseWhen(caseWhenExpr) + .build() + } } } @@ -809,12 +959,6 @@ object CometSize extends CometExpressionSerde[Size] { binding, (builder, unaryExpr) => builder.setIsNotNull(unaryExpr)) } - - private def createLiteralExprProto(legacySizeOfNull: Boolean): Option[ExprOuterClass.Expr] = { - val value = if (legacySizeOfNull) -1 else null - exprToProtoInternal(Literal(value, IntegerType), Seq.empty, binding = true) - } - } object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with ArraysBase { @@ -868,7 +1012,10 @@ object CometArraysZip extends CometExpressionSerde[ArraysZip] { return Unsupported(Some(s"Unsupported child data type: $dt")) } } - Compatible() + // `convert` puts every child in the null-check predicate as well as in the zip itself; + // Spark's ArraysZip evaluates all of its children rather than short-circuiting, so a single + // stateful child is enough to diverge. + NullGuard.supportLevel(expr.children: _*) } override def convert( diff --git a/spark/src/main/scala/org/apache/comet/serde/conditional.scala b/spark/src/main/scala/org/apache/comet/serde/conditional.scala index df94be8c93c..bdafbe273ac 100644 --- a/spark/src/main/scala/org/apache/comet/serde/conditional.scala +++ b/spark/src/main/scala/org/apache/comet/serde/conditional.scala @@ -22,10 +22,30 @@ package org.apache.comet.serde import scala.jdk.CollectionConverters._ import org.apache.spark.sql.catalyst.expressions.{Attribute, CaseWhen, Coalesce, Expression, If, IsNotNull} +import org.apache.spark.sql.types.NullType import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal +/** + * Native CASE merges the rows each branch produced with Arrow's `merge_n`, which builds the + * result through a `MutableArrayData` that carries a validity bitmap; a `NullArray` cannot hold + * one ("Arrays of type Null cannot contain a null bitmask"), so a CASE whose result type is + * `NullType` fails whenever more than one branch contributes rows. Spark normally folds such an + * expression away (`IF(c, NULL, NULL)`), but a `NullType`-typed non-foldable branch keeps it. + */ +private[serde] object NullTypeBranches { + def supportLevel(expr: Expression): SupportLevel = + if (expr.dataType == NullType) { + Unsupported(Some("native CASE cannot merge NullType branches")) + } else { + Compatible() + } +} + object CometIf extends CometExpressionSerde[If] { + + override def getSupportLevel(expr: If): SupportLevel = NullTypeBranches.supportLevel(expr) + override def convert( expr: If, inputs: Seq[Attribute], @@ -50,6 +70,10 @@ object CometIf extends CometExpressionSerde[If] { } object CometCaseWhen extends CometExpressionSerde[CaseWhen] { + + override def getSupportLevel(expr: CaseWhen): SupportLevel = + NullTypeBranches.supportLevel(expr) + override def convert( expr: CaseWhen, inputs: Seq[Attribute], @@ -89,6 +113,15 @@ object CometCaseWhen extends CometExpressionSerde[CaseWhen] { } object CometCoalesce extends CometExpressionSerde[Coalesce] { + + // Every child but the last is a guard; the last one is the ELSE, evaluated on the rows the + // guards left over. The result is a native CASE, so it shares that serde's NullType rule. + override def getSupportLevel(expr: Coalesce): SupportLevel = + NullTypeBranches.supportLevel(expr) match { + case _: Compatible => NullGuard.supportLevel(expr.children: _*) + case unsupported => unsupported + } + override def convert( expr: Coalesce, inputs: Seq[Attribute], diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index ee3e80059d5..9d955af560b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -20,7 +20,7 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, Sha1, Sha2, XxHash64} -import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, StringType, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, NullType, StringType, StructType} import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType} @@ -119,9 +119,15 @@ private object HashUtils { private val unsupportedDecimalReason = "`DecimalType` with precision > 18 is not supported (Spark hashes via Java `BigDecimal`)" private val unsupportedTimeTypeReason = "`TimeType` is not supported" + private val unsupportedNullTypeReason = + "`NullType` is not supported (the native hasher has no arm for it)" val unsupportedReasons: Seq[String] = - Seq(unsupportedDecimalReason, unsupportedTimeTypeReason, "Unsupported child data type") + Seq( + unsupportedDecimalReason, + unsupportedTimeTypeReason, + unsupportedNullTypeReason, + "Unsupported child data type") def supportLevelForChildren(expr: Expression): SupportLevel = { expr.children.iterator @@ -135,6 +141,7 @@ private object HashUtils { private def unsupportedReasonFor(dt: DataType): Option[String] = dt match { case d: DecimalType if d.precision > 18 => Some(unsupportedDecimalReason) + case _: NullType => Some(unsupportedNullTypeReason) case s: StructType => s.fields.iterator.flatMap(f => unsupportedReasonFor(f.dataType).iterator).toSeq.headOption case a: ArrayType => unsupportedReasonFor(a.elementType) diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala index 51fa428b543..40fb9b82586 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -159,11 +159,19 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { override def getCompatibleNotes(): Seq[String] = Seq(MapKeyDedupPolicySupport.nullKeyReason) + private val scalarSideReason: String = + "native map takes the first row of a scalar list where the other argument is per-row" + override def getSupportLevel(expr: MapFromArrays): SupportLevel = { if (MapKeyDedupPolicySupport.isLastWin) { Incompatible(Some(MapKeyDedupPolicySupport.incompatibleReason)) + } else if (expr.left.foldable != expr.right.foldable) { + // DataFusion's `map` reads a scalar list argument through its first row only, so a literal + // array beside a per-row one fails with "map requires key and value lists to have the same + // length" as soon as the batch holds more than one row. + Unsupported(Some(scalarSideReason)) } else { - Compatible(None) + NullGuard.supportLevel(expr.left, expr.right) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index bf5bd48673c..de88377b33b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -183,14 +183,12 @@ object Utils extends CometTypeShim with Logging { /** * Nullability to declare for a nested child (array element, struct field, map value) of type - * `dataType`. A `NullType` child is always declared nullable, whatever Spark's `containsNull` / - * `valueContainsNull` / `StructField.nullable` says: every value is null, so a non-nullable - * flag is a contradiction, and native kernels that rebuild a list around the input's actual - * child (DataFusion's `map_entries` and `array_repeat`, Comet's `spark_array_slice`, ...) - * compare that child's nullability with the one they assume and fail on the mismatch - * (`map_entries(map_filter(map(), ...))`, `slice(filter(array(), ...), 1, 1)`). Applied here - * and in `QueryPlanSerde.serializeDataType` so the JVM-exported field and the type declared to - * native agree. Map keys are not children in this sense: Arrow requires them non-nullable. + * `dataType`. A `NullType` child is always declared nullable, whatever Spark's flag says: every + * value is null, so a non-nullable flag is a contradiction, and native kernels that rebuild a + * list around the input's actual child compare that child's nullability with the one they + * assume and fail on the mismatch. Applied here and in `QueryPlanSerde.serializeDataType` so + * the JVM-exported field and the type declared to native agree. Map keys are not children in + * this sense: Arrow requires them non-nullable. */ def declaredChildNullability(dataType: DataType, nullable: Boolean): Boolean = nullable || dataType == NullType @@ -327,15 +325,8 @@ object Utils extends CometTypeShim with Logging { /** * The only supported way to build an `ArrowStreamWriter` in Comet; enforced by the scalastyle - * `arrowstreamwriter` rule. - * - * Arrow requires map keys to be non-nullable and rejects a stream whose schema violates that - * ("Map data key type should be a non-nullable"). Comet always declares keys non-nullable in - * `toArrowField`, but Arrow's `MinorType.NULL` factory discards the field it is handed and - * rebuilds a nullable one (`Types.java` returns `new NullVector(field.getName())` even though - * `NullVector(Field)` exists), so any `NullType` map key silently turns nullable once a vector - * exists for it. Repairing here, rather than at each call site, means a new IPC writer cannot - * reintroduce the bug by forgetting to ask. + * `arrowstreamwriter` rule. Repairs the declared schema with [[withNonNullableMapKeys]], so a + * new IPC writer cannot reintroduce the nullable `NullType` map key by forgetting to ask. * * Returns the writer together with the root it is bound to, which is `root` itself unless the * declared schema needed repairing. Callers must use the returned root: a writer serializes the @@ -536,9 +527,8 @@ object Utils extends CometTypeShim with Logging { // - Comet decodes dictionaries during execution, so a dictionary-encoded column // shouldn't happen. If it does, each partition can have a different dictionary, // and appending index vectors would silently mix incompatible dictionaries. - // - `VectorSchemaRootAppender` loops forever on a `NullVector` that is a direct - // child of a struct or of a map entry, e.g. `map(k, NULL)` or `map()`; see - // `hasNullDirectlyUnderStruct` for the Arrow mechanics. + // - `VectorSchemaRootAppender` cannot grow a `NullVector` directly under a struct + // (see `hasNullDirectlyUnderStruct`). val skipReason = if (!reader.getDictionaryVectors.isEmpty) { Some("unexpected dictionary-encoded column") diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala index 45f4bb22e0d..5161c03389c 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala @@ -151,11 +151,9 @@ private[python] trait CometArrowPythonRunnerBase new FieldType(false, ArrowType.Struct.INSTANCE, null), childFields.asJava) val structVec = structField.createVector(allocator).asInstanceOf[StructVector] - // Declare the root's schema from `structField`, not from `structVec.getField`: Arrow's - // `MinorType.NULL` factory builds a NullType map key as a nullable `NullVector`, so the - // live vector reports an invalid key field even when `structField` is valid. - // `Utils.newArrowStreamWriter` repairs the declared schema if needed and returns the root - // the writer is bound to, which is the one to close. + // Declare the root's schema from `structField`, not from the live vector, whose NullType + // map key reports nullable (see `Utils.withNonNullableMapKeys`); `newArrowStreamWriter` + // returns the root the writer is bound to, which is the one to close. val declaredRoot = new VectorSchemaRoot(Seq(structField).asJava, Seq[FieldVector](structVec).asJava, 0) val (boundRoot, writer) = @@ -209,10 +207,8 @@ private[python] trait CometArrowPythonRunnerBase // identical. val childNames = inputStructType.fieldNames streamFields = batchFields.zipWithIndex.map { case (field, i) => - // A NullType map key comes back from Arrow as a nullable `NullVector`, which - // `MapVector.initializeChildrenFromFields` rejects when `createVector` rebuilds the - // struct in `startWriter`. Repair the key nullability the same way - // `Utils.serializeBatches` does. + // `startWriter` rebuilds the struct from these fields, so repair the NullType map key + // nullability first (see `Utils.withNonNullableMapKeys`). Utils.withNonNullableMapKeys(renamed(field, childNames(i), forceNullable = true)) } startWriter(streamFields, dataOut) diff --git a/spark/src/test/resources/pyspark/test_pyarrow_udf.py b/spark/src/test/resources/pyspark/test_pyarrow_udf.py index d4f0b505549..15efda928a5 100644 --- a/spark/src/test/resources/pyspark/test_pyarrow_udf.py +++ b/spark/src/test/resources/pyspark/test_pyarrow_udf.py @@ -1094,10 +1094,8 @@ def _normalize(row): def test_map_in_arrow_null_typed_map_children(spark, tmp_path, accelerated): """ `transform_values(map(), ...)` yields MapType(NullType, LongType) and `map(id, NULL)` - yields MapType(LongType, NullType). Arrow's MinorType.NULL factory rebuilds a NullType map - key as *nullable*, which MapVector.initializeChildrenFromFields rejects ("Map data key type - should be a non-nullable"), so the accelerated runner has to repair the key field before it - builds its destination struct from the live Comet vectors. + yields MapType(LongType, NullType); the accelerated runner repairs the NullType map key + (see CometArrowPythonRunnerBase) before building its destination struct. pyarrow refuses to *declare* a non-nullable null-typed field (`A null type field may not be non-nullable`), so the NullType-key map can only ever be a UDF input, never part of the diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql index fa1ef356925..0fd189bc2a9 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql @@ -468,3 +468,9 @@ INSERT INTO cl_src_map VALUES query spark_answer_only SELECT size(collect_list(m)) FROM cl_src_map + +-- Native collect accumulates the input's own child array, so a NullType-bearing element keeps +-- its producer's nested nullability while the declared output widens it; the gate keeps such +-- inputs in Spark rather than failing the batch at runtime. +query expect_fallback(native collect_list/collect_set rebuilds a NullType-bearing element) +SELECT grp, sort_array(collect_list(named_struct('i', i, 'n', NULL))) FROM cl_src_int GROUP BY grp diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql index bf37b8a3b79..0e162ab422e 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql @@ -298,3 +298,7 @@ SELECT grp, sort_array(collect_set(DISTINCT i)) FROM cs_src_int GROUP BY grp ORD query SELECT grp, sort_array(collect_set(i)) FROM cs_src_int GROUP BY grp HAVING size(collect_set(i)) > 1 ORDER BY grp + +-- Same gate as collect_list: a NullType-bearing element stays in Spark. +query expect_fallback(native collect_list/collect_set rebuilds a NullType-bearing element) +SELECT grp, sort_array(collect_set(named_struct('i', i, 'n', NULL))) FROM cs_src_int GROUP BY grp diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_except.sql b/spark/src/test/resources/sql-tests/expressions/array/array_except.sql index afe79948e23..011d3699619 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_except.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_except.sql @@ -96,3 +96,11 @@ INSERT INTO test_except_flt_negzero VALUES query ignore(https://issues.apache.org/jira/browse/SPARK-54918) SELECT a, b, array_except(a, b) FROM test_except_flt_negzero + +-- The set-op kernel asserts identical element types, nested nullability included, and the two +-- sides can arrive with different nested nullability (a literal element is non-nullable, a lambda +-- variable over a list is not). Both sides are cast to a deeply-nullable element type first. +-- A nested list rather than a struct, because `isTypeSupported` declines struct elements before +-- `convert` runs. +query +SELECT array_except(transform(a, x -> array(1)), transform(b, x -> array(x))) FROM test_array_except diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_intersect.sql b/spark/src/test/resources/sql-tests/expressions/array/array_intersect.sql index fd5ac89a11e..b93c33e1dc8 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_intersect.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_intersect.sql @@ -77,6 +77,11 @@ query SELECT a, array_intersect(a, a) FROM test_intersect_dups -- empty array combinations +-- The empty operands are Null-typed literals. The NullType-element gate reports Unsupported for +-- them, so the JVM codegen dispatcher evaluates them instead: still inside Comet, and matching +-- Spark. A plain `query` pins both halves at once -- delete the gate and this file's +-- allowIncompatible=true hands these to the native kernel, whose NullType short-circuit returns +-- the other side's entries. query SELECT array_intersect(array(), array()), array_intersect(array(), array(1, 2)), array_intersect(array(1, 2), array()) @@ -257,3 +262,18 @@ SELECT array_intersect(array(1, NULL, 3), b) FROM test_array_intersect -- conditional (CASE WHEN) arrays query SELECT array_intersect(CASE WHEN a IS NOT NULL THEN a ELSE array(0) END, b) FROM test_array_intersect + +-- The set-op kernel asserts identical element types, nested nullability included, and the two +-- sides can arrive with different nested nullability (a literal field is non-nullable, a lambda +-- variable over a list is not). Both sides are cast to a deeply-nullable element type first. +query +SELECT array_intersect(transform(a, x -> named_struct('i', 1)), transform(b, x -> named_struct('i', x))) FROM test_array_intersect + +-- The set-op kernel short-circuits on a NullType-element side and returns the other side's +-- distinct entries, so `array_intersect(array(), array(NULL))` would come back as [NULL] where +-- Spark returns []. Only a side that Spark leaves as array reaches it (a typed sibling +-- makes Spark cast the Null side first). The gate reports Unsupported, which routes these to the +-- JVM codegen dispatcher at every setting -- including this file's allowIncompatible=true, where +-- dropping the gate would hand them to the kernel and turn the first case into [NULL]. +query +SELECT array_intersect(array(), array(NULL)), array_intersect(array(NULL), array()), array_intersect(array(NULL, NULL), array(NULL)) diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql b/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql index 6d58a62ac8f..10425d5cb56 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql @@ -123,3 +123,13 @@ SELECT array_repeat(CAST(NULL AS STRING), cnt) FROM test_array_repeat -- native declared nullable; the kernel's result must still match the planned type query SELECT array_repeat(filter(array(), x -> true), 2) FROM test_array_repeat + +-- map_entries produces a list whose non-NullType item is declared non-nullable; native +-- array_repeat rebuilds that item as nullable, so the composition stays in Spark. +query expect_fallback(native array_repeat rebuilds a non-nullable list item as nullable) +SELECT array_repeat(map_entries(map(coalesce(long_v, 0), NULL)), 2) FROM test_array_repeat + +-- Same mismatch with no NullType anywhere, so the guard is not a NullType-specific workaround: +-- any list whose item Spark declares non-nullable hits it. This one fails on main as well. +query expect_fallback(native array_repeat rebuilds a non-nullable list item as nullable) +SELECT array_repeat(map_entries(map(coalesce(long_v, 0), long_v)), 2) FROM test_array_repeat diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_union.sql b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql index 622606fc262..4c08823ad56 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_union.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql @@ -47,7 +47,8 @@ query SELECT a, b, array_union(a, b) FROM test_union_nulls -- empty array combinations -query +-- Both sides are Null-typed empty arrays, which the NullType-element gate keeps in Spark. +query expect_fallback(native array_union drops the entries of a NullType-element array) SELECT array_union(array(), array()) FROM test_union_nulls query @@ -56,7 +57,7 @@ SELECT array_union(array(), array(1, 2)) FROM test_union_nulls query SELECT array_union(array(1, 2), array()) FROM test_union_nulls -query +query expect_fallback(native array_union drops the entries of a NullType-element array) SELECT array_union(array(), array(NULL)) FROM test_union_nulls -- both-NULL arrays @@ -247,7 +248,14 @@ SELECT array_union(array(NULL, 99), b) FROM test_array_union query SELECT array_union(CASE WHEN a IS NOT NULL THEN a ELSE array(0) END, b) FROM test_array_union --- A NullType element built by the JVM codegen dispatcher (containsNull=false in Spark) reaches --- native declared nullable; the kernel's result must still match the planned type +-- DataFusion's set-op kernel treats a Null element type as "return distinct(other side)" and +-- drops the NULL entries the Null-typed list actually holds, so NullType-element unions stay +-- in Spark. +query expect_fallback(native array_union drops the entries of a NullType-element array) +SELECT array_union(transform(a, x -> NULL), array()) FROM test_array_union + +-- The set-op kernel asserts identical element types, nested nullability included, and the two +-- sides can arrive with different nested nullability (a literal field is non-nullable, a lambda +-- variable over a list is not). Both sides are cast to a deeply-nullable element type first. query -SELECT array_union(filter(array(), x -> true), filter(array(), x -> true)) FROM test_array_union +SELECT array_union(transform(a, x -> named_struct('i', 1)), transform(b, x -> named_struct('i', x))) FROM test_array_union diff --git a/spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql index e62e3842e8f..ae77d006b86 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql @@ -144,6 +144,12 @@ select arrays_zip(a, b) FROM test_arrays_zip query SELECT arrays_zip(a, b)['a'] FROM (SELECT array(1, 2, 3) as a, array(3, 4, 5) as b) +-- The serde guards every argument with CASE WHEN arg IS NOT NULL, and the zip inside the THEN +-- branch only runs on the rows the guards selected. A stateful argument then advances its +-- counter on those rows only, even when it is not the nullable one, so the pair stays in Spark. +query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's) +SELECT arrays_zip(transform(a, x -> named_struct('i', monotonically_increasing_id(), 'n', NULL)), IF(size(b) > 1, b, CAST(NULL AS array))) FROM test_arrays_zip + query SELECT arrays_zip(a, b)['b'] FROM (SELECT array(1, 2, 3) as a, array(3, 4, 5) as b) diff --git a/spark/src/test/resources/sql-tests/expressions/array/create_array.sql b/spark/src/test/resources/sql-tests/expressions/array/create_array.sql index 33e37b6765f..40aad48d789 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/create_array.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/create_array.sql @@ -140,3 +140,8 @@ SELECT array(named_struct('a', 1), named_struct('a', a)) FROM test_create_array -- Empty array cast to a nested array type: the element type has to survive with no children. query SELECT CAST(array() AS ARRAY>) + +-- A non-foldable all-NullType argument (built by the JVM codegen dispatcher) would make native +-- make_array collapse the whole batch into a single list row, so it stays in Spark. +query expect_fallback(native make_array builds a single row from a NullType batch) +SELECT array(aggregate(arr, NULL, (acc, x) -> NULL)) FROM test_create_array_complex diff --git a/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql index 30ab1bba587..a291d275312 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/element_at_ansi.sql @@ -113,3 +113,11 @@ SELECT id, element_at(IF(monotonically_increasing_id() % 2 = 0, CAST(NULL AS ARRAY), array(1)), 1) AS v1, element_at(IF(rand(7L) < 2, CAST(NULL AS ARRAY), array(1)), 1 + (id % (id - 2))) AS v2 FROM ansi_element_at_null + +-- non-deterministic collection under the ANSI null guard, built by a lambda +-- The same restriction as above reached through the JVM codegen dispatcher: the guard's two +-- copies of the dispatched `transform` share one kernel and its state, so it stays in Spark. +-- ============================================================================ + +query expect_fallback(nullable nondeterministic array or map operand) +SELECT element_at(transform(IF(monotonically_increasing_id() % 2 = 0, arr, CAST(NULL AS array)), x -> x + 1), 1) FROM ansi_element_at_oob diff --git a/spark/src/test/resources/sql-tests/expressions/array/slice.sql b/spark/src/test/resources/sql-tests/expressions/array/slice.sql index d3234ba391d..73327e04ef6 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/slice.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/slice.sql @@ -282,3 +282,9 @@ SELECT slice(array(map(k, v), map(k + 1, v)), 1, 1) FROM test_slice_map -- native declared nullable; the kernel's result must still match the planned type query SELECT slice(filter(array(), x -> true), 1, 1) FROM test_slice + +-- map_entries produces a list whose NullType-bearing struct item is declared non-nullable; +-- native spark_array_slice keeps that item where a nullable one is promised, so the +-- composition stays in Spark. +query expect_fallback(native spark_array_slice keeps a non-nullable list item where a nullable one is promised) +SELECT slice(map_entries(map(coalesce(start_idx, 0), NULL)), 1, 1) FROM test_slice diff --git a/spark/src/test/resources/sql-tests/expressions/array/transform.sql b/spark/src/test/resources/sql-tests/expressions/array/transform.sql index f7184a616ac..340cf23dc1e 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/transform.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/transform.sql @@ -85,13 +85,5 @@ SELECT transform(array(1, 2, 3), x -> x * x) query SELECT transform(array(), x -> x) --- Non-empty NullType results nested in list / struct outputs (codegen writes an Arrow NullVector --- child; `array()` above only covers the empty case) -query -SELECT transform(a, x -> NULL) FROM test_transform - -query -SELECT transform(a, x -> named_struct('v', x, 'n', NULL)) FROM test_transform - query SELECT transform(array(NULL), x -> x) diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql b/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql index 508571d728d..7e3b89a03d3 100644 --- a/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql +++ b/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql @@ -33,3 +33,8 @@ SELECT CASE WHEN s IS NULL THEN 'null_val' ELSE s END FROM test_case_when -- literal arguments query SELECT CASE WHEN i = 1 THEN s WHEN i = 2 THEN 'fixed' ELSE s END FROM test_case_when + +-- A NullType result stays in Spark: native CASE merges the rows of its branches through Arrow's +-- merge_n, which cannot build a NullArray with a validity bitmap. +query expect_fallback(native CASE cannot merge NullType branches) +SELECT CASE WHEN i = 1 THEN aggregate(array(i), NULL, (acc, x) -> NULL) END FROM test_case_when diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql index 6457768aedf..8ee23c8cb4c 100644 --- a/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql +++ b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql @@ -33,3 +33,14 @@ SELECT coalesce(a, 99) FROM test_coalesce -- literal arguments query SELECT coalesce(NULL, NULL, 99), coalesce(1, NULL, 99), coalesce(NULL) + +-- The serde guards every argument but the last with CASE WHEN arg IS NOT NULL THEN arg, and the +-- two copies of a non-deterministic argument advance their state independently: the THEN copy +-- can answer NULL for a row the guard selected, in a column declared non-nullable. +query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's) +SELECT coalesce(IF(monotonically_increasing_id() % 2 = 0, a, NULL), b, 0) FROM test_coalesce + +-- A NullType result stays in Spark: the serde builds a native CASE, which merges the rows of its +-- branches through Arrow's merge_n and cannot build a NullArray with a validity bitmap. +query expect_fallback(native CASE cannot merge NullType branches) +SELECT coalesce(aggregate(array(a), NULL, (acc, x) -> NULL), aggregate(array(b), NULL, (acc, x) -> NULL)) FROM test_coalesce diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql b/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql index 53b908ca793..51752b54a33 100644 --- a/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql +++ b/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql @@ -30,3 +30,9 @@ SELECT IF(a > 0, 'positive', 'non-positive') FROM test_if -- literal arguments query SELECT IF(true, 1, 2), IF(false, 1, 2), IF(NULL, 1, 2) + +-- A NullType result stays in Spark: native CASE merges the rows of its branches through Arrow's +-- merge_n, which cannot build a NullArray with a validity bitmap. Spark folds `IF(c, NULL, NULL)` +-- itself, so the branch has to be a non-foldable NullType expression. +query expect_fallback(native CASE cannot merge NullType branches) +SELECT IF(cond, aggregate(array(a), NULL, (acc, x) -> NULL), NULL) FROM test_if diff --git a/spark/src/test/resources/sql-tests/expressions/hash/hash.sql b/spark/src/test/resources/sql-tests/expressions/hash/hash.sql index fc93e9d4bf6..651c4fc83fc 100644 --- a/spark/src/test/resources/sql-tests/expressions/hash/hash.sql +++ b/spark/src/test/resources/sql-tests/expressions/hash/hash.sql @@ -30,3 +30,7 @@ SELECT md5(col), md5(cast(a as string)), md5(cast(b as string)), hash(col), hash -- native engine as scalar values rather than being folded away by Spark's optimizer. query SELECT md5('Spark SQL'), sha1('test'), sha2('test', 0), sha2('test', 256), sha2('test', 224), sha2('test', 384), sha2('test', 512), sha2('test', 128), sha2('test', -1), sha2(cast(null as string), 256), hash('test'), xxhash64('test') + +-- The native hasher has no arm for NullType, which a non-foldable struct or array can carry. +query expect_fallback(`NullType` is not supported) +SELECT hash(named_struct('a', a, 'b', NULL)), xxhash64(transform(array(a), x -> NULL)) FROM test diff --git a/spark/src/test/resources/sql-tests/expressions/map/create_map.sql b/spark/src/test/resources/sql-tests/expressions/map/create_map.sql index 4236c54d435..40ed8a0f3b3 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/create_map.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/create_map.sql @@ -52,9 +52,6 @@ SELECT map('a', NULL) query SELECT map(k, NULL) FROM test_create_map -query -SELECT id, map() FROM (SELECT explode(sequence(1, 3)) AS id) - query SELECT size(map()), size(map('a', NULL)) diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql b/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql index 11403fdfe65..f6b9ddcc01b 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql @@ -68,6 +68,3 @@ FROM test_map_entries_nested -- must reach DataFusion's map_entries declared nullable, or its ListArray build panics query SELECT map_entries(map_filter(map(), (k, v) -> true)) FROM test_map_entries - -query -SELECT map_entries(map('a', CAST(NULL AS int))) FROM test_map_entries diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql index 9ccfe34d573..202e9b562f4 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql @@ -63,3 +63,14 @@ SELECT map_from_arrays(NULL, NULL) -- empty arrays produce MapType(NullType, NullType) query SELECT map_from_arrays(array(), array()) + +-- The serde's null guard serializes both inputs twice, and a non-deterministic input would +-- advance differently in each copy, so it falls back. +query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's) +SELECT map_from_arrays(IF(monotonically_increasing_id() % 2 = 0, k, NULL), v) FROM test_map_from_arrays + +-- A literal array beside a per-row one: DataFusion's map kernel reads the scalar list through its +-- first row only and fails the length check, so the shape stays in Spark. Independent of NullType +-- (the typed literal fails the same way); found while probing the NullType flavour. +query expect_fallback(native map takes the first row of a scalar list where the other argument is per-row) +SELECT map_from_arrays(array(coalesce(k[0], 'x')), array(1)), map_from_arrays(array(coalesce(k[0], 'x')), array(NULL)) FROM test_map_from_arrays diff --git a/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql b/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql index d6353efa28c..ce1253a876a 100644 --- a/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql +++ b/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql @@ -39,8 +39,5 @@ SELECT named_struct('x', a, 'y', 'fixed_val', 'z', c) FROM test_named_struct query expect_fallback(duplicate field names) SELECT named_struct('x', a, 'x', b) FROM test_named_struct -query expect_fallback(duplicate struct field name) -SELECT transform(array(a), v -> named_struct('x', v, 'x', NULL)) FROM test_named_struct - query expect_fallback(duplicate struct field name) SELECT transform(array(a), v -> named_struct('x', v, 'x', v + 1)) FROM test_named_struct diff --git a/spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql b/spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql index c5d9e64a634..3afaf8b32bc 100644 --- a/spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql +++ b/spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql @@ -26,3 +26,9 @@ SELECT s.name, s.age, s.score FROM test_struct query SELECT s.name, s.age + 1, s.score * 2 FROM test_struct + +-- A struct row that a kernel grew through MutableArrayData (element_at on an out-of-range index) +-- can carry a Null-typed child with a validity bitmap; projecting that child rebuilds a clean +-- NullArray instead of failing validation. The typed sibling is the control. +query +SELECT try_element_at(array(named_struct('n', s.name, 'z', NULL)), CAST(s.age % 2 + 1 AS INT)).z, try_element_at(array(named_struct('n', s.name, 'z', NULL)), CAST(s.age % 2 + 1 AS INT)).n FROM test_struct diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index ad86dd15bc1..c2b50644108 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1099,6 +1099,33 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + test("size - non-deterministic child under the null guard") { + withParquetTable((0 until 16).map(i => Tuple1(i.toLong)), "t", withDictionary = false) { + // Non-legacy size wraps a nullable child in `CASE WHEN child IS NOT NULL`, which would + // evaluate a stateful child twice, so that shape stays in Spark; legacy mode builds no + // guard and keeps it native. + val nullableStateful = + "SELECT _1, size(IF(monotonically_increasing_id() % 2 = 0, array(_1), NULL)) FROM t" + withSQLConf(SQLConf.LEGACY_SIZE_OF_NULL.key -> "false") { + checkSparkAnswerAndFallbackReason( + nullableStateful, + "non-deterministic child under a null guard is evaluated on different rows than Spark's") + } + withSQLConf( + SQLConf.LEGACY_SIZE_OF_NULL.key -> "true", + SQLConf.ANSI_ENABLED.key -> "false") { + checkSparkAnswerAndOperator(nullableStateful) + } + // A non-nullable child gets no guard, so a stateful one whose length depends on the + // counter is evaluated once and matches Spark. The lambda runs through the JVM codegen + // dispatcher, where a guard's two copies would share one kernel and its counter. + withSQLConf(SQLConf.LEGACY_SIZE_OF_NULL.key -> "false") { + checkSparkAnswerAndOperator( + "SELECT _1, size(filter(array(_1, 1, 2), x -> x < monotonically_increasing_id())) FROM t") + } + } + } + // https://github.com/apache/datafusion-comet/issues/4560 test("array_size returns null for null input") { val table = "t1" diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index 088ebeb01e7..59197d08610 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala @@ -327,11 +327,8 @@ class CometCodegenSourceSuite extends AnyFunSuite { } test("canHandle rejects duplicate struct field names, in outputs and inputs") { - // Spark keeps `named_struct('a', x, 'a', y)` as two positional fields, but Arrow's - // StructVector keys children by name and collapses them on allocation, so the generated - // ordinal-based child casts hit a missing or differently typed vector. `CometCreateNamedStruct` - // rejects this at serde level; the gate must do the same for a struct buried inside a larger - // expression that is dispatched as a whole (e.g. a `transform` lambda body). + // See `CometBatchKernelCodegen.hasDuplicateStructFieldNames`; the gate must also catch a + // struct buried inside a larger expression dispatched as a whole (e.g. a `transform` body). val x = BoundReference(0, IntegerType, nullable = false) val dupNull = CreateNamedStruct(Seq(Literal("a"), x, Literal("a"), Literal(null, NullType))) val dupInt = CreateNamedStruct(Seq(Literal("a"), x, Literal("a"), Add(x, Literal(1)))) @@ -343,7 +340,7 @@ class CometCodegenSourceSuite extends AnyFunSuite { Seq(dupNull, dupInt, CreateArray(Seq(dupInt)), dupInput).foreach { expr => val reason = CometBatchKernelCodegen.canHandle(expr) assert( - reason.exists(_.contains("duplicate struct field name a")), + reason.exists(_.contains("duplicate struct field name")), s"expected canHandle to reject $expr; got: $reason") } val distinct = @@ -397,24 +394,11 @@ class CometCodegenSourceSuite extends AnyFunSuite { .getContainsNull) } - test("NullType output writes setNull without reading a source value") { - // A NullType leaf has no Arrow data buffer, so the emitted write must be `setNull` only. - val src = CometBatchKernelCodegen - .generateSource(Literal(null, NullType), IndexedSeq(nullableString)) - .body - assert( - src.contains("org.apache.arrow.vector.NullVector"), - s"expected the output vector to be a NullVector; got:\n$src") - assert( - !src.contains("getDataVector"), - s"expected no child-vector access for a scalar NullType output; got:\n$src") - } - test("nested NullType output casts the child vector and writes setNull into it") { - // The scalar case above cannot distinguish `emitWrite`'s NullType branch from `defaultBody`'s - // own `ev.isNull -> output.setNull(i)` short-circuit, which emits the same text. A NullType - // *value* child inside a map is only reachable through `emitWrite`, so assert on that: the - // child vector must be cast to NullVector and written through `setNull`. + // A scalar NullType output cannot distinguish `emitWrite`'s NullType branch from + // `defaultBody`'s own `ev.isNull -> output.setNull(i)` short-circuit, which emits the same + // text. A NullType *value* child inside a map is only reachable through `emitWrite`, so + // assert on that: the child vector must be cast to NullVector and written through `setNull`. val src = CometBatchKernelCodegen .generateSource( Literal.create(Map("a" -> null), MapType(StringType, NullType, valueContainsNull = true)), @@ -432,10 +416,7 @@ class CometCodegenSourceSuite extends AnyFunSuite { } test("gate and output emitters agree across the whole accepted type surface") { - // `CometBatchKernelCodegen.isSupportedDataType`, `outputVectorClass`, `emitWrite` and - // `emitSpecializedGetterExpr` each carry a doc comment saying they must stay in step, but - // nothing enforced it. Assert the implication directly: if `canHandle` greenlights an - // output type, generating the kernel for it must not throw. + // If `canHandle` greenlights an output type, generating the kernel for it must not throw. val leaves: Seq[DataType] = Seq( NullType, BooleanType, diff --git a/spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala index 6e50e7cb9ae..d3d69cf2e5d 100644 --- a/spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala @@ -23,12 +23,14 @@ import scala.jdk.CollectionConverters._ import scala.util.Random import org.apache.hadoop.fs.Path -import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.catalyst.expressions.StructsToCsv +import org.apache.spark.sql.comet.CometProjectExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions._ import org.apache.spark.sql.types.StringType +import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometCsvExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -65,6 +67,28 @@ class CometCsvExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + test("to_csv - a row that renders empty is NULL") { + // Spark hands the row to univocity's `writeRowToString` with `skipEmptyLines`, which turns an + // empty rendering (a lone null field under the default empty `nullValue`) into NULL rather + // than "". Spark 3.5+ crashes on that NULL in its own generated code + // (`nullSafeCodeGen` never marks the result null), so Spark's answer is only comparable on + // 3.4; the native answer is pinned on every version. + withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[StructsToCsv]) -> "true") { + val df = spark + .range(0, 4, 1, 1) + .select( + to_csv(struct(when(col("id") > 1, col("id")).as("c"))).as("lone"), + to_csv(struct(when(col("id") > 1, col("id")).as("c"), lit(null).as("n"))).as("pair")) + if (!isSpark35Plus) { + checkSparkAnswerAndOperator(df) + } else { + checkAnswer(df, Seq(Row(null, ","), Row(null, ","), Row("2", "2,"), Row("3", "3,"))) + assert(collect(df.queryExecution.executedPlan) { case p: CometProjectExec => p }.nonEmpty) + checkSparkAnswerAndOperator(df.select(col("pair"))) + } + } + } + test("to_csv - with configurable formatting options") { val table = "t1" withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[StructsToCsv]) -> "true") { diff --git a/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala b/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala new file mode 100644 index 00000000000..aa551e6ea2b --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala @@ -0,0 +1,728 @@ +/* + * 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.comet + +import scala.collection.mutable.ArrayBuffer +import scala.util.{Failure, Success, Try} + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.{Expression, JsonToStructs, RuntimeReplaceable, Sequence, StringToMap} +import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.comet.CometProjectExec +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.NullType + +import org.apache.comet.serde.{QueryPlanSerde, SupportLevel} + +/** + * Cross-product sweep of the `NullType` shapes the JVM codegen dispatcher admits against the + * expressions that consume them. + * + * Three sweeps share one driver, `sweep`: producers under consumers, producers under operators, + * and producers nested in a container and then put under a serializing operator. + */ +class CometNullTypeCompositionSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + /** + * Non-foldable `NullType`-bearing expressions over a live column. Each is a shape that the + * widened gate lets reach native execution; the constant-folded forms are literals and take a + * different path, so every producer here references `id`. + */ + private val arrayOfNull = Seq( + "transform(array(id), x -> NULL)", + "filter(array(CAST(NULL AS int)), x -> id IS NOT NULL)") + + private val mapWithNullValue = Seq("map(id, NULL)") + + private val mapWithNullKey = Seq("transform_values(map(), (k, v) -> id)") + + private val arrayOfStructWithNull = Seq("map_entries(map(id, NULL))") + + private val structWithNull = Seq("named_struct('a', id, 'b', NULL)") + + private val scalarNull = Seq("aggregate(array(id), NULL, (acc, x) -> NULL)") + + /** Consumers valid for a value of any type, written as templates over the producer `%s`. */ + private val anyTypeConsumers = Seq( + "to_json(struct(%s AS c))", + "to_csv(struct(%s AS c))", + "hash(%s)", + "xxhash64(%s)", + "CAST(%s AS string)", + "CASE WHEN id > 2 THEN %s END", + "IF(id > 2, %s, NULL)", + "coalesce(%s, %s)", + "%s IS NULL", + "%s = %s", + "%s <=> %s") + + /** Consumers valid for any `array`. */ + private val arrayConsumers = Seq( + "size(%s)", + "reverse(%s)", + "array_distinct(%s)", + "sort_array(%s)", + "array_sort(%s)", + "element_at(%s, 1)", + "slice(%s, 1, 1)", + "array_repeat(%s, 2)", + "array_union(%s, %s)", + "array_except(%s, %s)", + "array_intersect(%s, %s)", + // Set ops only diverge when a Null-typed side meets a side that actually holds entries, so + // the same-producer templates above cannot reach that branch on their own. + "array_union(%s, array(1))", + "array_union(array(1), %s)", + "array_except(%s, array(1))", + "array_intersect(%s, array(1))", + "arrays_overlap(%s, %s)", + "arrays_zip(%s, %s)", + "concat(%s, %s)", + "flatten(array(%s))", + "array_position(%s, NULL)", + "array_contains(%s, NULL)", + // Spark rejects a NullType needle, so the typed needle is what reaches the serde. + "array_contains(%s, %s[0])", + "array_remove(%s, NULL)", + "array_append(%s, NULL)", + "array_insert(%s, 1, NULL)", + "exists(%s, x -> x IS NULL)", + "forall(%s, x -> x IS NULL)", + "filter(%s, x -> x IS NULL)", + "transform(%s, x -> x)", + "aggregate(%s, 0, (acc, x) -> acc)", + "zip_with(%s, %s, (x, y) -> x)", + "map_from_entries(arrays_zip(%s, %s))", + "%s[0]", + "array_compact(%s)", + "array_max(%s)", + "array_min(%s)", + "array_join(%s, ',')", + // shuffle is non-deterministic, so only the sorted result compares. + "sort_array(shuffle(%s))", + "map_from_arrays(%s, array(id))") ++ anyTypeConsumers + + /** Consumers valid for any `map`. */ + private val mapConsumers = Seq( + "size(%s)", + "map_keys(%s)", + "map_values(%s)", + "map_entries(%s)", + "map_concat(%s, %s)", + "map_filter(%s, (k, v) -> k IS NOT NULL)", + "transform_values(%s, (k, v) -> v)", + "transform_keys(%s, (k, v) -> k)", + "map_zip_with(%s, %s, (k, v1, v2) -> v1)", + "map_from_entries(map_entries(%s))", + "element_at(%s, id)", + "%s[id]") ++ anyTypeConsumers + + private val structConsumers = Seq( + "%s.a", + "%s.b", + "struct(%s)", + "array(%s)", + "array(%s).b", + "size(array(%s))") ++ anyTypeConsumers + + private val scalarConsumers = Seq( + "array(%s)", + "map(id, %s)", + "named_struct('a', %s)", + "coalesce(%s, NULL)", + "abs(%s)", + // Spark's analyzer rejects a void scalar in an array position (no implicit cast), so these + // stay skipped; they are here so the sweep notices if that ever changes. + "array_union(%s, array(1))", + "array_distinct(%s)", + "array_contains(%s, 1)", + "CAST(%s AS array)") ++ anyTypeConsumers + + private def cases(wrap: String => String = identity): Seq[(String, String)] = + Seq( + arrayOfNull -> arrayConsumers, + arrayOfStructWithNull -> arrayConsumers, + mapWithNullValue -> mapConsumers, + mapWithNullKey -> mapConsumers, + structWithNull -> structConsumers, + scalarNull -> scalarConsumers).flatMap { case (producers, consumers) => + for (p <- producers; c <- consumers) + yield { + val wrapped = wrap(p) + (wrapped, substitute(c, wrapped, p)) + } + } + + /** + * Fills a consumer template, putting `first` in the leading placeholder and `rest` in any + * others. Wrapping only the first argument keeps a stateful producer to a single occurrence: + * two of them desynchronize inside Spark itself, whereas the first child is one that every + * consumer evaluates, so a divergence there belongs to Comet. + */ + private def substitute(template: String, first: String, rest: String): String = { + val at = template.indexOf("%s") + template.substring(0, at) + first + template.substring(at + 2).replace("%s", rest) + } + + /** + * Makes a producer nullable and non-deterministic, the only input that can tell apart the two + * copies a serde null guard serializes (see `NullGuard.doubleEvaluationReason`). + * + * Applied to the first argument only, so multi-argument consumers stay in the sweep: a serde + * that null-guards every child, as `CometArraysZip` does, needs just one stateful argument to + * diverge, and restricting the sweep to single-placeholder consumers hid exactly that case. + */ + private def nullableNondeterministic(producer: String): String = + s"IF(monotonically_increasing_id() % 2 = 0, $producer, NULL)" + + /** Makes a producer nullable while keeping it deterministic. */ + private def nullableDeterministic(producer: String): String = + s"IF(id % 2 = 0, $producer, NULL)" + + /** + * Non-nullable, non-deterministic, NullType-bearing producers that record the counter, either + * in their value or, for the `filter` one, in their length. Under a null guard whose other + * argument is nullable, the THEN branch evaluates them on the guard's filtered rows only, so + * the counter sequence differs from Spark's even though the stateful child itself is never + * null; and a lambda producer runs through the JVM codegen dispatcher, whose kernel cache makes + * the guard's two copies share one counter, so it diverges under a single-argument guard too + * (`size(%s)` on a head that guarded non-nullable children). Paired with a deterministic + * producer of the same type for the sibling slot. + */ + private val statefulProducers: Seq[(String, String, Seq[String])] = Seq( + ( + "transform(array(id), x -> named_struct('i', monotonically_increasing_id(), 'n', NULL))", + "transform(array(id), x -> named_struct('i', x, 'n', NULL))", + arrayConsumers), + ( + "filter(transform(array(id, 1, 2), x -> named_struct('i', x, 'n', NULL)), " + + "s -> s.i < monotonically_increasing_id())", + "filter(transform(array(id, 1, 2), x -> named_struct('i', x, 'n', NULL)), s -> s.i < id)", + arrayConsumers), + ("map(monotonically_increasing_id(), NULL)", "map(id, NULL)", mapConsumers), + ( + "named_struct('i', monotonically_increasing_id(), 'n', NULL)", + "named_struct('i', id, 'n', NULL)", + structConsumers)) + + /** + * Each consumer with the stateful producer in the first slot and either a nullable + * deterministic sibling or a plain one in the others. The stateful producer stays first because + * Spark's `BinaryExpression.eval` returns NULL on a null left operand without evaluating the + * right one, so a stateful right operand's counter is an artifact of Spark's short-circuit and + * not a contract Comet can match; the left operand is evaluated on every row by both engines. + */ + private def crossInputCases: Seq[(String, String)] = + for { + (stateful, deterministic, consumers) <- statefulProducers + c <- consumers + rest <- Seq(nullableDeterministic(deterministic), deterministic) + } yield (stateful, substitute(c, stateful, rest)) + + /** + * Query templates that put a producer under a different physical operator. The consumer sweep + * holds the operator fixed at a projection, so it never sees the paths that serialize and + * re-read a value rather than compute over it. + */ + private val operatorTemplates = Seq( + "project" -> "SELECT %s AS c FROM t", + "filter" -> "SELECT %s AS c FROM t WHERE id > 2", + "sort-by-id" -> "SELECT %s AS c FROM t ORDER BY id", + "sort-by-value" -> "SELECT %s AS c FROM t ORDER BY c", + "limit" -> "SELECT %s AS c FROM t LIMIT 3", + "take-ordered" -> "SELECT %s AS c FROM t ORDER BY id LIMIT 3", + "groupby-key" -> "SELECT %s AS c, count(*) FROM t GROUP BY c", + "groupby-value" -> "SELECT id, first(%s) AS c FROM t GROUP BY id", + "groupby-last" -> "SELECT id, last(%s) AS c FROM t GROUP BY id", + "collect-list" -> "SELECT collect_list(%s) AS c FROM t", + "collect-set" -> "SELECT collect_set(%s) AS c FROM t", + "max" -> "SELECT max(%s) AS c FROM t", + "min" -> "SELECT min(%s) AS c FROM t", + "count-distinct" -> "SELECT count(DISTINCT %s) AS c FROM t", + "distinct" -> "SELECT DISTINCT %s AS c FROM t", + "union-all" -> "SELECT %s AS c FROM t UNION ALL SELECT %s AS c FROM t", + "union-distinct" -> "SELECT %s AS c FROM t UNION SELECT %s AS c FROM t", + "window-order" -> "SELECT %s AS c, row_number() OVER (ORDER BY id) AS r FROM t", + "window-partition" -> "SELECT id, count(*) OVER (PARTITION BY %s) AS n FROM t", + // The producer is computed on the build side, so the value itself crosses the exchange + // (shuffle or broadcast) rather than being projected after the join. + "join-shuffle" -> + "SELECT a.id, b.c FROM t a JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id = b.bid", + "join-broadcast" -> + ("SELECT /*+ BROADCAST(b) */ a.id, b.c FROM t a " + + "JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id = b.bid"), + "join-nested-loop" -> + ("SELECT /*+ BROADCAST(b) */ a.id, b.c FROM t a " + + "JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id > b.bid"), + "expand-cube" -> "SELECT id, count(%s) AS n FROM t GROUP BY CUBE(id)", + "explode" -> "SELECT explode(%s) AS c FROM t", + "repartition" -> "SELECT /*+ REPARTITION(3) */ %s AS c FROM t", + // Above `spark.shuffle.sort.bypassMergeThreshold`, so the JVM shuffle takes its sort-based + // writer, which hands a whole destination partition to one native call. + "repartition-many" -> "SELECT /*+ REPARTITION(300) */ %s AS c FROM t", + "coalesce-partitions" -> "SELECT /*+ COALESCE(1) */ %s AS c FROM t", + // The value itself is the partitioning key, the join key and a scalar subquery result, so + // the hash partitioner, the join's key comparison and the subquery's scalar conversion see + // a NullType rather than only carrying one past. + "distribute-by-value" -> "SELECT %s AS c FROM t DISTRIBUTE BY c", + "join-on-value" -> + ("SELECT a.id, b.id FROM (SELECT id, %s AS c FROM t) a " + + "JOIN (SELECT id, %s AS c FROM t) b ON a.c <=> b.c AND a.id = b.id"), + "scalar-subquery" -> "SELECT id, (SELECT max(x) FROM (SELECT %s AS x FROM t)) AS c FROM t", + "subquery" -> "SELECT id FROM t WHERE id IN (SELECT id FROM t WHERE %s IS NOT NULL)", + "nested-project" -> "SELECT c FROM (SELECT %s AS c, id FROM t ORDER BY id) x") + + /** + * The subset of `operatorTemplates` that serializes the value rather than only computing over + * it. + */ + private val serializingOperators = { + val names = Set( + "project", + "sort-by-id", + "groupby-value", + "collect-list", + "collect-set", + "repartition", + "repartition-many", + "distribute-by-value", + "join-on-value", + "union-all", + "join-shuffle") + require(names.subsetOf(operatorTemplates.map(_._1).toSet), "unknown operator name") + operatorTemplates.filter { case (name, _) => names(name) } + } + + private val arrayOfStructWithNullField = Seq("array(named_struct('a', id, 'b', NULL))") + + private def allProducers: Seq[String] = + arrayOfNull ++ mapWithNullValue ++ mapWithNullKey ++ arrayOfStructWithNull ++ + structWithNull ++ scalarNull ++ arrayOfStructWithNullField + + /** + * Containers to wrap a producer in before putting it under a serializing operator. The nested + * nullability mismatches only show where a container nests the value and an operator re-reads + * the nesting: `collect_list(array(map(k, NULL)))` fails while both `array(map(k, NULL))` and + * `collect_list(map(k, NULL))` pass. + */ + private val nestingWrappers = Seq( + "array(%s)", + "array_repeat(%s, 2)", + "element_at(%s, 1)", + "slice(%s, 1, 1)", + "map(id, %s)", + "named_struct('s', %s)") + + private val excludedOptimizerRules = Seq( + "ConstantFolding", + "NullPropagation", + "SimplifyBinaryComparison", + "SimplifyConditionals").map("org.apache.spark.sql.catalyst.optimizer." + _).mkString(",") + + /** + * How rows are cut into batches and which path an exchange takes. The expression sweeps vary + * what is computed; these vary the machinery it runs through, because the defaults hide a whole + * class of defects: 16 rows in one batch never reuse a builder, `REPARTITION(3)` never leaves + * the JVM shuffle's bypass writer (one native call per batch), and a `spark.comet.batchSize` of + * 8192 never puts a batch boundary inside a stateful expression. + * + * `rows` is the size of `t`; the shuffle profiles need enough rows per destination partition to + * span several `spark.comet.shuffle.jvm.batchSize` batches inside one native call. + */ + private case class Profile(name: String, rows: Int, confs: Seq[(String, String)]) + + private val defaultProfile = Profile("default", 16, Seq.empty) + + /** Native batches of two rows, so every operator and stateful expression crosses a boundary. */ + private val smallBatchProfile = Profile( + "small-batches", + 200, + Seq( + CometConf.COMET_BATCH_SIZE.key -> "2", + CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2", + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true")) + + /** + * Every registered array, map, struct and any-type aggregate serde opted into its native + * kernel. An `Incompatible` serde routes through the JVM codegen dispatcher by default, so + * without this profile its native kernel never runs in the sweep; `allowIncompatible` is what + * lets a user reach it, and the kernel's NullType handling has to hold there too. + */ + private lazy val allowIncompatibleProfile = Profile( + "allow-incompatible", + 16, + registeredSerdes.toSeq + .map(cls => CometConf.getExprAllowIncompatConfigKey(cls) -> "true") + .sortBy(_._1)) + + /** + * Profiles for the sweeps whose queries end in a projection: batching and the kernel choice are + * what vary. + */ + private lazy val kernelProfiles = + Seq(defaultProfile, smallBatchProfile, allowIncompatibleProfile) + + /** + * Profiles for the sweeps that put a value through an operator. Chosen so every pair of + * settings below appears together at least once: JVM shuffle through the bypass writer + * (partition count under `spark.shuffle.sort.bypassMergeThreshold`) and through the sort-based + * writer (above it, one whole partition per native call), with and without forced spills, + * native shuffle, AQE on and off, native columnar-to-row on and off, and batch sizes of two. + */ + private val physicalProfiles = Seq( + defaultProfile, + Profile( + "default-native-c2r", + 16, + Seq(CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true")), + smallBatchProfile.copy(confs = smallBatchProfile.confs ++ Seq( + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")), + Profile( + "jvm-sort-writer", + 900, + Seq( + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.SHUFFLE_PARTITIONS.key -> "300", + CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2", + CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> "100000", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false")), + Profile( + "jvm-sort-writer-spills", + 900, + Seq( + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.SHUFFLE_PARTITIONS.key -> "300", + CometConf.COMET_BATCH_SIZE.key -> "2", + CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2", + CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> "10", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true")), + Profile( + "native-shuffle", + 900, + Seq( + CometConf.COMET_SHUFFLE_MODE.key -> "native", + SQLConf.SHUFFLE_PARTITIONS.key -> "300", + CometConf.COMET_BATCH_SIZE.key -> "2", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false"))) + + private def rowsOf( + query: String, + cometEnabled: Boolean, + ansi: Boolean, + profile: Profile): (Seq[String], Boolean) = { + val confs = Seq( + CometConf.COMET_ENABLED.key -> cometEnabled.toString, + CometConf.COMET_EXEC_ENABLED.key -> cometEnabled.toString, + SQLConf.ANSI_ENABLED.key -> ansi.toString, + // Keep the producers and their consumers out of the optimizer's hands so they are + // evaluated per row by the engine under test rather than folded at plan time: the + // producers are deterministic and non-nullable, so without this `p IS NULL` becomes + // `false`, `p = p` and `p <=> p` become `true`, and `coalesce(p, p)` / `IF(c, p, p)` + // collapse to `p`, leaving a projection of literals that proves nothing. + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedOptimizerRules) ++ profile.confs + // withSQLConf's body is typed `=> Unit` on the older supported Spark versions and generic + // only on the newer ones, so the result is captured out of band to stay portable. + var result: (Seq[String], Boolean) = (Seq.empty, false) + withSQLConf(confs: _*) { + val df = spark.sql(query) + val rows = df.collect().map(_.toString()).sorted.toSeq + // The sweep's queries put the NullType expression in a projection, so a native + // CometProjectExec is what distinguishes a case that actually exercised a native kernel + // from one that merely fell back. Without this, a sweep where everything falls back would + // still be green while proving nothing. + val nativeProject = + collectFirst(df.queryExecution.executedPlan) { case _: CometProjectExec => + () + }.isDefined + result = (rows, nativeProject) + } + result + } + + // ANSI is a dimension of the consumer sweeps because it changes which serdes wrap their child + // in a null guard and which kernels raise on out-of-range access; the operators below carry + // no ANSI semantics of their own. + for (ansi <- Seq(false, true); profile <- kernelProfiles) { + val tag = s"(ansi=$ansi ${profile.name})" + test(s"NullType producers survive every consumer that Spark accepts $tag") { + sweep( + "consumer", + cases().map { case (producer, expr) => (producer, s"SELECT $expr FROM t") }, + comparedFloor = 120, + nativeFloor = 100, + ansi = ansi, + profile = profile) + } + + // The plain producers are non-nullable, so this is the only sweep where a serde's null guard + // runs natively over a NullType-bearing column that is NULL on some rows and takes its ELSE + // branch; the non-deterministic sweep below makes the same serdes fall back instead. + test(s"nullable deterministic NullType producers survive every consumer $tag") { + sweep( + "nullable", + cases(nullableDeterministic).map { case (producer, expr) => + (producer, s"SELECT $expr FROM t") + }, + comparedFloor = 140, + nativeFloor = 110, + ansi = ansi, + profile = profile) + } + + test(s"stateful NullType producers survive guards filtered by a sibling $tag") { + sweep( + "cross-input", + crossInputCases.map { case (producer, expr) => (producer, s"SELECT $expr FROM t") }, + comparedFloor = 100, + nativeFloor = 60, + ansi = ansi, + profile = profile) + } + + test(s"nullable non-deterministic NullType producers survive every consumer $tag") { + sweep( + "non-deterministic", + cases(nullableNondeterministic).map { case (producer, expr) => + (producer, s"SELECT $expr FROM t") + }, + comparedFloor = 140, + nativeFloor = 110, + ansi = ansi, + profile = profile) + } + } + + // Every physical profile, because the operator sweep is the one whose queries actually cut + // rows into batches and push them through an exchange. + for (profile <- physicalProfiles) { + test(s"NullType producers survive every operator that Spark accepts (${profile.name})") { + sweep( + "operator", + for (producer <- allProducers; (op, template) <- operatorTemplates) + yield (op, template.replace("%s", producer)), + comparedFloor = 160, + nativeFloor = 100, + profile = profile) + } + } + + /** + * Registered array, map and struct serdes whose nested-typed argument no `NullType` producer + * can occupy, so no consumer template can reach them: `Sequence` takes integral or temporal + * bounds, and `StringToMap` and `JsonToStructs` take a string. + */ + private lazy val serdesWithoutNestedInput: Set[Class[_ <: Expression]] = + Set(classOf[Sequence], classOf[StringToMap], classOf[JsonToStructs]) + + /** Aggregates whose serde accepts any input type, so a `NullType` shape can reach them. */ + /** The serdes the sweep must reach: those with a nested-typed argument plus the aggregates. */ + private lazy val registeredSerdes: Set[Class[_]] = + (QueryPlanSerde.arrayExpressions.keySet ++ + QueryPlanSerde.mapExpressions.keySet ++ + QueryPlanSerde.structExpressions.keySet).toSet[Class[_]] -- + serdesWithoutNestedInput ++ anyTypeAggregates + + private lazy val anyTypeAggregates: Set[Class[_]] = Set( + classOf[CollectList], + classOf[CollectSet], + classOf[Count], + classOf[First], + classOf[Last], + classOf[Max], + classOf[Min]) + + /** + * Every expression class in the analyzed and optimized plans of `query`, or none if Spark + * rejects it. Both plans, because the optimizer is what inserts some expressions (`MapSort` + * under a map grouping key) and what expands `RuntimeReplaceable` ones. + */ + private def expressionClasses(query: String): Set[Class[_]] = + Try { + val qe = spark.sql(query).queryExecution + Seq(qe.analyzed, qe.optimizedPlan) + }.toOption.toSeq.flatten + .toSet[LogicalPlan] + .flatMap { plan => + plan.flatMap(_.expressions).flatMap { root => + root.collect { case e: Expression => e }.flatMap { + case r: RuntimeReplaceable => r +: r.replacement.collect { case e: Expression => e } + case e => Seq(e) + } + } + } + .map(_.getClass) + + // The consumer and operator lists are written by hand, so this is the check that a serde added + // to the registry later, or one forgotten now, does not silently stay outside the sweep. + test("the sweep reaches every registered array, map, struct and any-type aggregate serde") { + withTempView("t") { + spark.range(0, 8).createOrReplaceTempView("t") + // `registeredSerdes` is derived from QueryPlanSerde's registries and feeds two protections: + // this staleness check and `allowIncompatibleProfile`'s conf list. Were those registries + // renamed or emptied, both would degrade silently in the same direction -- `missing` becomes + // trivially empty and the profile collapses into a duplicate of `defaultProfile`. Pin the + // size so that cannot pass unnoticed. + assert( + registeredSerdes.size >= 40, + s"only ${registeredSerdes.size} serdes discovered; QueryPlanSerde's registries have moved " + + "and both this check and the allow-incompatible profile are now vacuous") + assert( + allowIncompatibleProfile.confs.size == registeredSerdes.size, + "the allow-incompatible profile must opt every serde the sweep reaches into its kernel") + val reached = + (cases().map { case (_, expr) => + s"SELECT $expr FROM t" + } ++ + (for (producer <- allProducers; (_, template) <- operatorTemplates) + yield template.replace("%s", producer))) + .flatMap(expressionClasses) + .toSet + val missing = registeredSerdes -- reached + assert( + missing.isEmpty, + s"registered serdes no sweep template reaches: ${missing + .map(_.getSimpleName) + .toSeq + .sorted + .mkString(", ")}") + } + } + + for (profile <- physicalProfiles) { + test(s"nested NullType values survive the operators that serialize them (${profile.name})") { + // Select producers by the type they actually have: `filter(array(CAST(NULL AS int)), ...)` + // is `array`, and nesting it only exercises a pre-existing nested-container limitation. + val nullBearingProducers = allProducers.filter { p => + Try(spark.range(0, 8).selectExpr(s"$p AS c").schema.head.dataType).toOption + .exists(SupportLevel.containsType(_, classOf[NullType])) + } + assert( + nullBearingProducers.size >= 5, + s"only ${nullBearingProducers.size} producers still carry a NullType; the producer list " + + "has drifted away from what this sweep is meant to cover") + val nested = + for (producer <- nullBearingProducers; wrapper <- nestingWrappers) + yield wrapper.replace("%s", producer) + sweep( + "nesting", + for (value <- nested.distinct; (op, template) <- serializingOperators) + yield (op, template.replace("%s", value)), + comparedFloor = 200, + nativeFloor = 100, + profile = profile) + } + } + + /** + * Runs every `(label, query)` twice, Comet off and on, and reports all divergences together. A + * query Spark itself rejects is skipped; one whose Comet arm throws or disagrees is a failure. + * + * The floors are floors rather than equalities because which cases Spark accepts varies across + * the supported Spark versions. `comparedFloor` fails a sweep whose templates have gone stale; + * `nativeFloor` fails one where Comet fell back almost everywhere, since falling back is a pass + * and such a sweep would prove nothing about the native kernels. + */ + private def sweep( + name: String, + queries: Seq[(String, String)], + comparedFloor: Int, + nativeFloor: Int, + ansi: Boolean = false, + profile: Profile = defaultProfile): Unit = { + var native = 0 + withTempPath { dir => + // One partition, so every batch holds several rows: a null guard over a non-deterministic + // child only diverges where the CASE sees both matching and non-matching rows in one batch. + spark.range(0, profile.rows, 1, 1).write.parquet(dir.getAbsolutePath) + withTempView("t") { + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t") + + val failures = ArrayBuffer.empty[String] + var compared = 0 + var skipped = 0 + + for ((label, query) <- queries) { + Try(rowsOf(query, cometEnabled = false, ansi, profile)).toOption match { + case None => + skipped += 1 + case Some((sparkRows, _)) => + compared += 1 + Try(rowsOf(query, cometEnabled = true, ansi, profile)) match { + case Failure(e) => + failures += s"[threw:$label] $query\n " + + s"${e.getClass.getSimpleName}: ${firstLine(causeText(e))}" + case Success((cometRows, _)) if cometRows != sparkRows => + failures += s"[mismatch:$label] $query\n" + + s" spark: ${preview(sparkRows)}\n" + + s" comet: ${preview(cometRows)}" + case Success((_, nativeProject)) => + if (nativeProject) native += 1 + } + } + } + + logWarning( + s"NullType $name sweep (ansi=$ansi, profile=${profile.name}): " + + s"compared=$compared native=$native skipped=$skipped total=${queries.length}") + assert( + compared >= comparedFloor, + s"$name sweep only compared $compared of ${queries.length} cases; " + + "its templates have gone stale") + assert( + native >= nativeFloor, + s"$name sweep only executed $native cases natively; Comet is falling back almost " + + "everywhere, so the comparison proves nothing about the native kernels") + if (failures.nonEmpty) { + // Name the profile: under `allow-incompatible` a mismatch can be a serde's *documented* + // difference (array_intersect's element order, array_except's null handling) rather than + // a NullType defect, and the two have to be told apart. + fail( + s"${failures.length} of $compared NullType $name cases diverge from Spark " + + s"(ansi=$ansi, profile=${profile.name}, $skipped invalid in Spark):\n" + + failures.mkString("\n")) + } + } + } + } + + private def causeText(e: Throwable): String = { + val builder = new StringBuilder + var current: Throwable = e + var depth = 0 + while (current != null && depth < 10) { + builder.append(Option(current.getMessage).getOrElse("")).append('\n') + current = current.getCause + depth += 1 + } + builder.toString() + } + + private def firstLine(message: String): String = + Option(message).map(_.linesIterator.next()).getOrElse("") + + private def preview(rows: Seq[String]): String = + rows.take(3).mkString(", ") + (if (rows.length > 3) s", ... (${rows.length} rows)" else "") +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala index 1acb28650bc..70a045901d8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -96,6 +96,35 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar checkShuffleAnswer(shuffled, 1) } + test("columnar shuffle spanning several native writer batches with NullType columns") { + // The sort-based writer (`CometUnsafeShuffleWriter`, chosen once the partition count exceeds + // `spark.shuffle.sort.bypassMergeThreshold`) hands a whole destination partition to + // `process_sorted_row_partition`, which cuts it into `spark.comet.shuffle.jvm.batchSize` + // batches while reusing its builders. A `NullBuilder` keeps its length across `finish`, so + // every Null-bearing shape used to fail or miscount from the second batch on. The bypass + // writer never sends more than one batch per call, hence the partition count here, and the + // spill threshold is lifted so the writer sees several rows per partition. + withSQLConf( + CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2", + CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> "100000", + CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.key -> "false") { + withParquetTable((0L until 2000L).map(Tuple1(_)), "tbl") { + val producers = Seq( + "named_struct('v', _1, 'n', NULL)", + "named_struct('s', named_struct('v', _1, 'n', NULL))", + "array(named_struct('v', _1, 'n', NULL))", + "element_at(transform(array(_1), x -> named_struct('v', x, 'n', NULL)), 1)", + "NULL", + "map(_1, NULL)", + "transform(array(_1), x -> NULL)") + producers.foreach { producer => + val df = sql(s"SELECT /*+ REPARTITION(300) */ _1, $producer AS c FROM tbl") + checkShuffleAnswer(df, 1) + } + } + } + } + test("columnar shuffle with Map[_, NullType] column") { val df = sql("SELECT id, map(id, null) AS m FROM VALUES (1), (2), (3) AS t(id)") val shuffled = df.repartition(2, $"id") diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index e72f4d10f74..36c667ae9b0 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1007,6 +1007,39 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache round-trips NullType columns and children") { + withNativeCache { + // Non-foldable NullType shapes (each references `id`), so the cached batches hold Arrow + // NullVectors at the top level and under a list, a map value and a struct field rather + // than literals folded away at plan time. The nullable IF shape puts a NULL row around a + // NullType-bearing struct. + val query = + """ + SELECT + id AS l, + aggregate(array(id), NULL, (acc, x) -> NULL) AS n, + transform(array(id), x -> NULL) AS an, + map(id, NULL) AS mn, + named_struct('a', id, 'b', NULL) AS sn, + IF(id % 2 = 0, named_struct('a', id, 'b', NULL), NULL) AS sno + FROM range(100) + """ + val expected = spark.sql(query).orderBy("l").collect() + + spark.sql(query).createOrReplaceTempView("null_types_cache") + spark.catalog.cacheTable("null_types_cache") + spark.table("null_types_cache").count() + + assert( + cachedBatchTypes("null_types_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val df = spark.sql("SELECT * FROM null_types_cache").orderBy("l") + assert(df.collect() === expected) + assert(df.queryExecution.executedPlan.toString().contains("CometInMemoryTableScan")) + } + } + test("Comet in-memory cache prunes only on columns that have bounds") { assume(isSpark40Plus, "collated string types require Spark 4.0+") withNativeCache { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 50dd9b97f4f..ddca23252cc 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -1387,8 +1387,7 @@ class CometJoinSuite extends CometTestBase { // map(k, NULL) and transform_values(map(), ...) leave NullType children in the map type // (a bare map() would be constant-folded into a literal the optimizer hoists above the // join). The build side goes through CometBroadcastExchangeExec's batch coalescing and an - // Arrow IPC round trip on the JVM, where a nested NullVector used to hang the appender and - // the IPC reader rejected the NullType map key (see #5525). + // Arrow IPC round trip on the JVM (see #5525). Seq(true, false).foreach { aqe => withSQLConf( CometConf.COMET_BATCH_SIZE.key -> "100", diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index 89c38aa0552..1f34781f6c9 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -27,9 +27,9 @@ import scala.concurrent.duration.DurationInt import scala.jdk.CollectionConverters._ import org.apache.arrow.c.CDataDictionaryProvider -import org.apache.arrow.vector.{FieldVector, IntVector, VectorSchemaRoot} -import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} -import org.apache.arrow.vector.ipc.ArrowStreamReader +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.{ListVector, MapVector} +import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} import org.apache.arrow.vector.types.pojo.{ArrowType, Field} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.catalyst.InternalRow @@ -214,31 +214,45 @@ class UtilsSuite extends CometTestBase { } /** - * One map column of `numRows` rows. With an `IntegerType` key every row is a single entry `i -> - * NULL` (a `NullVector` map value); with a `NullType` key every row is an empty map, as `map()` - * produces (a `NullVector` map key). Both nest a `NullVector` inside the entries struct. + * One `map` column of `numRows` empty maps, as `map()` produces (a `NullVector` + * key). */ - private def nullTypeMapBatch(numRows: Int, keyType: DataType): ColumnarBatch = { - val field = Utils.toArrowField("m", MapType(keyType, NullType), nullable = true, "UTC") + private def emptyNullKeyMapBatch(numRows: Int): ColumnarBatch = { + val field = Utils.toArrowField("m", MapType(NullType, NullType), nullable = true, "UTC") val vector = field.createVector(CometArrowAllocator).asInstanceOf[MapVector] vector.allocateNew() - val entries = vector.getDataVector.asInstanceOf[StructVector] (0 until numRows).foreach { i => vector.startNewValue(i) - keyType match { - case NullType => - vector.endValue(i, 0) - case _ => - entries.setIndexDefined(i) - entries.getChild(MapVector.KEY_NAME).asInstanceOf[IntVector].setSafe(i, i) - vector.endValue(i, 1) - } + vector.endValue(i, 0) } - entries.setValueCount(if (keyType == NullType) 0 else numRows) + vector.getDataVector.setValueCount(0) vector.setValueCount(numRows) new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector, null)), numRows) } + /** + * Writes `bound` through `writer` and asserts the stream reads back with a non-nullable key. + */ + private def assertRoundTrip( + bound: VectorSchemaRoot, + writer: ArrowStreamWriter, + out: ByteArrayOutputStream, + numRows: Int): Unit = { + bound.setRowCount(numRows) + writer.start() + writer.writeBatch() + writer.end() + val reader = + new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray), CometArrowAllocator) + try { + assert(reader.loadNextBatch()) + assert(reader.getVectorSchemaRoot.getRowCount == numRows) + assert(!mapKeyField(reader.getVectorSchemaRoot.getSchema.getFields.get(0)).isNullable) + } finally { + reader.close() + } + } + private def mapKeyField(field: Field): Field = field.getChildren.get(0).getChildren.get(0) /** One `array` column; row `i` holds `i` nulls. */ @@ -254,29 +268,11 @@ class UtilsSuite extends CometTestBase { new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector, null)), numRows) } - /** One `array>` column; every row holds one struct. */ - private def nullStructListBatch(numRows: Int): ColumnarBatch = { - val elementType = StructType(Seq(StructField("a", NullType))) - val field = Utils.toArrowField("l", ArrayType(elementType), nullable = true, "UTC") - val vector = field.createVector(CometArrowAllocator).asInstanceOf[ListVector] - vector.allocateNew() - val elements = vector.getDataVector.asInstanceOf[StructVector] - (0 until numRows).foreach { i => - vector.startNewValue(i) - elements.setIndexDefined(i) - vector.endValue(i, 1) - } - elements.setValueCount(numRows) - vector.setValueCount(numRows) - new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector, null)), numRows) - } - test("withNonNullableMapKeys restores the non-nullable key flag a NullVector drops") { - val batch = nullTypeMapBatch(2, NullType) + val batch = emptyNullKeyMapBatch(2) val field = batch.column(0).asInstanceOf[CometVector].getValueVector.getField - // `toArrowField` declares the key non-nullable, but Arrow's `MinorType.NULL` factory builds the - // key `NullVector` from the name alone, so the vector reports a nullable key. If this assertion - // starts failing, Arrow fixed that and `withNonNullableMapKeys` can go. + // The live vector reports a nullable key (see `Utils.withNonNullableMapKeys`). If this + // assertion starts failing, Arrow fixed that and the repair can go. assert(mapKeyField(field).isNullable) val repaired = Utils.withNonNullableMapKeys(field) @@ -291,95 +287,32 @@ class UtilsSuite extends CometTestBase { } test("newArrowStreamWriter keeps a root whose declared schema is already valid") { - val batch = nullTypeMapBatch(2, NullType) + val batch = emptyNullKeyMapBatch(2) val vector = batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector] val declared = Utils.withNonNullableMapKeys(vector.getField) - // The live vector still reports a nullable key, so a root declared from it would be swapped - // for a repaired copy. One declared from `declared` must be kept as-is: the row count is set - // only after the writer exists, and a swapped root would not see it. + // A root whose declared schema is already valid must be kept, not swapped for a copy. val root = new VectorSchemaRoot(Seq(declared).asJava, Seq(vector).asJava, 0) val out = new ByteArrayOutputStream() val (bound, writer) = Utils.newArrowStreamWriter(root, null, Channels.newChannel(out)) assert(bound eq root) - root.setRowCount(2) - writer.start() - writer.writeBatch() - writer.end() - - val reader = - new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray), CometArrowAllocator) - assert(reader.loadNextBatch()) - assert(reader.getVectorSchemaRoot.getRowCount == 2) - assert(!mapKeyField(reader.getVectorSchemaRoot.getSchema.getFields.get(0)).isNullable) - reader.close() + assertRoundTrip(bound, writer, out, numRows = 2) batch.close() } test("newArrowStreamWriter returns the root a later row count must be set on") { - val batch = nullTypeMapBatch(2, NullType) + val batch = emptyNullKeyMapBatch(2) val vector = batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector] - // Declared from the live vector, so the key is nullable and the root must be swapped. Setting - // the row count on the returned root has to reach the writer; setting it on the original one - // would ship an empty batch ("Array length did not match record batch length" downstream). + // Declared from the live vector, so the key is nullable and the root must be swapped. val root = new VectorSchemaRoot(Seq(vector).asJava) val out = new ByteArrayOutputStream() val (bound, writer) = Utils.newArrowStreamWriter(root, null, Channels.newChannel(out)) assert(bound ne root) - bound.setRowCount(2) - writer.start() - writer.writeBatch() - writer.end() - - val reader = - new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray), CometArrowAllocator) - assert(reader.loadNextBatch()) - assert(reader.getVectorSchemaRoot.getRowCount == 2) - assert(!mapKeyField(reader.getVectorSchemaRoot.getSchema.getFields.get(0)).isNullable) - reader.close() + assertRoundTrip(bound, writer, out, numRows = 2) batch.close() } - test("serializeBatches round-trips a NullType map key through Arrow IPC") { - // The IPC reader rebuilds a MapVector from the stream's schema and rejects a nullable key - // ("Map data key type should be a non-nullable"), which is exactly what a NullVector key - // reports unless the written schema is repaired. - val numRows = 3 - val batch = nullTypeMapBatch(numRows, NullType) - val (rowCount, buf) = Utils.serializeBatches(Iterator(batch)).next() - assert(rowCount == numRows) - - val decoded = Utils.decodeBatches(buf, "test").toSeq - assert(decoded.map(_.numRows()).sum == numRows) - decoded.foreach(_.close()) - } - - test("coalesceBroadcastBatches ships struct-nested NullType uncoalesced") { - // VectorSchemaRootAppender cannot grow a NullVector nested in a struct, including the map - // entries struct (NullVector.reAlloc is a no-op), so such buffers must be passed through, - // not appended. The list case pins that a struct below a list is still a struct. - val cases: Seq[(String, Int => ColumnarBatch)] = Seq( - "map" -> (nullTypeMapBatch(_, IntegerType)), - "map" -> (nullTypeMapBatch(_, NullType)), - "array>" -> nullStructListBatch) - cases.foreach { case (name, batch) => - val numRows = 4 - val numBatches = 3 - val batches = (0 until numBatches).map(_ => batch(numRows)) - val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq - - val (result, batchCount, totalRows) = Utils.coalesceBroadcastBatches(bufs.iterator) - // The pass-through signature: original buffers, nothing coalesced. - assert(batchCount == 0 && totalRows == 0, name) - assert(result.length == numBatches, name) - - val decoded = result.iterator.flatMap(b => Utils.decodeBatches(b, "test")).toSeq - assert(decoded.map(_.numRows()).sum == numRows.toLong * numBatches, name) - decoded.foreach(_.close()) - } - } - test("coalesceBroadcastBatches bypasses exactly the schemas with a NullType under a struct") { // Exhaustive over the shape space of the bypass rule: VectorAppender hangs only when a // NullVector is a *direct* child of a struct (see `Utils.hasNullDirectlyUnderStruct` for @@ -427,7 +360,9 @@ class UtilsSuite extends CometTestBase { CometArrowAllocator) .next() } - val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq + // Force serialization eagerly: on Scala 2.12 `toSeq` is a lazy Stream, and the inputs + // are closed on the next line. + val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toVector batches.foreach(_.close()) val (result, batchCount, totalRows) = Await.result( @@ -439,7 +374,8 @@ class UtilsSuite extends CometTestBase { (batchCount == 0) == expectBypass, s"$name: batchCount=$batchCount but bypass expected=$expectBypass") assert(result.length == (if (expectBypass) numBatches else 1), name) - if (!expectBypass) assert(totalRows == numRows.toLong * numBatches, name) + if (expectBypass) assert(totalRows == 0, name) + else assert(totalRows == numRows.toLong * numBatches, name) val decoded = result.iterator.flatMap(b => Utils.decodeBatches(b, "test")).toSeq assert(decoded.map(_.numRows()).sum == numRows.toLong * numBatches, name) @@ -448,14 +384,12 @@ class UtilsSuite extends CometTestBase { } test("coalesceBroadcastBatches keeps coalescing plain null lists") { - // A NullVector directly under a list is safe to append: the list's capacity loop only looks - // at its own offset and validity buffers, and ListVector.setValueCount sets the child's - // count from the last offset. Bypassing coalescing here would cost every consuming task one - // IPC stream per original buffer instead of one. + // A NullVector directly under a list appends fine (see `Utils.hasNullDirectlyUnderStruct`); + // bypassing here would cost every consuming task one IPC stream per original buffer. val numRows = 4 val numBatches = 3 val batches = (0 until numBatches).map(_ => nullListBatch(numRows)) - val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq + val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toVector val (result, batchCount, totalRows) = Utils.coalesceBroadcastBatches(bufs.iterator) assert(batchCount == numBatches) From 4be5fa0859267a3ee5ccbd6e4e6660a007bbe845 Mon Sep 17 00:00:00 2001 From: grorge Date: Mon, 7 Sep 2026 08:44:35 +0800 Subject: [PATCH 4/6] fix: keep NullType-under-struct broadcast build sides on Spark's broadcast Arrow's VectorSchemaRootAppender loops forever on a NullVector directly under a struct, so coalesceBroadcastBatches shipped such build sides uncoalesced and every consuming task opened one Arrow IPC stream per buffer (CometBatchRDD.compute). The NullType-columns microbenchmark measured that cost: a 64-buffer map build side over 8 consumer tasks opened 512 streams and ran 0.9x Spark where the fallback path ran 1.2x, and the cost grows with buffers x tasks. CometBroadcastExchangeExec now reports a build side with a NullType directly under a struct or map entry as Unsupported, naming the columns, so the broadcast stays on Spark as it did before the codegen dispatcher admitted NullType outputs. array is insulated by the list and still coalesces natively. The planner gate and the coalescer's now defensive bypass share one predicate, Utils.hasNullTypeUnderStruct, so both lift together once Arrow's appender can grow a NullVector under a struct; UtilsSuite runs the real appender over every shape and pins the predicate to exactly the hanging ones. CometJoinSuite asserts the fallback reason and that the plan carries Spark's BroadcastExchangeExec. Assisted-by: Claude Code (claude-fable-5) --- .../comet/CometBroadcastExchangeExec.scala | 22 ++++++++++++++- .../apache/spark/sql/comet/util/Utils.scala | 27 +++++++++++++++++++ .../apache/comet/exec/CometJoinSuite.scala | 22 +++++++++++---- .../spark/sql/comet/util/UtilsSuite.scala | 16 ++++------- 4 files changed, 70 insertions(+), 17 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala index 9ae92845c9c..2de985db5a7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala @@ -45,7 +45,7 @@ import org.apache.spark.util.io.ChunkedByteBuffer import com.google.common.base.Objects import org.apache.comet.{CometConf, ConfigEntry} -import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.{Compatible, OperatorOuterClass, SupportLevel, Unsupported} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.ShimCometBroadcastExchangeExec @@ -264,6 +264,26 @@ object CometBroadcastExchangeExec extends CometSink[BroadcastExchangeExec] { override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED) + /** + * A build side with a `NullType` directly under a struct or map entry stays on Spark's + * broadcast. `Utils.coalesceBroadcastBatches` cannot merge such batches (Arrow's appender + * cannot grow a `NullVector` under a struct) and would ship them uncoalesced, which costs every + * consuming task one IPC stream per broadcast buffer: measurably slower than Spark's broadcast + * once the build side spans many batches. See `Utils.hasNullTypeUnderStruct` for lifting this + * when Arrow is fixed. + */ + override def getSupportLevel(b: BroadcastExchangeExec): SupportLevel = { + val uncoalescable = b.child.output.filter(a => Utils.hasNullTypeUnderStruct(a.dataType)) + if (uncoalescable.isEmpty) { + Compatible(None) + } else { + Unsupported( + Some("NullType directly under a struct or map entry in the broadcast build side " + + s"(${uncoalescable.map(a => s"${a.name}: ${a.dataType.simpleString}").mkString(", ")}) " + + "cannot be coalesced for broadcast")) + } + } + override def createExec( nativeOp: OperatorOuterClass.Operator, b: BroadcastExchangeExec): CometNativeExec = { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index de88377b33b..3d1217053af 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -343,6 +343,30 @@ object Utils extends CometTypeShim with Logging { // scalastyle:on arrowstreamwriter } + /** + * Whether a `NullType` is a direct child of a struct (map keys and values included, as the + * entries struct's children) anywhere in `dataType`. This is the Spark-side twin of + * [[hasNullDirectlyUnderStruct]], for deciding at planning time what the broadcast coalescer + * would refuse: `CometBroadcastExchangeExec` keeps such build sides on Spark's broadcast, since + * shipping them uncoalesced costs every consuming task one IPC stream per broadcast buffer. + * + * When Comet moves to an Arrow Java release whose `VectorAppender` can grow a `NullVector` + * under a struct, drop this gate from `CometBroadcastExchangeExec.getSupportLevel` and the + * bypass in [[coalesceBroadcastBatches]]; `UtilsSuite` runs the real appender over every shape + * and will show which ones no longer need it. + */ + def hasNullTypeUnderStruct(dataType: DataType): Boolean = { + // A list insulates whatever is below it, so `inStruct` resets when descending into one. + def check(dt: DataType, inStruct: Boolean): Boolean = dt match { + case NullType => inStruct + case ArrayType(element, _) => check(element, inStruct = false) + case StructType(fields) => fields.exists(f => check(f.dataType, inStruct = true)) + case MapType(key, value, _) => check(key, inStruct = true) || check(value, inStruct = true) + case _ => false + } + check(dataType, inStruct = false) + } + /** * Whether an Arrow `Null` field is a direct child of a struct (map entries included) anywhere * in `schema`. Arrow's `VectorAppender` cannot grow such a column: a struct's capacity is the @@ -354,6 +378,9 @@ object Utils extends CometTypeShim with Logging { * validity buffers. So a `Null` under a list appends fine (`array(NULL)`), and so does a list * under a struct even when the list holds nulls (`map(k, array(NULL))`) - the struct only sees * the list's own, growable capacity. Top-level `NullVector`s are fine too. + * + * The planner refuses these schemas before they reach a Comet broadcast (see + * [[hasNullTypeUnderStruct]]); this check keeps the coalescer safe should one arrive anyway. */ private def hasNullDirectlyUnderStruct(schema: Schema): Boolean = { def check(field: Field): Boolean = { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index ddca23252cc..baaba71238c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -32,7 +32,8 @@ import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, IsNot import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, MetadataBuilder, NullType, StructField, StructType} @@ -1386,8 +1387,9 @@ class CometJoinSuite extends CometTestBase { test("Broadcast HashJoin with NullType map columns on the build side") { // map(k, NULL) and transform_values(map(), ...) leave NullType children in the map type // (a bare map() would be constant-folded into a literal the optimizer hoists above the - // join). The build side goes through CometBroadcastExchangeExec's batch coalescing and an - // Arrow IPC round trip on the JVM (see #5525). + // join). Such a build side cannot be coalesced for a Comet broadcast (see #5525 and + // `Utils.hasNullTypeUnderStruct`), so the exchange and the join stay on Spark with a + // reason that names the columns. Seq(true, false).foreach { aqe => withSQLConf( CometConf.COMET_BATCH_SIZE.key -> "100", @@ -1409,9 +1411,19 @@ class CometJoinSuite extends CometTestBase { // covering nothing. assert(df.schema("m1").dataType.asInstanceOf[MapType].valueType === NullType) assert(df.schema("m2").dataType.asInstanceOf[MapType].keyType === NullType) - checkSparkAnswerAndOperator( + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( df, - Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec])) + "NullType directly under a struct or map entry in the broadcast build side " + + "(m1: map, m2: map) cannot be coalesced for broadcast") + def operators(p: SparkPlan): Seq[SparkPlan] = p +: (p match { + case a: AdaptiveSparkPlanExec => operators(a.executedPlan) + case s: QueryStageExec => operators(s.plan) + case r: ReusedExchangeExec => operators(r.child) + case _ => p.children.flatMap(operators) + }) + val all = operators(cometPlan) + assert(all.exists(_.isInstanceOf[BroadcastExchangeExec]), cometPlan.treeString) + assert(!all.exists(_.isInstanceOf[CometBroadcastExchangeExec]), cometPlan.treeString) } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index 1f34781f6c9..b54616bf332 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -318,7 +318,8 @@ class UtilsSuite extends CometTestBase { // NullVector is a *direct* child of a struct (see `Utils.hasNullDirectlyUnderStruct` for // the Arrow mechanics). Each shape runs the real appender under a timeout, so a rule that // is too narrow shows up as a timeout on the hanging shapes instead of a hung build, and - // one that is too wide shows up as a needless bypass. + // one that is too wide shows up as a needless bypass. Once an Arrow release fixes the + // appender, the shapes reported as needless bypasses are the ones whose gate can go. val nullStruct = StructType(Seq(StructField("a", NullType))) val shapes: Seq[(DataType, Any)] = Seq( NullType -> null, @@ -335,15 +336,6 @@ class UtilsSuite extends CometTestBase { MapType(IntegerType, ArrayType(NullType)) -> ArrayBasedMapData(Array[Any](1), Array[Any](new GenericArrayData(Array[Any](null)))), MapType(NullType, NullType) -> ArrayBasedMapData(Array.empty[Any], Array.empty[Any])) - // A list insulates whatever is below it, so `inStruct` resets when descending into one. - def nullUnderStruct(dt: DataType, inStruct: Boolean): Boolean = dt match { - case NullType => inStruct - case ArrayType(element, _) => nullUnderStruct(element, inStruct = false) - case StructType(fields) => fields.exists(f => nullUnderStruct(f.dataType, inStruct = true)) - case MapType(k, v, _) => - nullUnderStruct(k, inStruct = true) || nullUnderStruct(v, inStruct = true) - case _ => false - } val numRows = 4 val numBatches = 3 @@ -369,7 +361,9 @@ class UtilsSuite extends CometTestBase { Future(Utils.coalesceBroadcastBatches(bufs.iterator))(ExecutionContext.global), 10.seconds) - val expectBypass = nullUnderStruct(dataType, inStruct = false) + // The planner's gate (`CometBroadcastExchangeExec.getSupportLevel`) uses this same + // predicate, so this also pins that it refuses exactly the shapes the appender cannot take. + val expectBypass = Utils.hasNullTypeUnderStruct(dataType) assert( (batchCount == 0) == expectBypass, s"$name: batchCount=$batchCount but bypass expected=$expectBypass") From 02f20d88635c91781b6805e68a4525f568605d2b Mon Sep 17 00:00:00 2001 From: grorge Date: Sun, 6 Sep 2026 20:43:52 +0800 Subject: [PATCH 5/6] test: add a NullType-columns microbenchmark Times NullType map and struct projections consumed by the codegen dispatcher, and broadcast joins whose build side carries a NullType column, against Spark and the prior fallback path (dispatcher off), each at the default and at a small batch size and with few and many broadcast build files. After timing, every case checks that all arms return the same rows and profiles one execution per arm: executed plan, JVM allocation, peak and retained heap, peak task execution memory, retained Arrow memory, and for the broadcast cases the buffer count, coalescing, consumer task count and resulting Arrow IPC stream count. Assisted-by: Claude Code (claude-fable-5) --- .../sql/benchmark/CometBenchmarkBase.scala | 119 +++-- .../CometNullTypeColumnsBenchmark.scala | 475 ++++++++++++++++++ 2 files changed, 558 insertions(+), 36 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometNullTypeColumnsBenchmark.scala diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala index 487e4797d12..d605c6d83a9 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala @@ -42,6 +42,8 @@ import org.apache.spark.sql.types.DecimalType import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions +import CometBenchmarkBase.BenchmarkArm + trait CometBenchmarkBase extends SqlBasedBenchmark with AdaptiveSparkPlanHelper @@ -62,6 +64,7 @@ trait CometBenchmarkBase .set( "spark.shuffle.manager", "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + extraSparkConf.foreach { case (key, value) => conf.set(key, value) } val sparkSession = SparkSession .builder() @@ -80,6 +83,12 @@ trait CometBenchmarkBase sparkSession } + /** + * Static Spark settings a benchmark needs in its context, applied on top of the defaults above + * before the context starts. Session-level SQL configs belong in the cases instead. + */ + protected def extraSparkConf: Map[String, String] = Map.empty + def runCometBenchmark(args: Array[String]): Unit override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { @@ -123,26 +132,25 @@ trait CometBenchmarkBase * SQL query to benchmark * @param extraCometConfigs * Additional configurations to apply for the Comet case (optional) + * @param extraArms + * Extra Comet cases, each overlaying its configs on the Comet case, e.g. a prior path with a + * dispatcher off (see [[CometBenchmarkBase.BenchmarkArm]]). + * @param expectNative + * Whether the Comet case is checked to be fully Comet native. False for a case that measures + * a deliberate fallback. */ final def runExpressionBenchmark( name: String, cardinality: Long, query: String, - extraCometConfigs: Map[String, String] = Map.empty): Unit = { + extraCometConfigs: Map[String, String] = Map.empty, + extraArms: Seq[BenchmarkArm] = Nil, + expectNative: Boolean = true): Unit = { val benchmark = new Benchmark(name, cardinality, output = output) - // Constant folding is excluded so that expressions over literal arguments are still evaluated - // per row. It must be excluded for both arms: if only Comet excludes it, Spark folds the - // expression away and does no per-row work, and the comparison is meaningless. - val noConstantFolding = - SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName) - - val sparkConfigs = Seq(noConstantFolding, CometConf.COMET_ENABLED.key -> "false") - - val cometConfigs = Seq( - noConstantFolding, - CometConf.COMET_ENABLED.key -> "true", - CometConf.COMET_EXEC_ENABLED.key -> "true") ++ extraCometConfigs + val arms = expressionBenchmarkArms(extraCometConfigs, extraArms) + .map(arm => if (arm.name == "Comet") arm.copy(expectNative = expectNative) else arm) + val sparkConfigs = arms.head.configs // Check that the benchmarked expression survives optimization into the Spark baseline plan. // The query is not executed: before execution `stripAQEPlan` yields the initial physical @@ -169,38 +177,64 @@ trait CometBenchmarkBase } } - // Check that the plan is fully Comet native before running the benchmark. Unlike the check - // above this one does execute the query, because AQE can re-plan a join after the first - // stage completes and the Comet operators are what we are checking for. - withSQLConf(cometConfigs: _*) { - val df = spark.sql(query) - df.noop() - val plan = stripAQEPlan(df.queryExecution.executedPlan) - findFirstNonCometOperator(plan).foreach { op => - warn( - benchmark, - s"""WARNING: the Comet plan is NOT fully Comet native, so the Comet case below is - |partly or wholly measuring Spark. - |First non-Comet operator: ${op.nodeName} - |Comet plan:""".stripMargin + "\n" + plan.treeString) + // Check that the plan of every arm expected to be native is fully Comet native before running + // the benchmark. Unlike the check above this one does execute the query, because AQE can + // re-plan a join after the first stage completes and the Comet operators are what we are + // checking for. + arms.filter(_.expectNative).foreach { arm => + withSQLConf(arm.configs: _*) { + val df = spark.sql(query) + df.noop() + val plan = stripAQEPlan(df.queryExecution.executedPlan) + findFirstNonCometOperator(plan).foreach { op => + warn( + benchmark, + s"""WARNING: the ${arm.name} plan is NOT fully Comet native, so that case below is + |partly or wholly measuring Spark. + |First non-Comet operator: ${op.nodeName} + |Comet plan:""".stripMargin + "\n" + plan.treeString) + } } } - benchmark.addCase("Spark") { _ => - withSQLConf(sparkConfigs: _*) { - spark.sql(query).noop() - } - } - - benchmark.addCase("Comet") { _ => - withSQLConf(cometConfigs: _*) { - spark.sql(query).noop() + arms.foreach { arm => + benchmark.addCase(arm.name) { _ => + withSQLConf(arm.configs: _*) { + spark.sql(query).noop() + } } } benchmark.run() } + /** + * The arms `runExpressionBenchmark` times, in order: the Spark baseline, the Comet case with + * `extraCometConfigs`, then each extra arm overlaid on the Comet case. Exposed so a benchmark + * can run its own checks (result equality, memory, plan) over exactly the timed configurations. + */ + protected def expressionBenchmarkArms( + extraCometConfigs: Map[String, String], + extraArms: Seq[BenchmarkArm]): Seq[BenchmarkArm] = { + // Constant folding is excluded so that expressions over literal arguments are still evaluated + // per row. It must be excluded for both arms: if only Comet excludes it, Spark folds the + // expression away and does no per-row work, and the comparison is meaningless. + val noConstantFolding = + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName) + + val sparkConfigs = Seq(noConstantFolding, CometConf.COMET_ENABLED.key -> "false") + + val cometConfigs = Seq( + noConstantFolding, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true") ++ extraCometConfigs + + Seq( + BenchmarkArm("Spark", sparkConfigs, expectNative = false), + BenchmarkArm("Comet", cometConfigs, expectNative = true)) ++ + extraArms.map(arm => arm.copy(configs = cometConfigs ++ arm.configs)) + } + /** * Returns the value of `spark.sql.optimizer.excludedRules` with `rule` appended, so that * benchmark-specific exclusions do not clobber exclusions already configured by the caller. @@ -382,6 +416,19 @@ trait CometBenchmarkBase object CometBenchmarkBase { + /** + * One timed case of `runExpressionBenchmark`. + * + * @param name + * the case name in the results table + * @param configs + * the SQL configs the case runs under; for an extra arm, these overlay the Comet case + * @param expectNative + * whether the plan is checked to be fully Comet native before timing. False for an arm that + * deliberately measures a fallback path. + */ + case class BenchmarkArm(name: String, configs: Seq[(String, String)], expectNative: Boolean) + /** * SplitMix64 finalizer, used to turn a row id into a well distributed pseudo-random value. * diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometNullTypeColumnsBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometNullTypeColumnsBenchmark.scala new file mode 100644 index 00000000000..3be38cf4a8a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometNullTypeColumnsBenchmark.scala @@ -0,0 +1,475 @@ +/* + * 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.benchmark + +import java.lang.management.{ManagementFactory, MemoryType} +import java.nio.charset.StandardCharsets + +import scala.collection.JavaConverters._ + +import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted, SparkListenerTaskEnd} +import org.apache.spark.sql.Row +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometPlan} +import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} +import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, QueryStageExec} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec} +import org.apache.spark.util.io.ChunkedByteBuffer + +import org.apache.comet.{CometArrowAllocator, CometConf} + +import CometBenchmarkBase.BenchmarkArm + +/** + * Benchmark for NullType-bearing columns on the two paths the NullType work touches: projections + * of NullType maps / structs consumed by the JVM codegen dispatcher, and broadcast joins whose + * build side carries a NullType column (a plain `array` coalesces through + * `Utils.coalesceBroadcastBatches`; a NullType directly under a struct or map entry cannot be + * coalesced, so the exchange and join stay on Spark). + * + * Every case is timed on five arms under the same session, warmup and iteration policy: Spark, + * Comet, Comet on the prior path (codegen dispatcher off, so the NullType projection falls back + * and takes its operator with it), and the two Comet arms again with a small + * `spark.comet.batchSize`. The build side is sized for few and for many broadcast buffers. + * + * After timing, each case checks that every arm returns the same rows and profiles one execution + * per arm: the executed plan (native or where it falls back), JVM allocation, peak and retained + * heap, Spark's peak task execution memory, the native operators' own memory metrics, retained + * Arrow memory, and for the broadcast cases the number of broadcast buffers, whether they were + * coalesced, the number of consuming tasks, and the resulting count of Arrow IPC streams opened + * (every consuming task decodes every buffer, see `CometBatchRDD.compute`). To run: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometNullTypeColumnsBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometNullTypeColumnsBenchmark-**results.txt". + */ +object CometNullTypeColumnsBenchmark extends CometBenchmarkBase { + + private val smallBatchSize = 1024 + + private val dispatchOff = CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false" + private val smallBatch = CometConf.COMET_BATCH_SIZE.key -> smallBatchSize.toString + + // The prior path: with the JVM codegen dispatcher off, a NullType-bearing projection falls + // back to Spark and takes its enclosing operator with it. Each Comet arm is also run with a + // small `spark.comet.batchSize`, the batch size of the native operators and of the JVM-side + // conversions. The native Parquet scan emits 8192-row batches whatever this is set to, so the + // broadcast buffer count is varied through the build size instead (see the profile columns). + private val extraArms: Seq[BenchmarkArm] = Seq( + BenchmarkArm("Comet (dispatch off, prior path)", Seq(dispatchOff), expectNative = false), + BenchmarkArm(s"Comet (batch $smallBatchSize)", Seq(smallBatch), expectNative = true), + BenchmarkArm( + s"Comet (dispatch off, prior path, batch $smallBatchSize)", + Seq(dispatchOff, smallBatch), + expectNative = false)) + + /** + * Times the case on every arm, then checks result equality and profiles each arm. + * `expectNative` is false for a case whose Comet plan is meant to fall back. + */ + private def runNullTypeCase( + name: String, + cardinality: Long, + query: String, + expectNative: Boolean = true): Unit = { + val arms = if (expectNative) extraArms else extraArms.map(_.copy(expectNative = false)) + runExpressionBenchmark( + name, + cardinality, + query, + extraArms = arms, + expectNative = expectNative) + verifyAndProfile(name, query, expressionBenchmarkArms(Map.empty, arms)) + } + + private def projectionBenchmarks(values: Int): Unit = { + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable(dir, spark.sql(s"SELECT value AS c1 FROM $tbl")) + + runBenchmark("NullType projection consumed") { + runNullTypeCase( + "map(c1, NULL) projected and consumed", + values, + "SELECT size(map_keys(map(c1, NULL))) FROM parquetV1Table") + // Control: a plain list, the shape that stays on every fast path. + runNullTypeCase( + "transform to array projected and consumed (control)", + values, + "SELECT size(transform(array(c1), x -> NULL)) FROM parquetV1Table") + } + } + } + } + + private def broadcastBuildBenchmarks(): Unit = { + val probeRows = 1024 * 1024 + + // Few and many broadcast buffers: the native scan emits 8192-row batches and the broadcast + // collects one buffer per batch, so the build size sets the buffer count. Files are split and + // packed by Spark's defaults; the hint forces the broadcast join. + for ((label, buildRows) <- Seq(("2 buffers", 2 * 8192), ("64 buffers", 64 * 8192))) { + withTempPath { dir => + withTempTable("probe", "build") { + spark + .range(probeRows) + .selectExpr("id AS k", "id % 100 AS v") + .write + .parquet(s"${dir.getAbsolutePath}/probe") + // Parquet cannot store NullType, so the build table holds a plain column and each case + // projects its collection column in the join subquery. + spark + .range(buildRows) + .selectExpr("id AS k") + .write + .parquet(s"${dir.getAbsolutePath}/build") + + spark.read.parquet(s"${dir.getAbsolutePath}/probe").createOrReplaceTempView("probe") + spark.read.parquet(s"${dir.getAbsolutePath}/build").createOrReplaceTempView("build") + + runBenchmark(s"BroadcastHashJoin build with NullType column ($label)") { + def joinQuery(buildProjection: String, consumed: String): String = + s"""SELECT /*+ BROADCAST(b) */ count(p.v + $consumed) + |FROM probe p JOIN (SELECT k, $buildProjection FROM build) b ON p.k = b.k + |""".stripMargin + + // Control: a plain list coalesces, since a NullVector under a list is safe to + // append. + runNullTypeCase( + "array column (control, coalesced)", + probeRows, + joinQuery("transform(array(k), x -> NULL) AS null_list", "size(b.null_list)")) + // A NullType directly under a struct or map entry cannot be coalesced (Arrow's + // appender cannot grow such a NullVector), so `CometBroadcastExchangeExec` refuses + // the build side and the exchange and join stay on Spark. + runNullTypeCase( + "map column (Spark broadcast)", + probeRows, + joinQuery("map(k, NULL) AS null_map", "size(map_keys(b.null_map))"), + expectNative = false) + // Consumed whole: a field access like `b.null_struct.a` is simplified to `b.k` by + // SimplifyExtractValueOps and the struct would never reach the broadcast. + runNullTypeCase( + "struct with NULL field column (Spark broadcast)", + probeRows, + joinQuery( + "named_struct('a', k, 'b', NULL) AS null_struct", + "size(array(b.null_struct))"), + expectNative = false) + } + } + } + } + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + runBenchmarkWithTable("NullType projections", 1024 * 1024 * 10) { values => + projectionBenchmarks(values) + } + broadcastBuildBenchmarks() + } + + // --------------------------------------------------------------------------------------------- + // Result equality and per-arm profiling + // --------------------------------------------------------------------------------------------- + + /** An order-independent digest of a result set: row count and the sum / xor of row hashes. */ + private case class Digest(rows: Long, sum: Long, xor: Long) + + private case class BroadcastProfile(buffers: Int, coalescedBatches: Long, consumerTasks: Int) { + def ipcStreams: Long = buffers.toLong * consumerTasks + } + + private case class Profile( + arm: String, + plan: String, + wallMs: Long, + jvmAllocatedBytes: Long, + heapPeakBytes: Long, + heapRetainedBytes: Long, + sparkPeakExecutionMemoryBytes: Long, + nativeMemoryBytes: Long, + arrowRetainedBytes: Long, + broadcast: Either[String, BroadcastProfile]) + + /** Records task and stage metrics of the jobs run while it is registered. */ + private class ProfileListener extends SparkListener { + var taskPeakExecutionMemory = 0L + var broadcastConsumerTasks = 0 + + override def onTaskEnd(event: SparkListenerTaskEnd): Unit = synchronized { + Option(event.taskMetrics).foreach { metrics => + taskPeakExecutionMemory = math.max(taskPeakExecutionMemory, metrics.peakExecutionMemory) + } + } + + // A stage that consumes a Comet broadcast has the broadcast's `CometBatchRDD` in its lineage, + // and every one of its tasks decodes every broadcast buffer. + override def onStageCompleted(event: SparkListenerStageCompleted): Unit = synchronized { + if (event.stageInfo.rddInfos.exists(_.name == "CometBatchRDD")) { + broadcastConsumerTasks += event.stageInfo.numTasks + } + } + } + + private def verifyAndProfile(name: String, query: String, arms: Seq[BenchmarkArm]): Unit = { + val digests = arms.map(arm => arm.name -> underConf(arm.configs)(digest(query))) + val (_, reference) = digests.head + val disagreeing = digests.filter(_._2 != reference) + if (disagreeing.nonEmpty) { + val detail = digests.map { case (arm, d) => s" $arm: $d" }.mkString("\n") + throw new IllegalStateException( + s"Result mismatch in '$name': arms disagree with ${digests.head._1}\n$detail") + } + emit(s"\n$name: results identical across ${arms.size} arms " + + s"(${reference.rows} rows, digest ${reference.sum.toHexString}/${reference.xor.toHexString})") + + val profiles = arms.map(arm => underConf(arm.configs)(profile(arm, query))) + emitProfiles(profiles) + } + + /** `withSQLConf` returns Unit; this variant returns the block's value. */ + private def underConf[T](configs: Seq[(String, String)])(f: => T): T = { + var result: Option[T] = None + withSQLConf(configs: _*) { result = Some(f) } + result.get + } + + private def digest(query: String): Digest = { + val (rows, sum, xor) = spark + .sql(query) + .rdd + .mapPartitions(NullTypeRowDigest.digestPartition) + .fold((0L, 0L, 0L)) { case ((r1, s1, x1), (r2, s2, x2)) => + (r1 + r2, s1 + s2, x1 ^ x2) + } + Digest(rows, sum, xor) + } + + private def profile(arm: BenchmarkArm, query: String): Profile = { + val heapPools = + ManagementFactory.getMemoryPoolMXBeans.asScala.filter(_.getType == MemoryType.HEAP) + val threads = ManagementFactory.getThreadMXBean.asInstanceOf[com.sun.management.ThreadMXBean] + def allocatedByAllThreads(): Long = + threads.getThreadAllocatedBytes(threads.getAllThreadIds).filter(_ > 0).sum + def usedHeapAfterGc(): Long = { + System.gc() + System.gc() + val runtime = Runtime.getRuntime + runtime.totalMemory - runtime.freeMemory + } + + val listener = new ProfileListener + spark.sparkContext.addSparkListener(listener) + try { + val heapBefore = usedHeapAfterGc() + heapPools.foreach(_.resetPeakUsage()) + val arrowBefore = CometArrowAllocator.getAllocatedMemory + val allocatedBefore = allocatedByAllThreads() + val start = System.nanoTime() + + // Execute the dataset's own physical plan rather than `noop()`: a write plans a separate + // query, whose operators would be the ones carrying the metrics and the final AQE plan. + val qe = spark.sql(query).queryExecution + SQLExecution.withNewExecutionId(qe) { + qe.toRdd.foreach(_ => ()) + } + + val wallMs = (System.nanoTime() - start) / 1000000 + val allocated = allocatedByAllThreads() - allocatedBefore + val heapPeak = heapPools.map(_.getPeakUsage.getUsed).sum + spark.sparkContext.listenerBus.waitUntilEmpty() + val arrowRetained = CometArrowAllocator.getAllocatedMemory - arrowBefore + val heapRetained = usedHeapAfterGc() - heapBefore + + val plan = stripAQEPlan(qe.executedPlan) + Profile( + arm.name, + describePlan(arm, plan), + wallMs, + allocated, + heapPeak, + heapRetained, + listener.taskPeakExecutionMemory, + nativeMemory(plan), + arrowRetained, + broadcastProfile(plan, listener.broadcastConsumerTasks)) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + /** + * Memory reported by the native operators' own metrics: the hash join's `build_mem_used` and + * the aggregate's `peak_mem_used`, summed over the plan and over tasks. Spark's task-level + * `peakExecutionMemory` only sees Spark operators, so the two columns are complementary. + */ + private def nativeMemory(plan: SparkPlan): Long = + operators(plan) + .collect { case p: CometPlan => p } + .flatMap { p => + p.metrics.collect { + case (name, metric) if name == "build_mem_used" || name == "peak_mem_used" => + metric.value + } + } + .sum + + /** + * Every operator of the executed plan in pre-order, descending into AQE query stages and reused + * exchanges, whose operators are only reachable through `plan` / `child` after execution. The + * original instances are returned (never a transformed copy: a copy carries fresh, unregistered + * metrics), each once, since a reused exchange is shared. + */ + private def operators(plan: SparkPlan): Seq[SparkPlan] = { + val seen = java.util.Collections.newSetFromMap( + new java.util.IdentityHashMap[SparkPlan, java.lang.Boolean]) + def walk(p: SparkPlan): Seq[SparkPlan] = + if (!seen.add(p)) { + Nil + } else { + p match { + case s: QueryStageExec => s +: walk(s.plan) + case r: ReusedExchangeExec => r +: walk(r.child) + case _ => p +: p.children.flatMap(walk) + } + } + walk(plan) + } + + private def describePlan(arm: BenchmarkArm, plan: SparkPlan): String = + if (arm.name == "Spark") { + "Spark" + } else { + // `findFirstNonCometOperator` walks a subtree and stops at query stage leaves, so ask it + // about each operator of the expanded plan in turn; a node is non-Comet when it is the + // first one reported for its own subtree. Stage wrappers, reused exchanges and AQE shuffle + // reads are plumbing Comet consumes natively. + val plumbing = + Seq(classOf[QueryStageExec], classOf[ReusedExchangeExec], classOf[AQEShuffleReadExec]) + operators(plan).find(op => + findFirstNonCometOperator(op, plumbing: _*).exists(_ eq op)) match { + case None => "fully Comet native" + case Some(op) => s"falls back at ${op.nodeName}" + } + } + + private def broadcastProfile( + plan: SparkPlan, + consumerTasks: Int): Either[String, BroadcastProfile] = { + val all = operators(plan) + all.collect { case b: CometBroadcastExchangeExec => b } match { + case Seq(b) => + val buffers = b.executeBroadcast[Array[ChunkedByteBuffer]]().value.length + Right(BroadcastProfile(buffers, b.metrics("numCoalescedBatches").value, consumerTasks)) + case Seq() => + val sparkBroadcast = all.exists(_.isInstanceOf[BroadcastExchangeExec]) + Left(if (sparkBroadcast) "Spark broadcast" else "no broadcast") + case many => Left(s"${many.size} Comet broadcasts") + } + } + + private def emitProfiles(profiles: Seq[Profile]): Unit = { + def mb(bytes: Long): String = f"${bytes / (1024.0 * 1024.0)}%.1f" + val header = Seq( + "arm", + "plan", + "wall ms", + "JVM alloc MB", + "heap peak MB", + "heap retained MB", + "Spark peak exec MB", + "native mem MB", + "Arrow retained MB", + "bcast buffers", + "coalesced", + "consumer tasks", + "IPC streams") + val rows = profiles.map { p => + val (buffers, coalesced, consumers, ipc) = p.broadcast match { + case Right(b) => + ( + b.buffers.toString, + b.coalescedBatches.toString, + b.consumerTasks.toString, + b.ipcStreams.toString) + case Left(reason) => (reason, "-", "-", "-") + } + Seq( + p.arm, + p.plan, + p.wallMs.toString, + mb(p.jvmAllocatedBytes), + mb(p.heapPeakBytes), + mb(p.heapRetainedBytes), + mb(p.sparkPeakExecutionMemoryBytes), + mb(p.nativeMemoryBytes), + mb(p.arrowRetainedBytes), + buffers, + coalesced, + consumers, + ipc) + } + val widths = (header +: rows).transpose.map(_.map(_.length).max) + def line(cells: Seq[String]): String = + cells + .zip(widths) + .zipWithIndex + .map { case ((cell, w), i) => + if (i < 2) cell.padTo(w, ' ') else cell.reverse.padTo(w, ' ').reverse + } + .mkString(" ") + emit( + "Per-arm profile of one execution after warmup (JVM alloc: bytes allocated by all live " + + "threads; Spark peak exec: max peakExecutionMemory over tasks, reported by Spark " + + "operators only; native mem: the native join build / aggregate peak memory metrics " + + "summed over operators and tasks; IPC streams = broadcast buffers x consumer tasks):") + emit(line(header)) + emit("-" * (widths.sum + 2 * (widths.size - 1))) + rows.foreach(row => emit(line(row))) + emit("") + } + + /** Writes a line to the console and to the results file, like the timing tables. */ + private def emit(text: String): Unit = { + // scalastyle:off println + println(text) + // scalastyle:on println + output.foreach(_.write((text + "\n").getBytes(StandardCharsets.UTF_8))) + } +} + +/** + * Row digest used from executor closures; kept outside the benchmark object so it captures + * nothing. + */ +private[benchmark] object NullTypeRowDigest { + def digestPartition(rows: Iterator[Row]): Iterator[(Long, Long, Long)] = { + var count = 0L + var sum = 0L + var xor = 0L + rows.foreach { row => + val h = row.mkString("\u0001").hashCode.toLong + count += 1 + sum += h + xor ^= h + } + Iterator((count, sum, xor)) + } +} From ba976f0e4bcb978ac9d74267e90497e769118ac4 Mon Sep 17 00:00:00 2001 From: grorge Date: Thu, 10 Sep 2026 21:00:40 +0800 Subject: [PATCH 6/6] fix: keep the NullType-gated shapes in the Comet pipeline and trim the composition sweep The NullType gates on CometArrayRepeat, CometArrayUnion, CometCreateArray, CometIf, CometCaseWhen and CometCoalesce reported Unsupported on serdes that did not mix in CodegenDispatchFallback, so the whole projection fell back to Spark instead of running the expression through the JVM codegen dispatcher like array_intersect and array_except already did. CometBatchKernelCodegen.canHandle accepts every output type these gates cover, so the serdes mix the trait in and the plan stays fully native. Enrolling coalesce exposed a dispatcher bug: a top-level struct output with a NullType field failed to allocate ("Unknown type: NULL"), because StructVector's constructor builds a NullableStructWriter over the field's children and that writer has no arm for a Null child; RenamedStructVector now adds the children after construction, the way a struct nested under a list already got them. The slice gate goes: after the rebase onto main, #5766's convert no longer serializes a return type, so the item nullability it guarded against cannot disagree and every shape it covered runs natively. The array_repeat gate stays (DataFusion's list repeat rebuilds the item as nullable and fails on a non-nullable input item such as map_entries' entry struct) but exempts a CreateArray child, since native make_array emits a nullable item whatever Spark's containsNull says, so array_repeat(array(c), 2) over a non-nullable c repeats natively. Every serde whose support level moved on this branch publishes its reason through getUnsupportedReasons, so the generated compatibility pages carry the restriction; the runtime note starts with the published text. The SQL file tests that asserted a fallback now assert the native plan and Spark's answer, and CometNullTypeCompositionSuite checks per gate that the reason is published, that the serde is enrolled, and that the query keeps its CometProjectExec. CometNullTypeCompositionSuite ran every consumer sweep under every kernel profile and every operator sweep under six physical profiles. Each sweep now runs only under the profiles that can change its outcome: the plain sweep keeps every kernel profile; the nullable sweep keeps the default and allow-incompatible profiles (a null guard is per row, so batching cannot change it, but the kernel under the guard can); the stateful sweep keeps the default and two-row batches (its divergence needs a batch boundary inside the producer); the non-deterministic sweep runs on the default profile (every guarding serde leaves its kernel). The columnar-to-row-only physical profile is folded into the sort-writer one, and Spark's reference rows are cached across the profiles that share a row count. 37 tests in 10 min 49 s become 28 in about 8 min on the same machine, with the shuffle and batching coverage of the operator and nesting sweeps intact. Assisted-by: Claude Code (claude-fable-5) --- .../user-guide/latest/scala_java_udfs.md | 1 + .../CometBatchKernelCodegenOutput.scala | 17 +- .../scala/org/apache/comet/serde/arrays.scala | 92 ++++---- .../org/apache/comet/serde/conditional.scala | 21 +- .../scala/org/apache/comet/serde/maps.scala | 2 + .../expressions/array/array_repeat.sql | 15 +- .../expressions/array/array_union.sql | 13 +- .../expressions/array/create_array.sql | 5 +- .../sql-tests/expressions/array/slice.sql | 6 +- .../expressions/conditional/case_when.sql | 7 +- .../expressions/conditional/coalesce.sql | 12 +- .../expressions/conditional/if_expr.sql | 9 +- .../comet/CometCodegenSourceSuite.scala | 24 +++ .../comet/CometNullTypeCompositionSuite.scala | 200 ++++++++++++++---- 14 files changed, 310 insertions(+), 114 deletions(-) diff --git a/docs/source/user-guide/latest/scala_java_udfs.md b/docs/source/user-guide/latest/scala_java_udfs.md index 3713ebf735d..3ab23481a32 100644 --- a/docs/source/user-guide/latest/scala_java_udfs.md +++ b/docs/source/user-guide/latest/scala_java_udfs.md @@ -47,6 +47,7 @@ This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabl - Hive `GenericUDF` and `SimpleUDF`. - `UserDefinedType` arguments and return types, and `NullType` arguments. 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`). A `NullType` _return_ type is supported: Comet writes an all-null Arrow vector for it. - 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. +- Struct types with duplicate field names (`named_struct('x', 1, 'x', 2)`) anywhere in an argument or return type. Arrow's `StructVector` addresses children by name, so the two fields would collapse into one; 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. diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index 68bdcd46ab5..36902ebee18 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -138,9 +138,22 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { override def getField: Field = exportField } + /** + * The superclass constructor builds a `NullableStructWriter` over `getField.getChildren`, and + * that writer has no arm for a Null child (`UnsupportedOperationException: Unknown type: NULL` + * for a `struct<..., x: null>` output). So the superclass is given a childless copy of the + * field, `getField` answers with it until construction completes (`constructed` is still the + * JVM default `false` while the superclass constructor runs, since this class's own fields are + * initialized afterwards), and the children come from `initializeChildrenFromFields` in + * [[allocateOutput]], the same way a struct nested under a list gets them. + */ private final class RenamedStructVector(exportField: Field, allocator: BufferAllocator) - extends StructVector(exportField, allocator, null) { - override def getField: Field = exportField + extends StructVector( + new Field(exportField.getName, exportField.getFieldType, null), + allocator, + null) { + private val constructed: Boolean = true + override def getField: Field = if (constructed) exportField else super.getField } /** diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 45a076caf6e..4cd81eef331 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -220,6 +220,8 @@ object CometArrayIntersect private val nullElementReason: String = "native array_intersect returns the other side's entries for a NullType-element array" + override def getUnsupportedReasons(): Seq[String] = Seq(nullElementReason) + override def getSupportLevel(expr: ArrayIntersect): SupportLevel = { // The native array_intersect dedups by raw bytes, which is wrong under non-default collations. // That is Incompatible rather than Unsupported because there is something real to opt into: a @@ -328,8 +330,14 @@ object CometArrayExcept private val incompatReason = "Null handling and ordering may differ from Spark" + private val elementTypeReason = + "native array_except supports only boolean, integral, floating-point, decimal, date, " + + "timestamp and string elements, or arrays of those" + override def getIncompatibleReasons(): Seq[String] = Seq(incompatReason) + override def getUnsupportedReasons(): Seq[String] = Seq(elementTypeReason) + override def getSupportLevel(expr: ArrayExcept): SupportLevel = { // Surface the native element-type restriction in EXPLAIN. Unsupported rather than // Incompatible for these types: the JVM codegen dispatcher evaluates them natively and does @@ -339,7 +347,7 @@ object CometArrayExcept // handling differences below remain a genuine opt-in. expr.children.map(_.dataType).find(dt => !isTypeSupported(dt)) match { case Some(dt) => - Unsupported(Some(s"native array_except does not support element type $dt")) + Unsupported(Some(s"$elementTypeReason ($dt)")) case None => Incompatible(Some(incompatReason)) } } @@ -502,25 +510,6 @@ object CometArrayInsert extends CometExpressionSerde[ArrayInsert] with ArraysBas } object CometSlice extends CometExpressionSerde[Slice] { - - override def getSupportLevel(expr: Slice): SupportLevel = { - expr.x.dataType match { - // Native spark_array_slice rebuilds the sliced list around the input's actual child, - // whose non-NullType item keeps Spark's containsNull, while `convert` promises a - // nullable item; for containsNull = false the two disagree and the native plan is - // rejected (e.g. slice(map_entries(map(k, NULL)), 1, 1)). A bare NullType item is - // declared nullable at the FFI boundary (Utils.declaredChildNullability) and slices fine. - case ArrayType(elementType, false) - if elementType != NullType && - SupportLevel.containsType(elementType, classOf[NullType]) => - Unsupported( - Some( - "native spark_array_slice keeps a non-nullable list item where a nullable " + - "one is promised")) - case _ => Compatible() - } - } - override def convert( expr: Slice, inputs: Seq[Attribute], @@ -577,11 +566,18 @@ private[serde] object NullElementSetOp { }) } -object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { +object CometArrayUnion extends CometExpressionSerde[ArrayUnion] with CodegenDispatchFallback { + + private val nullElementReason = + "native array_union drops the entries of a NullType-element array" + override def getUnsupportedReasons(): Seq[String] = Seq(nullElementReason) + + // A NullType-element side runs through the JVM codegen dispatcher (`CodegenDispatchFallback`) + // rather than the kernel that drops its entries; see `NullElementSetOp`. override def getSupportLevel(expr: ArrayUnion): SupportLevel = { if (NullElementSetOp.hasNullElementSide(expr)) { - Unsupported(Some("native array_union drops the entries of a NullType-element array")) + Unsupported(Some(nullElementReason)) } else { Compatible() } @@ -601,7 +597,14 @@ object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { } } -object CometCreateArray extends CometExpressionSerde[CreateArray] with ArraysBase { +object CometCreateArray + extends CometExpressionSerde[CreateArray] + with ArraysBase + with CodegenDispatchFallback { + + private val nullTypeBatchReason = "native make_array builds a single row from a NullType batch" + + override def getUnsupportedReasons(): Seq[String] = Seq(nullTypeBatchReason) override def getSupportLevel(expr: CreateArray): SupportLevel = { // DataFusion's make_array funnels an argument list that is entirely Null-typed into @@ -611,9 +614,10 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] with ArraysBas // `aggregate(arr, NULL, (acc, x) -> NULL)` admitted by the JVM codegen dispatcher) therefore // fails the scalar-function row-count check for batches with more than one row. All-literal // NULL arguments arrive as scalars and broadcast correctly, and empty `array()` never reaches - // make_array (`convert` emits a literal). + // make_array (`convert` emits a literal). `CodegenDispatchFallback` keeps the non-scalar + // case in the Comet pipeline through the JVM codegen dispatcher. if (expr.children.exists(c => c.dataType == NullType && !c.foldable)) { - Unsupported(Some("native make_array builds a single row from a NullType batch")) + Unsupported(Some(nullTypeBatchReason)) } else { Compatible() } @@ -658,18 +662,33 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] with ArraysBas } } -object CometArrayRepeat extends CometExpressionSerde[ArrayRepeat] { +object CometArrayRepeat extends CometExpressionSerde[ArrayRepeat] with CodegenDispatchFallback { + + private val nonNullableItemReason = + "native array_repeat rebuilds a non-nullable list item as nullable" + + override def getUnsupportedReasons(): Seq[String] = Seq(nonNullableItemReason) override def getSupportLevel(expr: ArrayRepeat): SupportLevel = { - expr.left.dataType match { - // DataFusion's list repeat rebuilds the repeated list's item field as nullable. Comet - // declares a non-NullType item with Spark's containsNull, so for containsNull = false - // the planned and produced types disagree and the native plan is rejected. A NullType - // item is declared nullable on the FFI boundary (Utils.declaredChildNullability) and - // matches the rebuild, so plain array stays native. - case ArrayType(elementType, false) if elementType != NullType => - Unsupported(Some("native array_repeat rebuilds a non-nullable list item as nullable")) - case _ => Compatible() + expr.left match { + // Native `make_array` always emits a nullable item, whatever Spark's containsNull says, so + // `array(c)` over a non-nullable `c` repeats natively. + case _: CreateArray => Compatible() + case left => + left.dataType match { + // DataFusion's list repeat rebuilds the repeated list's item field as nullable, and + // fails when the input's item is not ("ListArray expected data type List(non-null + // Struct(..)) got List(Struct(..))"). A non-NullType item reaches native with Spark's + // containsNull (`map_entries` makes a non-nullable entry struct; the JVM codegen + // dispatcher declares its list items with Spark's flag), so containsNull = false is + // the plan-time proxy for that input. `CodegenDispatchFallback` runs those through + // the dispatcher instead. A NullType item is declared nullable on the FFI boundary + // (Utils.declaredChildNullability) and matches the rebuild, so plain array + // stays native. + case ArrayType(elementType, false) if elementType != NullType => + Unsupported(Some(nonNullableItemReason)) + case _ => Compatible() + } } } @@ -990,7 +1009,8 @@ object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with Array object CometArraysZip extends CometExpressionSerde[ArraysZip] { override def getUnsupportedReasons(): Seq[String] = Seq( - "Not all input data types are supported; falls back to Spark for unsupported types") + "Not all input data types are supported; falls back to Spark for unsupported types", + NullGuard.reason) private def isTypeSupported(dt: DataType): Boolean = { import DataTypes._ diff --git a/spark/src/main/scala/org/apache/comet/serde/conditional.scala b/spark/src/main/scala/org/apache/comet/serde/conditional.scala index bdafbe273ac..7878e560757 100644 --- a/spark/src/main/scala/org/apache/comet/serde/conditional.scala +++ b/spark/src/main/scala/org/apache/comet/serde/conditional.scala @@ -31,18 +31,24 @@ import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal * result through a `MutableArrayData` that carries a validity bitmap; a `NullArray` cannot hold * one ("Arrays of type Null cannot contain a null bitmask"), so a CASE whose result type is * `NullType` fails whenever more than one branch contributes rows. Spark normally folds such an - * expression away (`IF(c, NULL, NULL)`), but a `NullType`-typed non-foldable branch keeps it. + * expression away (`IF(c, NULL, NULL)`), but a `NullType`-typed non-foldable branch keeps it. The + * three serdes below mix in `CodegenDispatchFallback`, so that shape runs through the JVM codegen + * dispatcher and the projection stays in the Comet pipeline. */ private[serde] object NullTypeBranches { + val reason = "native CASE cannot merge NullType branches" + def supportLevel(expr: Expression): SupportLevel = if (expr.dataType == NullType) { - Unsupported(Some("native CASE cannot merge NullType branches")) + Unsupported(Some(reason)) } else { Compatible() } } -object CometIf extends CometExpressionSerde[If] { +object CometIf extends CometExpressionSerde[If] with CodegenDispatchFallback { + + override def getUnsupportedReasons(): Seq[String] = Seq(NullTypeBranches.reason) override def getSupportLevel(expr: If): SupportLevel = NullTypeBranches.supportLevel(expr) @@ -69,7 +75,9 @@ object CometIf extends CometExpressionSerde[If] { } } -object CometCaseWhen extends CometExpressionSerde[CaseWhen] { +object CometCaseWhen extends CometExpressionSerde[CaseWhen] with CodegenDispatchFallback { + + override def getUnsupportedReasons(): Seq[String] = Seq(NullTypeBranches.reason) override def getSupportLevel(expr: CaseWhen): SupportLevel = NullTypeBranches.supportLevel(expr) @@ -112,7 +120,10 @@ object CometCaseWhen extends CometExpressionSerde[CaseWhen] { } } -object CometCoalesce extends CometExpressionSerde[Coalesce] { +object CometCoalesce extends CometExpressionSerde[Coalesce] with CodegenDispatchFallback { + + override def getUnsupportedReasons(): Seq[String] = + Seq(NullTypeBranches.reason, NullGuard.reason) // Every child but the last is a guard; the last one is the ELSE, evaluated on the rows the // guards left over. The result is a native CASE, so it shares that serde's NullType rule. diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala index 40fb9b82586..4d089de8850 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -162,6 +162,8 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { private val scalarSideReason: String = "native map takes the first row of a scalar list where the other argument is per-row" + override def getUnsupportedReasons(): Seq[String] = Seq(scalarSideReason, NullGuard.reason) + override def getSupportLevel(expr: MapFromArrays): SupportLevel = { if (MapKeyDedupPolicySupport.isLastWin) { Incompatible(Some(MapKeyDedupPolicySupport.incompatibleReason)) diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql b/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql index 10425d5cb56..7944a471990 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_repeat.sql @@ -125,11 +125,20 @@ query SELECT array_repeat(filter(array(), x -> true), 2) FROM test_array_repeat -- map_entries produces a list whose non-NullType item is declared non-nullable; native --- array_repeat rebuilds that item as nullable, so the composition stays in Spark. -query expect_fallback(native array_repeat rebuilds a non-nullable list item as nullable) +-- array_repeat rebuilds that item as nullable, so the serde declines it and the JVM codegen +-- dispatcher runs it inside the Comet pipeline (the plan stays fully native). +query SELECT array_repeat(map_entries(map(coalesce(long_v, 0), NULL)), 2) FROM test_array_repeat -- Same mismatch with no NullType anywhere, so the guard is not a NullType-specific workaround: -- any list whose item Spark declares non-nullable hits it. This one fails on main as well. -query expect_fallback(native array_repeat rebuilds a non-nullable list item as nullable) +query SELECT array_repeat(map_entries(map(coalesce(long_v, 0), long_v)), 2) FROM test_array_repeat + +-- array(c) over a non-nullable c has containsNull = false too, but native make_array emits a +-- nullable item whatever Spark says, so the serde keeps this one on the native kernel. +query +SELECT array_repeat(array(coalesce(int_v, 0)), 2) FROM test_array_repeat + +query +SELECT array_repeat(array(coalesce(int_v, 0), coalesce(cnt, 0)), cnt) FROM test_array_repeat diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_union.sql b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql index 4c08823ad56..2d7815328b1 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_union.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql @@ -47,8 +47,9 @@ query SELECT a, b, array_union(a, b) FROM test_union_nulls -- empty array combinations --- Both sides are Null-typed empty arrays, which the NullType-element gate keeps in Spark. -query expect_fallback(native array_union drops the entries of a NullType-element array) +-- Both sides are Null-typed empty arrays; the NullType-element gate hands them to the JVM +-- codegen dispatcher, which keeps the projection in the Comet pipeline. +query SELECT array_union(array(), array()) FROM test_union_nulls query @@ -57,7 +58,7 @@ SELECT array_union(array(), array(1, 2)) FROM test_union_nulls query SELECT array_union(array(1, 2), array()) FROM test_union_nulls -query expect_fallback(native array_union drops the entries of a NullType-element array) +query SELECT array_union(array(), array(NULL)) FROM test_union_nulls -- both-NULL arrays @@ -249,9 +250,9 @@ query SELECT array_union(CASE WHEN a IS NOT NULL THEN a ELSE array(0) END, b) FROM test_array_union -- DataFusion's set-op kernel treats a Null element type as "return distinct(other side)" and --- drops the NULL entries the Null-typed list actually holds, so NullType-element unions stay --- in Spark. -query expect_fallback(native array_union drops the entries of a NullType-element array) +-- drops the NULL entries the Null-typed list actually holds, so NullType-element unions run +-- through the JVM codegen dispatcher instead of the kernel. +query SELECT array_union(transform(a, x -> NULL), array()) FROM test_array_union -- The set-op kernel asserts identical element types, nested nullability included, and the two diff --git a/spark/src/test/resources/sql-tests/expressions/array/create_array.sql b/spark/src/test/resources/sql-tests/expressions/array/create_array.sql index 40aad48d789..c9601adea3c 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/create_array.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/create_array.sql @@ -142,6 +142,7 @@ query SELECT CAST(array() AS ARRAY>) -- A non-foldable all-NullType argument (built by the JVM codegen dispatcher) would make native --- make_array collapse the whole batch into a single list row, so it stays in Spark. -query expect_fallback(native make_array builds a single row from a NullType batch) +-- make_array collapse the whole batch into a single list row, so the whole expression runs +-- through the dispatcher instead. +query SELECT array(aggregate(arr, NULL, (acc, x) -> NULL)) FROM test_create_array_complex diff --git a/spark/src/test/resources/sql-tests/expressions/array/slice.sql b/spark/src/test/resources/sql-tests/expressions/array/slice.sql index 73327e04ef6..c10ef5e97a6 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/slice.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/slice.sql @@ -284,7 +284,7 @@ query SELECT slice(filter(array(), x -> true), 1, 1) FROM test_slice -- map_entries produces a list whose NullType-bearing struct item is declared non-nullable; --- native spark_array_slice keeps that item where a nullable one is promised, so the --- composition stays in Spark. -query expect_fallback(native spark_array_slice keeps a non-nullable list item where a nullable one is promised) +-- native spark_array_slice keeps that item, and since the serde no longer serializes a return +-- type (the output field comes from the input's), the two agree and it runs natively. +query SELECT slice(map_entries(map(coalesce(start_idx, 0), NULL)), 1, 1) FROM test_slice diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql b/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql index 7e3b89a03d3..b9817746e40 100644 --- a/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql +++ b/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql @@ -34,7 +34,8 @@ SELECT CASE WHEN s IS NULL THEN 'null_val' ELSE s END FROM test_case_when query SELECT CASE WHEN i = 1 THEN s WHEN i = 2 THEN 'fixed' ELSE s END FROM test_case_when --- A NullType result stays in Spark: native CASE merges the rows of its branches through Arrow's --- merge_n, which cannot build a NullArray with a validity bitmap. -query expect_fallback(native CASE cannot merge NullType branches) +-- A NullType result cannot run natively (native CASE merges the rows of its branches through +-- Arrow's merge_n, which cannot build a NullArray with a validity bitmap), so the JVM codegen +-- dispatcher runs it inside the Comet pipeline. +query SELECT CASE WHEN i = 1 THEN aggregate(array(i), NULL, (acc, x) -> NULL) END FROM test_case_when diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql index 8ee23c8cb4c..003d69e12cd 100644 --- a/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql +++ b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce.sql @@ -36,11 +36,13 @@ SELECT coalesce(NULL, NULL, 99), coalesce(1, NULL, 99), coalesce(NULL) -- The serde guards every argument but the last with CASE WHEN arg IS NOT NULL THEN arg, and the -- two copies of a non-deterministic argument advance their state independently: the THEN copy --- can answer NULL for a row the guard selected, in a column declared non-nullable. -query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's) +-- can answer NULL for a row the guard selected, in a column declared non-nullable. The JVM +-- codegen dispatcher evaluates Spark's own code once per row, so it runs there instead. +query SELECT coalesce(IF(monotonically_increasing_id() % 2 = 0, a, NULL), b, 0) FROM test_coalesce --- A NullType result stays in Spark: the serde builds a native CASE, which merges the rows of its --- branches through Arrow's merge_n and cannot build a NullArray with a validity bitmap. -query expect_fallback(native CASE cannot merge NullType branches) +-- A NullType result cannot run natively (the serde builds a native CASE, which merges the rows +-- of its branches through Arrow's merge_n and cannot build a NullArray with a validity bitmap), +-- so the JVM codegen dispatcher runs it inside the Comet pipeline. +query SELECT coalesce(aggregate(array(a), NULL, (acc, x) -> NULL), aggregate(array(b), NULL, (acc, x) -> NULL)) FROM test_coalesce diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql b/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql index 51752b54a33..87d4748593b 100644 --- a/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql +++ b/spark/src/test/resources/sql-tests/expressions/conditional/if_expr.sql @@ -31,8 +31,9 @@ SELECT IF(a > 0, 'positive', 'non-positive') FROM test_if query SELECT IF(true, 1, 2), IF(false, 1, 2), IF(NULL, 1, 2) --- A NullType result stays in Spark: native CASE merges the rows of its branches through Arrow's --- merge_n, which cannot build a NullArray with a validity bitmap. Spark folds `IF(c, NULL, NULL)` --- itself, so the branch has to be a non-foldable NullType expression. -query expect_fallback(native CASE cannot merge NullType branches) +-- A NullType result cannot run natively (native CASE merges the rows of its branches through +-- Arrow's merge_n, which cannot build a NullArray with a validity bitmap), so the JVM codegen +-- dispatcher runs it inside the Comet pipeline. Spark folds `IF(c, NULL, NULL)` itself, so the +-- branch has to be a non-foldable NullType expression. +query SELECT IF(cond, aggregate(array(a), NULL, (acc, x) -> NULL), NULL) FROM test_if diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index 59197d08610..dce8e01d6bb 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala @@ -394,6 +394,30 @@ class CometCodegenSourceSuite extends AnyFunSuite { .getContainsNull) } + test("top-level struct output with a NullType field allocates") { + // `StructVector`'s constructor builds a `NullableStructWriter` over the field's children, + // which throws `Unknown type: NULL` for a Null child; the children must be added after + // construction. Reachable from the JVM codegen dispatcher through e.g. + // `coalesce(named_struct('i', monotonically_increasing_id(), 'n', NULL), ...)`. + val dataType = StructType( + Seq( + StructField("i", LongType, nullable = false), + StructField("n", NullType), + StructField("l", ArrayType(IntegerType)))) + val field = CometBatchKernelCodegen.toFfiArrowField("out", dataType, nullable = true) + val vector = CometBatchKernelCodegen.allocateOutput(field, 4, 0) + try { + val struct = vector.asInstanceOf[org.apache.arrow.vector.complex.StructVector] + assert(struct.getChild("n").isInstanceOf[org.apache.arrow.vector.NullVector]) + assert(struct.getChildrenFromFields.size == 3) + // The export names survive: the list child is labelled `item`, not Arrow Java's `$data$`. + assert(struct.getField == field) + assert(struct.getField.getChildren.get(2).getChildren.get(0).getName == "item") + } finally { + vector.close() + } + } + test("nested NullType output casts the child vector and writes setNull into it") { // A scalar NullType output cannot distinguish `emitWrite`'s NullType branch from // `defaultBody`'s own `ev.isNull -> output.setNull(i)` short-circuit, which emits the same diff --git a/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala b/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala index aa551e6ea2b..aaff7cee8d3 100644 --- a/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala @@ -19,11 +19,12 @@ package org.apache.comet +import scala.collection.mutable import scala.collection.mutable.ArrayBuffer import scala.util.{Failure, Success, Try} import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.catalyst.expressions.{Expression, JsonToStructs, RuntimeReplaceable, Sequence, StringToMap} +import org.apache.spark.sql.catalyst.expressions.{ArrayExcept, ArrayIntersect, ArrayRepeat, ArraysZip, ArrayUnion, CaseWhen, Coalesce, CreateArray, Expression, If, JsonToStructs, MapFromArrays, RuntimeReplaceable, Sequence, StringToMap} import org.apache.spark.sql.catalyst.expressions.aggregate._ import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.comet.CometProjectExec @@ -31,7 +32,7 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.NullType -import org.apache.comet.serde.{QueryPlanSerde, SupportLevel} +import org.apache.comet.serde.{CodegenDispatchFallback, CometExpressionSerde, Compatible, QueryPlanSerde, SupportLevel, Unsupported} /** * Cross-product sweep of the `NullType` shapes the JVM codegen dispatcher admits against the @@ -382,18 +383,17 @@ class CometNullTypeCompositionSuite extends CometTestBase with AdaptiveSparkPlan Seq(defaultProfile, smallBatchProfile, allowIncompatibleProfile) /** - * Profiles for the sweeps that put a value through an operator. Chosen so every pair of - * settings below appears together at least once: JVM shuffle through the bypass writer - * (partition count under `spark.shuffle.sort.bypassMergeThreshold`) and through the sort-based - * writer (above it, one whole partition per native call), with and without forced spills, - * native shuffle, AQE on and off, native columnar-to-row on and off, and batch sizes of two. + * Profiles for the sweeps that put a value through an operator. Each shuffle path appears once: + * JVM shuffle through the bypass writer (partition count under + * `spark.shuffle.sort.bypassMergeThreshold`), through the sort-based writer (above it, one + * whole partition per native call) with and without forced spills, and native shuffle. Across + * them AQE, native columnar-to-row and the two-row batch size each take both values, and every + * pair of those settings appears together at least once except native columnar-to-row off with + * AQE off, whose two mechanisms do not interact. Five profiles rather than one per setting, + * because each one runs the whole operator and nesting sweeps. */ private val physicalProfiles = Seq( defaultProfile, - Profile( - "default-native-c2r", - 16, - Seq(CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true")), smallBatchProfile.copy(confs = smallBatchProfile.confs ++ Seq( CometConf.COMET_SHUFFLE_MODE.key -> "jvm", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")), @@ -406,7 +406,7 @@ class CometNullTypeCompositionSuite extends CometTestBase with AdaptiveSparkPlan CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE.key -> "2", CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> "100000", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", - CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false")), + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true")), Profile( "jvm-sort-writer-spills", 900, @@ -465,44 +465,61 @@ class CometNullTypeCompositionSuite extends CometTestBase with AdaptiveSparkPlan // ANSI is a dimension of the consumer sweeps because it changes which serdes wrap their child // in a null guard and which kernels raise on out-of-range access; the operators below carry // no ANSI semantics of their own. - for (ansi <- Seq(false, true); profile <- kernelProfiles) { - val tag = s"(ansi=$ansi ${profile.name})" - test(s"NullType producers survive every consumer that Spark accepts $tag") { - sweep( - "consumer", - cases().map { case (producer, expr) => (producer, s"SELECT $expr FROM t") }, - comparedFloor = 120, - nativeFloor = 100, - ansi = ansi, - profile = profile) + // Each consumer sweep runs under the kernel profiles that can change its outcome and no + // others, since every (ansi, profile) pair is a full pass over the producer x consumer product. + for (ansi <- Seq(false, true)) { + def tag(profile: Profile): String = s"(ansi=$ansi ${profile.name})" + + // The plain sweep is the one that exercises every kernel, so it takes every kernel profile. + for (profile <- kernelProfiles) { + test(s"NullType producers survive every consumer that Spark accepts ${tag(profile)}") { + sweep( + "consumer", + cases().map { case (producer, expr) => (producer, s"SELECT $expr FROM t") }, + comparedFloor = 120, + nativeFloor = 100, + ansi = ansi, + profile = profile) + } } // The plain producers are non-nullable, so this is the only sweep where a serde's null guard // runs natively over a NullType-bearing column that is NULL on some rows and takes its ELSE - // branch; the non-deterministic sweep below makes the same serdes fall back instead. - test(s"nullable deterministic NullType producers survive every consumer $tag") { - sweep( - "nullable", - cases(nullableDeterministic).map { case (producer, expr) => - (producer, s"SELECT $expr FROM t") - }, - comparedFloor = 140, - nativeFloor = 110, - ansi = ansi, - profile = profile) + // branch; the non-deterministic sweep below makes the same serdes leave the kernel instead. + // The guard is per row, so batching changes nothing here; which kernel sits under the guard + // does, so the allow-incompatible profile stays. + for (profile <- Seq(defaultProfile, allowIncompatibleProfile)) { + test(s"nullable deterministic NullType producers survive every consumer ${tag(profile)}") { + sweep( + "nullable", + cases(nullableDeterministic).map { case (producer, expr) => + (producer, s"SELECT $expr FROM t") + }, + comparedFloor = 140, + nativeFloor = 110, + ansi = ansi, + profile = profile) + } } - test(s"stateful NullType producers survive guards filtered by a sibling $tag") { - sweep( - "cross-input", - crossInputCases.map { case (producer, expr) => (producer, s"SELECT $expr FROM t") }, - comparedFloor = 100, - nativeFloor = 60, - ansi = ansi, - profile = profile) + // A stateful producer only diverges where a batch boundary falls inside it, so this sweep + // takes the two-row batches as well as the default; the kernel choice never sees the state. + for (profile <- Seq(defaultProfile, smallBatchProfile)) { + test(s"stateful NullType producers survive guards filtered by a sibling ${tag(profile)}") { + sweep( + "cross-input", + crossInputCases.map { case (producer, expr) => (producer, s"SELECT $expr FROM t") }, + comparedFloor = 100, + nativeFloor = 60, + ansi = ansi, + profile = profile) + } } - test(s"nullable non-deterministic NullType producers survive every consumer $tag") { + // A non-deterministic child makes every guarding serde decline its native kernel, so neither + // batching nor the kernel choice reaches the producer; ANSI still decides which serdes guard. + test( + s"nullable non-deterministic NullType producers survive every consumer ${tag(defaultProfile)}") { sweep( "non-deterministic", cases(nullableNondeterministic).map { case (producer, expr) => @@ -511,7 +528,90 @@ class CometNullTypeCompositionSuite extends CometTestBase with AdaptiveSparkPlan comparedFloor = 140, nativeFloor = 110, ansi = ansi, - profile = profile) + profile = defaultProfile) + } + } + + /** + * The gates the producers above trip. Each must publish its reason for the compatibility guide + * (`getUnsupportedReasons` is what `GenerateDocs` reads; the runtime note can add detail after + * it), and the ones whose whole shape Spark's own `doGenCode` can evaluate must enroll in the + * JVM codegen dispatcher so the projection stays in the Comet pipeline instead of falling back. + * Every query here is a shape from this suite's sweeps or the SQL file tests. + */ + test("NullType gates enroll in the JVM codegen dispatcher and publish their reasons") { + val nullArray = "transform(array(id), x -> NULL)" + val nullScalar = "aggregate(array(id), NULL, (acc, x) -> NULL)" + val dispatched: Seq[(String, Class[_ <: Expression])] = Seq( + "array_repeat(map_entries(map(id, NULL)), 2)" -> classOf[ArrayRepeat], + // Beside a typed side Spark's coercion casts the NullType one away, so both sides are NullType. + s"array_union($nullArray, array())" -> classOf[ArrayUnion], + s"array($nullScalar)" -> classOf[CreateArray], + s"IF(id > 2, $nullScalar, NULL)" -> classOf[If], + s"CASE WHEN id > 2 THEN $nullScalar END" -> classOf[CaseWhen], + s"coalesce($nullScalar, $nullScalar)" -> classOf[Coalesce], + "coalesce(IF(monotonically_increasing_id() % 2 = 0, id, NULL), id)" -> classOf[Coalesce]) + // Declined for a reason the dispatcher shares (a NullType input it cannot read, a kernel + // that only runs when the user opts in) or on a serde outside this PR's enrollment; these + // only have to publish their reason. + val documentedOnly: Seq[(String, Class[_ <: Expression])] = Seq( + s"array_intersect($nullArray, array())" -> classOf[ArrayIntersect], + "array_except(array(named_struct('a', id)), array(named_struct('a', id)))" -> + classOf[ArrayExcept], + "arrays_zip(array(monotonically_increasing_id()), array(id))" -> classOf[ArraysZip], + "map_from_arrays(array(id), array(1))" -> classOf[MapFromArrays]) + // Shapes a gate must leave on the native kernel: the gate's plan-time proxy would match, but + // the native producer is known to be safe. + val nativeShapes: Seq[(String, Class[_ <: Expression])] = Seq( + // Native make_array emits a nullable item whatever Spark's containsNull says. + "array_repeat(array(id), 2)" -> classOf[ArrayRepeat]) + + withTempView("t") { + spark.range(0, 8).createOrReplaceTempView("t") + def serdeAndNode( + expr: String, + cls: Class[_ <: Expression]): (CometExpressionSerde[Expression], Expression) = { + val query = s"SELECT $expr AS c FROM t" + val node = spark + .sql(query) + .queryExecution + .analyzed + .flatMap(_.expressions) + .flatMap(_.collect { case e if e.getClass == cls => e }) + .headOption + .getOrElse(fail(s"$query has no ${cls.getSimpleName} in its analyzed plan")) + (QueryPlanSerde.exprSerdeMap(cls).asInstanceOf[CometExpressionSerde[Expression]], node) + } + for ((expr, cls) <- dispatched ++ documentedOnly) { + val (serde, node) = serdeAndNode(expr, cls) + val reason = serde.getSupportLevel(node) match { + case Unsupported(Some(reason)) => reason + case other => + fail(s"${cls.getSimpleName} reports $other for $expr, expected Unsupported") + } + assert( + serde.getUnsupportedReasons().exists(reason.startsWith), + s"${cls.getSimpleName}'s reason for $expr is not in getUnsupportedReasons: $reason") + } + for ((expr, cls) <- nativeShapes) { + val (serde, node) = serdeAndNode(expr, cls) + assert( + serde.getSupportLevel(node).isInstanceOf[Compatible], + s"${cls.getSimpleName} must keep $expr on the native kernel") + } + for ((expr, cls) <- dispatched ++ nativeShapes) { + val query = s"SELECT $expr AS c FROM t" + assert( + QueryPlanSerde.exprSerdeMap(cls).isInstanceOf[CodegenDispatchFallback], + s"${cls.getSimpleName} must mix in CodegenDispatchFallback") + val (sparkRows, _) = rowsOf(query, cometEnabled = false, ansi = false, defaultProfile) + val (cometRows, nativeProject) = + rowsOf(query, cometEnabled = true, ansi = false, defaultProfile) + assert( + cometRows == sparkRows, + s"$query: comet ${preview(cometRows)}, spark ${preview(sparkRows)}") + assert(nativeProject, s"$query left the Comet pipeline instead of dispatching") + } } } @@ -645,6 +745,13 @@ class CometNullTypeCompositionSuite extends CometTestBase with AdaptiveSparkPlan * `nativeFloor` fails one where Comet fell back almost everywhere, since falling back is a pass * and such a sweep would prove nothing about the native kernels. */ + /** + * Spark's answer to a query depends on the ANSI flag and the row count but on no Comet setting, + * so profiles that share both reuse it rather than run the reference side once per profile. + * `None` records a query Spark rejects. + */ + private val sparkRowsCache = mutable.HashMap.empty[(String, Boolean, Int), Option[Seq[String]]] + private def sweep( name: String, queries: Seq[(String, String)], @@ -665,10 +772,13 @@ class CometNullTypeCompositionSuite extends CometTestBase with AdaptiveSparkPlan var skipped = 0 for ((label, query) <- queries) { - Try(rowsOf(query, cometEnabled = false, ansi, profile)).toOption match { + val sparkRows = sparkRowsCache.getOrElseUpdate( + (query, ansi, profile.rows), + Try(rowsOf(query, cometEnabled = false, ansi, profile)).toOption.map(_._1)) + sparkRows match { case None => skipped += 1 - case Some((sparkRows, _)) => + case Some(sparkRows) => compared += 1 Try(rowsOf(query, cometEnabled = true, ansi, profile)) match { case Failure(e) =>