fix: do not resolve DPP subqueries when AQE reads a native scan's partitioning

[dwsmith1983]

Which issue does this PR close?

Part of #6133. This fixes the q5 failure. The q64 row loss has a separate cause, tracked in #6264.

Rationale for this change

AQE applies Spark's own query stage rules before Comet's. When CoalesceShufflePartitions changes a stage, optimizeQueryStage runs ValidateRequirements, which reads outputPartitioning on every node in the stage. CometPlanAdaptiveDynamicPruningFilters has not run yet at that point, so a scan's DPP subquery is still Spark's SubqueryAdaptiveBroadcastExec placeholder.

For a non-bucketed scan, CometNativeScanExec.outputPartitioning returned UnknownPartitioning(perPartitionData.length). Computing that length runs InSubqueryExec.updateResult on the placeholder, which throws SubqueryAdaptiveBroadcastExec does not support the execute() code path. That is the stack trace in the issue.

Coalescing skips shuffle reads that share a subtree with a scan, and CoalesceShufflePartitions only splits a stage into separate groups at Spark's UnionExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec and CartesianProductExec. So the failure needs one of those in the scan's stage. q5 has a union and hits it when Comet's union is not used.

FileSourceScanExec does not have this problem: it reports UnknownPartitioning(0) for a non-bucketed scan and only resolves DPP when it builds its RDD.

CometIcebergNativeScanExec fails the same way. Its outputPartitioning counted the scan's tasks, which goes through BatchScanExec.inputRDD and translateRuntimeFilterV2 and throws Can't translate ... IN dynamicpruning ... no subquery result because the subquery has no result yet. On Spark 3.5 and later, BatchScanExec without a key-grouped partitioning reports UnknownPartitioning(0).

What changes are included in this PR?

• CometNativeScanExec.outputPartitioning reports UnknownPartitioning(0) for a non-bucketed scan, as FileSourceScanExec does. A bucketed scan still reports its bucket layout.
• CometIcebergNativeScanExec.outputPartitioning reports UnknownPartitioning(0).
• CometNativeExec sizes the native RDD of a block whose leaf is a CometScanWithPlanData or a CometIcebergNativeScanExec from perPartitionData.length, at execution, after findAllPlanData has resolved the scan's subqueries. It used to read that count from outputPartitioning. A bucketed scan has one entry per bucket, so its count is unchanged.

Planning now sees 0 partitions for a non-bucketed native scan and for an Iceberg native scan, the same as for Spark's scans. The one exception is Spark 3.4, where BatchScanExec reports SinglePartition for a single input partition and the Iceberg native scan reports UnknownPartitioning(0). In particular a single split scan no longer satisfies AllTuples while planning, which also matches Spark. The execution paths that need a count take it from the RDD.

How are these changes tested?

Two tests in CometExecSuite, next to the existing AQE DPP tests:

• "AQE DPP: shuffle coalescing in the scan's stage keeps DPP working" joins a partitioned fact table to a filtered dimension and unions it with an aggregate, with Comet's union disabled so Spark's UnionExec keeps the aggregate's shuffle read in the scan's stage. On main it fails with the stack trace from the issue. With the change it matches Spark, and on Spark 3.5 and later it checks that the DPP scan and a coalesced shuffle read share the final stage, and that DPP still prunes to 1 of 10 partitions.
• "AQE DPP: bucketed scan keeps its partitioning for a sort-merge join" checks that a bucketed scan with DPP still feeds a sort-merge join without a shuffle. It passes on main and guards the bucketed branch.

CometIcebergNativeSuite gets the same shuffle coalescing test for the Iceberg native scan. Before the Iceberg change it fails with the translateRuntimeFilterV2 error on Spark 3.5 and 4.1. With it, the answer matches Spark, the Iceberg scan and a coalesced shuffle read share the final stage, and the scan plans 1 of the 10 file tasks.

CometExecSuite, CometScanWithPlanDataSuite, CometTopKSuite, CometJoinSuite, CometIcebergNativeSuite and the DPP fallback suites pass on Spark 3.5. The AQE DPP tests, the Iceberg DPP tests and CometTopKSuite pass on Spark 4.0 and 4.1, and the AQE DPP and Iceberg DPP tests pass on Spark 3.4.

[sunchao]

Summary

• Prior state and problem: AQE could read scan partitioning before DPP subqueries were executable, causing native Parquet and Iceberg scans to fail during stage validation.
• Design approach: Report UnknownPartitioning(0) during planning and obtain execution partition counts from resolved perPartitionData.
• Correctness / compatibility analysis: Checked Spark sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0, plus standalone scans, fused execution, broadcast sizing, shuffle inputs and limit consumers. The change matches non-bucketed V1 scans and ordinary V2 scans on Spark 3.5+.
• Key design decisions: Preserve bucketed partitioning and resolve subqueries before reading execution data. Reusing the existing lazy scan data adds no per-row work or unnecessary abstraction.
• Implementation sketch: Update both scan partitioning overrides, add scan-specific sizing cases in buildNativeContext, and cover coalescing and bucketed joins with three regression tests.
• Behavioral changes worth calling out: A single-split native scan no longer advertises an AllTuples guarantee through its partition count. Spark 3.4 V2 scans can report SinglePartition, so the Iceberg override is deliberately more conservative there.
• Suggested improvements: No introduced P1/P2 issues found within this review.

Reviewed both commits and all five changed files in the full diff from base 605051ad239ef704f5f25d67910a446a6b6d7c70 to head a3e1ea43e3073c841565c4ed3ec604679c7b4e90. The PR is non-draft. No existing reviews, comments or threads were present. Routed skills: review-comet-pr and review-comet-shuffle-pr for partitioning and execution sizing.

Exact-head CI: the label check passed. Comet CI, CodeQL and the title workflow remain action_required; Comet CI has no executed jobs or build/test verdict.

Validation: make core passed. The Spark 4.1.3 reactor build and 53 focused DPP/plan-data tests passed, including all three new regressions, with no failures or cancellations. git diff --check passed. Other Spark profiles were checked against source only. Full Spark SQL/Iceberg upstream suites and platform matrices were not run; ScalaStyle and Spotless were skipped. The working tree remains clean.

[dwsmith1983]

@andygrove could you approve a CI run on a88941e2e? CI has not run on this PR yet. It now also covers the native Iceberg scan, which failed the same way.

[dwsmith1983]

@sunchao could you approve a CI run on 588c029f7? CI has not run on this PR yet, and it fixes the q5 failure from #6133.

[Neuw84]

We tested this PR together with #6270 on TPC-DS at SF1000: Parquet on S3, a Spark 4.1 cluster on EKS, 8 executors x 13 cores on m5.4xlarge, one availability zone, AQE on with 300 shuffle partitions.

The build was main b58b2f3a with #6268 (588c029f) and #6270 (b3f4f058) merged on top (HEAD 0f22d064), with the native library built for x86-64-v3. Comet ran with native scan, exec and shuffle enabled.

• q5 now completes. 23.1 s, 100 rows, the same checksum as vanilla Spark, with 77/77 operators on Comet (11x CometNativeScanExec). Before this fix, q5 failed at this scale with native scans.
• The whole run passes. All 103 queries complete. Row counts match vanilla Spark on every query, and checksums match on all but q65, which has ties in its ORDER BY and differs the same way for every engine we compare.

Thanks for the fix.

[andygrove]

I would like to do a deeper review of this PR before it is merged.

[sunchao]

Summary

• Prior state and problem: AQE requirement validation could force unresolved DPP subqueries through native scan partition counting, causing Parquet and Iceberg queries to fail.
• Design approach: Report UnknownPartitioning(0) during planning and obtain execution partition counts from resolved scan data.
• Correctness / compatibility analysis: Checked Spark sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. Traced standalone scans, fused execution, broadcast alignment, shuffle inputs, unions and limit consumers. Execution retains the actual task count.
• Key design decisions: Preserve bucketed Parquet partitioning and resolve subqueries before accessing scan data. Reusing existing lazy data adds no per-row work or unnecessary abstraction.
• Implementation sketch: Update both scan partitioning overrides, add scan-specific sizing in buildNativeContext, and add three regression tests covering coalescing, pruning and bucketed joins.
• Behavioral changes worth calling out: Compared with branch-1.1 at 992c806a7e38c2e88bd018aa5774164b0850e1fa, the intended changes avoid premature DPP execution and stop advertising an AllTuples guarantee for single-split non-bucketed scans. Spark 3.4 V2 scans can report SinglePartition, so the Iceberg override is deliberately more conservative there.
• Suggested improvements: No introduced P1/P2 issues found within this review.

Reviewed all five files in the full diff from base fef94f6cd78b18151dff57b7a936798385356de5 to head 6a074c6b3fd62295a6b2dd802a2c26f2959a43bd. The PR is non-draft. Read the supplied reviews, issue comments and empty inline/thread collections. No substantiated existing P1/P2 concern remains unresolved. Routed skills: review-comet-pr, review-comet-shuffle-pr and review-comet-ffi-pr.

Exact-head CI: the label check passed. Comet CI and CodeQL report action_required, providing no build/test verdict.

Validation: a fresh native build and Spark 4.1.3 reactor run passed all 32 TopK and native-shuffle input tests. Verified existing exact-head reactor results for 57 DPP and scan-data tests, including all three new regressions. No failures or cancellations. git diff --check passed and the working tree is clean. Other Spark profiles were checked against source only. Full upstream Spark SQL/Iceberg suites and platform matrices were not run. ScalaStyle and Spotless were skipped.

[andygrove]

I went through this in more depth as I said I would. Without the change, the new coalescing tests fail on 3.5 and 4.1 with the errors from #6133, and the Iceberg one fails on 3.4 too. With it they pass on 3.4, 3.5 and 4.1, and so do Spark's DynamicPartitionPruningSuite and the bucketed read suites with Comet on. Nothing at execution still takes the partition count from outputPartitioning. EnsureRequirements never sees a Comet scan, so the only planning reader is ValidateRequirements, and this can only make that check stricter.

CI hasn't run at the current head, and the Spark SQL and Iceberg suites have never run on this PR. Since this changes what every native scan reports, I've added run-spark-4.1-tests, run-iceberg-tests and run-all-spark-profiles.

While checking the Spark 3.4 SinglePartition case in the description, I found an older bug nearby that this PR doesn't cause. I filed it as #6667.

[andygrove]

Now that buildNativeContext sizes every CometScanWithPlanData from perPartitionData, could this doc also say that outputPartitioning must not read perPartitionData, or anything else that runs the scan's DPP subqueries? AQE reads it while optimizing a stage, before CometPlanAdaptiveDynamicPruningFilters has converted the placeholders, and that's the q5 crash. The Delta scan in #5365 copied the old pattern and works around it with hasUnevaluableSubqueryFilter. With the rule written here it could just report UnknownPartitioning(0) like CometNativeScanExec does now.

[dwsmith1983]

could this doc also say that outputPartitioning must not read perPartitionData, or anything else that runs the scan's DPP subqueries?

Added in 2eb4ff6. The CometScanWithPlanData doc now says perPartitionData.length is what sizes the native RDD at execution, that outputPartitioning must not read perPartitionData or anything that evaluates the scan's DPP subqueries because AQE calls it while optimizing a stage before CometPlanAdaptiveDynamicPruningFilters has converted the placeholders, and that a scan that cannot report a real partitioning without them should report UnknownPartitioning(0), as CometNativeScanExec does for a non-bucketed scan.

I pushed this while the label runs were still going so the doc change would not wait on them. It only touches that comment, so the code those runs are testing is unchanged.

[dwsmith1983]

The failure on every profile was CometTaskMetricsSuite "coalesced native scan reports rows and bytes from every partition it reads", which #6547 added after this branch last merged main. Its precondition read outputPartitioning.numPartitions on the executed scan and expected 4, and this PR makes a non-bucketed scan report UnknownPartitioning(0). e4bdd80 has the test read perPartitionData.length, the count the native RDD is sized from, which is what the check meant. The suite passes with the change on the Spark 4.1 profile.

@andygrove can you add it back to the merge queue? It should be fine now.