diff --git a/dev/scalastyle-config.xml b/dev/scalastyle-config.xml index 9de6df51ef7..edd3de30730 100644 --- a/dev/scalastyle-config.xml +++ b/dev/scalastyle-config.xml @@ -193,6 +193,19 @@ This file is divided into 3 sections: ]]> + + new ArrowStreamWriter + + + (\.toUpperCase|\.toLowerCase)(?!(\(|\(Locale.ROOT\))) `MAX_ROUNDED_ARRAY_LENGTH` throws at runtime. Wired as `CometScalarFunction("array_repeat")` against `datafusion-spark`'s `SparkArrayRepeat`, which returns NULL for NULL count and repeats NULL elements (matching Spark). +- Spark 3.5.8 (audited 2026-05-27): baseline. `ArrayRepeat(left, right) extends BinaryExpression with ExpectsInputTypes`; `inputTypes = Seq(AnyDataType, IntegerType)`. NULL count yields NULL; count <= 0 yields empty array; count > `MAX_ROUNDED_ARRAY_LENGTH` throws at runtime. Wired through `CometArrayRepeat` to `datafusion-spark`'s `SparkArrayRepeat` (an array argument whose item Spark declares non-nullable, other than a `CreateArray` or one with NullType elements, routes through the JVM codegen dispatcher, since the kernel rebuilds the item as nullable), which returns NULL for NULL count and repeats NULL elements (matching Spark). Spark's `eval` evaluates the count first and skips the element when it is NULL, while its generated code evaluates both, as the native kernel does; where Spark would use `eval` (`QueryPlanSerde.sparkEvaluatesInterpreted`), a nullable count runs through the codegen dispatcher, which then evaluates the tree through `eval` too. - Spark 4.0.1 (audited 2026-05-27): error message uses `createArrayWithElementsExceedLimitError(prettyName, count)`; semantics unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. ## array_union - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `ArrayUnion(left, right) extends ArrayBinaryLike with ComplexTypeMergingExpression`; result is left-side distinct elements followed by new right-side elements. Wired as `CometScalarFunction("array_union")`. +- Spark 3.5.8 (audited 2026-05-27): baseline. `ArrayUnion(left, right) extends ArrayBinaryLike with ComplexTypeMergingExpression`; result is left-side distinct elements followed by new right-side elements. Wired through `CometArrayUnion` to the native `array_union`, with both sides cast to a deeply-nullable element type; a `NullType`-element side routes through the JVM codegen dispatcher. - Spark 4.0.1 (audited 2026-05-27): `nullIntolerant = true` moves into `ArrayBinaryLike`; overflow path uses `arrayFunctionWithElementsExceedLimitError`. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. - Float/double element types follow the same Spark 4.2.0-only gate as `array_distinct`. Result element ordering matches Spark (left-side distinct elements followed by new right-side elements). @@ -153,7 +153,7 @@ ## arrays_zip - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `ArraysZip(children, names)`; returns an array of structs, padding shorter inputs with NULL. Comet routes via `CometArraysZip` and rejects unsupported child element types (anything outside primitives, decimals, dates/timestamps, strings, binary, and nested arrays/structs of those). +- Spark 3.5.8 (audited 2026-05-27): baseline. `ArraysZip(children, names)`; returns an array of structs, padding shorter inputs with NULL. Comet routes via `CometArraysZip` and rejects unsupported child element types (anything outside primitives, decimals, dates/timestamps, strings, binary, and nested arrays/structs of those). The inputs are wrapped in `CASE WHEN a IS NULL THEN NULL WHEN b IS NULL THEN NULL ... ELSE arrays_zip(...) END`: Spark's generated code stops evaluating the arguments at the first NULL one, and one WHEN per argument keeps native from evaluating a later argument on those rows (a single `AND` of the null checks can, which raised a division by zero under ANSI). Spark's interpreted `eval` evaluates every argument, so where Spark would use it (`spark.sql.codegen.factoryMode=NO_CODEGEN`, beneath a non-leaf `CodegenFallback` expression such as a higher-order function, as the input of an imperative aggregate, or under a generator outside a whole-stage stage, all of which `QueryPlanSerde.sparkEvaluatesInterpreted` tracks) `arrays_zip` over several arguments falls back to Spark, and the codegen dispatcher evaluates a tree it runs through `eval`. `arrays_zip()` with no arguments falls back to Spark. - Spark 4.0.1 (audited 2026-05-27): the length-mismatch error switches from `IllegalArgumentException` to `SparkIllegalArgumentException("_LEGACY_ERROR_TEMP_3235")`; runtime unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. diff --git a/docs/source/contributor-guide/expression-audits/collection_funcs.md b/docs/source/contributor-guide/expression-audits/collection_funcs.md index bcdebc165fc..51e4b84cd01 100644 --- a/docs/source/contributor-guide/expression-audits/collection_funcs.md +++ b/docs/source/contributor-guide/expression-audits/collection_funcs.md @@ -44,10 +44,10 @@ ## size - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Size(child, legacySizeOfNull) extends UnaryExpression with ExpectsInputTypes`; `inputTypes = Seq(TypeCollection(ArrayType, MapType)) -> IntegerType`. `legacySizeOfNull=true` returns `-1` for NULL input; `false` returns NULL. Comet routes via `CometSize`, which emits a `CaseWhen(isNotNull(child), size_scalar(child), Literal(legacySizeOfNull))`. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Size(child, legacySizeOfNull) extends UnaryExpression with ExpectsInputTypes`; `inputTypes = Seq(TypeCollection(ArrayType, MapType)) -> IntegerType`. `legacySizeOfNull=true` returns `-1` for NULL input; `false` returns NULL. Comet routes via `CometSize`, which emits `size_scalar(child)` directly when `legacySizeOfNull` is true or the child is non-nullable (native `size` already answers `-1` for NULL), and otherwise `CaseWhen(isNotNull(child), size_scalar(child), NULL)`; a nullable non-deterministic child falls back to Spark, since the guard would evaluate it twice. - Spark 4.0.1 (audited 2026-05-27): byte-for-byte identical to 3.5.8. - Spark 4.1.1 (audited 2026-05-27): byte-for-byte identical to 3.5.8. -- Both `ArrayType` and `MapType` inputs are `Compatible` and run natively; every other child type is `Unsupported`. +- `ArrayType` and `MapType` inputs are `Compatible` and run natively, except for the nullable non-deterministic child above; every other child type is `Unsupported`. - Performance (tuned 2026-07-10, PR [#4877](https://github.com/apache/datafusion-comet/pull/4877)): compute list row sizes from the offset buffer instead of allocating a sliced `ArrayRef` per row via `list_array.value(i)`, removing one heap allocation per row. ~94% faster. Benchmark: `benches/array_size.rs`. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/contributor-guide/expression-audits/conditional_funcs.md b/docs/source/contributor-guide/expression-audits/conditional_funcs.md index 29caf1493bc..0744a89192d 100644 --- a/docs/source/contributor-guide/expression-audits/conditional_funcs.md +++ b/docs/source/contributor-guide/expression-audits/conditional_funcs.md @@ -24,7 +24,7 @@ ## coalesce - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Coalesce(children) extends Expression with ComplexTypeMergingExpression`; returns the first non-null child, evaluated left-to-right with short-circuit. Result type is the merged child type. Comet routes via `CometCoalesce`, which serialises as nested `CaseWhen(IsNotNull(c1) -> c1, IsNotNull(c2) -> c2, ..., else cN)` so the native engine preserves the short-circuit semantics. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Coalesce(children) extends Expression with ComplexTypeMergingExpression`; returns the first non-null child, evaluated left-to-right with short-circuit. Result type is the merged child type. Comet routes via `CometCoalesce`, which serialises as nested `CaseWhen(IsNotNull(c1) -> c1, IsNotNull(c2) -> c2, ..., else cN)` so the native engine preserves the short-circuit semantics. A non-deterministic argument other than the last would be serialized twice, so `CometCoalesce` sends such a coalesce to the JVM codegen dispatcher, which runs it unless it calls a non-deterministic user-defined function. - Spark 4.0.1 (audited 2026-05-27): byte-for-byte identical to 3.5.8. - Spark 4.1.1 (audited 2026-05-27): byte-for-byte identical to 3.5.8. - Performance (tuned 2026-09-28, PR [#6350](https://github.com/apache/datafusion-comet/pull/6350)): runs through the `CASE WHEN` evaluation described under `when`. 12-13x faster for two columns. Benchmark: `benches/conditional.rs`. diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index 11014f34909..d25d9348f0b 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -54,6 +54,11 @@ JVM shuffle (`CometColumnarExchange`) is used instead of native shuffle (`CometE [Supported partition key types](native_shuffle.md#when-native-shuffle-is-used) for the exact rules. Complex types are fully supported as data columns in both implementations. +4. **Round-robin over a hashed `NullType` column**: native round-robin partitioning hashes the + row's columns (the first `maxHashColumns`, or all of them for `0`, unless positional placement + applies), and the native hasher cannot hash a `NullType` column or field, so such a shuffle + uses JVM columnar shuffle when the shuffle mode and shuffle manager allow it. + ## Input Handling ### Spark Row-Based Input diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index f852dd497a1..d73b08c4fe0 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -54,7 +54,9 @@ Native shuffle (`CometExchange`) is selected when all of the following condition - `SinglePartition` - `RoundRobinPartitioning`, disabled by default via `spark.comet.shuffle.native.partitioning.roundrobin.enabled`, because Comet's hash-based - assignment puts unsorted rows in different partitions than Spark does + assignment puts unsorted rows in different partitions than Spark does. It is also refused + when a column it hashes (the first `maxHashColumns`, or all of them for `0`, unless + positional placement applies) contains a `NullType`, which the native hasher cannot hash 4. **Supported partition key types**: The rule differs by partitioning, and neither restricts data columns. Complex types are fully supported as data columns in native shuffle. diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index 9bae4312078..c881f66860e 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -52,7 +52,10 @@ Some Spark expressions have two implementations in Comet: expression inside Comet's native pipeline (via the Arrow-direct codegen dispatcher). This produces byte-exact Spark results at the cost of one JNI round-trip per batch. It is gated globally by `spark.comet.exec.scalaUDF.codegen.enabled` (enabled by default); when the - dispatcher is disabled, these expressions fall back to Spark. + dispatcher is disabled, these expressions fall back to Spark. The dispatcher also cannot read + an input column whose type contains `NullType` (for example an `array` column produced + below an exchange), so an expression that needs the dispatcher over such a column falls back to + Spark as well. - A **native** (Rust / DataFusion) implementation that avoids the JNI round-trip but has known semantic differences from Spark for some inputs or patterns. diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 657a32de532..64c146285d9 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -44,6 +44,26 @@ than Comet operators can be slower than Spark's cache. With Kryo and `spark.kryo.registrationRequired=true`, Comet needs its Kryo registrator whether or not the cache is enabled; see [Kryo serialization](../installation.md#kryo-serialization). +## Broadcast Exchange + +A broadcast build side that has a `NullType` column directly under a struct field or a map entry +(for example `map(k, NULL)` or `named_struct('a', k, 'b', NULL)`) stays on Spark's broadcast +exchange, and the join above it falls back to Spark. Comet coalesces the build side's batches into +one Arrow buffer before broadcasting, and Arrow Java's appender cannot grow a `NullVector` in that +position. A `NullType` element of a list (`array`) is not affected. + +## Non-Deterministic Expressions Beside NullType Values + +Spark's generated code skips an argument on rows where another argument decides the result (for +example a `NULL` operand of `+`, or a `NULL` divisor), while a native expression evaluates every +argument over the whole batch, and some expressions evaluate an argument twice. A stateful +non-deterministic expression such as `monotonically_increasing_id()` or `rand()` then advances on +rows Spark does not evaluate. An operator that evaluates a non-deterministic expression, and that +runs a JVM codegen dispatcher kernel computing a `NullType` value (as its result or inside it) +itself or in the native plan below it up to a shuffle (for example `map(k, NULL)`, including below a broadcast or a union), falls back to +Spark. A non-deterministic column that the dispatcher runs as a whole, such as +`map(monotonically_increasing_id(), NULL)`, is not affected. + ## Sampling Comet runs `SampleExec` natively when sampling is performed without replacement, which covers @@ -136,3 +156,9 @@ achieves the same semantic goals: The only difference is that Comet's partition assignments will differ from Spark's. When results are sorted, they will be identical to Spark. Unsorted results may have different row ordering. + +A hashed column whose type contains `NullType` (a `NULL` column, or a struct, array or map with a `NULL`-typed +part) cannot be hashed natively, so such a shuffle falls back to the JVM columnar shuffle when +`spark.comet.shuffle.mode` and the shuffle manager allow it, and to Spark's shuffle otherwise. This does not apply when rows are placed +by position (`spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled`) or when the column is past +the first `spark.comet.shuffle.native.partitioning.roundrobin.maxHashColumns` columns. diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index b5abfb88935..51b4cf50985 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -37,7 +37,10 @@ for how Comet chooses. Most expressions can also be disabled with `spark.comet.expression.EXPRNAME.enabled=false`, where `EXPRNAME` is the Spark expression class name (for example `Length` or `StartsWith`). See the -[Comet Configuration Guide](configs.md) for the full list. +[Comet Configuration Guide](configs.md) for the full list. Comet adds `Cast` expressions of its +own to give nested values one type (the arguments of `array_union`, `array_intersect` and +`array_except` over nested elements, and the branches of a `CASE WHEN` or `coalesce` whose type +contains a struct), so disabling `Cast` makes those fall back to Spark too. ## Status legend @@ -94,7 +97,7 @@ The tables below list every Spark built-in expression with its current status. | `bit_xor` | ✅ | Native | | | `bool_and` | ✅ | — | | | `bool_or` | ✅ | — | | -| `collect_list` | ✅ | Native | | +| `collect_list` | ✅ | Native | `NullType` input, including one nested in a complex type, falls back to Spark | | `collect_set` | ✅ | Native | | | `corr` | ✅ | Native | | | `count` | ✅ | Native | | @@ -152,24 +155,24 @@ The tables below list every Spark built-in expression with its current status. | Function | Status | Implementation | Notes | | --- | --- | --- | --- | -| `array` | ✅ | Native | | -| `array_append` | ✅ | Native | | +| `array` | ✅ | Hybrid | A non-literal NullType element routes through the JVM codegen dispatcher | +| `array_append` | ✅ | Native | On Spark 3.x, a non-deterministic array, or a non-deterministic item beside an array that can be NULL, falls back to Spark (Spark 4.0+ rewrites it to `array_insert`) | | `array_compact` | ✅ | — | | | `array_contains` | ✅ | Native | Float/double element arrays route through the JVM codegen dispatcher by default; the native path is opt-in via allowIncompatible | | `array_distinct` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)) | -| `array_except` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible ([details](compatibility/expressions/array.md)) | +| `array_except` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible, except for element types it does not support (for example structs), which never take the native path ([details](compatibility/expressions/array.md)) | | `array_insert` | ✅ | Native | | -| `array_intersect` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible ([details](compatibility/expressions/array.md)) | +| `array_intersect` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible, except for a `NullType`-element side, which never takes the native path ([details](compatibility/expressions/array.md)) | | `array_join` | ✅ | Hybrid | Native for literal or column delimiter and null replacement; other cases and non-UTF8_BINARY collations use the JVM codegen dispatcher ([details](compatibility/expressions/array.md)) | | `array_max` | ✅ | Hybrid | Native Spark-compatible floating-point and nested ordering; non-default string collations use the JVM codegen dispatcher ([details](compatibility/expressions/array.md)) | | `array_min` | ✅ | Hybrid | Native Spark-compatible floating-point and nested ordering; non-default string collations use the JVM codegen dispatcher ([details](compatibility/expressions/array.md)) | | `array_position` | ✅ | Native | Binary/struct/map/null elements fall back | | `array_prepend` | ✅ | — | | | `array_remove` | ✅ | Native | | -| `array_repeat` | ✅ | Native | | -| `array_union` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)) | +| `array_repeat` | ✅ | Hybrid | Repeating an array whose elements Spark declares non-nullable routes through the JVM codegen dispatcher, unless the array is built by `array(...)` or its elements are NULL-typed. So does a nullable count where Spark would evaluate the expression without generated code (see `arrays_zip`), since Spark's interpreted path skips the element when the count is NULL | +| `array_union` | ✅ | Hybrid | A NullType-element side routes through the JVM codegen dispatcher; floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)) | | `arrays_overlap` | ✅ | Native | | -| `arrays_zip` | ✅ | Native | | +| `arrays_zip` | ✅ | Native | A non-deterministic argument, or no argument at all, falls back to Spark. So do several arguments where Spark would evaluate them without generated code (`spark.sql.codegen.factoryMode=NO_CODEGEN`, on Spark 3.4 only with whole-stage codegen off, under an expression with no generated code such as a higher-order function, as the input of an imperative aggregate such as `collect_list`, or under an `explode` that Spark leaves out of whole-stage codegen), since Spark's interpreted path evaluates every argument and its generated code stops at the first NULL one | | `element_at` | ✅ | Native | | | `flatten` | ✅ | Native | Binary/struct/map elements fall back | | `get` | ✅ | — | | @@ -204,10 +207,10 @@ The tables below list every Spark built-in expression with its current status. | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `array_size` | ✅ | — | | -| `cardinality` | ✅ | Native | | +| `cardinality` | ✅ | Native | When NULL input returns NULL (`spark.sql.legacy.sizeOfNull=false` or ANSI mode), a nullable non-deterministic argument falls back to Spark | | `concat` | ✅ | Hybrid | Binary/array children and non-UTF8_BINARY collations route through the JVM codegen dispatcher | | `reverse` | ✅ | Hybrid | Binary input (Spark 4.2), arrays with binary, struct, or map elements, and collated strings, route through the JVM codegen dispatcher ([details](compatibility/expressions/array.md)) | -| `size` | ✅ | Native | | +| `size` | ✅ | Native | When NULL input returns NULL (`spark.sql.legacy.sizeOfNull=false` or ANSI mode), a nullable non-deterministic argument falls back to Spark | --- @@ -215,7 +218,7 @@ The tables below list every Spark built-in expression with its current status. | Function | Status | Implementation | Notes | | --- | --- | --- | --- | -| `coalesce` | ✅ | Native | | +| `coalesce` | ✅ | Hybrid | A non-deterministic argument other than the last routes through the JVM codegen dispatcher (or falls back when it calls a non-deterministic user-defined function) | | `if` | ✅ | Native | | | `ifnull` | ✅ | — | | | `nanvl` | ✅ | Codegen dispatch | | @@ -265,7 +268,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `date_from_unix_date` | ✅ | Native | | | `date_part` | ✅ | — | | | `date_sub` | ✅ | Native | | -| `date_trunc` | ✅ | Hybrid | | +| `date_trunc` | ✅ | Hybrid | A non-literal format runs through the JVM codegen dispatcher where Spark would evaluate the expression without generated code (see `arrays_zip`), since Spark's interpreted path returns NULL for an invalid format without evaluating the value | | `dateadd` | ✅ | Native | The 2-argument form is native; the `dateadd(UNIT, n, ts)` form parses to `timestampadd` and runs through codegen dispatch | | `datediff` | ✅ | Native | The 2-argument form is native; the `datediff(UNIT, start, end)` form parses to `timestampdiff` and runs through codegen dispatch | | `datepart` | ✅ | — | | @@ -312,7 +315,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `to_timestamp_ntz` | ✅ | — | Rewrites to `to_timestamp` (`TimestampNTZType`) | | `to_unix_timestamp` | ✅ | Hybrid | | | `to_utc_timestamp` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default (handles all timezone forms); the native path is opt-in via allowIncompatible ([details](compatibility/expressions/datetime.md)) | -| `trunc` | ✅ | Hybrid | | +| `trunc` | ✅ | Hybrid | A non-literal format runs through the JVM codegen dispatcher where Spark would evaluate the expression without generated code (see `arrays_zip`), since Spark's interpreted path returns NULL for an invalid format without evaluating the value | | `try_make_interval` | ✅ | — | Rewrites to `MakeInterval`; same support as `make_interval` (Spark 4.0+) | | `try_make_timestamp` | ✅ | — | | | `try_make_timestamp_ltz` | ✅ | — | Same support as `try_make_timestamp` (Spark 4.0+) | @@ -357,12 +360,12 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `crc32` | ✅ | Native | | -| `hash` | ✅ | Native | | +| `hash` | ✅ | Native | `NullType` input, including one nested in a complex type, falls back to Spark | | `md5` | ✅ | Native | | | `sha` | ✅ | Native | | | `sha1` | ✅ | Native | | | `sha2` | ✅ | Native | | -| `xxhash64` | ✅ | Native | | +| `xxhash64` | ✅ | Native | `NullType` input, including one nested in a complex type, falls back to Spark | --- @@ -408,7 +411,7 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | `map_concat` | ✅ | Codegen dispatch | | | `map_contains_key` | ✅ | — | | | `map_entries` | ✅ | Native | | -| `map_from_arrays` | ✅ | Native | | +| `map_from_arrays` | ✅ | Native | A non-deterministic argument falls back to Spark | | `map_from_entries` | ✅ | Hybrid | BinaryType keys/values and `spark.sql.mapKeyDedupPolicy=LAST_WIN` route through the JVM codegen dispatcher ([details](compatibility/expressions/map.md)) | | `map_keys` | ✅ | Native | | | `map_values` | ✅ | Native | | diff --git a/docs/source/user-guide/latest/scala_java_udfs.md b/docs/source/user-guide/latest/scala_java_udfs.md index 43a30374578..ec73b3f43b8 100644 --- a/docs/source/user-guide/latest/scala_java_udfs.md +++ b/docs/source/user-guide/latest/scala_java_udfs.md @@ -45,14 +45,17 @@ 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`. -- `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. +- UDFs marked non-deterministic (`asNondeterministic()`). Spark keeps such a function's state in one object that all of its calls share and advances it row by row, which Comet's per-batch evaluation cannot reproduce, so the operator falls back to Spark. +- A subquery inside the UDF's argument tree (`my_udf((SELECT max(x) FROM t))`). A subquery beside the UDF (`my_udf(x) + (SELECT max(x) FROM t)`) is not affected. +- Struct types with duplicate field names (`named_struct('x', 1, 'x', 2)`) in a column the UDF tree reads or in its return type, the values that cross into Arrow. 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. A struct the UDF's arguments build from other columns (`my_udf(named_struct('x', a, 'x', b))`) is evaluated inside the generated code and is not affected. When a UDF is rejected, the reason surfaces through Comet's standard fallback diagnostics; the query still runs on Spark. ## Behavior -- Non-deterministic expressions referenced from the argument tree (`rand`, `uuid`, `monotonically_increasing_id`) produce per-partition sequences consistent with Spark. +- Non-deterministic expressions referenced from the argument tree (`rand`, `uuid`, `monotonically_increasing_id`) produce per-partition sequences consistent with Spark, and two identical calls each keep their own sequence, as in Spark. - `TaskContext.get()` inside the user function returns the driving Spark task's context. - The Spark task thread's context ClassLoader is propagated to the thread that runs the user function, so functions defined in jars supplied with `--jars` / `spark.jars` resolve the same way they do under Spark's own execution. - The user function must be closure-serializable; the same function that works with Spark's executor execution works here. diff --git a/native/common/src/struct_nulls.rs b/native/common/src/struct_nulls.rs index f3390707702..ad03e8a7f14 100644 --- a/native/common/src/struct_nulls.rs +++ b/native/common/src/struct_nulls.rs @@ -28,8 +28,9 @@ //! whatever sits in the child slot //! (). -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,25 @@ 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/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index de458eca329..7094ef87695 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -6570,6 +6570,28 @@ mod tests { Ok(()) } + #[test] + fn test_literal_to_list_with_only_empty_children() -> Result<(), DataFusionError> { + // [[], null]: no child contributes values, so the values array is built empty. It must + // still have the child list's element type, as it does when some child is non-empty. + let data = ListLiteral { + list_values: vec![ListLiteral::default(), ListLiteral::default()], + null_mask: vec![true, false], + ..Default::default() + }; + let inner = DataType::List(Arc::new(Field::new("item", DataType::Int64, true))); + let nested_type = DataType::List(Arc::new(Field::new("item", inner.clone(), true))); + + let array = literal_to_array_ref(nested_type, data)?; + assert_eq!(array.data_type(), &inner); + let list = array.as_any().downcast_ref::().unwrap(); + assert_eq!(list.len(), 2); + assert!(list.value(0).is_empty()); + assert!(list.is_null(1)); + + Ok(()) + } + /// Test that reproduces the "Cast error: Casting from Int8 to Date32 not supported" error /// that occurs when performing date subtraction with Int8 (TINYINT) values. /// This corresponds to the Scala test "date_sub with int arrays" in CometExpressionSuite. diff --git a/native/shuffle/src/remote_schema_tests.rs b/native/shuffle/src/remote_schema_tests.rs index 9d445a0241c..13cd566c6ba 100644 --- a/native/shuffle/src/remote_schema_tests.rs +++ b/native/shuffle/src/remote_schema_tests.rs @@ -19,9 +19,9 @@ use crate::ShuffleCodecContext; 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,53 @@ 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] +#[cfg_attr(miri, ignore)] // The codec matrix calls ZSTD_createCCtx. +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 1918ce3b18c..6378f010962 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!( @@ -1432,6 +1440,18 @@ pub fn process_sorted_row_partition( output_data.write_all(&frozen)?; current_row += n; + drop(batch); + + // A finished `NullBuilder` keeps its length; see `recreate_null_type_builders`. Done once + // the batch is written and dropped, and only when another batch follows. + if current_row < row_num { + recreate_null_type_builders( + &mut data_builders, + schema, + batch_size, + prefer_dictionary_ratio, + )?; + } } Ok(( @@ -1441,6 +1461,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, @@ -1504,10 +1555,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/arrays_overlap.rs b/native/spark-expr/src/array_funcs/arrays_overlap.rs index 13af6636dab..7707a1b21d3 100644 --- a/native/spark-expr/src/array_funcs/arrays_overlap.rs +++ b/native/spark-expr/src/array_funcs/arrays_overlap.rs @@ -131,9 +131,15 @@ impl ScalarUDFImpl for SparkArraysOverlap { Ok(ColumnarValue::Array(result)) } (left, right) => { - // Handle scalar inputs by converting to arrays - let left_arr = left.to_array(1)?; - let right_arr = right.to_array(1)?; + // At least one side is a scalar. A scalar beside a column is expanded to the batch + // length so every row is compared; only two scalars produce a scalar result. + let both_scalar = matches!( + (left, right), + (ColumnarValue::Scalar(_), ColumnarValue::Scalar(_)) + ); + let rows = if both_scalar { 1 } else { args.number_rows }; + let left_arr = left.to_array(rows)?; + let right_arr = right.to_array(rows)?; let result = match (left_arr.data_type(), right_arr.data_type()) { (DataType::List(_), DataType::List(_)) => arrays_overlap_list::( left_arr.as_any().downcast_ref().unwrap(), @@ -149,8 +155,13 @@ impl ScalarUDFImpl for SparkArraysOverlap { ) } }; - let scalar = ScalarValue::try_from_array(&result, 0)?; - Ok(ColumnarValue::Scalar(scalar)) + if both_scalar { + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } else { + Ok(ColumnarValue::Array(result)) + } } } } 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..c6f9ee35dda 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,10 +144,52 @@ 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 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:?}") + } + } + } + } // 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 diff --git a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala index f36487080d3..b088af12e72 100644 --- a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala +++ b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala @@ -77,7 +77,13 @@ object GenerateDocs { /** Build the documentation notes for a single expression serde. */ private def exprNotes(cls: Class[_], serde: CometExpressionSerde[_]): ExprNotes = { - val optIn = serde.isInstanceOf[NativeOptInAvailable] + // A `CodegenDispatchFallback` serde that enrolls only its `Unsupported` cases leaves its + // `Incompatible` cases falling back to Spark, so they are documented as such. + val optIn = serde match { + case d: CodegenDispatchFallback => d.dispatchesIncompatible + case _: NativeOptInAvailable => true + case _ => false + } val key = serde match { case n: NativeOptInAvailable => n.nativeOptInConfigKeyOverride.getOrElse(CometConf.getExprAllowIncompatConfigKey(cls)) @@ -439,8 +445,10 @@ object GenerateDocs { } if (n.unsupportedReasons.nonEmpty) { val header = if (n.codegenDispatchFallback) { - "\nThe following cases have no native implementation and always run in the JVM using" + - " Spark's code-generated implementation (inside the Comet pipeline):\n\n" + "\nThe following cases have no native implementation. They run in the JVM using" + + " Spark's own implementation (inside the Comet pipeline) when the codegen" + + " dispatcher is enabled and accepts the expression, and fall back to Spark" + + " otherwise:\n\n" } else { "\nThe following cases are not supported by Comet and always fall back to Spark," + " regardless of any `allowIncompatible` setting:\n\n" 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 eadda24da70..39ef0830762 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -19,12 +19,16 @@ package org.apache.comet.codegen +import scala.util.control.NonFatal + import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.types.pojo.Field import org.apache.spark.internal.Logging -import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, Unevaluable} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{ApplyFunctionExpression, BoundReference, CallMethodViaReflection, Expression, Literal, Nondeterministic, PlanExpression, ScalaUDF, UnaryExpression, Unevaluable} import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -82,11 +86,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. * * The `StringType` case admits non-default collations (Spark 4+), on purpose. Collation is * carried by the expression, not by the value: the kernel runs Spark's own `doGenCode` against @@ -102,8 +103,15 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come * operator against the Catalyst `DataType`, which keeps its collation, and does not depend on * what this predicate admits. Rejecting collated strings here would force a full Spark fallback * for a route that is already correct. + * + * `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 @@ -111,16 +119,43 @@ 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 => st.fields.forall(f => isSupportedDataType(f.dataType, allowNullType)) + case mt: MapType => + isSupportedDataType(mt.keyType, allowNullType) && + isSupportedDataType(mt.valueType, allowNullType) + case _ => false + } + + /** + * 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 hasDuplicateStructFieldNames(dt: DataType): Boolean = dt match { 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) + 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`. @@ -138,17 +173,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 [[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)) { - return Some(s"codegen dispatch: unsupported 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 // operator falls back to Spark cleanly rather than tripping a Janino compile failure - // mid-execution (Comet has no recovery for that). + // mid-execution: the interpreted fallback in `compile` only replaces the expression's own + // code, and still generates the same typed Arrow accessors. val maxFields = SQLConf.get.wholeStageMaxNumFields val totalFields = numOfNestedFields(boundExpr.dataType) + boundExpr.collect { case b: BoundReference => numOfNestedFields(b.dataType) }.sum @@ -167,15 +204,14 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come // `CometScalaUDFCodegen.kernelCache` prevents concurrent partitions from racing on shared // state inside the expression. // - // Nondeterministic / stateful expressions are accepted: each cache entry holds one kernel + // Nondeterministic / stateful built-ins are accepted: each cache entry holds one kernel // instance with a single `init(partitionIndex)` call, so `Rand` / `MonotonicallyIncreasingID` - // state advances correctly across batches. + // state advances correctly across batches, and each occurrence gets its own entry + // (`DispatchOccurrence`). A non-deterministic user function is refused (`userStateReason`). // - // `ExecSubqueryExpression` (`ScalarSubquery`, `InSubqueryExec`) is accepted: the surrounding - // Comet operator's inherited `SparkPlan.waitForSubqueries` populates the subquery's - // `result` field before evaluation. The closure serializer captures that value into the - // arg-0 bytes, and the dispatcher keys its compile cache on those bytes, so distinct subquery - // results produce distinct cache entries. + // A subquery is refused (`subqueryReason`): the plan is serialized before the operator's + // `prepare` starts its subqueries, so the shipped copy never receives the result and its + // evaluation fails with "has not finished". // // `Unevaluable`: rejected by default. `isCodegenInertUnevaluable` exempts version-specific // leaves that are `Unevaluable` but never invoked by codegen (e.g. Spark 4.0's @@ -193,12 +229,30 @@ 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 + if (boundExpr.exists(_.isInstanceOf[PlanExpression[_]])) { + return Some(subqueryReason) + } + if (boundExpr.exists(DispatchOccurrence.holdsSharedState)) { + return Some(userStateReason) } - badRef.map(b => - s"codegen dispatch: unsupported input type ${b.dataType} at ordinal ${b.ordinal}") + boundExpr.collectFirst(Function.unlift(inputRejection)) + } + + val userStateReason: String = + "codegen dispatch: a non-deterministic user-defined function or reflected method keeps " + + "state that Spark shares between all of its calls and advances row by row, which the " + + "dispatcher, evaluating each expression over a whole batch, cannot reproduce" + + val subqueryReason: String = + "codegen dispatch: a subquery inside the expression would be evaluated before its result " + + "is available" + + /** Why `expr` cannot be read as a codegen input, if it cannot. */ + private def inputRejection(expr: Expression): Option[String] = expr match { + case b: BoundReference => + typeRejection(b.dataType, allowNullType = false, "input") + .map(reason => s"$reason at ordinal ${b.ordinal}") + case _ => None } /** @@ -215,7 +269,27 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come def toFfiArrowField(name: String, dataType: DataType, nullable: Boolean): Field = CometBatchKernelCodegenOutput.toFfiArrowField(name, dataType, nullable) - def compile(boundExpr: Expression, inputSchema: Seq[ArrowColumnSpec]): CompiledKernel = { + def compile(boundExpr: Expression, inputSchema: Seq[ArrowColumnSpec]): CompiledKernel = + try { + compileAs(boundExpr, inputSchema) + } catch { + // Spark's own `doGenCode` does not always compile: `ElementAt` on an in-bounds literal + // index is non-nullable, yet with ANSI off its generated code still assigns an undeclared + // `isNull` on the out-of-bounds branch. Spark's whole-stage codegen falls back to the + // interpreted path on a compile failure, and so does the kernel: the expression is wrapped + // in `InterpretedKernelExpr`, a `CodegenFallback`, so only the Arrow reads and writes stay + // generated and the expression itself runs through its interpreted `eval`. + case NonFatal(t) if !boundExpr.isInstanceOf[InterpretedKernelExpr] => + logWarning( + s"CometBatchKernelCodegen: generated code for ${boundExpr.getClass.getSimpleName} " + + "did not compile; evaluating it through its interpreted eval instead", + t) + compileAs(InterpretedKernelExpr(boundExpr), inputSchema) + } + + private def compileAs( + boundExpr: Expression, + inputSchema: Seq[ArrowColumnSpec]): CompiledKernel = { val src = generateSource(boundExpr, inputSchema) val (clazz, _) = try { @@ -622,3 +696,75 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come } } } + +/** + * Runs `child` through its interpreted `eval` inside a dispatcher kernel. Used when the code + * Spark generates for `child` does not compile (see `CometBatchKernelCodegen.compile`). + * `CodegenFallback.doGenCode` emits one `eval(row)` call against the kernel's typed Arrow getters + * and registers the partition initialization of any `Nondeterministic` node under it. + */ +private[comet] case class InterpretedKernelExpr(child: Expression) + extends UnaryExpression + with CodegenFallback { + override def dataType: DataType = child.dataType + override def nullable: Boolean = child.nullable + override def eval(input: InternalRow): Any = child.eval(input) + override protected def withNewChildInternal(newChild: Expression): InterpretedKernelExpr = + copy(child = newChild) +} + +/** + * Ships one occurrence of a non-deterministic `child` to the dispatcher. The dispatcher caches a + * kernel instance per serialized expression, and a non-deterministic kernel keeps state across + * batches (`monotonically_increasing_id`'s counter, `rand`'s generator). Two identical + * occurrences in one plan serialize to the same bytes, so without a distinct `occurrence` the + * second would continue the first one's state, where Spark gives each its own. The dispatcher + * unwraps this before it compiles `child`, so it never reaches generated code. + * + * Only the state of Catalyst's `Nondeterministic` nodes is per occurrence in Spark. A + * non-deterministic user function (`ScalaUDF`, `Invoke`, `StaticInvoke`, + * `ApplyFunctionExpression`) keeps its state in one object every call references, and a reflected + * method (`CallMethodViaReflection`) can keep it in static fields, which no arrangement of + * kernels reproduces, so `CometBatchKernelCodegen.canHandle` refuses them (see + * [[DispatchOccurrence.holdsSharedState]]). + */ +private[comet] case class DispatchOccurrence(child: Expression, occurrence: Long) + extends UnaryExpression + with Unevaluable { + override def dataType: DataType = child.dataType + override def nullable: Boolean = child.nullable + override protected def withNewChildInternal(newChild: Expression): DispatchOccurrence = + copy(child = newChild) +} + +private[comet] object DispatchOccurrence { + private val nextOccurrence = new java.util.concurrent.atomic.AtomicLong() + + /** `expr` as the dispatcher should receive it: tagged when it keeps per-occurrence state. */ + def tag(expr: Expression): Expression = + if (expr.exists(_.isInstanceOf[Nondeterministic])) { + DispatchOccurrence(expr, nextOccurrence.getAndIncrement()) + } else { + expr + } + + /** + * Non-deterministic in itself, whatever its children are, without being a Catalyst + * `Nondeterministic`: its state lives in an object the plan shares between its calls. + */ + def holdsSharedState(e: Expression): Boolean = e match { + case _: CallMethodViaReflection => true + case _: Nondeterministic | _: PlanExpression[_] => false + case u: ScalaUDF => !u.udfDeterministic + case i: Invoke => !i.isDeterministic + case s: StaticInvoke => !s.isDeterministic + case f: ApplyFunctionExpression => !f.function.isDeterministic + case other => !other.deterministic && other.children.forall(_.deterministic) + } + + /** The expression to compile from what the dispatcher received. */ + def untag(expr: Expression): Expression = expr match { + case DispatchOccurrence(child, _) => child + case other => other + } +} 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..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 } /** @@ -179,6 +192,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 +223,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 +426,10 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { */ private def emitSpecializedGetterExpr(target: String, idx: String, elemType: DataType): String = elemType match { + // 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)" case ShortType => s"$target.getShort($idx)" diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala index e14fb086b12..27d00641330 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala @@ -28,7 +28,7 @@ import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withFallback import org.apache.comet.DataTypeSupport.isComplexType import org.apache.comet.serde.{CodegenDispatchFallback, CometExpressionSerde, CometTimeZone, Compatible, ExprOuterClass, Incompatible, SupportLevel, Unsupported} import org.apache.comet.serde.ExprOuterClass.Expr -import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProtoInternal, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProtoInternal, liftFallbackReasons, serializeDataType} import org.apache.comet.shims.{CometExprShim, CometTypeShim} object CometCast @@ -110,7 +110,16 @@ object CometCast } else if (CometTimeZone.nativeId(cast.timeZoneId).isEmpty) { CometTimeZone.supportLevel(cast.timeZoneId) } else { - isSupported(cast.child.dataType, cast.dataType, cast.timeZoneId, evalMode(cast)) + // A relabel-only cast runs natively, but main dispatched some of them (`array` with + // different nested nullability). The dispatcher evaluates the child the way Spark's + // generated code does, skipping an argument an earlier one decides, and a native child + // evaluates every argument, so a non-deterministic child keeps main's route. + isSupported( + cast.child.dataType, + cast.dataType, + cast.timeZoneId, + evalMode(cast), + relabel = cast.child.deterministic) } } @@ -121,7 +130,12 @@ object CometCast val cometEvalMode = evalMode(cast) cast.child match { case _: Literal => - exprToProtoInternal(Literal.create(cast.eval(), cast.dataType), inputs, binding) + // The folded literal is not in the plan, so a reason it records is lifted onto `cast`. + val folded = Literal.create(cast.eval(), cast.dataType) + exprToProtoInternal(folded, inputs, binding).orElse { + liftFallbackReasons(folded, cast) + None + } case _ => if (isAlwaysCastToNull(cast.child.dataType, cast.dataType, cometEvalMode)) { exprToProtoInternal(Literal.create(null, cast.dataType), inputs, binding) @@ -136,6 +150,23 @@ object CometCast } } + /** + * True when `fromType` and `toType` differ only in nested field names and nullability. Spark + * and the native cast both match struct fields by position, so such a cast changes no value and + * only relabels the type; a set op adds one to give its two sides the same type. + */ + private def isRelabel(fromType: DataType, toType: DataType): Boolean = + (fromType, toType) match { + case (ArrayType(fromElement, _), ArrayType(toElement, _)) => + isRelabel(fromElement, toElement) + case (MapType(fromKey, fromValue, _), MapType(toKey, toValue, _)) => + isRelabel(fromKey, toKey) && isRelabel(fromValue, toValue) + case (StructType(fromFields), StructType(toFields)) => + fromFields.length == toFields.length && + fromFields.zip(toFields).forall { case (f, t) => isRelabel(f.dataType, t.dataType) } + case _ => fromType == toType + } + // Some casts like date -> int/byte / long are always null. Terminate early in planning private def isAlwaysCastToNull( fromType: DataType, @@ -191,7 +222,8 @@ object CometCast fromType: DataType, toType: DataType, timeZoneId: Option[String], - evalMode: CometEvalMode.Value): SupportLevel = { + evalMode: CometEvalMode.Value, + relabel: Boolean = true): SupportLevel = { // Spark 4's `VariantType` (SPARK-45827) has no native counterpart in Comet, and the codegen // dispatcher also cannot serialize `VariantType` in the data args or return type. The @@ -219,7 +251,7 @@ object CometCast return Unsupported(Some(nonDefaultCollationReason)) } - if (fromType == toType) { + if (fromType == toType || (relabel && isRelabel(fromType, toType))) { return Compatible() } @@ -235,9 +267,9 @@ object CometCast if toElementType != DataTypes.IntegerType && toElementType != DataTypes.StringType => unsupported(fromType, toType) case (dt: ArrayType, DataTypes.StringType) => - isSupported(dt.elementType, DataTypes.StringType, timeZoneId, evalMode) + isSupported(dt.elementType, DataTypes.StringType, timeZoneId, evalMode, relabel) case (dt: ArrayType, dt1: ArrayType) => - isSupported(dt.elementType, dt1.elementType, timeZoneId, evalMode) + isSupported(dt.elementType, dt1.elementType, timeZoneId, evalMode, relabel) case (dt: DataType, _) if dt.typeName == "timestamp_ntz" => toType match { case DataTypes.StringType => Compatible() @@ -279,7 +311,7 @@ object CometCast if (isAlwaysCastToNull(a.dataType, b.dataType, evalMode)) { return unsupported(fromType, toType) } - isSupported(a.dataType, b.dataType, timeZoneId, evalMode) match { + isSupported(a.dataType, b.dataType, timeZoneId, evalMode, relabel) match { case Compatible(_, _) => // all good case other => @@ -295,9 +327,9 @@ object CometCast isAlwaysCastToNull(from_map.valueType, to_map.valueType, evalMode)) { unsupported(fromType, toType) } else { - isSupported(from_map.keyType, to_map.keyType, timeZoneId, evalMode) match { + isSupported(from_map.keyType, to_map.keyType, timeZoneId, evalMode, relabel) match { case Compatible(_, _) => - isSupported(from_map.valueType, to_map.valueType, timeZoneId, evalMode) + isSupported(from_map.valueType, to_map.valueType, timeZoneId, evalMode, relabel) case other => other } } diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 623c603af81..fd37c9d4470 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -19,11 +19,15 @@ package org.apache.comet.rules +import java.nio.ByteBuffer + import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal +import org.apache.spark.SparkEnv import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, InputFileBlockLength, InputFileBlockStart, InputFileName, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Nondeterministic, Remainder, SparkPartitionID} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.rules.Rule @@ -53,14 +57,19 @@ import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ +import com.google.protobuf.Descriptors.FieldDescriptor +import com.google.protobuf.Message + import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometConf.{COMET_SPARK_TO_ARROW_ENABLED, COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST} import org.apache.comet.CometSparkSessionExtensions._ +import org.apache.comet.codegen.DispatchOccurrence import org.apache.comet.rules.CometExecRule.allExecs import org.apache.comet.serde._ import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.operator._ import org.apache.comet.shims.{CometTypeShim, ShimCometStreaming, ShimCometWindowGroupLimit, ShimSubqueryBroadcast} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen object CometExecRule { @@ -73,6 +82,167 @@ object CometExecRule { scan.relation.cacheBuilder.serializer.isInstanceOf[ArrowCachedBatchSerializer] && ArrowCachedBatchSerializer.supportsSchema(scan.relation.output) + val nondeterministicBesideNullTypeReason: String = + "a non-deterministic expression that does not run entirely in the codegen dispatcher, in " + + "an operator whose native plan also runs a dispatcher kernel with a NullType result: " + + "native evaluation computes every argument over the whole batch and can evaluate an " + + "argument twice, while Spark skips arguments that other arguments decide, so a stateful " + + "expression advances on rows Spark does not evaluate" + + /** + * True when `op` evaluates a non-deterministic expression outside the codegen dispatcher, and + * `op`'s native plan runs a dispatcher kernel whose result type contains NullType. + * + * Spark's generated code skips an argument when another one decides the result (a NULL operand, + * a NULL divisor), while native evaluation computes every argument over the whole batch, and + * some serdes serialize an argument twice (`ln`'s domain check), so a stateful expression + * (`monotonically_increasing_id`, `rand`) advances on rows Spark does not evaluate. That + * happens on main for any native operator. But the dispatcher used to refuse NullType results, + * and an operator with one fell back to Spark together with the native operators above it up to + * a shuffle. Every operator this check refuses therefore ran in Spark on main, so refusing all + * of them costs nothing against main and does not depend on modelling each expression's + * evaluation order. A top-level expression the dispatcher runs as a whole is exempt: its kernel + * evaluates Spark's own generated code, with its own state per occurrence. The exemption reads + * the dispatcher's tag on `op`'s expressions, which a rebuilt tree (a decimal expression + * wrapped in `CheckOverflow`) does not carry; that only makes the operator fall back. + * + * The kernels are read from the serialized plans rather than from `op`'s expressions, since + * serialization can dispatch a tree it built itself (a folded map literal expanded into + * `CreateMap`, a decimal addition wrapped in `CheckOverflow`). + */ + private[rules] def nondeterministicBesideNullTypeKernel( + op: SparkPlan, + nativeOp: OperatorOuterClass.Operator): Boolean = + op.expressions.exists(evaluatesNondeterministicNatively) && + (runsNullTypeKernel(nativeOp) || projectionRunsNullTypeKernel(op) || + op.children.exists(nullTypeKernelBelow)) + + private def evaluatesNondeterministicNatively(expr: Expression): Boolean = expr match { + case e if !e.exists(isStateful) => false + case e if e.getTagValue(CometExplainInfo.DISPATCHED_SELF).isDefined => false + case alias: Alias => evaluatesNondeterministicNatively(alias.child) + case _ => true + } + + /** + * A node whose value depends on how many rows were evaluated before it: a Catalyst + * `Nondeterministic` node or a user function that keeps state, whatever its arguments. The + * partition id and the input file only depend on the partition, however many rows Spark skips. + */ + private def isStateful(e: Expression): Boolean = e match { + case _: SparkPartitionID | _: InputFileName | _: InputFileBlockStart | + _: InputFileBlockLength => + false + case _: Nondeterministic => true + case other => DispatchOccurrence.holdsSharedState(other) + } + + /** + * `TakeOrderedAndProjectExec` serializes as a sink, a leaf scan, and builds its native sort and + * projection only when it executes, so those expressions are serialized again here. + */ + private def projectionRunsNullTypeKernel(op: SparkPlan): Boolean = op match { + case top: TakeOrderedAndProjectExec => + deferredRunsNullTypeKernel(top.sortOrder ++ top.projectList, top.child.output) + case _ => false + } + + private def deferredRunsNullTypeKernel( + exprs: Seq[Expression], + inputs: Seq[Attribute]): Boolean = + exprs.exists(e => QueryPlanSerde.exprToProto(e, inputs).exists(runsNullTypeKernel)) + + /** + * The native plans below `op` up to a shuffle. A sink is a leaf scan of the serialized plan + * above it, and a broadcast, union or coalesce sink exists only when the plan below it is + * native, so on main a NullType kernel there kept the operators above it in Spark as well; a + * shuffle took a Spark child, so only its own partitioning expressions count. + */ + private def nullTypeKernelBelow(plan: SparkPlan): Boolean = plan match { + // A Comet shuffle's partitioning expressions are serialized when it runs. Main refused a + // shuffle keyed on a NullType kernel, and then the operators above it ran in Spark too. + case shuffle: CometShuffleExchangeExec => + shuffle.outputPartitioning match { + case partitioning: Expression => + deferredRunsNullTypeKernel(partitioning.children, shuffle.child.output) + case _ => false + } + case stage: ShuffleQueryStageExec => nullTypeKernelBelow(stage.plan) + case read: AQEShuffleReadExec => nullTypeKernelBelow(read.child) + case _: ShuffleExchangeLike => false + case native: CometNativeExec => + runsNullTypeKernel(native.nativeOp) || native.children.exists(nullTypeKernelBelow) + case stage: BroadcastQueryStageExec => nullTypeKernelBelow(stage.plan) + case reused: ReusedExchangeExec => nullTypeKernelBelow(reused.child) + case top: CometTakeOrderedAndProjectExec => + deferredRunsNullTypeKernel(top.sortOrder ++ top.projectList, top.child.output) || + nullTypeKernelBelow(top.child) + case comet: CometPlan => comet.children.exists(nullTypeKernelBelow) + case _ => false + } + + private val dispatcherClassName = classOf[CometScalaUDFCodegen].getName + + /** + * Walks a serialized plan or expression for a dispatcher kernel that computes a NullType value: + * as its result, or anywhere in the tree it runs (a whole `coalesce` dispatched with an INT + * result can hold a `transform` with an `array` result, which main refused). + */ + private def runsNullTypeKernel(message: Message): Boolean = message match { + case udf: ExprOuterClass.JvmScalarUdf + if udf.getClassName == dispatcherClassName && + (protoContainsNullType(udf.getReturnType) || kernelComputesNullType(udf)) => + true + case _ => + message.getAllFields.asScala.exists { + case (field, value) if field.getJavaType == FieldDescriptor.JavaType.MESSAGE => + value match { + case values: java.util.List[_] => + values.asScala.exists(v => runsNullTypeKernel(v.asInstanceOf[Message])) + case child: Message => runsNullTypeKernel(child) + case _ => false + } + case _ => false + } + } + + /** + * The kernel's first argument is its closure-serialized expression (see + * `CometScalaUDF.emitJvmCodegenDispatch`). One that cannot be read counts as computing + * NullType. + */ + private def kernelComputesNullType(udf: ExprOuterClass.JvmScalarUdf): Boolean = + udf.getArgsCount == 0 || !udf.getArgs(0).hasLiteral || { + val bytes = udf.getArgs(0).getLiteral.getBytesVal.toByteArray + try { + val loader = Option(Thread.currentThread().getContextClassLoader) + .getOrElse(classOf[Expression].getClassLoader) + SparkEnv.get.closureSerializer + .newInstance() + .deserialize[Expression](ByteBuffer.wrap(bytes), loader) + .exists(e => containsNullType(e.dataType)) + } catch { + case NonFatal(_) => true + } + } + + private def containsNullType(dt: DataType): Boolean = dt match { + case NullType => true + case ArrayType(element, _) => containsNullType(element) + case MapType(key, value, _) => containsNullType(key) || containsNullType(value) + case StructType(fields) => fields.exists(f => containsNullType(f.dataType)) + case _ => false + } + + private def protoContainsNullType(dt: Types.DataType): Boolean = { + val info = dt.getTypeInfo + dt.getTypeId == Types.DataType.DataTypeId.NULL || + (info.hasList && protoContainsNullType(info.getList.getElementType)) || + (info.hasMap && (protoContainsNullType(info.getMap.getKeyType) || + protoContainsNullType(info.getMap.getValueType))) || + (info.hasStruct && info.getStruct.getFieldDatatypesList.asScala.exists(protoContainsNullType)) + } + private[rules] def removePlaceholders(plan: SparkPlan): SparkPlan = plan.transformUp { // revertUnsafePartialAggregates re-runs transform over already wrapped query stages, which // can produce CometSinkPlaceHolder(CometSinkPlaceHolder(stage)). Remove sinks bottom-up. @@ -769,7 +939,10 @@ case class CometExecRule(session: SparkSession) } case other => other } - if (cometChild.isInstanceOf[CometNativeExec]) { + // The join's own exchange stays on Spark when the broadcast gate refuses the build + // side, so this one must too: a Comet broadcast here would not be reused by it. + val gated = CometBroadcastExchangeExec.getSupportLevel(b).isInstanceOf[Unsupported] + if (cometChild.isInstanceOf[CometNativeExec] && !gated) { logInfo( "Converting SubqueryBroadcastExec to " + "CometSubqueryBroadcastExec for DPP exchange reuse") @@ -1099,24 +1272,37 @@ case class CometExecRule(session: SparkSession) childOp.foreach(builder.addChildren) return serde .convert(op, builder, childOp: _*) - .map { nativeOp => - val exec = serde.createExec(nativeOp, op) - rollUpInfoMessages(op, exec) - exec - } + .flatMap(nativeOp => createExec(serde, nativeOp, op)) } else { return serde .convert(op, builder) - .map { nativeOp => - val exec = serde.createExec(nativeOp, op) - rollUpInfoMessages(op, exec) - exec - } + .flatMap(nativeOp => createExec(serde, nativeOp, op)) } } None } + /** + * Builds the exec for a converted operator, unless the conversion put a non-deterministic + * expression next to a NullType dispatcher result + * ([[CometExecRule.nondeterministicBesideNullTypeKernel]]). The check runs after conversion, so + * it sees which kernels the plan runs, and an operator its serde refused keeps the serde's own + * reason. + */ + private def createExec( + serde: CometOperatorSerde[SparkPlan], + nativeOp: OperatorOuterClass.Operator, + op: SparkPlan): Option[SparkPlan] = { + if (CometExecRule.nondeterministicBesideNullTypeKernel(op, nativeOp)) { + withFallbackReason(op, CometExecRule.nondeterministicBesideNullTypeReason) + None + } else { + val exec = serde.createExec(nativeOp, op) + rollUpInfoMessages(op, exec) + Some(exec) + } + } + /** * Lift fallback reasons recorded on `op`'s expression trees onto `op` itself. * diff --git a/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala b/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala index 206bd7c8207..bff2862ae3a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala @@ -122,7 +122,15 @@ trait CometExpressionSerde[T <: Expression] { * dispatcher. Every other `Incompatible` expression falls back to Spark, and every other * `Unsupported` expression falls back to Spark. */ -trait CodegenDispatchFallback extends NativeOptInAvailable { self: CometExpressionSerde[_] => } +trait CodegenDispatchFallback extends NativeOptInAvailable { self: CometExpressionSerde[_] => + + /** + * Whether `Incompatible` cases also route through the dispatcher. A serde that enrolls only for + * its `Unsupported` cases sets this to false, so its `Incompatible` cases fall back to Spark + * unless the user opts into the native path, as for a serde without this mixin. + */ + def dispatchesIncompatible: Boolean = true +} /** * Marker for serdes that have a native implementation the user can opt into. Normally these diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 58caedcff14..2208a0504a9 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -28,9 +28,10 @@ import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf import org.apache.comet.CometExplainInfo import org.apache.comet.CometSparkSessionExtensions.{withCodegenDispatchExpr, withFallbackReason} -import org.apache.comet.codegen.CometBatchKernelCodegen +import org.apache.comet.DataTypeSupport +import org.apache.comet.codegen.{CometBatchKernelCodegen, DispatchOccurrence, InterpretedKernelExpr} import org.apache.comet.serde.ExprOuterClass.Expr -import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType, sparkEvaluatesInterpreted} import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** @@ -116,6 +117,13 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { case None => } + // Where Spark would evaluate this tree through `eval` (see `sparkEvaluatesInterpreted`), the + // kernel does too: some expressions' `eval` evaluates arguments their generated code skips + // (`ArraysZip` evaluates every argument before it looks for a NULL one), and an error one of + // them raises must reach the query as it does in Spark. + val kernelExpr = + if (sparkEvaluatesInterpreted) InterpretedKernelExpr(boundExpr) else boundExpr + // Serialize via Spark's closure serializer: respects the task context classloader (so user // UDF jars are visible) and matches Spark's wire format. The bytes become arg 0 of the // JvmScalarUdf proto and self-describe the expression so this works in cluster mode without @@ -130,7 +138,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { val bytes = try { val serializer = SparkEnv.get.closureSerializer.newInstance() - val buffer = serializer.serialize(boundExpr) + val buffer = serializer.serialize(DispatchOccurrence.tag(kernelExpr)) val serialized = new Array[Byte](buffer.remaining()) buffer.get(serialized) serialized @@ -159,12 +167,18 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { return None } } - val returnTypeProto = serializeDataType(expr.dataType).getOrElse { - withFallbackReason( - expr, - s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}") - return None - } + // Declared deep-nullable, like the Arrow field the kernel exports + // (`CometScalaUDFCodegen.lookupOrCompile`): native constructors (`make_array`, + // `named_struct`) and Parquet columns produce nested fields nullable, and a kernel that + // compares two inputs' types (set ops, CASE, list compare) rejects a dispatched side that + // keeps Spark's non-null flags. + val returnTypeProto = + serializeDataType(DataTypeSupport.deepNullable(expr.dataType)).getOrElse { + withFallbackReason( + expr, + s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}") + return None + } val udfBuilder = ExprOuterClass.JvmScalarUdf .newBuilder() 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 4a4b9eedbe7..234321a0ca8 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -28,9 +28,11 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback 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 @@ -38,7 +40,7 @@ import org.apache.spark.sql.types._ import org.apache.comet.CometConf import org.apache.comet.CometExplainInfo -import org.apache.comet.CometSparkSessionExtensions.{appendTagValues, withFallbackReason, withFallbackReasons, withInfo, withNativeExpr} +import org.apache.comet.CometSparkSessionExtensions.{appendTagValues, isSpark35Plus, withFallbackReason, withFallbackReasons, withInfo, withNativeExpr} import org.apache.comet.DataTypeSupport import org.apache.comet.expressions._ import org.apache.comet.parquet.CometParquetUtils @@ -73,7 +75,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, @@ -715,7 +717,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val info = DataTypeInfo.newBuilder() val list = ListInfo.newBuilder() list.setElementType(elementType.get) - list.setContainsNull(a.containsNull) + list.setContainsNull(Utils.declaredChildNullability(a.elementType, a.containsNull)) nestedParquetFieldId(parentField, elementPath, includeFieldIds) .foreach(list.setElementFieldId) @@ -739,7 +741,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) @@ -756,7 +758,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 @@ -822,7 +825,22 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { aggExpr: AggregateExpression, inputs: Seq[Attribute], binding: Boolean, - conf: SQLConf): Option[AggExpr] = { + conf: SQLConf): Option[AggExpr] = + aggExpr.aggregateFunction match { + // Spark feeds an imperative aggregate through `update`, which evaluates its inputs with + // `eval` under every codegen setting. + case _: ImperativeAggregate => + val enclosing = underInterpretedAncestor + withInterpretedEvaluation(aggExprToProto0(aggExpr, inputs, binding, conf, enclosing)) + case _ => aggExprToProto0(aggExpr, inputs, binding, conf, underInterpretedAncestor) + } + + private def aggExprToProto0( + aggExpr: AggregateExpression, + inputs: Seq[Attribute], + binding: Boolean, + conf: SQLConf, + filterInterpreted: Boolean): Option[AggExpr] = { // Distinct aggregates with a single column are supported (e.g., COUNT(DISTINCT x), // SUM(DISTINCT x), AVG(DISTINCT x)). The multi-stage plan generated by Spark @@ -904,7 +922,11 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { s"FILTER (WHERE ...) is not supported for aggregate mode ${aggExpr.mode}") return None } - val filterProto = exprToProto(aggExpr.filter.get, inputs, binding) + // Spark builds the FILTER predicate with `Predicate.create`, not through the aggregate's + // `update`, so it keeps the evaluation of the scope enclosing the aggregate. + val filterProto = interpretedAncestor.withValue(filterInterpreted) { + exprToProto(aggExpr.filter.get, inputs, binding) + } if (filterProto.isEmpty) { return None } @@ -1020,6 +1042,50 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * is not supported. */ def exprToProtoInternal( + expr: Expression, + inputs: Seq[Attribute], + binding: Boolean): Option[Expr] = expr match { + case _: CodegenFallback if !expr.isInstanceOf[LeafExpression] && !underInterpretedAncestor => + interpretedAncestor.withValue(true)(exprToProtoInternal0(expr, inputs, binding)) + case _ => exprToProtoInternal0(expr, inputs, binding) + } + + /** + * Set while the subtree of a non-leaf `CodegenFallback` expression is serialized (or an input + * Spark evaluates through `eval` for another reason, see [[withInterpretedEvaluation]]). Spark + * evaluates that subtree through `eval` even inside generated code, and an expression whose + * `eval` evaluates more of its arguments than its generated code does (`ArraysZip`) can raise + * an error there that the native path, which follows the generated code, does not. + */ + private val interpretedAncestor = new scala.util.DynamicVariable[Boolean](false) + + def underInterpretedAncestor: Boolean = interpretedAncestor.value + + /** + * Serializes `f` as if beneath a `CodegenFallback` expression: for an operator or input that + * Spark evaluates through `eval` (an imperative aggregate's inputs, a generator outside a + * whole-stage stage). + */ + def withInterpretedEvaluation[T](f: => T): T = interpretedAncestor.withValue(true)(f) + + /** + * True when Spark would evaluate the expression being serialized through `eval` rather than + * generated code: beneath a non-leaf `CodegenFallback` expression, or in a projection Spark + * builds without generated code (codegen factory mode NO_CODEGEN, which also turns whole-stage + * codegen off from Spark 3.5, SPARK-44236). The native path and the codegen dispatcher follow + * the generated code, which can evaluate fewer arguments than `eval` does. On Spark 3.4 a + * NO_CODEGEN plan can still interpret a projection outside a whole-stage stage (too many + * fields, a sort's ordering); NO_CODEGEN is a test setting, and this does not model those. + */ + def sparkEvaluatesInterpreted: Boolean = underInterpretedAncestor || { + val conf = SQLConf.get + val noCodegen = conf + .getConfString(SQLConf.CODEGEN_FACTORY_MODE.key) + .equalsIgnoreCase(CodegenObjectFactoryMode.NO_CODEGEN.toString) + noCodegen && (isSpark35Plus || !conf.wholeStageEnabled) + } + + private def exprToProtoInternal0( expr: Expression, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { @@ -1066,13 +1132,20 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { // result through the JVM codegen dispatcher (Spark's own `doGenCode` inside the Comet // pipeline) so the projection stays native while still matching Spark. Everything else // falls back to Spark. Falling back is also the result when the dispatcher cannot - // handle the expression. - val dispatched = dispatchIfFallback(handler, expr, inputs, binding).map { - case (h, proto) => - val key = h.nativeOptInConfigKeyOverride - .getOrElse(CometConf.getExprAllowIncompatConfigKey(exprConfName)) - withInfo(expr, NativeOptIn.message(exprConfName, key)) - proto + // handle the expression, or when the serde enrolls only its `Unsupported` cases + // (`dispatchesIncompatible`). + val dispatchesIncompatible = handler match { + case h: CodegenDispatchFallback => h.dispatchesIncompatible + case _ => false + } + val dispatchable = + if (dispatchesIncompatible) dispatchIfFallback(handler, expr, inputs, binding) + else None + val dispatched = dispatchable.map { case (h, proto) => + val key = h.nativeOptInConfigKeyOverride + .getOrElse(CometConf.getExprAllowIncompatConfigKey(exprConfName)) + withInfo(expr, NativeOptIn.message(exprConfName, key)) + proto } dispatched.orElse { val optionalNotes = notes.map(str => s" ($str)").getOrElse("") 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..f8dbda700ad 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,30 @@ 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. + * Each copy gets its own instance, natively or in the JVM codegen dispatcher (see + * `DispatchOccurrence`), 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. + * + * 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 8928e40cbeb..7393f2d5be7 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, import org.apache.spark.sql.catalyst.util.ArrayData import org.apache.spark.sql.comet.CometExecUtils 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.{isSpark41Plus, isSpark42Plus, withFallbackReason} @@ -1169,14 +1169,12 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { } } + private val respectNullsReason = + "`collect_set` with `RESPECT NULLS` falls back to Spark, since the native " + + "implementation always drops null inputs." + override def getUnsupportedReasons(): Seq[String] = - if (isSpark42Plus) { - Seq( - "`collect_set` with `RESPECT NULLS` falls back to Spark, since the native " + - "implementation always drops null inputs.") - } else { - Nil - } + if (isSpark42Plus) Seq(respectNullsReason) else Nil override def getSupportLevel(expr: CollectSet): SupportLevel = { // The native path always drops null inputs. Spark 4.2 added an `ignoreNulls` field to @@ -1185,7 +1183,7 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { // through 4.1 the field does not exist, `RESPECT NULLS`/`IGNORE NULLS` are rejected at // analysis time, and CometCollectShim.ignoreNulls hardcodes true, making this a no-op. if (!CometCollectShim.ignoreNulls(expr)) { - Unsupported(Some("collect_set with RESPECT NULLS (ignoreNulls = false) is not supported")) + Unsupported(Some(respectNullsReason)) } else if (isSpark42Plus) { Compatible() } else { @@ -1239,8 +1237,46 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { } } +/** + * NullType gate for collect_list. Without a grouping key, the final aggregation merges the + * partial states as they arrive from the shuffle, and a nested element there still carries the + * non-nullable fields its producer gave it (`named_struct('a', id)`, `array(id)`, `map_entries`), + * while the declared output has them nullable: the planner's nullability cast on the collect + * child only covers the partial side. DataFusion then rejects the batch ("column types must match + * schema types"). A grouped collect_list, and collect_set in either form, are not affected; the + * serde cannot see the grouping keys, so the gate covers grouped collect_list as well. + * + * This is not specific to NullType: `SELECT collect_list(array(id)) FROM t` over more than one + * partition fails the same way on main. For NullType-bearing input main fails even sooner: a + * nested element such as `array(named_struct('a', id, 'b', NULL))` raises this error from a + * single partition. The gate trades the NullType shapes that do run on main (a struct with a NULL + * field, and any grouped collect_list) for keeping all of them off that path; it can go once the + * final merge normalizes the state as well. + */ +object CometCollectListNullType { + + val reason: String = + "collect_list over a NullType-bearing input falls back to Spark: without a grouping key, " + + "the native final merge rejects an element whose producer declares non-nullable fields, " + + "and the grouping keys are not known when the aggregate is planned" + + def supportLevel(dt: DataType): Option[SupportLevel] = + if (SupportLevel.containsType(dt, classOf[NullType])) { + Some(Unsupported(Some(reason))) + } else { + None + } +} + object CometCollectList extends CometAggregateExpressionSerde[CollectList] { + private val respectNullsReason = + "`collect_list` with `RESPECT NULLS` falls back to Spark, since the native " + + "implementation always drops null inputs." + + override def getUnsupportedReasons(): Seq[String] = + (if (isSpark42Plus) Seq(respectNullsReason) else Nil) :+ CometCollectListNullType.reason + override def getSupportLevel(expr: CollectList): SupportLevel = { // The native path delegates to SparkCollectList, which always drops null inputs. Spark 4.2 // added an `ignoreNulls` field to CollectList that `RESPECT NULLS` sets to false, preserving @@ -1249,9 +1285,9 @@ object CometCollectList extends CometAggregateExpressionSerde[CollectList] { // are rejected at analysis time, and CometCollectShim.ignoreNulls hardcodes true, making this // a no-op. if (!CometCollectShim.ignoreNulls(expr)) { - Unsupported(Some("collect_list with RESPECT NULLS (ignoreNulls = false) is not supported")) + Unsupported(Some(respectNullsReason)) } else { - Compatible() + CometCollectListNullType.supportLevel(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 b2a75f84782..0d2d2f76cfb 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -23,13 +23,13 @@ import scala.annotation.tailrec import scala.jdk.CollectionConverters._ import org.apache.spark.SPARK_VERSION -import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, 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.{ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, 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.comet.util.Utils import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, withFallbackReason} import org.apache.comet.DataTypeSupport.{deepNullable, isComplexType} import org.apache.comet.serde.QueryPlanSerde._ import org.apache.comet.shims.{CometExprShim, CometTypeShim} @@ -62,6 +62,16 @@ object CometArrayRemove object CometArrayAppend extends CometExpressionSerde[ArrayAppend] with ArraysBase { + override def getUnsupportedReasons(): Seq[String] = Seq(NullGuard.reason) + + // The item is serialized once, inside THEN, but it is guarded too when the array can be NULL: + // Spark's generated code evaluates it on every row (only the interpreted `eval` skips it for a + // NULL array), while native THEN evaluates it on the selected rows only. With a non-nullable + // array every row is selected, so the item runs on every row in both engines. + override def getSupportLevel(expr: ArrayAppend): SupportLevel = + if (expr.left.nullable) NullGuard.supportLevel(expr.children: _*) + else NullGuard.supportLevel(expr.left) + override def convert( expr: ArrayAppend, inputs: Seq[Attribute], @@ -228,13 +238,29 @@ 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 getUnsupportedReasons(): Seq[String] = + Seq(nullElementReason, SetOpArguments.argumentReason("array_intersect")) + 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)) @@ -245,12 +271,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) - - val arraysIntersectScalarExpr = - scalarFunctionExprToProto("array_intersect", leftArrayExprProto, rightArrayExprProto) - arraysIntersectScalarExpr + // 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 + } + SetOpArguments.serialize(expr, inputs, binding).flatMap { case Seq(left, right) => + scalarFunctionExprToProto("array_intersect", Some(left), Some(right)) + } } } @@ -348,18 +378,27 @@ 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, SetOpArguments.argumentReason("array_except")) + 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"$elementTypeReason ($dt)")) + case None => Incompatible(Some(incompatReason)) } - Incompatible(Some(reason)) } @tailrec @@ -381,21 +420,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") + withFallbackReason(expr, s"$elementTypeReason ($dt)") return None case None => } - val leftArrayExprProto = exprToProtoInternal(expr.left, inputs, binding) - val rightArrayExprProto = exprToProtoInternal(expr.right, inputs, binding) - - val arrayExceptScalarExpr = - scalarFunctionExprToProto("array_except", leftArrayExprProto, rightArrayExprProto) - arrayExceptScalarExpr + SetOpArguments.serialize(expr, inputs, binding).flatMap { case Seq(left, right) => + scalarFunctionExprToProto("array_except", Some(left), Some(right)) + } } } @@ -571,26 +607,136 @@ object CometArrayDistinct extends CometScalarFunction[ArrayDistinct]("array_dist ArraySetSupport.supportLevel(expr.dataType) } -object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { +/** + * 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). Equal Spark + * types are no evidence of equal native types either: the JVM codegen dispatcher and the native + * constructors declare every nested field nullable, while `map_entries` keeps its key field + * non-nullable. Field names can differ too: with case-insensitive analysis Spark accepts + * `struct` beside `struct` without a cast, and compares struct values by position, and a + * native `CASE` names a merged struct's fields after its ELSE branch where Spark names them after + * the first branch. Casting both sides to one target, the set op's own element type made deeply + * nullable, only relabels metadata (struct casts are positional) and never changes values, so the + * kernel always sees one type. Every side is cast, even one whose Spark type already is the + * target, since its native type need not be; the native cast is a no-op when it is. A primitive + * element type needs nothing, since only nested fields take part in the comparison. + */ +private[serde] object SetOpArguments { + def unify(expr: Expression): Seq[Expression] = + expr.dataType match { + case ArrayType(et, _) if isComplexType(et) => + val unified = ArrayType(deepNullable(et), containsNull = true) + expr.children.map(Cast(_, unified)) + case _ => expr.children + } + + /** + * Serializes `expr`'s arguments through [[unify]]. The casts it adds are not in the operator's + * tree, so a reason recorded under one would never reach `CometExecRule`'s roll-up; when an + * argument does not serialize, the reasons found under it are recorded on `expr` itself. + */ + def serialize( + expr: Expression, + inputs: Seq[Attribute], + binding: Boolean): Option[Seq[ExprOuterClass.Expr]] = { + val unified = unify(expr) + val protos = unified.map(exprToProtoInternal(_, inputs, binding)) + if (protos.forall(_.isDefined)) { + Some(protos.flatten) + } else { + unified.zip(protos).collect { case (arg, None) => arg }.foreach { arg => + QueryPlanSerde.liftFallbackReasons(arg, expr) + } + if (!hasFallbackReason(expr)) { + withFallbackReason(expr, argumentReason(expr.prettyName)) + } + None + } + } + + /** The reason recorded when an argument fails without a reason of its own. */ + def argumentReason(name: String): String = s"$name: an argument could not be serialized" +} + +/** + * 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] with CodegenDispatchFallback { override def getIncompatibleReasons(): Seq[String] = Seq(ArraySetSupport.floatingPointReason) - override def getSupportLevel(expr: ArrayUnion): SupportLevel = - ArraySetSupport.supportLevel(expr.dataType) + // Only the NullType-element case runs through the dispatcher. A floating-point element type + // still falls back to Spark unless the user opts in, as `CometArrayDistinct` does. + override def dispatchesIncompatible: Boolean = false + + private val nullElementReason = + "native array_union drops the entries of a NullType-element array" + + override def getUnsupportedReasons(): Seq[String] = + Seq(nullElementReason, SetOpArguments.argumentReason("array_union")) + + // 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(nullElementReason)) + } else { + ArraySetSupport.supportLevel(expr.dataType) + } + } 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 arraysUnionScalarExpr = - scalarFunctionExprToProto("array_union", leftArrayExprProto, rightArrayExprProto) - arraysUnionScalarExpr + SetOpArguments.serialize(expr, inputs, binding).flatMap { case Seq(left, right) => + scalarFunctionExprToProto("array_union", Some(left), Some(right)) + } } } -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 + // 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. Only a + // literal is sure to arrive as a scalar: a foldable argument that constant folding did not + // reduce (`element_at(array(NULL), 1)` with `ConstantFolding` excluded) is still evaluated + // per row. Empty `array()` never reaches 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.isInstanceOf[Literal])) { + Unsupported(Some(nullTypeBatchReason)) + } else { + Compatible() + } + } + override def convert( expr: CreateArray, inputs: Seq[Attribute], @@ -630,6 +776,58 @@ object CometCreateArray extends CometExpressionSerde[CreateArray] with ArraysBas } } +object CometArrayRepeat extends CometExpressionSerde[ArrayRepeat] with CodegenDispatchFallback { + + private val nonNullableItemReason = + "native array_repeat rebuilds a non-nullable list item as nullable" + + private val interpretedReason = + "`array_repeat` with a nullable count, where Spark evaluates it without generated code, " + + "runs through the codegen dispatcher: its interpreted eval skips the element when the " + + "count is NULL, where the native kernel evaluates both" + + override def getUnsupportedReasons(): Seq[String] = + Seq(nonNullableItemReason, interpretedReason) + + override def getSupportLevel(expr: ArrayRepeat): SupportLevel = { + // Spark's generated code evaluates the element and the count; its `eval` evaluates the count + // first and skips the element when the count is NULL, so an element that raises an error + // raises natively only. The dispatcher evaluates the tree through `eval` in these contexts. + if (sparkEvaluatesInterpreted && expr.right.nullable) { + return Unsupported(Some(interpretedReason)) + } + 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), so containsNull = + // false is the plan-time proxy for that input. It is broader than the defect: a + // dispatched producer now declares its items nullable, and those still go through the + // dispatcher here. `CodegenDispatchFallback` runs them there instead of the kernel. + // 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() + } + } + } + + 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( @@ -848,12 +1046,23 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] { object CometSize extends CometExpressionSerde[Size] { + private val childTypeReason = "Only array and map inputs are supported" + + override def getUnsupportedReasons(): Seq[String] = Seq(childTypeReason, NullGuard.reason) + 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")) + Unsupported(Some(s"$childTypeReason ($other)")) } } @@ -862,21 +1071,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() + } } } @@ -891,12 +1109,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 { @@ -927,8 +1139,20 @@ 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") + private val inputTypeReason = + "Not all input data types are supported; falls back to Spark for unsupported types" + + private val noArgumentsReason = "`arrays_zip()` with no arguments falls back to Spark" + + private val interpretedReason = + "`arrays_zip` over several arguments falls back to Spark where Spark evaluates it without " + + "generated code (codegen factory mode NO_CODEGEN, under a CodegenFallback expression, as " + + "the input of an imperative aggregate, or under a generator outside a whole-stage " + + "stage): its interpreted eval evaluates every argument, where its generated code, which " + + "the native null guard follows, stops at the first NULL one" + + override def getUnsupportedReasons(): Seq[String] = + Seq(inputTypeReason, noArgumentsReason, interpretedReason, NullGuard.reason) private def isTypeSupported(dt: DataType): Boolean = { import DataTypes._ @@ -944,13 +1168,26 @@ object CometArraysZip extends CometExpressionSerde[ArraysZip] { } override def getSupportLevel(expr: ArraysZip): SupportLevel = { + // The native kernel needs at least one argument, and so does `convert`'s CASE (one WHEN per + // argument); Spark returns an empty array. + if (expr.children.isEmpty) { + return Unsupported(Some(noArgumentsReason)) + } val inputTypes = expr.children.map(_.dataType).toSet for (dt <- inputTypes) { if (!isTypeSupported(dt)) { - return Unsupported(Some(s"Unsupported child data type: $dt")) + return Unsupported(Some(s"$inputTypeReason ($dt)")) } } - Compatible() + // The guard follows Spark's generated code, which stops at the first NULL argument. Spark's + // interpreted `eval` evaluates them all first, so an argument that raises an error after a + // NULL one raises in Spark there and not natively. + if (expr.children.size > 1 && sparkEvaluatesInterpreted) { + return Unsupported(Some(interpretedReason)) + } + // `convert` puts every child in a null check as well as in the zip itself, so a single + // stateful child is enough to diverge. + NullGuard.supportLevel(expr.children: _*) } override def convert( @@ -962,25 +1199,30 @@ object CometArraysZip extends CometExpressionSerde[ArraysZip] { expr.children.map(exprToProtoInternal(_, inputs, binding)) val names: Seq[Any] = expr.names.map(_.eval(EmptyRow)) - // mimic Spark's ArraysZip behavior: returns NULL if any argument is NULL - val combinedNullCheck = expr.children.map(child => IsNotNull(child)).reduce(And) - val isNotNullExpr = exprToProtoInternal(combinedNullCheck, inputs, binding) + // Mimic Spark's ArraysZip: NULL if any argument is NULL. One WHEN per argument, in order: + // Spark's generated code stops evaluating the arguments at the first NULL one, and native + // CASE evaluates each later WHEN and the ELSE only on the rows no earlier WHEN matched. A + // single AND of the null checks can evaluate a later argument on rows Spark skips, raising + // an error there (a division by zero under ANSI). + val isNullExprs = + expr.children.map(child => exprToProtoInternal(IsNull(child), inputs, binding)) val nullLiteralProto = exprToProtoInternal(Literal(null, expr.dataType), Seq.empty, binding = true) - if (exprChildren.forall( - _.isDefined) && isNotNullExpr.isDefined && nullLiteralProto.isDefined) { + if (exprChildren.forall(_.isDefined) && isNullExprs.forall( + _.isDefined) && nullLiteralProto.isDefined) { val arraysZip: ExprOuterClass.ArraysZip = ExprOuterClass.ArraysZip .newBuilder() .addAllValues(exprChildren.map(_.get).asJava) .addAllNames(names.map(_.toString).asJava) .build() - val caseWhenExpr = ExprOuterClass.CaseWhen - .newBuilder() - .addWhen(isNotNullExpr.get) - .addThen(ExprOuterClass.Expr.newBuilder().setArraysZip(arraysZip).build()) - .setElseExpr(nullLiteralProto.get) + val caseWhenBuilder = ExprOuterClass.CaseWhen.newBuilder() + isNullExprs.foreach { isNull => + caseWhenBuilder.addWhen(isNull.get).addThen(nullLiteralProto.get) + } + val caseWhenExpr = caseWhenBuilder + .setElseExpr(ExprOuterClass.Expr.newBuilder().setArraysZip(arraysZip).build()) .build() Some( ExprOuterClass.Expr 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..e4c2eb51733 100644 --- a/spark/src/main/scala/org/apache/comet/serde/conditional.scala +++ b/spark/src/main/scala/org/apache/comet/serde/conditional.scala @@ -21,9 +21,46 @@ 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.catalyst.expressions.{Attribute, CaseWhen, Cast, Coalesce, Expression, If, IsNotNull} +import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType} -import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal +import org.apache.comet.DataTypeSupport.deepNullable +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, liftFallbackReasons} + +/** + * Serializes a CASE WHEN or coalesce result branch. A native CASE coerces its branches to one + * type, folding from the ELSE branch, so a merged struct takes the ELSE branch's field names, + * where Spark names it after the first branch; with case-insensitive analysis the two can differ + * (`s` beside `named_struct('A', id)`). A consumer that compares types exactly (`array(...)`, the + * set ops) then meets a struct named differently from its sibling. A branch of a struct-bearing + * type is cast to the expression's own type, made deeply nullable, so every branch has Spark's + * names and the coercion has nothing left to change; struct casts are positional, so values are + * kept. + */ +private[serde] object CaseBranch { + def serialize( + expr: Expression, + branch: Expression, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + if (!hasStruct(expr.dataType)) { + return exprToProtoInternal(branch, inputs, binding) + } + // The cast is not in the plan, so a reason recorded under it is lifted onto `expr`. + val aligned = Cast(branch, deepNullable(expr.dataType)) + exprToProtoInternal(aligned, inputs, binding).orElse { + liftFallbackReasons(aligned, expr) + None + } + } + + private def hasStruct(dt: DataType): Boolean = dt match { + case _: StructType => true + case ArrayType(element, _) => hasStruct(element) + case MapType(key, value, _) => hasStruct(key) || hasStruct(value) + case _ => false + } +} object CometIf extends CometExpressionSerde[If] { override def convert( @@ -61,7 +98,7 @@ object CometCaseWhen extends CometExpressionSerde[CaseWhen] { }) val thenSeq = expr.branches.map(elements => { allBranches = allBranches :+ elements._2 - exprToProtoInternal(elements._2, inputs, binding) + CaseBranch.serialize(expr, elements._2, inputs, binding) }) assert(whenSeq.length == thenSeq.length) if (whenSeq.forall(_.isDefined) && thenSeq.forall(_.isDefined)) { @@ -69,8 +106,7 @@ object CometCaseWhen extends CometExpressionSerde[CaseWhen] { builder.addAllWhen(whenSeq.map(_.get).asJava) builder.addAllThen(thenSeq.map(_.get).asJava) if (expr.elseValue.isDefined) { - val elseValueExpr = - exprToProtoInternal(expr.elseValue.get, inputs, binding) + val elseValueExpr = CaseBranch.serialize(expr, expr.elseValue.get, inputs, binding) if (elseValueExpr.isDefined) { builder.setElseExpr(elseValueExpr.get) } else { @@ -88,11 +124,27 @@ object CometCaseWhen extends CometExpressionSerde[CaseWhen] { } } -object CometCoalesce extends CometExpressionSerde[Coalesce] { +object CometCoalesce extends CometExpressionSerde[Coalesce] with CodegenDispatchFallback { + + override def getUnsupportedReasons(): Seq[String] = Seq(NullGuard.reason) + + // Every child but the last is a guard; the last one is the ELSE, evaluated on the rows the + // guards left over. Only the guarded children are serialized twice (predicate and THEN), so + // only they need the single-evaluation check; the ELSE is serialized once, as Spark evaluates + // it once. + override def getSupportLevel(expr: Coalesce): SupportLevel = + NullGuard.supportLevel(expr.children.dropRight(1): _*) + override def convert( expr: Coalesce, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { + // The optimizer normally reduces a one-argument coalesce to its argument, but not when + // `NullPropagation` and `SimplifyConditionals` are excluded, and a CASE over no guarded + // children would have no WHEN clause, which native CASE rejects. + if (expr.children.size == 1) { + return exprToProtoInternal(expr.children.head, inputs, binding) + } val branches = expr.children.dropRight(1).map { child => (IsNotNull(child), child) } @@ -101,14 +153,14 @@ object CometCoalesce extends CometExpressionSerde[Coalesce] { exprToProtoInternal(elements._1, inputs, binding) }) val thenSeq = branches.map(elements => { - exprToProtoInternal(elements._2, inputs, binding) + CaseBranch.serialize(expr, elements._2, inputs, binding) }) assert(whenSeq.length == thenSeq.length) if (whenSeq.forall(_.isDefined) && thenSeq.forall(_.isDefined)) { val builder = ExprOuterClass.CaseWhen.newBuilder() builder.addAllWhen(whenSeq.map(_.get).asJava) builder.addAllThen(thenSeq.map(_.get).asJava) - val elseValueExpr = exprToProtoInternal(elseValue, inputs, binding) + val elseValueExpr = CaseBranch.serialize(expr, elseValue, inputs, binding) if (elseValueExpr.isDefined) { builder.setElseExpr(elseValueExpr.get) } else { diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index bf27ce2c552..7cb1ead0f98 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -511,6 +511,25 @@ object CometUnixDate extends CometExpressionSerde[UnixDate] { } } +/** + * Spark's `eval` for `trunc` / `date_trunc` reads a non-literal format first and returns NULL for + * an invalid one without evaluating the date or timestamp, while its generated code, like the + * native kernel, evaluates both. Where Spark would use `eval` + * ([[QueryPlanSerde.sparkEvaluatesInterpreted]]) the expression runs through the codegen + * dispatcher, which then evaluates it through `eval` as well. + */ +private object TruncFormat { + + /** Whether Spark would evaluate this format first and possibly skip the value. */ + def skipsValue(format: Expression): Boolean = + !format.isInstanceOf[Literal] && sparkEvaluatesInterpreted + + def interpretedReason(name: String): String = + s"`$name` with a non-literal format runs through the codegen dispatcher where Spark " + + "evaluates it without generated code: its interpreted eval returns NULL for an invalid " + + "format without evaluating the value, which the native kernel evaluates" +} + object CometTruncDate extends CometExpressionSerde[TruncDate] with CodegenDispatchFallback { val supportedFormats: Seq[String] = @@ -529,10 +548,15 @@ object CometTruncDate extends CometExpressionSerde[TruncDate] with CodegenDispat Seq(nonLiteralFormatIncompatReason) ++ DatetimeCollation.incompatibleReasons("trunc") override def getUnsupportedReasons(): Seq[String] = Seq( - "Only the following formats are supported: " + supportedFormats.mkString(", ")) + "Only the following formats are supported: " + supportedFormats.mkString(", "), + TruncFormat.interpretedReason("trunc")) override def getSupportLevel(expr: TruncDate): SupportLevel = { - if (DatetimeCollation.hasNonDefaultCollation(expr)) { + // Ahead of every opt-in branch (collation included), so `allowIncompatible` cannot route a + // non-literal format to the native kernel where Spark evaluates through `eval`. + if (TruncFormat.skipsValue(expr.format)) { + Unsupported(Some(TruncFormat.interpretedReason(expr.prettyName))) + } else if (DatetimeCollation.hasNonDefaultCollation(expr)) { Incompatible(Some(collationReason)) } else { expr.format match { @@ -605,10 +629,14 @@ object CometTruncTimestamp DatetimeCollation.incompatibleReasons("date_trunc") override def getUnsupportedReasons(): Seq[String] = Seq( - "Only the following formats are supported: " + supportedFormats.mkString(", ")) + "Only the following formats are supported: " + supportedFormats.mkString(", "), + TruncFormat.interpretedReason("date_trunc")) override def getSupportLevel(expr: TruncTimestamp): SupportLevel = { - if (DatetimeCollation.hasNonDefaultCollation(expr)) { + // Ahead of every opt-in branch; see CometTruncDate. + if (TruncFormat.skipsValue(expr.format)) { + Unsupported(Some(TruncFormat.interpretedReason(expr.prettyName))) + } else if (DatetimeCollation.hasNonDefaultCollation(expr)) { Incompatible(Some(collationReason)) } else if (CometTimeZone.nativeId(expr.timeZoneId).isEmpty) { CometTimeZone.supportLevel(expr.timeZoneId) 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 f9d3e2dbc00..d86bd852665 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -159,13 +159,17 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { override def getCompatibleNotes(): Seq[String] = Seq(MapKeyDedupPolicySupport.nullKeyReason) - override def getSupportLevel(expr: MapFromArrays): SupportLevel = { - if (MapKeyDedupPolicySupport.isLastWin) { - Incompatible(Some(MapKeyDedupPolicySupport.incompatibleReason)) - } else { - Compatible(None) + override def getUnsupportedReasons(): Seq[String] = Seq(NullGuard.reason) + + // The null guard comes before the LAST_WIN branch: `Incompatible` hands the expression to + // `convert` once the user opts in, so a shape the guard cannot run has to be refused first or + // the opt-in would reach it. + override def getSupportLevel(expr: MapFromArrays): SupportLevel = + NullGuard.supportLevel(expr.left, expr.right) match { + case _: Compatible if MapKeyDedupPolicySupport.isLastWin => + Incompatible(Some(MapKeyDedupPolicySupport.incompatibleReason)) + case level => level } - } override def convert( expr: MapFromArrays, diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala index 1cab882df93..90beabf8a72 100644 --- a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala @@ -35,7 +35,8 @@ import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.types.{BinaryType, DataType, StringType} -import org.apache.comet.codegen.{CometBatchKernel, CometBatchKernelCodegen} +import org.apache.comet.DataTypeSupport +import org.apache.comet.codegen.{CometBatchKernel, CometBatchKernelCodegen, DispatchOccurrence} import org.apache.comet.codegen.CometBatchKernelCodegen.{ArrayColumnSpec, ArrowColumnSpec, MapColumnSpec, ScalarColumnSpec, StructColumnSpec, StructFieldSpec} import org.apache.comet.udf.CometUDF @@ -81,7 +82,9 @@ class CometScalaUDFCodegen extends CometUDF with Logging { * Per-task cache keyed on serialized expression bytes plus per-column specs. The deserialized * `boundExpr` carries mutable state (`NamedLambdaVariable.value` for HOFs, `Rand`'s * `XORShiftRandom`) that must not be shared across concurrent tasks running the same query; - * keeping the cache per-task gives each task its own copy. Guarded by `this.synchronized`. + * keeping the cache per-task gives each task its own copy. Within a task, two occurrences of + * one non-deterministic expression get their own entries too, since the serde ships each with + * its own [[DispatchOccurrence]]. Guarded by `this.synchronized`. */ private val kernelCache : mutable.Map[CometScalaUDFCodegen.CacheKey, CometScalaUDFCodegen.CacheEntry] = @@ -155,9 +158,10 @@ class CometScalaUDFCodegen extends CometUDF with Logging { .getOrElse(classOf[Expression].getClassLoader) val boundExpr = try { - SparkEnv.get.closureSerializer - .newInstance() - .deserialize[Expression](ByteBuffer.wrap(bytes), loader) + DispatchOccurrence.untag( + SparkEnv.get.closureSerializer + .newInstance() + .deserialize[Expression](ByteBuffer.wrap(bytes), loader)) } catch { case NonFatal(t) => logError( @@ -169,9 +173,13 @@ class CometScalaUDFCodegen extends CometUDF with Logging { val compiled = CometBatchKernelCodegen.compile(boundExpr, specs) val kernel = compiled.newInstance() kernel.init(CometScalaUDFCodegen.currentPartitionIndex()) + // Exported deep-nullable, matching the return type the serde declares + // (`CometScalaUDF`): nested nullability is Arrow metadata only, so the kernel still + // writes with Spark's flags, but native consumers see the same field types a native + // producer of this Spark type would give them. val outputField = CometBatchKernelCodegen.toFfiArrowField( "codegen_result", - boundExpr.dataType, + DataTypeSupport.deepNullable(boundExpr.dataType), boundExpr.nullable) val entry = CometScalaUDFCodegen.CacheEntry(compiled, kernel, boundExpr.dataType, outputField) 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 615f7bada26..e59ffc412fc 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/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala index 19532cefa14..0d8cb5c599b 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 @@ -201,6 +201,10 @@ object CometArrowStream extends Logging { * a null and crash native validation. Both inputs are borrowed and unchanged; the returned * field owns no buffers. Missing dictionary metadata or provider entries propagate the shared * accessor's named column error. + * + * 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/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 7a0d231a141..517838ec8cf 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -41,7 +41,7 @@ import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType} +import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.MutablePair import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordComparator} @@ -670,6 +670,20 @@ object CometShuffleExchangeExec if (!config.get(conf)) { reasons += s"${config.key} is disabled" } + // Unless it places rows by position, native round-robin hashes the leading + // `maxHashColumns` columns (all of them for 0), and the native hasher has no NullType + // arm, at the top level or nested. + val hashed = + if (positionalRoundRobinSpec(partitioning, s.child).isDefined) { + Seq.empty + } else { + val maxHashColumns = + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get(conf) + if (maxHashColumns <= 0) inputs else inputs.take(maxHashColumns) + } + if (hashed.exists(a => containsNullType(a.dataType))) { + reasons += nullTypeRoundRobinReason + } case _ => reasons += s"unsupported Spark partitioning for native shuffle: ${partitioning.getClass.getName}" @@ -677,6 +691,18 @@ object CometShuffleExchangeExec reasons.toSeq } + private val nullTypeRoundRobinReason = + "native round-robin shuffle hashes the row's columns, and the native hasher does not " + + "support NullType columns or fields" + + private def containsNullType(dt: DataType): Boolean = dt match { + case NullType => true + case ArrayType(element, _) => containsNullType(element) + case MapType(key, value, _) => containsNullType(key) || containsNullType(value) + case StructType(fields) => fields.exists(f => containsNullType(f.dataType)) + case _ => false + } + /** * Reasons the columnar shuffle path cannot handle this shuffle. Empty means columnar is * supported. Pure: does not tag the node. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 91881c2c145..0b1a908bc45 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1798,7 +1798,20 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { val childExpr = op.generator.children.head - val childExprProto = exprToProto(childExpr, op.child.output) + // Outside a whole-stage stage, GenerateExec calls the generator's `eval`, which evaluates its + // child the same way, so the child is serialized as Spark interprets it. These are the + // conditions under which CollapseCodegenStages leaves the Generate out of a stage. + val interpreted = !op.conf.wholeStageEnabled || !op.generator.supportCodegen || + op.expressions.exists(_.exists(e => + e.isInstanceOf[CodegenFallback] && !e.isInstanceOf[LeafExpression])) || + WholeStageCodegenExec.isTooManyFields(op.conf, op.schema) || + op.children.exists(child => WholeStageCodegenExec.isTooManyFields(op.conf, child.schema)) + val childExprProto = + if (interpreted) { + QueryPlanSerde.withInterpretedEvaluation(exprToProto(childExpr, op.child.output)) + } else { + exprToProto(childExpr, op.child.output) + } if (childExprProto.isEmpty) { return None 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 c884344ef15..d87c12d7970 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,11 +21,12 @@ 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._ import org.apache.arrow.c.CDataDictionaryProvider +import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} @@ -184,6 +185,18 @@ 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 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 + /** Maps field from Spark to Arrow. NOTE: timeZoneId required for TimestampType */ def toArrowField(name: String, dt: DataType, nullable: Boolean, timeZoneId: String): Field = { dt match { @@ -192,7 +205,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( @@ -200,24 +218,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) @@ -233,6 +257,145 @@ 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. 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 + * 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 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` cannot show when that happens, since the + * bypassed shapes never reach the appender; check with a direct appender probe first. + */ + 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 + * 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. + * + * 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 = { + 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 @@ -267,10 +430,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) { @@ -335,15 +498,19 @@ object Utils extends CometTypeShim with Logging { * allocation) for each one. With coalescing, we decode and append all batches into one * VectorSchemaRoot on the driver, then re-serialize once. Each consumer task then deserializes * a single Arrow IPC stream. + * + * `parentAllocator` is only overridden by tests, to bound what a runaway append can take. */ def coalesceBroadcastBatches( - input: Iterator[ChunkedByteBuffer]): (Array[ChunkedByteBuffer], Long, Long) = { + input: Iterator[ChunkedByteBuffer], + parentAllocator: BufferAllocator = org.apache.comet.CometArrowAllocator) + : (Array[ChunkedByteBuffer], Long, Long) = { val buffers = input.filterNot(_.size == 0).toArray if (buffers.isEmpty) { return (Array.empty, 0L, 0L) } - val allocator = org.apache.comet.CometArrowAllocator + val allocator = parentAllocator .newChildAllocator("broadcast-coalesce", 0, Long.MaxValue) try { var targetRoot: VectorSchemaRoot = null @@ -358,19 +525,23 @@ 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` cannot grow a `NullVector` directly under a struct + // (see `hasNullDirectlyUnderStruct`). + 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()) { @@ -417,8 +588,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 cf2ac1ab0c7..294886a696b 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 @@ -183,8 +183,15 @@ 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 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) = + Utils.newArrowStreamWriter(declaredRoot, null, Channels.newChannel(dataOut)) + writeRoot = boundRoot + arrowWriter = writer arrowWriter.start() } @@ -239,7 +246,9 @@ private[python] trait CometArrowPythonRunnerBase // is identical. val childNames = inputStructType.fieldNames streamFields = batchFields.zipWithIndex.map { case (field, i) => - renamed(field, childNames(i), forceNullable = true) + // `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 683045743e3..0c5438002bb 100644 --- a/spark/src/test/resources/pyspark/test_pyarrow_udf.py +++ b/spark/src/test/resources/pyspark/test_pyarrow_udf.py @@ -1235,6 +1235,89 @@ 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); 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 + 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/aggregate/collect_list.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql index fa1ef356925..6d456102cce 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,12 @@ INSERT INTO cl_src_map VALUES query spark_answer_only SELECT size(collect_list(m)) FROM cl_src_map + +-- Without a grouping key the native final merge can see a nested element with the non-nullable +-- fields its producer declared, and fails the batch; the serde cannot see the grouping keys, so a +-- NullType-bearing element stays in Spark either way. +query expect_fallback(collect_list over a NullType-bearing input falls back to Spark) +SELECT grp, sort_array(collect_list(named_struct('i', i, 'n', NULL))) FROM cl_src_int GROUP BY grp + +query expect_fallback(collect_list over a NullType-bearing input falls back to Spark) +SELECT sort_array(collect_list(named_struct('i', i, 'n', NULL))) FROM cl_src_int diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_respect_nulls.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_respect_nulls.sql index 5218e020fc6..d7ba2376925 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_respect_nulls.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_respect_nulls.sql @@ -45,21 +45,21 @@ INSERT INTO crn_src VALUES -- collect_list RESPECT NULLS: nulls are kept (global aggregate) -- ============================================================ -query expect_fallback(collect_list with RESPECT NULLS (ignoreNulls = false) is not supported) +query expect_fallback(`collect_list` with `RESPECT NULLS` falls back to Spark) SELECT sort_array(collect_list(v) RESPECT NULLS) FROM crn_src -- ============================================================ -- collect_list RESPECT NULLS: per group, including an all-null group -- ============================================================ -query expect_fallback(collect_list with RESPECT NULLS (ignoreNulls = false) is not supported) +query expect_fallback(`collect_list` with `RESPECT NULLS` falls back to Spark) SELECT grp, sort_array(collect_list(v) RESPECT NULLS) FROM crn_src GROUP BY grp ORDER BY grp -- ============================================================ -- collect_set RESPECT NULLS: keeps a single null alongside distinct values -- ============================================================ -query expect_fallback(collect_set with RESPECT NULLS (ignoreNulls = false) is not supported) +query expect_fallback(`collect_set` with `RESPECT NULLS` falls back to Spark) SELECT grp, sort_array(collect_set(v) RESPECT NULLS) FROM crn_src GROUP BY grp ORDER 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 1fca2925de3..2e2400a90c7 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 @@ -261,3 +261,11 @@ 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 + +-- A NullType-bearing element runs natively, grouped or not. The order of collect_set is not +-- defined, so the result is sorted. +query +SELECT grp, sort_array(collect_set(named_struct('i', i, 'n', NULL))) FROM cs_src_int GROUP BY grp + +query +SELECT sort_array(collect_set(named_struct('i', i, 'n', NULL))) FROM cs_src_int diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_append_null_guard.sql b/spark/src/test/resources/sql-tests/expressions/array/array_append_null_guard.sql new file mode 100644 index 00000000000..d2ed0b43f23 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_append_null_guard.sql @@ -0,0 +1,41 @@ +-- 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. + +-- On Spark 4.0 and later array_append rewrites to array_insert, which has its own serde. +-- MaxSparkVersion: 3.5 + +statement +CREATE TABLE test_array_append_guard(arr array, val int) USING parquet + +statement +INSERT INTO test_array_append_guard VALUES (NULL, 1), (array(1), 2), (array(2), 3), (array(1, 2), NULL) + +-- The serde guards the array with CASE WHEN arr IS NOT NULL and serializes it twice, so a +-- non-deterministic array stays in Spark. +query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's) +SELECT array_append(IF(monotonically_increasing_id() % 2 = 0, arr, NULL), val) FROM test_array_append_guard + +-- The item is serialized once, inside the THEN branch, so native code evaluates it only for a +-- non-null array. Spark's generated code evaluates it on every row, so a stateful item would +-- read a different counter after the NULL array in the first row; it stays in Spark as well. +query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's) +SELECT array_append(arr, IF(monotonically_increasing_id() % 2 = 0, val, NULL)) FROM test_array_append_guard + +-- A non-nullable array selects every row, so the item is evaluated on every row in both engines +-- and a stateful item stays native. +query expect_native(array_append) +SELECT array_append(array(val), monotonically_increasing_id()) FROM test_array_append_guard 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 6ec4c275d11..bb4ad1caf1e 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 @@ -71,3 +71,17 @@ INSERT INTO test_except_float VALUES query SELECT a, b, array_except(a, b) FROM test_except_float + +-- Coverage for nested-list elements from two native producers, `map_keys` on one side and a +-- lambda on the other; the sides are cast to one deeply-nullable element type like the other set +-- ops. Not a regression witness: both producers already declare a nullable list item, so the +-- except kernel accepts the pair without the cast. A nested list rather than a struct, because +-- `isTypeSupported` declines struct elements before `convert` runs. +query expect_native(array_except) +SELECT array_except(array(map_keys(map(coalesce(b[0], 0), 1))), transform(array(coalesce(b[0], 0)), x -> array(x))) FROM test_array_except + +-- Struct elements have no native kernel. The branch reports Unsupported, which the JVM codegen +-- dispatcher runs even with this file's allowIncompatible=true (the Incompatible branch would +-- hand them to the kernel instead). +query expect_dispatch(array_except) +SELECT array_except(transform(a, x -> named_struct('i', x)), transform(b, x -> named_struct('i', 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 6279c104f34..056057a5c17 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 @@ -226,3 +226,21 @@ 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 two sides +-- of one Spark type can still reach it with different nested nullability: `map_entries` keeps +-- its key field non-nullable, while the JVM codegen dispatcher and the native constructors +-- declare every nested field nullable. Both sides are cast to a deeply-nullable element type +-- first; without that the kernel's type assertion fails the query. +query expect_native(array_intersect) +SELECT array_intersect(map_entries(map(coalesce(b[0], 0), 1)), transform(array(coalesce(b[0], 0)), x -> named_struct('key', x, 'value', 1))) 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 expect_dispatch(array_intersect) +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 58272d164f1..4262f6e2472 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,27 @@ 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 + +-- map_entries produces a list whose non-NullType item is declared non-nullable; native +-- 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 expect_dispatch(array_repeat) +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_dispatch(array_repeat) +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 expect_native(array_repeat) +SELECT array_repeat(array(coalesce(int_v, 0)), 2) FROM test_array_repeat + +query expect_native(array_repeat) +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_repeat_interpreted.sql b/spark/src/test/resources/sql-tests/expressions/array/array_repeat_interpreted.sql new file mode 100644 index 00000000000..92c1c3e9ae3 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_repeat_interpreted.sql @@ -0,0 +1,48 @@ +-- 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. + +-- Spark's generated code for `array_repeat` evaluates the element and the count, but its `eval` +-- evaluates the count first and skips the element when the count is NULL. Beneath an expression +-- with no generated code (here the `filter` that `array_compact` becomes) Spark uses `eval`, so +-- an element that raises an error must not raise on the rows with a NULL count. A nullable count +-- runs through the codegen dispatcher there, which evaluates the tree the same way. + +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE test_array_repeat_interpreted(id bigint) USING parquet + +-- One partition, so the four rows share one batch. +statement +INSERT INTO test_array_repeat_interpreted SELECT id FROM range(0, 4, 1, 1) + +-- Even rows have a NULL count and a zero divisor in the element. +query expect_dispatch(array_repeat) +SELECT id, array_compact(array_repeat(10 div (id % 2), IF(id % 2 = 0, CAST(NULL AS INT), 1))) FROM test_array_repeat_interpreted + +-- Outside such an expression Spark uses the generated code, which evaluates the element. +query expect_error(DIVIDE_BY_ZERO) +SELECT id, array_repeat(10 div (id % 2), IF(id % 2 = 0, CAST(NULL AS INT), 1)) FROM test_array_repeat_interpreted + +-- An imperative aggregate evaluates its input through `eval`, but its FILTER predicate is +-- generated code, so the element is evaluated there and raises, as in Spark. +query expect_error(DIVIDE_BY_ZERO) +SELECT sort_array(collect_list(id) FILTER (WHERE size(array_repeat(10 div (id % 2), IF(id % 2 = 0, CAST(NULL AS INT), 1))) > 0)) FROM test_array_repeat_interpreted + +-- The same shape over valid input, to show the aggregate and its FILTER run natively. +query expect_native(array_repeat) +SELECT sort_array(collect_list(id) FILTER (WHERE size(array_repeat(id, IF(id % 2 = 0, CAST(NULL AS INT), 1))) > 0)) FROM test_array_repeat_interpreted 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/array_union.sql b/spark/src/test/resources/sql-tests/expressions/array/array_union.sql index 27f11bc4abf..56aacd9e02e 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 @@ -50,7 +50,9 @@ query SELECT a, b, array_union(a, b) FROM test_union_nulls -- empty array combinations -query +-- 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 expect_dispatch(array_union) SELECT array_union(array(), array()) FROM test_union_nulls query @@ -59,7 +61,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_dispatch(array_union) SELECT array_union(array(), array(NULL)) FROM test_union_nulls -- both-NULL arrays @@ -227,3 +229,41 @@ 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 + +-- 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 run +-- through the JVM codegen dispatcher instead of the kernel. +query expect_dispatch(array_union) +SELECT array_union(transform(a, x -> NULL), array()) FROM test_array_union + +-- The set-op kernel asserts identical element types, nested nullability included, and two sides +-- of one Spark type can still reach it with different nested nullability: `map_entries` keeps +-- its key field non-nullable, while the JVM codegen dispatcher and the native constructors +-- declare every nested field nullable. Both sides are cast to a deeply-nullable element type +-- first; without that the kernel's type assertion fails the query. +query expect_native(array_union) +SELECT array_union(map_entries(map(coalesce(b[0], 0), 1)), transform(array(coalesce(b[0], 0)), x -> named_struct('key', x, 'value', 1))) FROM test_array_union + +-- With case-insensitive analysis Spark accepts struct sides whose field names differ only in +-- case, without a cast, and compares the structs by position. The set-op kernel's type +-- assertion compares field names too, so both sides are cast to the set op's own element type. +query expect_native(array_union) +SELECT array_union(array(named_struct('a', b[0])), array(named_struct('A', b[0]))) FROM test_array_union + +query expect_native(array_union) +SELECT array_union(transform(b, x -> named_struct('a', x, 'n', NULL)), transform(b, x -> named_struct('A', x, 'n', NULL))) FROM test_array_union + +-- Spark names a merged IF / CASE struct's fields after its first branch, a native CASE after +-- its ELSE branch. Here the IF side's Spark type already equals the set op's element type, so +-- it is cast anyway: its native type does not. +statement +CREATE TABLE test_union_branch_names(id bigint, s struct) USING parquet + +statement +INSERT INTO test_union_branch_names VALUES (0, named_struct('a', 1)), (1, named_struct('a', 2)), (2, NULL) + +query expect_native(array_union) +SELECT array_union(array(IF(id < 0, s, named_struct('A', id))), array(s)) FROM test_union_branch_names + +query expect_native(array_union) +SELECT array_union(array(CASE WHEN id < 0 THEN s ELSE named_struct('A', id) END), array(s)) FROM test_union_branch_names diff --git a/spark/src/test/resources/sql-tests/expressions/array/arrays_overlap.sql b/spark/src/test/resources/sql-tests/expressions/array/arrays_overlap.sql index 1ddfe14dbcd..39315c348d8 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/arrays_overlap.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/arrays_overlap.sql @@ -285,3 +285,16 @@ SELECT arrays_overlap(array(NULL, 99), b) FROM test_arrays_overlap -- conditional (CASE WHEN) arrays query SELECT arrays_overlap(CASE WHEN a IS NOT NULL THEN a ELSE array(0) END, b) FROM test_arrays_overlap + +-- A folded literal array reaches native as a scalar. Beside a column it has to be expanded to +-- the batch length: as a one-row list it ran past its end when the literal came second, and when +-- it came first only the first row was compared and that answer returned for the whole batch. +statement +CREATE TABLE test_overlap_batch(id bigint) USING parquet + +-- One file, so the rows share a batch. +statement +INSERT INTO test_overlap_batch SELECT id FROM range(0, 4, 1, 1) + +query +SELECT id, arrays_overlap(array(id), array(1L)), arrays_overlap(array(1L), array(id)) FROM test_overlap_batch 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..3d705b7fd48 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,16 @@ 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 + +-- No arguments: Spark returns an empty array per row; the native kernel needs at least one. +query expect_fallback(`arrays_zip()` with no arguments falls back to Spark) +SELECT arrays_zip() 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/arrays_zip_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip_ansi.sql new file mode 100644 index 00000000000..b29bfa55575 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip_ansi.sql @@ -0,0 +1,57 @@ +-- 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. + +-- Spark's generated code stops evaluating the arguments at the first NULL one, so an expression +-- that can only fail after it (here a division by zero under ANSI) must not fail natively either. + +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE test_arrays_zip_ansi(id bigint) USING parquet + +-- One partition, so the four rows share one file and one batch; with a row per batch the guard +-- would skip the whole batch of a NULL row and never reach the division. +statement +INSERT INTO test_arrays_zip_ansi SELECT id FROM range(0, 4, 1, 1) + +-- Even rows have a NULL first argument and a zero divisor in the second. +query expect_native(arrays_zip) +SELECT id, arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), array(10 div (id % 2))) FROM test_arrays_zip_ansi + +-- Spark evaluates the subtree of a CodegenFallback expression through `eval`, even inside +-- generated code; here that is the `filter` that `array_compact` becomes. The interpreted +-- `ArraysZip.eval` evaluates every argument before it looks for a NULL one, so the division +-- raises in Spark, and `arrays_zip` falls back to Spark under such an expression. +query expect_error(DIVIDE_BY_ZERO) +SELECT id, array_compact(arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), array(10 div (id % 2)))) FROM test_arrays_zip_ansi + +query expect_fallback(falls back to Spark where Spark evaluates it without generated code) +SELECT id, array_compact(arrays_zip(array(id), array(id + 1))) FROM test_arrays_zip_ansi + +-- An imperative aggregate (`collect_list`) evaluates its input through `eval` under every codegen +-- setting, so the division raises in Spark here too. +query expect_error(DIVIDE_BY_ZERO) +SELECT collect_list(arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), array(10 div (id % 2)))) FROM test_arrays_zip_ansi + +query expect_fallback(falls back to Spark where Spark evaluates it without generated code) +SELECT collect_list(arrays_zip(array(id), array(id + 1))) FROM test_arrays_zip_ansi + +-- A Generate whose expressions hold an expression with no generated code (here `transform`) is +-- left out of a whole-stage stage, and then evaluates its generator through `eval`, so the +-- `arrays_zip` under the `explode_outer` evaluates every argument even with codegen on. +query expect_error(DIVIDE_BY_ZERO) +SELECT id, explode_outer(arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), transform(array(id), x -> 10 div (x % 2)))) FROM test_arrays_zip_ansi diff --git a/spark/src/test/resources/sql-tests/expressions/array/arrays_zip_generate.sql b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip_generate.sql new file mode 100644 index 00000000000..9f4a4273807 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip_generate.sql @@ -0,0 +1,41 @@ +-- 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. + +-- Outside a whole-stage stage, GenerateExec evaluates its generator through `eval`, and so the +-- `arrays_zip` under it; projections still use generated code. + +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.sql.codegen.wholeStage=false + +statement +CREATE TABLE test_arrays_zip_generate(id bigint) USING parquet + +-- One partition, so the four rows share one batch. +statement +INSERT INTO test_arrays_zip_generate SELECT id FROM range(0, 4, 1, 1) + +-- Even rows have a NULL first argument and a zero divisor in the second. +query expect_error(DIVIDE_BY_ZERO) +SELECT id, explode(arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), array(10 div (id % 2)))) FROM test_arrays_zip_generate + +query expect_fallback(falls back to Spark where Spark evaluates it without generated code) +SELECT id, explode(arrays_zip(array(id), array(id + 1))) FROM test_arrays_zip_generate + +-- A projection is still generated code without whole-stage codegen, so it stays native and the +-- division is skipped on the rows with a NULL first argument, as in Spark. +query expect_native(arrays_zip) +SELECT id, arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), array(10 div (id % 2))) FROM test_arrays_zip_generate diff --git a/spark/src/test/resources/sql-tests/expressions/array/arrays_zip_interpreted.sql b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip_interpreted.sql new file mode 100644 index 00000000000..93cbfcd888e --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/arrays_zip_interpreted.sql @@ -0,0 +1,53 @@ +-- 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. + +-- Without generated code, Spark evaluates `arrays_zip` through its interpreted `eval`, which +-- evaluates every argument before it looks for a NULL one. The native null guard follows the +-- generated code, which stops at the first NULL argument, so `arrays_zip` over several +-- arguments falls back to Spark here; otherwise the division below would raise in Spark only. + +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.sql.codegen.wholeStage=false +-- Config: spark.sql.codegen.factoryMode=NO_CODEGEN + +statement +CREATE TABLE test_arrays_zip_interpreted(id bigint) USING parquet + +-- One partition, so the four rows share one batch. +statement +INSERT INTO test_arrays_zip_interpreted SELECT id FROM range(0, 4, 1, 1) + +-- Even rows have a NULL first argument and a zero divisor in the second. +query expect_error(DIVIDE_BY_ZERO) +SELECT id, arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), array(10 div (id % 2))) FROM test_arrays_zip_interpreted + +query expect_fallback(falls back to Spark where Spark evaluates it without generated code) +SELECT id, arrays_zip(array(id), array(id + 1)) FROM test_arrays_zip_interpreted + +-- A single argument has nothing to skip, so it stays native. +query expect_native(arrays_zip) +SELECT id, arrays_zip(array(id)) FROM test_arrays_zip_interpreted + +-- A NullType-bearing map runs through the codegen dispatcher as a whole, so the `arrays_zip` +-- inside it never reaches its own gate. The dispatcher's kernel evaluates the tree through +-- `eval` here, as Spark does, so the division raises in both. +query expect_error(DIVIDE_BY_ZERO) +SELECT id, map(0, arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), array(10 div (id % 2)), array(NULL))) FROM test_arrays_zip_interpreted + +-- The same shape over valid input, to show the map really is dispatched. +query expect_dispatch(map) +SELECT id, map(0, arrays_zip(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), array(id), array(NULL))) FROM test_arrays_zip_interpreted 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..aa616e2e880 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,26 @@ 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 the whole expression runs +-- through the dispatcher instead. +query expect_dispatch(array) +SELECT array(aggregate(arr, NULL, (acc, x) -> NULL)) FROM test_create_array_complex + +-- A foldable NullType argument is not necessarily a scalar at native: with constant folding +-- excluded (as in these fixtures), `element_at(array(NULL), 1)` and a dispatched `map(...)` chain +-- arrive as one value per row, and make_array would return one list row for the whole batch. Only +-- a literal is guaranteed to arrive as a scalar, so anything else goes through the dispatcher. +statement +CREATE TABLE test_create_array_batch(id bigint) USING parquet + +-- One file, so the rows share a batch. +statement +INSERT INTO test_create_array_batch SELECT id FROM range(0, 4, 1, 1) + +query expect_dispatch(array) +SELECT id, array(element_at(array(NULL), 1)) FROM test_create_array_batch + +query expect_dispatch(array) +SELECT id, array(map_values(map(1, NULL))[0]) FROM test_create_array_batch 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 89293ccefd9..5709fc0ccf0 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 @@ -124,3 +124,9 @@ 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 + +-- A non-deterministic collection built by a lambda, which the JVM codegen dispatcher would +-- evaluate, is declined the same way. + +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/routing_arrays_opt_in.sql b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_opt_in.sql index 64038315a83..62f3c0b28d2 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_opt_in.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_opt_in.sql @@ -56,7 +56,9 @@ SELECT sort_array(array(make_dt_interval(x), make_dt_interval(y))) FROM routing_ query expect_native(array_except) SELECT sort_array(array_except(a, array(1))) FROM routing_arrays -query expect_fallback(data type not supported) +-- A binary element has no native kernel; the serde reports it Unsupported, and with the +-- dispatcher disabled the expression falls back. +query expect_fallback(native array_except supports only) SELECT array_except(array(unhex('41')), array(unhex('42'))) query expect_fallback(child data type not supported) diff --git a/spark/src/test/resources/sql-tests/expressions/array/size.sql b/spark/src/test/resources/sql-tests/expressions/array/size.sql index fb2bedef1e3..7e4396d88f4 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/size.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/size.sql @@ -41,3 +41,12 @@ SELECT size(cast(NULL as map)) -- cardinality is a SQL alias for size query SELECT cardinality(arr), cardinality(m) FROM test_size + +-- Without the legacy behavior, the serde guards a nullable argument with CASE WHEN arg IS NOT +-- NULL and serializes it twice, so a non-deterministic one stays in Spark. The legacy path builds +-- no guard (native size already returns -1 for NULL). This SET applies to both matrix runs. +statement +SET spark.sql.legacy.sizeOfNull=false + +query expect_fallback(non-deterministic child under a null guard is evaluated on different rows than Spark's) +SELECT size(IF(monotonically_increasing_id() % 2 = 0, arr, NULL)) FROM test_size 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..6ce58ecdd65 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,14 @@ 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 + +-- map_entries produces a list whose NullType-bearing struct item is declared non-nullable; +-- 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 expect_native(slice) +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 46490917619..340cf23dc1e 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,10 @@ 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) + +query +SELECT transform(array(NULL), 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/conditional/case_when.sql b/spark/src/test/resources/sql-tests/expressions/conditional/case_when.sql index 508571d728d..93d153a7a3f 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,7 @@ 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 whose rows come from both the THEN and the implicit ELSE runs natively. +query expect_native(casewhen) +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..e95e9d5f553 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,19 @@ 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. The JVM +-- codegen dispatcher evaluates Spark's own code once per row, so it runs there instead. +query expect_dispatch(coalesce) +SELECT coalesce(IF(monotonically_increasing_id() % 2 = 0, a, NULL), b, 0) FROM test_coalesce + +-- A non-deterministic last argument is the ELSE, serialized and evaluated once, so it stays on +-- the native CASE. +query expect_native(coalesce) +SELECT coalesce(a, monotonically_increasing_id()) FROM test_coalesce + +-- A NullType result runs on the native CASE the serde builds. +query expect_native(coalesce) +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..0e7d2054c95 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,54 @@ 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 whose rows come from both branches runs natively. Spark folds +-- `IF(c, NULL, NULL)` itself, so the branch has to be a non-foldable NullType expression. +query expect_native(if) +SELECT IF(cond, aggregate(array(a), NULL, (acc, x) -> NULL), NULL) FROM test_if + +-- Both branches have the same Spark type (Spark's IF ignores nullability when it compares them), +-- but native `named_struct` declares a field built from a literal non-nullable and one built from +-- a column nullable. When no row takes the THEN branch, the result is the ELSE as it is unless the +-- planner coerces the branches to one type, which it does for IF and CASE WHEN alike. +statement +CREATE TABLE test_if_branch_types(id bigint) USING parquet + +-- One file, so the rows share a batch. +statement +INSERT INTO test_if_branch_types SELECT id FROM range(0, 4, 1, 1) + +query expect_native(if) +SELECT IF(id = -1, named_struct('i', 1L, 'n', NULL), named_struct('i', id, 'n', NULL)) FROM test_if_branch_types + +query expect_native(if) +SELECT IF(id = -1, named_struct('i', 1L), named_struct('i', id)) FROM test_if_branch_types + +query expect_native(casewhen) +SELECT CASE WHEN id = -1 THEN named_struct('i', 1L, 'n', NULL) ELSE named_struct('i', id, 'n', NULL) END FROM test_if_branch_types + +-- Spark names a merged struct's fields after the first branch, a native CASE after its ELSE +-- branch, and with case-insensitive analysis the two can differ. Each CASE WHEN or coalesce +-- branch is cast to the expression's own type (native IF already takes the THEN branch's +-- names), so a consumer that compares types exactly (here `array(...)`) sees Spark's names on +-- both of its arguments. +statement +CREATE TABLE test_if_branch_names(id bigint, s struct) USING parquet + +statement +INSERT INTO test_if_branch_names SELECT id, named_struct('a', id + 10) FROM range(0, 4, 1, 1) + +query expect_native(if) +SELECT array(IF(id >= 0, s, named_struct('A', id)), s) FROM test_if_branch_names + +query expect_native(casewhen) +SELECT array(CASE WHEN id >= 0 THEN s ELSE named_struct('A', id) END, s) FROM test_if_branch_names + +query expect_native(coalesce) +SELECT array(coalesce(s, named_struct('A', id)), s) FROM test_if_branch_names + +-- `array_append` (`array_insert` on Spark 4.0+) widens its array and item to one element type +-- only when their Spark types differ; here they match, so a CASE item named after its ELSE +-- branch reached the kernel's exact type check. +query expect_native(casewhen) +SELECT array_append(array(s), CASE WHEN id >= 0 THEN named_struct('a', id) ELSE named_struct('A', id) END) FROM test_if_branch_names diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_interpreted.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_interpreted.sql new file mode 100644 index 00000000000..fba8ddd4492 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_interpreted.sql @@ -0,0 +1,40 @@ +-- 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. + +-- Spark's `eval` for `trunc` and `date_trunc` reads a non-literal format first and returns NULL +-- for an invalid one without evaluating the date or timestamp, while its generated code, like the +-- native kernel, evaluates both. An imperative aggregate (`collect_list`) evaluates its input +-- through `eval`, so there the invalid value on the row with an invalid format must not raise. +-- The expression runs through the codegen dispatcher in that context, which evaluates it the +-- same way. + +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.comet.expression.TruncDate.allowIncompatible=true +-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true + +statement +CREATE TABLE test_trunc_interpreted(s string, fmt string) USING parquet + +-- One partition, so the rows share one batch. +statement +INSERT INTO test_trunc_interpreted SELECT * FROM VALUES ('not-a-date', 'bogus'), ('2024-05-17', 'year') AS t(s, fmt) DISTRIBUTE BY 1 + +query expect_dispatch(trunc) +SELECT collect_list(trunc(CAST(s AS date), fmt)) FROM test_trunc_interpreted + +query expect_dispatch(date_trunc) +SELECT collect_list(date_trunc(fmt, CAST(s AS timestamp))) FROM test_trunc_interpreted diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_interpreted_collation.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_interpreted_collation.sql new file mode 100644 index 00000000000..c981cf126ce --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_interpreted_collation.sql @@ -0,0 +1,33 @@ +-- 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. + +-- The interpreted-evaluation check for `trunc` (see trunc_interpreted.sql) runs ahead of the +-- collation opt-in, so a collated non-literal format cannot route the expression to the native +-- kernel through `allowIncompatible` where Spark evaluates it through `eval`. + +-- MinSparkVersion: 4.0 +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.comet.expression.TruncDate.allowIncompatible=true + +statement +CREATE TABLE test_trunc_interpreted_collation(s string, fmt string) USING parquet + +statement +INSERT INTO test_trunc_interpreted_collation SELECT * FROM VALUES ('not-a-date', 'bogus'), ('2024-05-17', 'year') AS t(s, fmt) DISTRIBUTE BY 1 + +query expect_dispatch(trunc) +SELECT collect_list(trunc(CAST(s AS date), CAST(fmt AS STRING COLLATE UTF8_LCASE))) FROM test_trunc_interpreted_collation 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 cbf90b18459..a241f5b4d4e 100644 --- a/spark/src/test/resources/sql-tests/expressions/hash/hash.sql +++ b/spark/src/test/resources/sql-tests/expressions/hash/hash.sql @@ -53,3 +53,10 @@ SELECT hash(array(-d)), xxhash64(array(-d)), hash(named_struct('a', -d, 'b', -f) -- so the null row is filtered out. query SELECT hash(map(-d, -f)), xxhash64(map(-d, -f)) FROM test_nan WHERE d IS NOT NULL + +-- 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)) FROM test + +query expect_fallback(`NullType` is not supported) +SELECT 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 8e479412453..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 @@ -39,3 +39,31 @@ 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 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()) + +-- 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/create_map_interpreted.sql b/spark/src/test/resources/sql-tests/expressions/map/create_map_interpreted.sql new file mode 100644 index 00000000000..cf0b74b7d04 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/create_map_interpreted.sql @@ -0,0 +1,38 @@ +-- 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. + +-- `element_at` on an in-bounds literal index is non-nullable, yet with ANSI off the code Spark +-- generates for it assigns an undeclared null flag, so it does not compile. Spark never falls +-- back while testing, so these queries run without generated code (NO_CODEGEN, whole-stage off) +-- for Spark to have an answer; the dispatcher then evaluates its kernel through `eval` from the +-- start, as Spark does. The compile-failure retry itself is covered in CometCodegenSuite. + +-- Config: spark.sql.codegen.wholeStage=false +-- Config: spark.sql.codegen.factoryMode=NO_CODEGEN + +statement +CREATE TABLE test_create_map_interpreted(id bigint) USING parquet + +-- One file, so the rows share a batch. +statement +INSERT INTO test_create_map_interpreted SELECT id FROM range(0, 4, 1, 1) + +query expect_dispatch(map) +SELECT id, map(0, named_struct('i', id), 1, element_at(array(named_struct('i', id)), 1)) FROM test_create_map_interpreted + +query expect_dispatch(map) +SELECT id, map(0, named_struct('i', id, 'n', NULL), 1, element_at(array(named_struct('i', id, 'n', NULL)), 1)) FROM test_create_map_interpreted 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_entries.sql b/spark/src/test/resources/sql-tests/expressions/map/map_entries.sql index 2d172f436f8..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 @@ -63,3 +63,8 @@ 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 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_from_arrays.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql index 178c07f432a..427e092aa55 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,17 @@ 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()) + +-- 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: the map kernel expands the scalar list to every row. +query expect_native(map_from_arrays) +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/map/map_from_arrays_ansi.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_ansi.sql new file mode 100644 index 00000000000..ce20db42f30 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_ansi.sql @@ -0,0 +1,35 @@ +-- 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. + +-- Spark does not evaluate the values of a row whose keys are NULL, so an expression that can +-- only fail there (here a division by zero under ANSI) must not fail natively either. +-- `CometMapExpressionSuite` covers native values; this covers values produced by a lambda the JVM +-- codegen dispatcher evaluates, whose struct result carries a NullType field. + +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE test_map_from_arrays_ansi(id bigint) USING parquet + +-- One partition, so the four rows share one file and one batch; with a row per batch the guard +-- would skip the whole batch of a NULL-key row and never reach the division. +statement +INSERT INTO test_map_from_arrays_ansi SELECT id FROM range(0, 4, 1, 1) + +-- Even rows have NULL keys and a zero divisor; odd rows have both. +query expect_native(map_from_arrays) +SELECT id, map_from_arrays(IF(id % 2 = 0, CAST(NULL AS ARRAY), array(id)), transform(array(id), x -> named_struct('q', 10 div (x % 2), 'n', NULL))) FROM test_map_from_arrays_ansi diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_last_win_opt_in.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_last_win_opt_in.sql new file mode 100644 index 00000000000..05fb9c3df91 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_last_win_opt_in.sql @@ -0,0 +1,43 @@ +-- 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. + +-- Under LAST_WIN the dedup difference is an opt-in (`allowIncompatible`), and the opt-in hands the +-- expression straight to the native kernel. The kernel's own gate (a stateful input under the +-- null guard) has nothing to do with dedup, so it must still refuse that shape with the opt-in on. + +-- Config: spark.sql.mapKeyDedupPolicy=LAST_WIN +-- Config: spark.comet.expression.MapFromArrays.allowIncompatible=true + +statement +CREATE TABLE test_map_from_arrays_lw(id bigint, k array, v array) USING parquet + +statement +INSERT INTO test_map_from_arrays_lw VALUES + (0, array(0), array(10)), (1, array(1), array(11)), (2, array(2), array(12)), (3, array(3), NULL) + +-- No duplicate keys, so the opt-in takes the native kernel. +query expect_native(map_from_arrays) +SELECT map_from_arrays(k, v) FROM test_map_from_arrays_lw + +-- The null guard would evaluate the stateful key array twice, so it stays in Spark, where row 2 +-- keeps {2: NULL}. +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, array(id), NULL), transform(array(id), x -> NULL)) FROM test_map_from_arrays_lw + +-- A literal array beside a per-row one takes the native kernel too. +query expect_native(map_from_arrays) +SELECT map_from_arrays(k, array(1)) FROM test_map_from_arrays_lw 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/resources/sql-tests/expressions/misc/nulltype_stateful_argument.sql b/spark/src/test/resources/sql-tests/expressions/misc/nulltype_stateful_argument.sql new file mode 100644 index 00000000000..aa697219277 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/nulltype_stateful_argument.sql @@ -0,0 +1,109 @@ +-- 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. + +-- Spark's generated code skips an argument on rows another argument decides (a NULL operand, a +-- NULL divisor), and native evaluation computes every argument over the whole batch, so a stateful +-- expression advances on rows Spark skips. An operator whose native plan runs a codegen dispatcher +-- kernel with a NullType result used to fall back to Spark, together with the native operators +-- above it up to a shuffle, and stays there when it evaluates a non-deterministic expression +-- outside the dispatcher. + +statement +CREATE TABLE test_nulltype_stateful(id bigint) USING parquet + +-- One file, so the rows share a partition and the counter. +statement +INSERT INTO test_nulltype_stateful SELECT id FROM range(0, 4, 1, 1) + +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, arrays_overlap(IF(id = 0, CAST(NULL AS ARRAY), transform(array(id), x -> NULL)), IF(monotonically_increasing_id() = 0, array(), array(NULL))) FROM test_nulltype_stateful + +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, arrays_overlap(IF(id = 0, CAST(NULL AS ARRAY), transform(array(id), x -> NULL)), array_repeat(NULL, CAST(monotonically_increasing_id() AS INT))) FROM test_nulltype_stateful + +-- The NullType value and the stateful argument in different columns of one projection. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, map(id, NULL), IF(id = 0, NULL, id) + monotonically_increasing_id() FROM test_nulltype_stateful + +-- The NullType value computed by a native operator below the one with the stateful argument, in +-- the same native block (SORT BY keeps the two projections apart without an exchange). +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, m, IF(id = 0, NULL, id) + monotonically_increasing_id() FROM (SELECT id, map(id, NULL) AS m FROM test_nulltype_stateful SORT BY id) + +-- The NullType value computed below a broadcast or a union, which exist only over a native plan. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT /*+ BROADCAST(b) */ t.id, IF(t.id = 0, NULL, t.id) + monotonically_increasing_id(), b.a FROM test_nulltype_stateful t JOIN (SELECT id, transform(array(id), x -> NULL) AS a FROM test_nulltype_stateful) b ON t.id = b.id + +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, IF(id = 0, NULL, id) + monotonically_increasing_id(), a FROM (SELECT id, transform(array(id), x -> NULL) AS a FROM test_nulltype_stateful UNION ALL SELECT id, transform(array(id), x -> NULL) FROM test_nulltype_stateful) + +-- Spark evaluates a divisor first and skips the dividend when the divisor is NULL. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, map(id, NULL), monotonically_increasing_id() DIV IF(id = 0, CAST(NULL AS BIGINT), 1L) FROM test_nulltype_stateful + +-- ln serializes its argument twice, for its domain check. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, map(id, NULL), ln(CAST(monotonically_increasing_id() AS DOUBLE)) FROM test_nulltype_stateful + +-- ORDER BY ... LIMIT becomes TakeOrderedAndProject, which builds its projection natively when it +-- runs. The optimizer collapses the second query's two projections into that projection too. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, map(id, NULL), IF(id = 0, CAST(NULL AS BIGINT), id) + monotonically_increasing_id() FROM (SELECT id FROM test_nulltype_stateful ORDER BY id) LIMIT 4 + +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, m, IF(id = 0, CAST(NULL AS BIGINT), id) + monotonically_increasing_id() FROM (SELECT id, map(id, NULL) AS m FROM test_nulltype_stateful ORDER BY id) LIMIT 4 + +-- The kernel in TakeOrderedAndProject's sort keys, and in the projection of one below the operator. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, IF(id = 0, CAST(NULL AS BIGINT), id) + monotonically_increasing_id() FROM (SELECT id FROM test_nulltype_stateful ORDER BY size(transform(array(id), x -> NULL)), id) LIMIT 4 + +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, m, IF(id = 0, CAST(NULL AS BIGINT), id) + monotonically_increasing_id() FROM (SELECT id, map(id, NULL) AS m FROM (SELECT id FROM test_nulltype_stateful ORDER BY id) LIMIT 4) + +-- A kernel with a typed result can run a NullType value inside it: the whole coalesce goes to the +-- dispatcher (its guarded argument is non-deterministic), and holds a transform main refused. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, coalesce(IF(monotonically_increasing_id() >= 0, size(transform(array(id), x -> NULL)), NULL), 0) AS a, IF(id = 0, CAST(NULL AS BIGINT), id) + monotonically_increasing_id() AS v FROM test_nulltype_stateful + +-- The kernel as the key of a shuffle below the operator: main refused the shuffle, so the +-- operator above it ran in Spark. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, IF(id = 0, NULL, id) + monotonically_increasing_id() FROM (SELECT id FROM test_nulltype_stateful DISTRIBUTE BY transform(array(id), x -> NULL)) + +-- A non-deterministic column the dispatcher runs as a whole evaluates Spark's own generated code, +-- so it stays native. +query +SELECT id, map(monotonically_increasing_id(), NULL) FROM test_nulltype_stateful + +-- Any other non-deterministic expression falls back, even one native evaluation gets right: main +-- ran all of these in Spark. +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, map(id, NULL), IF(id = 0, NULL, monotonically_increasing_id()) FROM test_nulltype_stateful + +-- A NULL literal is not a dispatcher kernel, so this ran natively on main and still does. +query +SELECT id, NULL AS n, 1L + monotonically_increasing_id() FROM test_nulltype_stateful + +-- A cast that only relabels nested nullability runs natively, but main ran this one (an +-- array field) through the dispatcher, which evaluates the child as Spark's generated code +-- does. A non-deterministic child keeps that route. +query expect_dispatch(cast) +SELECT id, CAST(named_struct('a', array(DATE '2020-01-01'), 'b', IF(id = 0, NULL, id) + monotonically_increasing_id()) AS struct, b: bigint>) FROM test_nulltype_stateful + +-- The partition id does not depend on the rows evaluated before it. +query +SELECT id, array_union(array_repeat(NULL, CAST(id AS INT)), array_repeat(NULL, CAST(id AS INT))), spark_partition_id() FROM test_nulltype_stateful diff --git a/spark/src/test/resources/sql-tests/expressions/misc/nulltype_stateful_shuffle.sql b/spark/src/test/resources/sql-tests/expressions/misc/nulltype_stateful_shuffle.sql new file mode 100644 index 00000000000..11c35f8aff3 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/nulltype_stateful_shuffle.sql @@ -0,0 +1,33 @@ +-- 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. + +-- MinSparkVersion: 4.0 + +-- shuffle keeps its own random state whatever its argument holds. Its argument here is +-- non-deterministic only through spark_partition_id(), which holds no row state, and the operator +-- runs a codegen dispatcher kernel with a NullType result, so it falls back to Spark as on main. +-- shuffle takes a seed from Spark 4.0. + +statement +CREATE TABLE test_nulltype_stateful_shuffle(id bigint) USING parquet + +-- One file, so the rows share a partition and the random state. +statement +INSERT INTO test_nulltype_stateful_shuffle SELECT id FROM range(0, 64, 1, 1) + +query expect_fallback(a non-deterministic expression that does not run entirely in the codegen dispatcher) +SELECT id, IF(id = 0, CAST(NULL AS BIGINT), id) + element_at(shuffle(array(id, CAST(spark_partition_id() AS BIGINT), id + 10), 42), 1) AS v, transform(array(id), x -> NULL) AS nulls FROM test_nulltype_stateful_shuffle 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..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 @@ -33,3 +33,11 @@ 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', 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 5aabafe4022..da4889074b9 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -569,6 +569,40 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + test("array_union: a folded struct literal side falls back with its reason") { + // With constant folding on, `array(named_struct('i', 1L, 'n', NULL))` becomes a struct-element + // Literal, which Comet cannot serialize. It sits under the cast `SetOpArguments` adds, which is + // not in the operator's tree, so its reason used to reach no node the roll-up could see + // (the strict fallback check in these suites failed the query instead of falling back). + withTable("union_struct_lit") { + spark.range(0, 4, 1, 1).toDF("id").write.mode("overwrite").saveAsTable("union_struct_lit") + checkSparkAnswerAndFallbackReason( + "SELECT array_union(array(named_struct('i', id, 'n', NULL)), " + + "array(named_struct('i', 1L, 'n', NULL))) FROM union_struct_lit", + "Unsupported data type") + } + } + + test("nested array literal whose inner arrays are all empty") { + // Constant folding turns `array(cast(array() AS array))` into a nested list literal + // with no values, which native built one list level deeper than its declared type. Consumers + // that compare their inputs' types then failed. A deeper literal that mixes an empty child + // with a populated one failed to build at all. + withTable("empty_nested_lit") { + spark.range(0, 4, 1, 1).toDF("id").write.saveAsTable("empty_nested_lit") + Seq( + "arrays_overlap(transform(array(id), x -> array()), array(array()))", + "arrays_overlap(array(array(id)), array(cast(array() AS array)))", + "array_union(array(array(id)), array(cast(array() AS array)))", + "IF(id > 1, array(array(id)), array(cast(array() AS array)))", + "coalesce(IF(id > 1, array(array(id)), NULL), array(cast(array() AS array)))", + "array(array(id)) = array(cast(array() AS array))", + "array(array(array()), array(array(array(1))))").foreach { e => + checkSparkAnswerAndOperator(s"SELECT id, $e FROM empty_nested_lit") + } + } + } + test("array_max") { Seq(true, false).foreach { dictionaryEnabled => withTempDir { dir => @@ -1293,6 +1327,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 THEN copy would count over fewer rows than Spark's. + 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 400ec6fcd2f..dce8e01d6bb 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 @@ -287,6 +290,203 @@ 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("canHandle rejects duplicate struct field names, in outputs and inputs") { + // 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)))) + 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")), + 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("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 + // 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") { + // 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, diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index fd3384a3527..f224b94b680 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.{SparkConf, SparkEnv, TaskContext} import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.api.java.UDF1 import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, ElementAt, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF} import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.util.GenericArrayData @@ -43,6 +43,7 @@ import org.apache.spark.unsafe.types.{ByteArray, UTF8String} import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec +import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.{CometInvokeTargets, CometScalaUDF, QueryPlanSerde} import org.apache.comet.serde.ExprOuterClass.Expr.ExprStructCase import org.apache.comet.udf.codegen.CometScalaUDFCodegen @@ -267,6 +268,37 @@ class CometCodegenSuite } } + test("a kernel whose generated code does not compile evaluates the expression instead") { + // `element_at` over an in-bounds literal index is non-nullable, yet with ANSI off the code + // Spark generates for it assigns an undeclared null flag, so it does not compile. The kernel + // retries over the expression's `eval`, as Spark's whole-stage codegen falls back. + withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { + val expr = + ElementAt(CreateArray(Seq(BoundReference(0, LongType, nullable = false))), Literal(1)) + assert(!expr.nullable) + val input = new BigIntVector("in", CometArrowAllocator) + val field = CometBatchKernelCodegen.toFfiArrowField("out", LongType, nullable = false) + val output = CometBatchKernelCodegen.allocateOutput(field, 2, 0) + try { + input.allocateNew() + input.setSafe(0, 7L) + input.setSafe(1, -3L) + input.setValueCount(2) + val spec = ArrowColumnSpec(classOf[BigIntVector], nullable = false) + val kernel = CometBatchKernelCodegen.compile(expr, IndexedSeq(spec)).newInstance() + kernel.init(0) + kernel.process(Array(input), output, 2) + output.setValueCount(2) + val result = output.asInstanceOf[BigIntVector] + assert(result.get(0) == 7L) + assert(result.get(1) == -3L) + } finally { + output.close() + input.close() + } + } + } + test("ScalaUDF over concat(c1, c2) suppresses the null short-circuit") { // Concat is not NullIntolerant. The dispatcher's short-circuit guard inspects every node in // the bound tree and must skip the whole-tree null short-circuit because one child is @@ -912,6 +944,106 @@ class CometCodegenSuite } } + test("identical non-deterministic dispatched expressions keep their own state") { + // Two identical occurrences serialize to the same bytes. Spark gives each its own counter, + // so both columns are 0..n-1; sharing one cached kernel made the second continue the first + // (n..2n-1). The coalesce goes to the dispatcher because its guarded argument is + // non-deterministic (`NullGuard`), and a map with a NullType value always does. Neither + // carries an expression ID, which would make the two occurrences' bytes differ anyway, as a + // ScalaUDF's encoders or a lambda variable do. Batches of 8 make each occurrence carry its + // counter across batches, which per-batch state would not. + withTempPath { dir => + spark.range(0, 64, 1, numPartitions = 1).write.parquet(dir.getCanonicalPath) + withTable("t") { + sql(s"CREATE TABLE t USING parquet LOCATION '${dir.getCanonicalPath}'") + val guarded = "coalesce(IF(id >= 0, monotonically_increasing_id(), " + + "CAST(NULL AS BIGINT)), -1L)" + Seq( + s"SELECT id, $guarded AS a, $guarded AS b FROM t", + "SELECT id, map(monotonically_increasing_id(), NULL) AS a, " + + "map(monotonically_increasing_id(), NULL) AS b FROM t").foreach { q => + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "8") { + assertCodegenRan { + checkSparkAnswerAndOperator(sql(q)) + } + } + } + } + } + } + + test("a non-deterministic user function stays in Spark") { + // Spark keeps such a function's state in one object every call references and advances it + // row by row. Kernels per expression either copy the object (two calls count from zero + // each) or share it across columns batch by batch, so the expression falls back. The Java + // UDF has no encoders, so identical calls serialize to identical bytes; the queries cover + // identical calls, different calls, a call inside a NullType map, and calls beside + // monotonically_increasing_id. + spark.udf.register( + "nextCount", + org.apache.spark.sql.functions + .udf(new CountingJavaUdf, org.apache.spark.sql.types.LongType) + .asNondeterministic()) + withTempPath { dir => + spark.range(0, 8, 1, numPartitions = 1).write.parquet(dir.getCanonicalPath) + withTable("t") { + sql(s"CREATE TABLE t USING parquet LOCATION '${dir.getCanonicalPath}'") + val guarded = "coalesce(IF(id >= 0, monotonically_increasing_id(), " + + "CAST(NULL AS BIGINT)), nextCount('x'))" + val keyed = "map(monotonically_increasing_id() + nextCount(CAST(id AS STRING)), NULL)" + val counted = "nextCount(CAST(monotonically_increasing_id() AS STRING))" + Seq( + "SELECT id, nextCount('x') AS a, nextCount('x') AS b FROM t", + "SELECT id, nextCount('x') AS a, nextCount('y') AS b FROM t", + "SELECT id, map(nextCount('x'), NULL) AS a, map(nextCount('y'), NULL) AS b FROM t", + s"SELECT id, $guarded AS a, $guarded AS b FROM t", + s"SELECT id, $keyed AS a, $keyed AS b FROM t", + s"SELECT id, $counted AS a, $counted AS b FROM t", + // A reflected method can keep state in static fields Spark shares the same way; this + // one does not, so the answers can be compared, but it is refused all the same. + "SELECT id, map(reflect('java.lang.String', 'valueOf', id), NULL) AS a FROM t") + .foreach { q => + checkSparkAnswerAndFallbackReason(q, CometBatchKernelCodegen.userStateReason) + } + } + } + } + + test("a subquery inside a dispatched expression stays in Spark") { + // The plan is serialized before the operator starts its subqueries, so a dispatched copy + // would never see the result. A subquery beside the dispatched expression is unaffected + // (see "ScalaUDF composed with reused scalar subquery across projection and filter"). + spark.udf.register("plusOne", (x: Long) => x + 1) + withSubjects("only") { + Seq( + "SELECT plusOne((SELECT max(id) FROM range(0, 2, 1, 1))) FROM t", + "SELECT map((SELECT max(id) FROM range(0, 2, 1, 1)), NULL) FROM t").foreach { q => + checkSparkAnswerAndFallbackReason(q, CometBatchKernelCodegen.subqueryReason) + } + } + } + + test("a non-deterministic expression beside a NullType dispatcher result stays in Spark") { + // Spark skips the right operand of `+` on a row whose left operand is NULL, while native + // evaluation computes it over the whole batch, so the counter advances on a row Spark skips. + // The projection fell back to Spark when the dispatcher refused NullType results. Here + // serialization dispatches a tree it builds itself: the folded `map('k', NULL)` literal is + // expanded into a `CreateMap`, and the decimal addition is wrapped in `CheckOverflow`. + withTempPath { dir => + spark.range(0, 4, 1, numPartitions = 1).write.parquet(dir.getCanonicalPath) + withTable("t") { + sql(s"CREATE TABLE t USING parquet LOCATION '${dir.getCanonicalPath}'") + val stateful = "IF(id = 0, CAST(NULL AS BIGINT), id) + monotonically_increasing_id()" + Seq( + s"SELECT id, map('k', NULL) AS m, $stateful AS v FROM t", + "SELECT id, map(CAST(id AS DECIMAL(10, 0)) + CAST(1 AS DECIMAL(10, 0)), NULL) AS m, " + + s"$stateful AS v FROM t").foreach { q => + checkSparkAnswerAndFallbackReason(q, CometExecRule.nondeterministicBesideNullTypeReason) + } + } + } + } + test("per-task cache isolates UDF state across sequential task runs in one session") { // Regression guard for the cache-scoping invariant on CometUdfBridge: instances live for // exactly one Spark task and are dropped on task completion, so a stateful kernel sees a @@ -1340,13 +1472,12 @@ class CometCodegenSuite } test("zero-column ScalaUDF produces one row per input row") { - // Non-deterministic (so Spark doesn't constant-fold) with a deterministic body (so - // Spark-vs-Comet comparison stays honest). The expression has no `AttributeReference`, - // so the serde produces an empty data-arg list and the dispatcher has no data column to - // read the batch size from. Guards the `numRows` path through the JNI bridge. + // A ScalaUDF is never constant-folded, so the call reaches Comet. The expression has no + // `AttributeReference`, so the serde produces an empty data-arg list and the dispatcher has + // no data column to read the batch size from. Guards the `numRows` path through the JNI + // bridge. import org.apache.spark.sql.functions.udf - val alwaysHello = udf(() => "hello").asNondeterministic() - spark.udf.register("helloU", alwaysHello) + spark.udf.register("helloU", udf(() => "hello")) withSubjects("a", "b", null, "c") { assertCodegenRan { checkSparkAnswerAndOperator(sql("SELECT helloU() FROM t")) @@ -1472,13 +1603,10 @@ class CometCodegenSuite test("ScalaUDF composed with reused scalar subquery across projection and filter") { // The same scalar subquery appears in two sites: the projection (which the dispatcher - // compiles into a fused kernel) and the filter (a separate operator). Each site holds its - // own `ScalarSubquery` expression instance with its own `@volatile result` field. Each - // surrounding operator's inherited `SparkPlan.waitForSubqueries` populates its instance's - // `result` before the dispatcher's bridge serializes the expression. The populated value - // travels through closure serialization into the cache key's bytes, so different subquery - // values compile distinct kernels. Exercises the full subquery-correctness invariant - // documented on `CometBatchKernelCodegen.canHandle`. + // compiles into a fused kernel) and the filter (a separate operator). In both sites the + // subquery is a sibling of the dispatched `addOne(x)`, evaluated natively, so neither + // dispatched tree carries it; a subquery inside a dispatched tree is refused (see "a + // subquery inside a dispatched expression stays in Spark"). spark.udf.register("addOne", (i: Int) => i + 1) withTable("t", "t2") { sql("CREATE TABLE t (x INT) USING parquet") @@ -2562,3 +2690,13 @@ private case class XyPair(x: Int, y: String) /** Element type for the `Array>` dynamically-sized output case. */ private case class IntStr(a: Int, b: String) + +/** Returns how many times it has been called before, ignoring its argument. */ +class CountingJavaUdf extends org.apache.spark.sql.api.java.UDF1[String, java.lang.Long] { + private var calls = 0L + override def call(ignored: String): java.lang.Long = { + val result = calls + calls += 1 + java.lang.Long.valueOf(result) + } +} 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/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index af8aa7f0b9e..a2da07975a1 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -24,8 +24,8 @@ import scala.util.Random import org.apache.hadoop.fs.Path import org.apache.spark.SparkThrowable import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row} -import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, FromUnixTime, In, InSet, Literal, StructsToJson, TruncDate, TruncTimestamp} -import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, OptimizeIn, SimplifyExtractValueOps} +import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Coalesce, FromUnixTime, In, InSet, Literal, StructsToJson, TruncDate, TruncTimestamp} +import org.apache.spark.sql.catalyst.optimizer.{ConstantFolding, ConvertToLocalRelation, NullPropagation, OptimizeIn, SimplifyConditionals, SimplifyExtractValueOps} import org.apache.spark.sql.comet.{CometFilterExec, CometProjectExec, CometSortExec, CometTakeOrderedAndProjectExec} import org.apache.spark.sql.execution.{LocalTableScanExec, ProjectExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper @@ -1796,6 +1796,36 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("one-argument coalesce kept by the optimizer") { + // `NullPropagation` and `SimplifyConditionals` normally reduce a one-argument coalesce to its + // argument. With them excluded it reaches the serde, which must emit the argument itself: a + // CASE over no guarded children has no WHEN clause, which native CASE rejects. (A SQL file + // test cannot set this: its harness overwrites `excludedRules` with `ConstantFolding`.) + val excluded = Seq(NullPropagation, SimplifyConditionals, ConstantFolding).map(_.ruleName) + withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excluded.mkString(",")) { + withTable("t") { + sql("CREATE TABLE t(id bigint) USING parquet") + // One file, so the rows share a batch and the counter runs over all of them. + sql("INSERT INTO t SELECT IF(id = 1, NULL, id) FROM range(0, 4, 1, 1)") + Seq( + "SELECT coalesce(id) FROM t", + // The single argument is evaluated once, so a non-deterministic one stays native. + "SELECT coalesce(IF(monotonically_increasing_id() % 2 = 0, id, NULL)) FROM t") + .foreach { query => + val df = sql(query) + assert( + df.queryExecution.optimizedPlan.expressions.exists(_.exists { + case c: Coalesce => c.children.size == 1 + case _ => false + }), + s"Expected a one-argument coalesce to survive the optimizer:\n" + + df.queryExecution.optimizedPlan.treeString) + checkSparkAnswerAndOperator(df) + } + } + } + } + test("scalar decimal overflow - legacy mode produces null") { // 1.1e19 * 1.1e19 = 1.21e38 fits in i128 (max ~1.7e38) but exceeds DECIMAL(38,0)'s // max of 10^38-1, so CheckOverflow nulls the result in legacy (non-ANSI) mode. diff --git a/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala b/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala index 4db950e73c9..ab203a4b666 100644 --- a/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala @@ -221,6 +221,24 @@ class CometNativeCastSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("a cast that only relabels nested field names and nullability is compatible") { + // A set op casts both sides to one element type; when only names or nullability change, the + // cast is a relabel, which must stay native rather than go through the codegen dispatcher. + val dates = ArrayType(ArrayType(DateType, containsNull = false), containsNull = false) + val nullableDates = ArrayType(ArrayType(DateType, containsNull = true), containsNull = true) + val lower = ArrayType(StructType(Seq(StructField("a", DateType, nullable = false)))) + val upper = ArrayType(StructType(Seq(StructField("A", DateType, nullable = true)))) + val mapOf = (t: DataType) => MapType(IntegerType, t, valueContainsNull = false) + Seq((dates, nullableDates), (lower, upper), (mapOf(lower), mapOf(upper))).foreach { + case (from, to) => + assert(CometCast.isSupported(from, to, None, CometEvalMode.LEGACY) == Compatible()) + } + // A change of element type is not a relabel and keeps its own rule. + assert( + CometCast.isSupported(dates, ArrayType(ArrayType(LongType)), None, CometEvalMode.LEGACY) + != Compatible()) + } + test("cast BooleanType to DecimalType(14,4)") { castTest(generateBools(), DataTypes.createDecimalType(14, 4)) } 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 c6a6805c72f..9d36c1da3c0 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -120,12 +120,63 @@ 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. The suite's + // `local[5]` session writes the table as five files, so each of the five map tasks sends each + // of the 300 destinations about 13 rows, several batches of two; with too few rows per + // destination no call would reach a second batch, and the test would pass with the reset + // removed. + 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 20000L).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") checkShuffleAnswer(shuffled, 1) } + test("columnar shuffle with Map[NullType, _] column") { + // map() leaves a NullType key. The shuffle rows go through the native builders, which declare + // the key non-nullable themselves, so this is coverage for the shape rather than a witness for + // the Arrow Java IPC key repair (`Utils.withNonNullableMapKeys`); `UtilsSuite` and the pyspark + // Arrow UDF tests cover that. 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/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 3e7cf698f53..4d5250a0f2c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -410,6 +410,66 @@ class CometExecSuite extends CometTestBase { } } + test("non-AQE DPP: a build side the broadcast gate refuses keeps Spark's DPP subquery") { + // `map(store_id, NULL)` puts a NullType under a map entry, so the join's broadcast stays on + // Spark (`CometBroadcastExchangeExec.getSupportLevel`). The DPP subquery must stay with it: + // a Comet broadcast there would not be reused by the join's exchange, so the build side + // would run twice. + withTempDir { dir => + val path = s"${dir.getAbsolutePath}/data" + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + spark + .range(100) + .selectExpr("id % 10 as store_id", "cast(id * 2 as int) as date_id") + .write + .partitionBy("store_id") + .parquet(s"$path/fact") + spark + .range(10) + .selectExpr("cast(id as int) as store_id", "cast(id as string) as country") + .write + .parquet(s"$path/dim") + } + + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") { + spark.read.parquet(s"$path/fact").createOrReplaceTempView("fact_null_map") + spark.read.parquet(s"$path/dim").createOrReplaceTempView("dim_null_map") + + val df = spark.sql("""SELECT f.date_id, f.store_id, d.m + |FROM fact_null_map f JOIN ( + | SELECT store_id, country, map(store_id, NULL) AS m FROM dim_null_map) d + |ON f.store_id = d.store_id + |WHERE d.country = 'DE'""".stripMargin) + val (_, cometPlan) = checkSparkAnswer(df) + + // Guard the premise: DPP planned a broadcast subquery for this join. + val subqueries = collectWithSubqueries(cometPlan) { + case s: SubqueryBroadcastExec => s + case s: CometSubqueryBroadcastExec => s + } + assert(subqueries.nonEmpty, s"Expected a DPP subquery:\n${cometPlan.treeString}") + // And that the build side under Spark's broadcast is Comet native: the rewrite only + // considers a native build side, so with a Spark one the next assertion would hold for + // the wrong reason. + val sparkBroadcasts = collectWithSubqueries(cometPlan) { case e: BroadcastExchangeExec => + e + } + assert( + sparkBroadcasts.exists(_.child.exists(_.isInstanceOf[CometNativeExec])), + s"Expected a Comet native build side under Spark's broadcast:\n${cometPlan.treeString}") + assert( + collectWithSubqueries(cometPlan) { case e: CometBroadcastExchangeExec => e }.isEmpty, + s"Expected no Comet broadcast:\n${cometPlan.treeString}") + assert( + collectWithSubqueries(cometPlan) { case e: ReusedExchangeExec => e }.nonEmpty, + s"Expected the DPP subquery to reuse the join's exchange:\n${cometPlan.treeString}") + } + } + } + test("non-AQE DPP: non-atomic type (struct/array) join key") { withTempDir { dir => val path = s"${dir.getAbsolutePath}/data" 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 b2c3764fd6d..3fff321ba93 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1206,6 +1206,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 on collated string columns") { 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 eec16205a11..7d41d02865d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -33,11 +33,11 @@ import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, Dynam 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.{LocalTableScanExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException -import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeLike} 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 import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus @@ -1823,4 +1823,50 @@ 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). 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", + 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) + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + df, + "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/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index b817ed7a503..80df7fabb52 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -1350,6 +1350,34 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + test("native shuffle: round robin over a hashed NullType-bearing column falls back") { + // Without positional placement, native round-robin hashes every column, and the native + // hasher has no NullType arm; `map(_1, NULL)` reaches the exchange from the codegen + // dispatcher, which runs NullType-bearing outputs. + withSQLConf( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable((0 until 4).map(i => (i.toLong, i)), "tbl") { + checkSparkAnswerAndFallbackReason( + "SELECT /*+ REPARTITION(2) */ map(_1, NULL) AS m FROM tbl", + "the native hasher does not support NullType") + // A column the hasher never reads keeps the native shuffle: past `maxHashColumns`, or + // with every row placed by position. + withSQLConf( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.key -> "1") { + checkSparkAnswerAndOperator( + sql("SELECT /*+ REPARTITION(2) */ _1, map(_1, NULL) AS m FROM tbl")) + } + withSQLConf( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "true") { + checkSparkAnswerAndOperator( + sql("SELECT /*+ REPARTITION(2) */ map(_1, NULL) AS m FROM tbl")) + } + } + } + } + test("native shuffle: round robin deterministic behavior") { // Test that round robin produces consistent results across multiple executions withSQLConf(CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true") { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala index bfab3727cab..93a4264d3e0 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.{CometTestBase, DataFrame} import org.apache.spark.sql.catalyst.plans.physical.RoundRobinPartitioning import org.apache.spark.sql.execution.SQLExecution import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.functions.{col, lit, rand, udf} +import org.apache.spark.sql.functions.{col, lit, monotonically_increasing_id, rand} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -113,15 +113,16 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp test("a nondeterministic projection or filter keeps content-hash placement") { // Evaluated row by row, but nothing bounds what a nondeterministic expression does between // attempts: one that reorders or re-filters rows on a retry sends them to different reducers. - val sameId = udf((id: Long) => id) withPositionalRoundRobin() { withParquetTable(1000) { t => val df = spark.table(t) - assert(!isPositional( - df.select(sameId.asNondeterministic()(col("id")).as("id")).repartition(numPartitions))) + // The same values either way; only the second projection is non-deterministic. + assert(isPositional(df.select((col("id") + lit(0L)).as("id")).repartition(numPartitions))) + assert( + !isPositional( + df.select((col("id") + monotonically_increasing_id() * 0).as("id")) + .repartition(numPartitions))) assert(!isPositional(df.filter(rand(42) < 0.5).repartition(numPartitions))) - // The same UDF marked deterministic is admitted, so it is the flag that excludes it. - assert(isPositional(df.select(sameId(col("id")).as("id")).repartition(numPartitions))) } } } 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..b6e95771194 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, 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 +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,203 @@ class UtilsSuite extends CometTestBase { } } } + + /** + * One `map` column of `numRows` empty maps, as `map()` produces (a `NullVector` + * key). + */ + 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() + (0 until numRows).foreach { i => + vector.startNewValue(i) + vector.endValue(i, 0) + } + 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. */ + 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) + } + + test("withNonNullableMapKeys restores the non-nullable key flag a NullVector drops") { + val batch = emptyNullKeyMapBatch(2) + val field = batch.column(0).asInstanceOf[CometVector].getValueVector.getField + // 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) + 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 = emptyNullKeyMapBatch(2) + val vector = + batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector] + val declared = Utils.withNonNullableMapKeys(vector.getField) + // 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) + assertRoundTrip(bound, writer, out, numRows = 2) + batch.close() + } + + test("newArrowStreamWriter returns the root a later row count must be set on") { + 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. + 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) + assertRoundTrip(bound, writer, out, numRows = 2) + batch.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 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. The bypassed shapes never reach the appender, so this test + // cannot tell when an Arrow release fixes it; that needs a direct appender probe. + 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])) + + 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() + } + // 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()) + + // A hanging append keeps growing its vectors after the timeout gives up on it; the bound + // turns that into an allocation failure in its own thread instead of exhausting the + // test JVM's direct memory. + val bounded = CometArrowAllocator.newChildAllocator(s"coalesce $name", 0, 64L << 20) + val (result, batchCount, totalRows) = Await.result( + Future(Utils.coalesceBroadcastBatches(bufs.iterator, bounded))(ExecutionContext.global), + 10.seconds) + bounded.close() + + // 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") + assert(result.length == (if (expectBypass) numBatches else 1), 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) + decoded.foreach(_.close()) + } + } + + test("coalesceBroadcastBatches keeps coalescing plain null lists") { + // 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).toVector + + 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") + } }