Conversation
8ebda6f to
bd2c151
Compare
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
This addresses repeated protobuf decoding and physical planning when many Spark tasks execute the same native block. The implementation separates an executor-local decoded-plan cache from a stage-attempt-scoped registry of shared DataFusion operators. Both flags default to false.
The identity combines a driver-generated block UUID, stage ID, stage-attempt number, serialized plan, configuration, batch size, partition count and task CPU count. That agrees with the maintained Spark 3.5 and 4.0 task and scheduler contracts: speculative copies and retries have increasing attempt numbers, while a stage retry gets a separate stage-attempt identity. Attempts above zero use private plans, and a shared tree claims each partition only once.
The execution boundary is narrow. JVM readers, streams, task contexts and memory pools stay in each attempt. SharedInputExec resolves them through that attempt's DataFusion task context. Admission rejects partition-dependent expressions, UDFs, subqueries, native file scans, shuffle operators, Top-K and positive sort offsets. Partitioned joins allocate their build state per execution. Full sorts preserve partitions, and final aggregates become FinalPartitioned when needed. I found no verified result or cross-task ownership defect in those admitted paths.
One integration issue needs correction: the appended JNI parameter conflicts with the existing shuffle signature test. The inline P2 identifies the exact failing assertion. CI at merge 4136f5ea, with parents 5ca14992 and bd2c151a, reports 1,626 native tests passing, including all 11 decoded-cache and 22 shared-tree tests. The Spark 4.1 execution job passed 980 tests. The shuffle job passed 500 and failed one, so Required Checks remains failed. Maintained Spark 3.4 and 4.1 branches were unavailable for source comparison. I did not run a local product suite.
Performance
The PR reports a matched off/on microbenchmark with 1,024 partitions and 32 rows per partition. With both flags enabled, it reports lower task CPU and cumulative requested allocation, and query elapsed time 4.0% and 12.6% lower for the 16-column and 64-column projection/filter cases. Those are author-reported results. They do not isolate physical sharing from decoded caching, and they do not establish lower peak memory or a benefit for wide aggregates.
The second inline P2 concerns stage-size scaling in metrics reporting. A task clones all accumulated partition metrics before selecting its own. A focused mechanism benchmark measured median snapshot/filter time of 0.19 µs for one partition and 1.10 ms for 32,768 retained partitions, with eight metrics per partition. This uses the exact selection branch and DataFusion's snapshot representation with small dependency stand-ins. It is not an end-to-end Spark result. Please qualify the repair with a matched cache-only versus sharing benchmark that keeps one tree alive across many task waves.
Design
The two caches have different purposes and lifetimes. The decoded cache retains immutable protobuf definitions between task waves and uses per-key initialization outside its map lock. The physical registry retains weak references, which prevents an idle tree from keeping task resources or operator metrics alive. It also means a gap with no active attempt ends reuse and the next wave recompiles the tree.
Physical construction currently holds the process-wide registry mutex. The source and a focused lock-scope probe confirm that an unrelated warm lookup waits behind a cold build. This is an explicit tradeoff in the proposed design. The reported single-plan benchmark does not measure its effect with concurrent unrelated plans. The key limits bound the registry, while operator metrics remain live for the lifetime of an active shared tree, which is the reason the reporting-cost finding matters.
Abstraction & complexity
SharedInputExec keeps the execution-time boundary small and reuses upstream operators instead of introducing parallel implementations of their kernels. Rebuilding the tree once to install input boundaries and preserving the Spark-to-native metric mapping are reasonable responsibilities for this layer. The conservative expression allowlist also gives future extensions a clear audit boundary.
The tests exercise actual operator identity, sparse partitions, cancellation, retry isolation, aggregate modes, join build sides, spill cleanup and metric ownership. The main improvements to address in this version are reconciling the JNI signature test and making task metric reporting independent of previously completed partitions. The performance evidence should then include the large-stage control described above.
| classLoader: ClassLoader, | ||
| sharedPlanScope: String = ""): Long |
There was a problem hiding this comment.
Correctness
[P2] Could you reconcile the existing signature test with this API change? CometNativeShuffleSuite still asserts that createPlan.getParameterTypes.last is ClassLoader. Appending sharedPlanScope makes it String, even with both new flags disabled. The current Spark 4.1 shuffle job fails at CometNativeShuffleSuite.scala:109 with exactly that mismatch, leaving Required Checks failed. Please update the test to the intended signature, or pass the scope through the existing configuration argument if preserving the signature remains the contract.
There was a problem hiding this comment.
Fixed in c2e48d2. I missed this existing assertion when adding sharedPlanScope. The test now checks the trailing TaskContext, ClassLoader, and String parameters and has a name that reflects the updated signature. The full CometNativeShuffleSuite passed locally on Spark 4.1: 57 tests, 0 failures.
| plan.metrics().map(|metrics| { | ||
| let mut selected = MetricsSet::new(); | ||
| for metric in metrics | ||
| .iter() | ||
| .filter(|m| m.partition() == Some(self.partition)) |
There was a problem hiding this comment.
Performance
[P2] Could each attempt retain or access only its own metric handles instead of snapshotting the entire shared tree here? DataFusion 55.1.0's plan.metrics() clones its Vec<Arc<Metric>> under the operator's mutex before this filter runs. Every executed partition appends metrics, and releasePlan calls this path for every task. When overlapping task waves or one long task keep the tree alive, the nth task therefore clones and scans metrics for all n partitions. Total completion-reporting work becomes quadratic in the partitions processed by that tree, whereas private plans only inspect their own metrics.
A focused benchmark of this branch and snapshot representation, using dependency stand-ins and eight metrics per partition, measured median flush times of 0.19 µs, 26.53 µs, 233.53 µs and 1.10 ms at 1, 1,024, 8,192 and 32,768 retained partitions. These are mechanism timings, not Spark query timings. The 64-entry and encoded-key limits do not bound the metrics in one live tree. Please make the reporting cost independent of completed partitions and add a matched large-stage benchmark with a long-lived task plus a cache-only control.
There was a problem hiding this comment.
Agreed. I reproduced this with actual DataFusion operators in a release build, with partition 0's stream kept alive while the remaining partitions execute and complete. For a 16-expression projection/filter tree, the median cumulative snapshot/filter time over three runs was 51.3 ms at 1,024 partitions and 3.90 s at 8,192 partitions. The matched decoded-cache-only control took 0.142 ms and 1.11 ms, respectively. These are native mechanism measurements, not Spark query timings.
The current DataFusion metrics API returns a full snapshot, so merely caching each task's handles after an initial full scan would move the quadratic work rather than remove it. I am evaluating options without modifying DataFusion and am adding the matched Spark cache-only versus sharing comparison with a long-lived task. This issue remains unresolved; the existing benchmark results do not justify the current reporting path.
There was a problem hiding this comment.
Thanks for identifying this. After investigating the alternatives, we plan to propose an upstream DataFusion PR to add partition-specific metrics access. The intent is to support indexed retrieval in ExecutionPlanMetricsSet and expose it through ExecutionPlan, so Comet can retrieve a task's metrics without cloning and scanning metrics from all previously executed partitions. A convenience API that still filters a full snapshot would not address the underlying cost.
We also experimented with bounding the number of partitions sharing each tree. That removed the large-stage regression in our local benchmark, but it sacrifices plan reuse and introduces scheduling-dependent tradeoffs. We would prefer to address the metrics API in DataFusion rather than make that workaround part of Comet's design.
Both new configs are currently disabled by default, so this does not affect users running with the defaults. Users explicitly enabling physical-plan sharing can still encounter this issue; it remains unresolved until the upstream API improvement is available and integrated. We will follow up with the DataFusion PR and rerun the matched large-stage/long-lived-task benchmark against the cache-only control after integration.
|
Review follow-up:
These are pooled medians from two sequential JVMs on M4 Max, Spark 4.1.3, The separate release-mode native probe also completed three repetitions per configuration with real operators and a held partition-0 stream. Cumulative snapshot/filter time in sharing mode was 51.3 ms, 3.90 s and 133.6 s at 1,024, 8,192 and 32,768 partitions. Cache-only was 0.142 ms, 1.11 ms and 4.84 ms. These are mechanism timings over the completed short partitions, not Spark query timings. The metrics finding remains open. The long-task Spark result confirms a regression, while the native probe establishes the snapshot scaling mechanism. The short-task gains do not justify merging the current reporting path. A bound on the number of partitions served by each physical tree is a possible approach without changing DataFusion, but its reuse/construction tradeoff still needs validation. I have not applied that change or treated this issue as resolved. Per-key registry build slots address a separate lock-contention issue. |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Follow-up at c2e48d2a (base 5ca14992): the P2 JNI-signature test finding is fixed. The assertion now checks the trailing TaskContext, ClassLoader, String parameters, and the current shuffle CI job passes that test. The revision changes only this test; production code is unchanged from bd2c151a.
The P2 metrics-scaling finding remains. Each task still snapshots all metrics accumulated on the shared tree before filtering its partition. A long-running task keeps that tree alive across task waves, and final releasePlan reporting bypasses the periodic update throttle. The new matched measurements include the needed decoded-cache-only control: the author reports elapsed time 35.5% higher and task CPU 21.6% higher with sharing in the 8,192-partition long-task case, alongside lower elapsed times in the reported 1,024-partition cases. I have not independently rerun those benchmarks. Bounding partitions per tree or making metrics retrieval partition-indexed still needs an implementation and the same matched validation.
No new findings. Current checks show 23 successes and 14 skips. I verified 11 cache and 22 sharing Rust cases, 980 execution tests, and 501 shuffle tests in CI logs. Those jobs actually checked out merge e1f2ab0c, whose first parent is 09b44ad6, rather than the authoritative base above. Eleven of the 14 authored files match the reviewed head exactly; I also inspected the inherited tracing/config differences in the other three. No local full suite was run. Maintained Spark 3.5/4.0 sources were checked; maintained 3.4/4.1 branches were unavailable. Keeping this review at COMMENT for the existing metrics P2.
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
Thanks for the thorough follow-up on the earlier rounds. The partition claims and stage-attempt scoping are careful, and the native tests cover a lot of lifecycle ground. I don't want to repeat the open metrics-scaling discussion, so this review covers other things. Inline comments cover specific code. Two broader points:
Applicability with default settings. spark.comet.shuffle.directRead.enabled defaults to true, so any stage reading a Comet shuffle gets a ShuffleScan input (CometSink.shouldUseShuffleScan). Admission rejects that, and native file scans too. That seems to leave only JVM-input stages whose expressions are limited to arithmetic and comparisons, with no Cast, Divide, CaseWhen, In or string functions. Could you report how many native blocks are admitted in TPC-H and TPC-DS with default settings? That would help us judge whether the added complexity pays off for real workloads, since the current benchmarks use spark.range with 32 rows per partition.
Scope. Could the decoded-plan cache and physical-plan sharing be split into two PRs? The cache is small and independent. Sharing is much larger, still has the open metrics-scaling issue, and relies on DataFusion operators keeping per-partition execution independent, which isn't an upstream contract. Your numbers compare sharing against cache-only but never cache-only against off, so we can't yet see what the cache gives on its own.
| )?; | ||
| let (scans, shuffle_scans, root_op) = | ||
| if let Some(key) = &exec_context.shared_plan_key { | ||
| let shared = super::shared_pipeline::get_or_build( |
There was a problem hiding this comment.
If the shared build fails, for example when convert_tree hits Unexpected operator in shared tree, the ? here fails the task. Could we fall back to planner.create_plan in that case and log it? Admission and the planner agree today, but a later planner change that adds a wrapper operator would otherwise break queries whenever the flag is on.
There was a problem hiding this comment.
Implemented in ffdbbd2. Shared-tree construction/conversion errors now log a warning and fall back to private planning. The regression test introduces an unsupported wrapper around a real DataFusion tree, verifies that conversion fails without leaving a registry entry, and verifies the private plan's output. Binding and execution errors still propagate: once task-owned input streams may have been imported, retrying them would not be safe.
| static PLAN_CACHE: LazyLock<PlanCache<Operator>> = | ||
| LazyLock::new(|| PlanCache::new(MAX_ENTRIES, MAX_ENCODED_BYTES)); | ||
|
|
||
| pub(super) fn decode_plan( |
There was a problem hiding this comment.
CometExecRDD.compute injects per-partition planning data for native scans (PlanDataInjector.injectPlanData), so every task in those stages sends different bytes. With the cache on, each of those tasks misses, copies its full plan into a new entry, and evicts entries that other stages could have reused. Could we skip the cache when per-partition plan data was injected, or key on the base plan instead? A test showing that a stage with distinct per-task bytes doesn't evict a reusable entry would help guard this.
There was a problem hiding this comment.
Addressed in ffdbbd2 by removing the decoded-plan cache and its config entirely. The standalone measurements showed only approximately 0.2%–1.4% lower elapsed time than both flags disabled, which does not justify maintaining this cache. Injected per-task plan bytes are now decoded without entering a native definition cache, so they cannot evict reusable definitions. The physical registry still rejects native file scans.
| } | ||
| } | ||
|
|
||
| test("shared DataFusion stateful operators across Spark partitions") { |
There was a problem hiding this comment.
These tests use checkSparkAnswerAndOperator, which passes whether or not a shared tree was used. I think the post-shuffle Final aggregates and joins here read through ShuffleScan, so they probably run on private plans even with the flag on. Could we expose a small counter or test hook for shared-tree binds and assert on it? Asserting the expected fallback cases would be useful too.
There was a problem hiding this comment.
Added shared_plan_tasks SQL metrics and explicit path assertions in ffdbbd2. Projection/filter tests assert the exact task count; positive sort, post-shuffle Final aggregate, and broadcast hash join tests arrange eligible JVM inputs and assert nonzero counts only when sharing is enabled, with AQE both on and off. A separate test verifies that the executed block actually contains a serialized ShuffleScan and asserts zero shared tasks. Native-scan and stateful-expression fallback cases also assert zero.
This metric counts tasks bound to a shared tree, not registry hits; native tests continue to verify actual operator identity. The broader output-equivalence tests have been renamed so their names do not imply that every case was admitted. I also found that the direct-read setting alone is not enough to infer the input representation: some exchange plans still serialize as ordinary Scan inputs, so the fallback test checks the serialized plan itself.
| withSQLConf( | ||
| CometConf.COMET_SHUFFLE_ENABLED.key -> "true", | ||
| CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "true", | ||
| CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> "true") { |
There was a problem hiding this comment.
This override (and the matching one in CometExecIteratorLifecycleSuite) turns on both flags for every existing test in the suite. The flag-off path, which is what users run by default, loses that coverage here. Could we enable the flags only in the new tests, or run the affected tests in both modes?
There was a problem hiding this comment.
Fixed in ffdbbd2. The suite-wide sharing/cache overrides are removed, restoring default-path coverage for existing tests. The cache flag itself is gone. New sharing tests enable the feature locally, and the lifecycle regressions run with sharing disabled and enabled. The lifecycle test's synthetic TaskContext now carries SQLConf through local properties, as real Spark tasks do; otherwise its empty properties silently selected the default setting.
| "lifecycle-test") | ||
| val limitOp = | ||
| CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get | ||
| val scanOp = |
There was a problem hiding this comment.
Could you explain the change from the limit plan to a bare scan plan here? It looks like the original limit case is no longer covered by this lifecycle test.
There was a problem hiding this comment.
You are right: I replaced the limit with a scan to enter an admitted sharing path, but that should have been an additional case. In ffdbbd2, the limit case is restored and tested with sharing off/on. A separate scan case is also tested in both modes, and inspects the final native metrics report to assert whether a shared tree was actually bound before the injected metrics-update failure. All nine lifecycle tests pass.
Remove the decoded-plan cache and its config after the standalone measurements showed little benefit. Shared construction failures now log and use private planning before any task input streams are imported. Report shared_plan_tasks through SQL metrics and assert sharing and fallback paths in JVM tests. Restore default-mode suite coverage and the original limit lifecycle case; exercise scan cleanup separately with real shared binding. Synthetic TaskContexts now receive SQLConf through task local properties.
viirya
added
the
run-spark-4.1-tests
|
Review follow-up in ffdbbd2:
Still outstanding: sharing-only versus disabled measurements after cache removal; TPC-H/TPC-DS admission and actual-bind coverage with normal defaults preserved apart from enabling sharing; and documenting/validating the execution-state assumptions for admitted DataFusion operators. The targeted positive-path tests deliberately configure eligible inputs and are not evidence of default-workload coverage. The metrics-scaling finding remains open. We plan to address it through partition-specific metrics access in DataFusion, integrate that API, and rerun the matched large-stage/long-lived-task benchmark. The remaining sharing config defaults to false, which limits default-user exposure but does not resolve the performance issue for users who enable it. I updated the PR description to distinguish current validation from historical measurements and outstanding work. |
3 participants
Which issue does this PR close?
Closes #1204.
Rationale for this change
Comet repeats physical plan construction for tasks executing the same native block. This PR allows eligible tasks within the same executor, Spark stage attempt, and native block to execute the same DataFusion physical plan using their partition IDs.
This remains experimental and disabled by default. The open metrics-scaling finding must be addressed: current DataFusion metrics access snapshots all retained partition metrics before Comet selects the current partition. We plan to address that through upstream partition-specific metrics access, then integrate and validate it here. Default-workload applicability and the execution-state assumptions for admitted operators also need further validation.
What changes are included in this PR?
shared_plan_tasksSQL metrics to distinguish shared binds from private execution. These count task binds, not registry hits or the number of concurrent users of a tree.spark.comet.exec.sharedPlan.enabledinternal and false by default. The decoded-plan cache and its config have been removed; protobuf plans are decoded per task.The registry is limited to 64 weak entries and 8 MiB of encoded keys. Last-owner release reclaims each tree and its metrics; idle gaps may cause rebuilding. Initial construction holds the registry mutex, so unrelated lookups can wait behind a cold build. These limits do not bound metrics accumulated within a live tree.
How are these changes tested?
Current revision:
--lib --tests -- -D warnings), JVM formatting/style checks, and CI suite-registration checks passed.run-spark-4.1-testslabel has been applied to request the Spark SQL gate for this revision; its result is pending.Performance evidence collected before removal of the decoded cache is recorded in the review discussion. The matched sharing-versus-cache-only measurements showed benefits for short-task projection/filter workloads, but also a 35.5% elapsed-time regression in the 8,192-partition long-task case. Cache-only versus both disabled showed only approximately 0.2%–1.4% elapsed-time reductions, motivating removal of that cache.
Those measurements do not validate this revised sharing-only implementation. Sharing-only versus disabled measurements, TPC-H/TPC-DS admission and actual-bind coverage with normal defaults, and revalidation after the upstream metrics improvement remain pending. No lower peak-memory or general production speedup is claimed.