Repository navigation
fix: [branch-1.1] initialize dispatched kernels with the partition the native plan computes (#6700) - #6876
Conversation
… computes (apache#6700) * fix: initialize dispatched kernels with the partition the native plan computes The JVM codegen dispatcher seeded kernels from `TaskContext.partitionId()`, which differs from the partition being computed under a union, a coalesce and a cartesian product. It also kept one kernel per task, so the parent partitions of a coalesce shared random state. The native planner now passes its partition index and plan id through the UDF bridge. Nondeterministic kernels are cached per plan and initialized with that index. Deterministic kernels never read it and stay shared across the plans of a task. * fix: drop a closed plan's dispatched kernels and check every node for determinism A coalesce kept the nondeterministic kernel of every parent partition until the task ended. `CometExecIterator.close` now calls `CometUdfBridge.releasePlan`, which hands the closed plan to a new `CometUDF.releasePlan` hook, and the dispatcher drops that plan's kernels there. Kernels shared by every plan stay until the task ends. The dispatcher shared a kernel whenever the root expression reported deterministic, but `Invoke` leaves its target object out of `deterministic`. So `make_valid_utf8` over `spark_partition_id()` still shared one kernel across the plans of a coalesce. It now checks every node. (cherry picked from commit 8d0bb01)
github-actions
Bot
added
bug
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Summary
- Prior state and problem: Dispatched kernels used the task partition ID and shared nondeterministic state across parent partitions, producing incorrect results under union, coalesce, and Cartesian execution.
- Design approach: Forward the computed partition index and native plan ID through JNI, isolate nondeterministic kernels by plan, and release them when that plan closes.
- Correctness: Spark sources confirm the partition-initialization semantics. Standalone probes passed for partition IDs, random sequences, monotonic IDs, nested
Invokenondeterminism, deterministic sharing, cleanup, and legacy overload delegation. No introduced P1/P2 correctness issue was identified. - Compatibility analysis: Relevant initialization and expression behavior were checked against Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0 sources. Default
CometUDFmethods preserve existing implementations. Runtime probes covered Spark 4.1.3 only. - Key design decisions: Deterministic kernels remain task-shared. Descendant nondeterminism is inspected because
Invoke.deterministiccan omit its target object. Nondeterministic entries follow native-plan lifetime. - Implementation sketch:
PhysicalPlannersupplies both identifiers toJvmScalarUdfExpr, the bridge forwards them, andCometExecIterator.close()invokes cache cleanup after native teardown. JNI argument order and reconstructed expressions preserve both identifiers. - Performance: The new shared-then-plan lookup hashes large nondeterministic closures twice per batch. A bounded dispatcher benchmark measured 2.28 ms per 4,096-row batch at head versus 1.26 ms with the base dispatcher for a captured 1 MiB array. This is the P2 below, also reported in the upstream PR.
- Design: The partition and lifetime model is straightforward and fits the existing dispatcher. The cache lookup needs adjustment to avoid the measured regression.
- Abstraction & complexity: The default evaluation overload and cleanup hook are appropriately scoped. Reusing a content hash across the two keys would preserve this design without introducing another management layer.
- Behavioral changes worth calling out: Partition-dependent expressions now initialize separately for each computed parent partition. Completed-plan kernels are removed while deterministic kernels remain available until task completion.
- Suggested improvements: Compute the closure/schema hash once and reuse it across shared and per-plan probes. Preserve the state-isolation and cleanup behavior.
Reviewed the entire ten-file diff from e9efd9f764ee0a59b7898ff028d6985d4a7a28e1 to 6bdf823283c88794b0321961d72fd813a7b622bd. The PR remains non-draft. Its reviews, conversation comments, inline comments, and review threads are empty. The linked upstream discussion was also inspected, including its hashing concern.
Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr, and review-comet-memory-pr.
Exact-head CI: 15 successful checks, four skipped, and ten running, with no failures observed. Native builds, Rust tests, and Spark SQL/Iceberg preparation remain pending. CI is not yet a completed validation verdict.
Validation: Seven standalone JVM probes passed using freshly compiled current dispatcher/codegen sources, cached supporting classes, and minimal allocator/Arrow adapters on Spark 4.1.3/JDK 21. Expression-crate cargo check and core cargo check --no-default-features passed. Default core checking failed because JNI headers were unavailable. Full Maven compilation stopped at an unwritable Zinc cache. Native integration tests and the full runtime version matrix were not run locally. Evidence is in /tmp/pr6876-review-6bdf8232/. Project source remains unchanged.
| // key serves every plan in the task. Only a kernel with a nondeterministic node is stored per | ||
| // plan. | ||
| val sharedKey = key.copy(planId = CometScalaUDFCodegen.NoPlan) | ||
| kernelCache.get(sharedKey).orElse(kernelCache.get(key)) match { |
There was a problem hiding this comment.
[P2] Reuse the closure hash across these two lookups. For a nondeterministic dispatched UDF with a large captured array, get(sharedKey) always misses before get(key) succeeds. Each case-class hash calls ByteBuffer.hashCode(), scanning the entire serialized closure again on every batch. The base performs one scan. With a captured 1 MiB lookup array, the current dispatcher took 2.28 ms per 4,096-row batch versus 1.26 ms with the base dispatcher, despite identical results within one unchanged partition. Precompute the closure/schema hash once and reuse it for both keys while preserving plan isolation.
Evidence: Reproduced with freshly compiled head and base dispatchers on Spark 4.1.3, Scala 2.13.17, and JDK 21. Define serializable Lookup(val data: Array[Byte]) extends (Int => Int) with apply(i) = data(Math.floorMod(i, data.length)).toInt. Serialize ScalaUDF(new Lookup(Array.fill[Byte](1048576)(7)), IntegerType, Seq(SparkPartitionID())) into the dispatcher’s argument-zero VarBinaryVector, then repeatedly evaluate 4,096 rows at partition 0, plan 10. After 150 warmups per implementation, nine alternating samples of 200 batches produced medians of 1.25946982 ms for base and 2.281142645 ms for head. The serialized expression was 1,051,614 bytes. Both used the same current codegen and supporting classes. A separate extracted-cache probe corroborated the extra hashing cost. Source and results: /tmp/pr6876-review-6bdf8232/DirectProbe.scala, direct.log, CacheProbe.scala, and validation.json. This is a dispatcher benchmark, not an end-to-end query timing.
There was a problem hiding this comment.
Does this issue also exist in main branch? This PR is backporting an already approved commit to main
3 participants
Which issue does this PR close?
Closes #6570.
Backport of #6700 to
branch-1.1.Rationale for this change
The JVM codegen dispatcher initialized kernels from the Spark task partition rather than the partition computed by the native plan. Those indexes can differ under union, coalesce, and cartesian product, producing incorrect results for partition-seeded expressions and sharing random state across coalesced parent partitions.
What changes are included in this PR?
Cherry-picks #6700 with
-x. The native planner now passes its partition index and plan ID through the UDF bridge, nondeterministic kernels are cached per native plan, and closing a plan releases its kernels. Deterministic kernels remain shared within a task.The 1.1 branch does not contain a later Parquet JVM-UDF test touched only to update a constructor call, so that inapplicable test-only hunk is omitted. The backport also adds the existing
isSpark40Plushelper import needed by the new regression test on this branch.How are these changes tested?
Ran:
Result: 3 suites completed, 171 tests passed.