feat: support NullType output types in codegen dispatch

[grorge123]

Which issue does this PR close?

Closes #5525.

Rationale for this change

Untyped constructors such as map(), map('a', NULL) and array() leave NullType children in their output type. CometBatchKernelCodegen.canHandle applied the same isSupportedDataType predicate to the output type and to every BoundReference input, and that predicate had no NullType case, so any expression whose output type contained NullType fell back to Spark and took the whole operator with it.

Only the input side needs that restriction: CometScalaUDFCodegen.specFor cannot build an ArrowColumnSpec for a NullVector, so a NullType input would throw at execute time, after the point where a fallback is still possible. On the output side the kernel just has to emit an all-null Arrow NullVector, which serde (NullType type id) and Utils.toArrowField (ArrowType.Null) already understand. No native changes are needed.

Note on scope: expressions that Spark can constant-fold (map(), map('a', NULL), array() with no column references) become a Literal before planning, and CometLiteral still rejects a complex literal containing NullType, so those forms keep falling back. The gate change helps the non-foldable forms: map(k, NULL) over a column, and higher-order functions such as transform_values(map(), ...).

What changes are included in this PR?

Codegen gate

• CometBatchKernelCodegen: isSupportedDataType gains a private allowNullType overload that recurses through array / struct / map children. canHandle checks the output type with allowNullType = true and keeps checking BoundReference inputs with the unchanged public predicate, so NullType inputs still fall back at plan time. The doc comment states the asymmetry and why.
• CometBatchKernelCodegen: canHandle also rejects duplicate struct field names, found recursively through the output type and through every BoundReference input (duplicateStructFieldNames). Spark keeps duplicates as distinct positional fields, but Arrow's StructVector keys children by name and collapses them, so the generated ordinal-based child casts would hit a missing or differently typed vector. CometCreateNamedStruct already refuses this at the serde level, but whole-expression dispatch never consults that rule for a named_struct nested inside e.g. a transform lambda, so the gate re-checks it. generateSource re-asserts the same condition.
• CometBatchKernelCodegenOutput: outputVectorClass maps NullType to NullVector; emitWrite gets a NullType branch that only calls setNull and never reads the source value; emitSpecializedGetterExpr returns "null" for NullType, keeping the emitter's type surface in sync with the gate.

Carrying NullVector columns through the JVM IPC paths (review follow-up)

Once a NullVector can sit inside a map or struct, two Arrow Java 18.3.0 behaviours bite the JVM-side IPC paths (broadcast, getByteArrayRdd, the PyArrow UDF runner):

• Arrow's MinorType.NULL factory discards the Field it is handed and rebuilds a nullable NullVector, so a NullType map key that Comet declared non-nullable comes back nullable, and MapVector.initializeChildrenFromFields rejects the schema on read ("Map data key type should be a non-nullable"). Utils.withNonNullableMapKeys repairs the key flag, and Utils.newArrowStreamWriter is now the only way to build an ArrowStreamWriter in Comet (enforced by a scalastyle rule), so every IPC writer gets the repair. It returns the root the writer is actually bound to, so a caller that sets the row count per batch (the PyArrow UDF runner) cannot set it on a superseded root. CometArrowStream.actualFieldOf applies the same repair to the stream schema handed to native. CometArrowPythonRunnerBase.startWriter repairs the child fields before createVector, which would otherwise fail on the same check.
• VectorSchemaRootAppender loops forever on a NullVector that is a direct child of a struct or of a map entry: a struct's capacity is the minimum over its direct children, and NullVector.reAlloc() is a no-op whose capacity equals its value count, so the struct's capacity loop never terminates. Utils.coalesceBroadcastBatches ships exactly those schemas uncoalesced, the same way it already handles dictionary-encoded columns. A list breaks that chain from both sides, because ListVector overrides BaseRepeatedValueVector.getValueCapacity with one that only reads its own offset and validity buffers - so a top-level NullType column, array(NULL), and a list under a struct such as map(k, array(NULL)) all keep coalescing.
• A NullType child is always declared nullable in the schema handed to native (Utils.declaredChildNullability, applied in Utils.toArrowField and in QueryPlanSerde.serializeDataType). Arrow rebuilds every NullVector as nullable, so a containsNull = false / valueContainsNull = false declaration that Spark derives for an untyped constructor would not match the vector native receives.

Docs

• docs/source/user-guide/latest/scala_java_udfs.md: NullType arguments stay in the unsupported list, NullType return types are documented as supported. While editing that bullet I also removed CalendarIntervalType from the unsupported list: the list dates from feat(experimental): ScalaUDF and Java UDF support via Janino codegen #4267, and feat: add CalendarIntervalType support #4898 later added CalendarIntervalType to the codegen gate and both emitters without touching this doc, so the entry has been stale since then.

Tests

• CometCodegenSourceSuite: seven new unit tests — canHandle accepts NullType outputs (top-level and nested), rejects NullType inputs (top-level and nested), rejects duplicate struct field names in both outputs and inputs, NullType children are declared nullable whatever the Spark flags say, the NullType output path writes setNull without reading a source, nested NullType output casts the child vector, and the gate and output emitters agree across the whole accepted type surface.
• UtilsSuite: withNonNullableMapKeys restores the key flag (and pins the Arrow behaviour it works around, so the workaround can be dropped when Arrow fixes it), serializeBatches round-trips a NullType map key, and newArrowStreamWriter keeps / returns the bound root so a row count set after construction is serialized. For the broadcast bypass, one test walks the whole shape space (top-level NullType, array<null>, nested lists, struct<a: null>, array<struct<a: null>>, struct<l: array<null>>, map<int, null>, map<int, array<null>>, map<null, null>) and asserts the rule fires on exactly the dangerous ones, running the real appender under a timeout so a rule that is too narrow fails instead of hanging the build; two further tests pin that struct-nested NullType ships uncoalesced and that plain null lists keep coalescing with their per-row lengths intact.
• CometJoinSuite: broadcast hash join with map(_1, NULL) and transform_values(map(), ...) on the build side, AQE on and off, with spark.sql.legacy.createEmptyCollectionUsingStringType=false pinned and the NullType premise asserted.
• CometColumnarShuffleSuite: JVM columnar shuffle with Map[_, NullType] and Map[NullType, _] columns.
• test_pyarrow_udf.py: mapInArrow over Map[NullType, _] / Map[_, NullType] input columns in both accelerated and fallback modes (pyarrow refuses to declare a non-nullable null-typed field, so the NullType-key map can only be a UDF input).
• SQL file tests: 9 new queries in map/create_map.sql (map(), map('a', NULL), map(k, NULL) over a table, nesting in array / struct / map, size / map_keys / map_values over them, and a map() column carried through ORDER BY), map_from_arrays(array(), array()) in map/map_from_arrays.sql, plus one untyped empty-constructor query in each of array_sort_comparator.sql, transform.sql, zip_with.sql, map_concat.sql, map_filter.sql, map_zip_with.sql, transform_keys.sql, transform_values.sql, and duplicate-name / declared-nullability coverage in struct/create_named_struct.sql, array/array_repeat.sql, array/array_union.sql, array/slice.sql and map/map_entries.sql. CometSqlFileTestSuite runs with constant folding disabled, so these exercise the codegen path directly.

How are these changes tested?

Run locally against the final commit.

spark4.1_2.13 profile:

• make — BUILD SUCCESS.
• make test-jvm — Tests: succeeded 2585, failed 0, canceled 5, ignored 16 (the 5 canceled are pre-existing assume skips: Spark-4-only behaviour, SPARK-54220, and the PySpark end-to-end case that needs PYSPARK_PYTHON). The three MinIO-backed suites (CometS3CredentialBridgeSuite, IcebergReadFromS3Suite, ParquetReadFromS3Suite) abort on this machine without AWS_REGION; re-run separately with AWS_REGION=us-east-1: Tests: succeeded 16, failed 0.
• ./mvnw test -pl common,spark -Dsuites="org.apache.spark.sql.comet.util.UtilsSuite,org.apache.comet.exec.CometJoinSuite,org.apache.comet.exec.CometShuffleSuite,org.apache.comet.exec.DisableAQECometShuffleSuite,org.apache.comet.CometSqlFileTestSuite" — Tests: succeeded 592, failed 0.
• pytest spark/src/test/resources/pyspark/test_pyarrow_udf.py with Spark 4.1.3 / pyspark 4.1.3 — 83 passed.
• spotless, scalastyle, prettier --check "**/*.md", dev/ci/check-suites.py, dev/ci/check-benchmark-runner.py — clean.

-Pspark-3.5 (Scala 2.12) profile:

• ./mvnw test -Pspark-3.5 -pl common,spark -Dsuites="org.apache.spark.sql.comet.util.UtilsSuite,org.apache.comet.exec.CometJoinSuite" — Tests: succeeded 40, failed 0.
• ./mvnw test -Pspark-3.5 -pl common,spark -Dsuites="org.apache.comet.exec.CometShuffleSuite,org.apache.comet.exec.DisableAQECometShuffleSuite,org.apache.comet.CometSqlFileTestSuite" — Tests: succeeded 552, failed 0.

Red checks, to confirm the new tests catch the failures they target: with only the map-key repair reverted, UtilsSuite and CometJoinSuite fail with IllegalArgumentException: Map data key type should be a non-nullable; with the Python runner repair reverted, test_map_in_arrow_null_typed_map_children[accelerated] fails with the same exception.

All suites touched here are already registered in the CI suite matrices, so no workflow change is needed.

Assisted-by: Claude Code (claude-fable-5)

[sunchao]

Source review of this pinned change found the two newly admitted broadcast failure paths below. Neither example was executed locally. The supplied snapshot has four action-required workflows and no passing checks; the author-reported local test results were not independently verified.

[sunchao]

Re-reviewed db1768fa2aea540821391c7aef4cff2a4b8c43fe. The two earlier witnesses are fixed in source. Two source-derived P2s remain; neither was executed here.

[P2] Reject duplicate names in newly admitted struct outputs. The output gate now admits transform(array(id), x -> named_struct('a', x, 'a', NULL)) over an INT input. Whole-expression ArrayTransform dispatch skips the inner CreateNamedStruct duplicate-name fallback. Spark retains both fields, but Arrow 18.3.0's default struct policy replaces the first a with the later NullVector. Generated output setup still casts child ordinal 0 to IntVector, so the query fails before rows are written. BASE rejected this Null-bearing output. Please retain fallback for duplicate names recursively, or allocate children without losing their ordinals.

[P2] Keep coalescing safe plain null lists. The new broadcast bypass also rejects plain array<null>, although its ListVector appender grows offsets/validity independently of the Null child and finalizes the child count. An already supported build payload such as IF(rand(17L) < 0.5, array(NULL), array(NULL,NULL)), selected after a join over a repartitioned broadcast build, therefore loses coalescing. Returning B original buffers makes each of P consuming tasks open B compression/IPC streams and schema roots instead of one: B×P setups rather than P. This is an operation-count consequence, not a measured speed ratio. Keep the struct/map protection, but allow the safe plain-list case to coalesce.

[grorge123]

We re-applied the representation-level serde restriction inside canHandle.
Besides that, we found another problem: filter(array(), ...) and map_filter(map(), ...) leave containsNull = false on a NullType child, and downstream native kernels that rebuild a list around that input (map_entries, array_repeat, slice) panic on the nullability mismatch. We fixed it by always declaring a NullType child nullable on both sides of the FFI boundary.

[sunchao]

Re-reviewed 1d0a6f493bbd60453234748d548c85d77348a98c. The previous duplicate-field, map-key schema, and broadcast-bypass issues are addressed in the inspected source. The five additional P2 cases below remain.

Could you add a small microbenchmark for consumed map(id, NULL) projections and the remaining struct/map broadcast bypass, with a plain array<null> control? Compare this head with Spark and safe prior paths using matched versions, settings, and warmup. Please confirm result equality and the executed native/fallback plans, and report throughput, allocations, peak/retained memory, and IPC setup counts for small/default batches and few/many build buffers.

[sunchao]

[P2] Preserve single evaluation of nullable stateful outputs

With codegen dispatch and ANSI enabled, consider element_at(transform(IF(monotonically_increasing_id() % 2 = 0, array(id), CAST(NULL AS array<bigint>)), x -> named_struct('id', x, 'n', NULL)), 1) over a native LONG batch containing id=0,1,2,3. Spark evaluates the transform once per row, retaining the struct for id=2. This gate now admits that Null-bearing output, but the existing ElementAt ANSI conversion serializes its left subtree into both the CASE predicate and the lookup. Native CASE tests all four rows, then reevaluates the transform on the two selected rows. The second selected row receives an odd counter and becomes NULL. Please materialize the left value once or retain Spark fallback for this composition. Current BASE rejects this output. The previous reviewed head admitted it but lacked the duplicated-left CASE, so the newly failing combination is in this increment. Source-derived witness, not executed.

[grorge123]

Fixed: CometElementAt keeps a non-deterministic nullable collection in Spark under ANSI, since the null guard's two copies each hold their own kernel state and the THEN copy only sees the CASE-selected rows. Reproduced the witness first (Comet returned [null] where Spark keeps [[2,null]]), and it now matches. The same guard is built by CometSize (non-legacy mode), CometArrayAppend, CometMapFromArrays and CometCoalesce, so they share the gate (NullGuard); CometSize drops the guard in legacy mode, where native already answers -1. The coalesce case is reachable on main without NullType (coalesce(IF(monotonically_increasing_id() % 2 = 0, array(id), NULL), array(id)) NPEs in columnar-to-row because the result is declared non-nullable); it is included since the sweep found it and the fix is one line on the shared guard.

Tests: expect_fallback queries in element_at_ansi.sql, map_from_arrays.sql and coalesce.sql, plus the non-deterministic dimension of CometNullTypeCompositionSuite.

[sunchao]

[P2] Residual stateful guard case across different arguments

At 79fc84008a7dca16a4c24f28036649c0af2603b0, the original nullable element_at witness is fixed, but the shared guard still misses a cross-input case. This is source-derived, not an executed reproduction. With a native LONG input t containing ordered ids 0, 1, 2, 3 in one partition/batch and codegen dispatch enabled, consider SELECT id, arrays_zip(transform(array(id), x -> named_struct('i', monotonically_increasing_id(), 'n', NULL)), IF(id % 2 = 0, array(id), CAST(NULL AS ARRAY<BIGINT>))) AS z FROM t. The first array is non-nullable and nondeterministic, while the second is nullable and deterministic, so neither matches this predicate. Spark evaluates the first array even on odd rows where the zip is NULL. CometArraysZip.convert instead puts the value-producing kernel inside a CASE whose THEN batch contains only even rows, changing the stateful child's evaluation sequence. The base rejected this NullType-bearing producer. This guard therefore misses a composition newly admitted by the PR. Could the guard preserve one evaluation before filtering or retain fallback when a nullable sibling filters a stateful child, with coverage for nullability and nondeterminism occurring in different arguments?

[grorge123]

Fixed: reproduced the witness first (Comet returned [[4,null],0] and [[5,null],2] where Spark keeps [[0,null],0] and [[2,null],2]), and the non-NullType flavour diverges the same way, so this is the guard's rule being too narrow rather than a NullType path. The witness's stateful child is a lambda, which runs through the JVM codegen dispatcher; its kernel cache is keyed by the serialized expression, so the guard's predicate copy and THEN branch run one kernel instance and share its counter: the predicate consumes it for the whole batch and the THEN branch continues from there. A natively evaluated stateful child gets its own instance per copy, but native CASE evaluates the THEN branch on the rows the predicate selected, so it diverges as soon as the guard filters. NullGuard now refuses any non-deterministic child inside a guard, whichever argument is the nullable one, and CometArraysZip, CometElementAt, CometArrayAppend, CometMapFromArrays, CometCoalesce and CometSize share it; CometSize builds no guard at all for a non-nullable child (or in legacy mode), so size(filter(arr, x -> x < monotonically_increasing_id())) is evaluated once and stays native.

Tests: expect_fallback witness in arrays_zip.sql; CometNullTypeCompositionSuite gains a cross-input sweep that puts a non-nullable stateful producer (including one whose length records the counter) beside a nullable deterministic sibling under every consumer, and a nullable deterministic sweep so the guards' ELSE branches run natively; the size witnesses are in CometArrayExpressionSuite.

[grorge123]

Beyond the five cases above, this revision adds CometNullTypeCompositionSuite, a differential sweep of the admitted shape space: every non-foldable NullType producer (transform(array(id), x -> NULL), map(id, NULL), transform_values(map(), ...), map_entries(map(id, NULL)), named_struct(..., NULL), an all-NULL aggregate) under every consumer, operator and nesting container Spark accepts, across ANSI on/off, a nullable non-deterministic wrapper, and native columnar-to-row on/off, comparing Comet with Spark.

The non-NullType flavour of the same item-field problem (slice(map_entries(map(k, v)), ...), slice(reverse(arrays_zip(...))), the #4789 nested-nullability contract) reproduces on main and is outside this PR's admission. The one native gap the sweep tolerates is also pre-existing: Comet's row shuffle writer has no case for a Null struct field (Unsupported data type of struct field: Null), matched by signature so every other failure stays red.

Assisted-by: Claude Code (claude-fable-5)

[sunchao]

The five earlier concrete witnesses are addressed in the inspected source at 79fc84008a7dca16a4c24f28036649c0af2603b0. A residual cross-input stateful-guard case remains in the existing discussion, and one newly reachable shuffle failure is detailed inline.

This pass used source analysis only. I did not run a Spark/Comet query, test or runtime reproduction. The current-head snapshot has no check runs or workflow runs, so CI validation is unavailable.

[sunchao]

[P2] Do not suppress shuffle regressions newly admitted by this PR

Source-derived, not executed. Over nonempty primitive Parquet t, use SELECT /*+ REPARTITION(3, c) */ c FROM (SELECT transform(array(id), x -> named_struct('v', x, 'n', NULL)) AS c FROM t) with native scan/codegen dispatch, CometShuffleManager, Comet shuffle enabled, shuffle mode auto, spark.comet.shuffle.convertFromSparkPlan.enabled=false, AQE off and JVM columnar-to-row. At the base the producer falls back, and the disabled conversion leaves its exchange in Spark. This PR admits a CometProject, so the array partition key rejects native shuffle but falls through to JVM columnar shuffle. The unsafe-row list writer then sends its struct's Null field to append_field, which reaches exactly the panic waived here. The writer defect is inherited, but its reachability under these settings is new. Please retain fallback for that schema or support Null struct fields, and do not waive failures that passed at the base under the same settings. With conversion enabled the base can also fail, so that is not a valid regression control.

[grorge123]

Fixed: the exact query above does not fail at this head (reproduced under the listed settings: CometProject into CometColumnarShuffle, results match), because a list of structs goes through the writer's row-major append_field, which has a Null arm. The shape that did panic is a top-level struct column with a Null field, e.g. REPARTITION(3) named_struct('v', id, 'n', NULL), which takes the field-major paths that lacked the arm; that was the failure the sweep tolerated. Both field-major paths now handle Null struct fields (every row null, as the row-major path does), and the sweep tolerates nothing: every producer under repartition, shuffle join, group-by and sort compares against Spark.

Tests: CometNullTypeCompositionSuite operator and nesting sweeps with the waiver removed.

[sunchao]

You're right. Correction to my [P2] example: at 79fc8400, the array-of-struct path reaches append_field, which already has a DataType::Null arm. My assertion that this path reaches the unsupported-struct-field panic was incorrect, and I withdraw that specific witness. Thanks for checking it.

I am reviewing the two field-major fixes and removal of the sweep waiver separately at 496c3f7f. This correction is based on the pinned source, not an executed SQL reproduction.

[sunchao]

Re-reviewed 496c3f7f against ef62b463. The broader nondeterminism guard addresses the cross-input case, and the new field-major NullType arms pass the single-batch controls. I found one additional P2 in builder reuse across batches, detailed inline.

I independently reran the exact-source native component harness: 8 tests passed and 3 failed on the same NullBuilder reuse defect. This is not a full Spark, JNI, or shuffle execution. The planner route is source-traced. Four workflows remain action_required, with no head or merge check results in the current snapshot. No fresh benchmark was run.

[sunchao]

[P2] Reset NullType builders before reusing them for another batch

These new arms handle the first batch, but process_sorted_row_partition creates its builders outside the batch loop and builder_to_array calls finish() without replacing them. In the pinned Arrow 58.4.0, NullBuilder::finish() leaves its length unchanged. With the exact current conversion functions, two consecutive two-row batches of struct<v:bigint,n:void> panic on the second finish because the struct has length 2 while its Null child has length 4. The single-batch and typed-null two-batch controls pass. A nested struct fails the same way.

SpillSorter sends an entire destination partition to this native call, so one call can exceed spark.comet.shuffle.jvm.batchSize. The PR now admits producers such as element_at(transform(array(id), x -> named_struct('v', x, 'n', NULL)), 1) over native primitive input. With CometShuffleManager, JVM shuffle mode, spark.comet.shuffle.convertFromSparkPlan.enabled=false, and AQE off, source tracing shows their exchange can enter this writer where the base retained Spark fallback.

Could we reset or recreate Null-containing builders between batches, or retain fallback, and add coverage spanning more than one writer batch? I reproduced the builder failure in an isolated harness using production conversion source and pinned dependencies. The SQL/planner route was source-traced, not executed end to end.

[grorge123]

Fixed: reproduced first with the sort-based writer (spark.shuffle.sort.bypassMergeThreshold=0, spark.comet.shuffle.jvm.batchSize=2, jvm shuffle mode, convertFromSparkPlan.enabled=false, AQE off, 16 rows into 2 partitions): named_struct('v', id, 'n', NULL) panics in StructBuilder::finish with (2 != 4) on the second batch, as you describe, and so do a nested struct and element_at(transform(array(id), x -> named_struct('v', x, 'n', NULL)), 1). The same reuse defect hits every other Null-bearing shape: a top-level NULL column fails the batch's row-count check, map(id, NULL) fails "keys and values have unequal length", and array<null> comes back with the wrong row count. Base control at ef62b463 under the same settings: the top-level NULL column and array(named_struct('v', id, 'n', NULL)) fail there too (row-count check and the same (2 != 4) panic through the row-major arm), the struct shapes hit the field-major unreachable instead, and map(id, NULL), array<null> and the element_at shape stayed in Spark on the base, so those three are admitted by this PR. NullBuilder::finish keeps its length in arrow 58.4.0 (a NullArray owns no buffers to hand over), so process_sorted_row_partition now recreates every builder whose type holds a Null anywhere (contains_null_type) after each batch; the other builders reset on finish and are kept. The bypass writer sends at most one batch per native call, which is why the single-batch controls and the sweep's REPARTITION(3) never saw it.

Tests: null_type_builders_start_every_batch_empty in row.rs drives two batches through the production make_builders / builder_to_array / recreate path for a Null field, a nested one and a top-level Null column; CometColumnarShuffleSuite "columnar shuffle spanning several native writer batches with NullType columns" runs the seven shapes above through REPARTITION(300) with jvm.batchSize=2 and the spill threshold lifted, under both AQE settings, comparing with Spark.