feat: route arrays_zip on map element types through the codegen dispatcher - #5688
feat: route arrays_zip on map element types through the codegen dispatcher#5688kazantsev-maksim wants to merge 80 commits into
arrays_zip on map element types through the codegen dispatcher#5688Conversation
This reverts commit 768b3e9.
arrays_zip on map element types through the codegen dispatcher
sunchao
left a comment
There was a problem hiding this comment.
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 { |
There was a problem hiding this comment.
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.
Which issue does this PR close?
Closes #5583
Rationale for this change
arrays_zipover an array whose element type is a map (array<map<string,int>>) is not executed natively:CometArraysZip's type gate returnsUnsupportedforMapTypeelement arrays, so the entire projection falls back to Spark, even though Spark supports it.The JVM codegen dispatcher already handles map types:
CometBatchKernelCodegen.isSupportedDataTypeadmitsMapTyperecursively, andArraysZipis a plain codegen expression, so dispatching it runs Spark's owndoGenCodeand matches Spark exactly by construction.What changes are included in this PR?
CometArraysZipnow mixes inCodegenDispatchFallback: arrays of maps are routed through the JVM codegen dispatcher, while scalar/array/struct element arrays keep the existing native path.getUnsupportedReasons(): it now surfaces only when the dispatcher is disabled or rejects the tree.NullTypeelements are not dispatchable —NullTypeis absent fromisSupportedDataType— soarrays_zipoverarray<null>continues to fall back to Spark.ArraysZipfails to serialize, the serde still returnsNoneand the projection falls back (the dispatcher cannot see that state).arrays_zipis now hybrid: map elements via dispatch, other elements natively).How are these changes tested?
New sql test cases added