Skip to content

feat: route arrays_zip on map element types through the codegen dispatcher - #5688

Open
kazantsev-maksim wants to merge 80 commits into
apache:mainfrom
kazantsev-maksim:arrays_zip_codegen
Open

feat: route arrays_zip on map element types through the codegen dispatcher#5688
kazantsev-maksim wants to merge 80 commits into
apache:mainfrom
kazantsev-maksim:arrays_zip_codegen

Conversation

@kazantsev-maksim

@kazantsev-maksim kazantsev-maksim commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #5583

Rationale for this change

arrays_zip over an array whose element type is a map (array<map<string,int>>) is not executed natively: CometArraysZip's type gate returns Unsupported for MapType element arrays, so the entire projection falls back to Spark, even though Spark supports it.

The JVM codegen dispatcher already handles map types: CometBatchKernelCodegen.isSupportedDataType admits MapType recursively, and ArraysZip is a plain codegen expression, so dispatching it runs Spark's own doGenCode and matches Spark exactly by construction.

What changes are included in this PR?

  • CometArraysZip now mixes in CodegenDispatchFallback: arrays of maps are routed through the JVM codegen dispatcher, while scalar/array/struct element arrays keep the existing native path.
  • Updated getUnsupportedReasons(): it now surfaces only when the dispatcher is disabled or rejects the tree.
  • Carve-out (unchanged behavior, documented): NullType elements are not dispatchable — NullType is absent from isSupportedDataType — so arrays_zip over array<null> continues to fall back to Spark.
  • Convert-side declines are also unchanged: when a child of ArraysZip fails to serialize, the serde still returns None and the projection falls back (the dispatcher cannot see that state).
  • Regenerated expression compatibility docs (arrays_zip is now hybrid: map elements via dispatch, other elements natively).

How are these changes tested?

New sql test cases added

@kazantsev-maksim kazantsev-maksim changed the title Arrays zip codegen feat: route arrays_zip on map element types through the codegen dispatcher Sep 4, 2026

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correctness

Reviewed head 74e92e84e1d99957c514500d5d9cc6db0f4fc311 against base 2da32915c87909d7dbf8d62adb02f8349c36bca0. I found no verified correctness defect in the changed behavior. There is a performance-validation request below.

The maintained Spark 3.5 and 4.0 implementations return null when an input array is null, pad shorter arrays with null fields, and retain the input-derived struct field names. Adding CodegenDispatchFallback lets the existing dispatcher evaluate map-containing ArraysZip trees using Spark's code generation. I checked the recursive map readers and writers as well as the expression semantics. The result keeps non-null struct elements, nullable zipped fields, and the nested map key/value types. The exported Arrow list and map field names agree with Comet's native type serialization.

The existing scalar, array and struct native type gate and convert implementation are unchanged. The per-expression disable switch still acts before dispatch. Disabling spark.comet.exec.scalaUDF.codegen.enabled, rejection by canHandle, or failure to serialize a dispatched input attribute or return type still declines the conversion. A native-path child conversion returning None is not retried through dispatch. The new marker does not alter that path for NullType, and the dispatcher itself rejects NullType outputs.

Local component validation compiled the exact Comet recursive kernel, input and output implementations against Spark 4.0.4, Scala 2.13.16 and Arrow 18.3.0 on JDK 17. All 224 row-and-batch comparisons with Spark expression evaluation passed across seven map shapes with ANSI mode on and off. These covered null and empty arrays, null map elements and values, unequal lengths, nested arrays/structs/maps, and array/struct map keys. The checks included Arrow C Data input/output round trips, exported schema names, and zero outstanding allocator bytes. The probe substitutes the scalar vector adapter and allocator. It does not run the full Comet planner, Rust consumer, JNI query path or SQL fixture suite. Maintained source comparison was limited to Spark 3.5 and 4.0, and no 3.4/4.1 runtime coverage is claimed.

The new SQL queries use table columns and the harness requires Comet execution plus Spark answer equality. They are useful regression coverage, but I have not run them here. At the September 4, 19:52 UTC snapshot, all six head workflow records were action_required, with no head or merge check runs. That is not passing CI evidence.

Performance

Supported native inputs retain the same execution path, so this does not add a per-row dispatcher cost to those cases. Map-containing trees now pay the existing bridge cost: input export/import per batch, expression/kernel cache lookup, Spark row-wise evaluation, and allocation and population of the nested output vector. The cache avoids recompiling the expression for each batch, but zipping maps still materializes result structures and copies their contents into Arrow output buffers.

Keeping surrounding expressions in Comet can avoid a projection fallback, but that benefit needs to be weighed against the bridge and nested-output costs for this newly enabled case. Neither the PR nor the correctness probe measures that tradeoff. The inline P2 asks for a matched microbenchmark before enabling this path by default. This is a performance-validation gap, not a claim that a regression has been measured.

Design

The change uses the existing dispatch decision point and preserves the native implementation for types it already accepts. That keeps map support aligned with Spark without adding another map-specific native implementation or changing the native compatibility gate. The two SQL cases exercise a map result and a mixed projection with scalar arithmetic. The explicit hybrid entry in the expression list reflects this division of execution paths.

Abstraction & complexity

The existing marker is an appropriate extension point for this change. No new dispatcher, map adapter, protocol representation or runtime cache is introduced. Type admission, binding, transport and output allocation remain centralized in the existing implementation. I found no actionable added abstraction or complexity concern in this diff.

}

object CometArraysZip extends CometExpressionSerde[ArraysZip] {
object CometArraysZip extends CometExpressionSerde[ArraysZip] with CodegenDispatchFallback {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Performance

[P2] Could you add a matched microbenchmark for the newly dispatched arrays-of-maps path before enabling it by default? This marker changes existing queries from Spark projection fallback to Spark-generated evaluation inside the JVM bridge. That path still evaluates ArraysZip row by row and additionally imports the input and materializes the nested Arrow output. Avoiding the surrounding projection fallback can pay for those costs, but the two correctness queries do not establish when it does.

Please compare dispatch enabled with the existing dispatcher-disabled fallback on identical nonconstant column inputs, including a map-only projection and the mixed native projection, small and multi-batch inputs, and small versus larger maps/arrays. Verify the same results and include warm steady-state timings with the batch size and data shape. I am not asserting a measured slowdown. The missing comparison is the evidence needed to assess the default-path performance tradeoff.

@andygrove andygrove added enhancement New feature or request area:expressions Expression evaluation array expressions labels Sep 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:expressions Expression evaluation array expressions enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

arrays_zip falls back to Spark for map element types

3 participants