Repository navigation
fix: declare required child distribution on aggregate, window and join operators - #7231
Conversation
…n operators AQE's OptimizeSkewInRebalancePartitions splits a skewed shuffle partition unless an operator above needs ClusteredDistribution, which ValidateRequirements checks. Comet's aggregate, window and shuffled join operators declared no requirement, so a REBALANCE(k) followed by GROUP BY k or PARTITION BY k returned one row per piece of a skewed key. Return the original Spark operator's requiredChildDistribution.
…roup limit Move the five identical requiredChildDistribution overrides into the CometDelegatesRequiredChildDistribution trait and mix it into CometWindowGroupLimitExec as well, whose Final mode needs its partition spec clustered. Extend the test with a window group limit and a project/filter below the aggregate, and drop its temp view afterwards.
andygrove
added
the
run-spark-4.1-tests
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Delegating to originalPlan.requiredChildDistribution matches Spark's contract, including Unspecified for skew joins, and it takes effect because Comet converts the plan in AQE's stage preparation rules, before ValidateRequirements checks each read rule's result. One point on the join cases in the test.
| checkNotSplit(window) | ||
| checkNotSplit(s"SELECT * FROM ($window) WHERE rn <= 2") | ||
| Seq("true", "false").foreach { preferSmj => | ||
| withSQLConf(SQLConf.PREFER_SORTMERGEJOIN.key -> preferSmj) { |
There was a problem hiding this comment.
I don't think these two cases reach the join changes. The GROUP BY a.k puts both aggregates above the join in the same stage, and the join passes its children's partitioning through, so the aggregate's ClusteredDistribution(a.k) already rejects the split. The cases would pass with only the aggregate change.
Also, PREFER_SORTMERGEJOIN=false alone does not plan a shuffled hash join here. With AUTO_BROADCASTJOIN_THRESHOLD=-1, canBuildLocalHashMapBySize compares sizeInBytes with -1 * numShufflePartitions, and the two sides of a self-join are the same size, so muchSmaller is false too. Both iterations plan a sort-merge join, and CometHashJoinExec is never reached.
Could the join cases drop the grouping, for example SELECT a.k, a.v, b.v FROM $rebalanced a JOIN $rebalanced b ON a.k = b.k, force the hash join with "spark.sql.join.forceApplyShuffledHashJoin" -> "true" as CometJoinSuite does, and assert the expected Comet join class is in the plan? checkSparkAnswer does not check which operators ran natively, so a fallback would also pass. I have not run this.
There was a problem hiding this comment.
I completed five independent review passes on ef4199bd443a4968a8f82f59786d2a3f4fabc094 and agree that the join regression should isolate and assert each intended Comet join. I found no additional implementation defect. Spark's and Comet's defaults do not switch the second iteration to shuffled hash join, and the aggregate requirement can mask either missing join override.
One practical addition to the suggested test change: could we use a smaller dedicated skewed fixture when removing the grouping? With the current data, that ungrouped self-join produces 17,440,000 rows: 4000 * 4000 + 900 * 40 * 40. checkSparkAnswer collects both results, so simply dropping GROUP BY would make this regression unnecessarily expensive. A smaller fixture with appropriately lowered skew thresholds can still exercise the split.
Please explicitly select and assert the native join class, and verify each case fails with only that join's distribution delegation removed while the aggregate fix stays enabled. I am leaving approval pending this existing coverage issue. The current CI regression passes, but that does not establish that it protects both join changes.
There was a problem hiding this comment.
Thanks both, you were right: the aggregate above the join rejected the split by itself, and PREFER_SORTMERGEJOIN=false never planned a shuffled hash join here. Fixed in 0107cbb.
The join cases now drop the GROUP BY and join the skewed side to a one-row-per-key side, so the output is about 40k rows instead of 17M and no separate fixture was needed. The first case uses the default sort-merge join and asserts CometSortMergeJoinExec is in the plan. The second sets spark.sql.join.forceApplyShuffledHashJoin with sort-merge preferred off, and asserts CometHashJoinExec.
I checked each join on its own by removing CometDelegatesRequiredChildDistribution from one join at a time, with the aggregate change still in place. Both mutations fail the test, with PartitioningCollection requires all of its partitionings have the same numPartitions thrown from that join's outputPartitioning, which is the split reaching the join. With the real implementation the test passes.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed ef4199bd443a4968a8f82f59786d2a3f4fabc094 with five independent passes covering AQE contracts, aggregate/window semantics, joins, regression tests, and design/release behavior. I found no introduced implementation defect. Delegating to Spark preserves the partial/global operator requirements and the exception for already optimized skew joins.
The remaining issue is the focused join regression raised in the existing discussion. Both iterations ordinarily select sort-merge join. Disabling the preference does not overcome the negative broadcast-size threshold or the equal-size inputs, and AQE's default shuffled-hash local-map threshold is 0, below this test's 2k advisory size. The aggregate above the join can also reject the split independently of either join override.
Could the join cases remove the grouping, explicitly force each intended join implementation, and assert the corresponding Comet join class? Please also check that each case fails when its join's distribution delegation is removed while the aggregate change remains. That would protect both join changes independently. I am leaving approval pending this existing test issue, without opening a duplicate inline thread.
Validation: inspected the implementation and pinned Spark contracts, and checked current CI logs. The new regression ran in the execution job, with 1,774 tests passing. The additional Spark 4.1 SQL jobs also passed. CI executed merge commit a003d8b47a3bb4124a1f6d2815493347aaf3af45, combining this head with base c4a0b1f2bce44ffa2bbeab3bbc1f715853734969. I did not rerun runtime tests locally.
Drop the aggregate above the join, which rejected the split on its own, and join the skewed side to a one-row-per-key side so the output stays small. Force the shuffled hash join explicitly and assert that the Comet join operator is in the plan. Each case fails when only that join's distribution delegation is removed.
…one (apache#7424) * fix: report SinglePartition from a native Iceberg scan that replaces one On Spark 3.4, BatchScanExec reports SinglePartition for a scan of one input partition, and Spark plans no exchange above it. The native Iceberg scan reported UnknownPartitioning(0). Since the aggregates declare their required distribution (apache#7231), EnsureRequirements, which AQE runs again over a DPP subquery's converted plan, inserted an exchange inside a native block that was already serialized, and the native scan lost its tasks. Report SinglePartition when the replaced scan does. * test: run the SPARK-34637 Iceberg DPP test on Spark 3.4 again
…one (apache#7424) * fix: report SinglePartition from a native Iceberg scan that replaces one On Spark 3.4, BatchScanExec reports SinglePartition for a scan of one input partition, and Spark plans no exchange above it. The native Iceberg scan reported UnknownPartitioning(0). Since the aggregates declare their required distribution (apache#7231), EnsureRequirements, which AQE runs again over a DPP subquery's converted plan, inserted an exchange inside a native block that was already serialized, and the native scan lost its tasks. Report SinglePartition when the replaced scan does. * test: run the SPARK-34637 Iceberg DPP test on Spark 3.4 again (cherry picked from commit 891008e)
Which issue does this PR close?
Closes #7208.
Rationale for this change
With AQE on,
REBALANCE(k)followed byGROUP BY kreturns duplicate groups when the data is skewed, and aPARTITION BY kwindow restartsrow_number()inside a key. Spark'sOptimizeSkewInRebalancePartitionssplits the skewed shuffle partition. In plain Spark that is safe, becauseValidateRequirementsrejects the split when an operator above needsClusteredDistribution(k). Comet's aggregate, window and join operators declared norequiredChildDistribution, so the check always passed and the native operators ran over the pieces of a key separately.What changes are included in this PR?
A new trait,
CometDelegatesRequiredChildDistribution, returnsoriginalPlan.requiredChildDistribution. It is mixed intoCometBaseAggregateExec,CometWindowExec,CometWindowGroupLimitExec,CometHashJoinExecandCometSortMergeJoinExec, so each declares the same requirement as the Spark operator it replaces.requiredChildOrderingand the broadcast hash join are unchanged.How are these changes tested?
New test in
CometExecSuite,AQE does not split a skewed rebalance partition under operators that need clustering. It uses the data from the issue and checks theGROUP BYquery (also with a project and filter below it), therow_number()window query, a window group limit (rn <= 2), and a self-join with sort-merge join on and off. Each query is compared with Spark, and the plan is asserted to have no skewedAQEShuffleReadExec.The test fails without the operator changes and passes with them.
CometExecSuite,CometWindowExecSuite,CometJoinSuiteandCometAggregateSuitepass on the default profile (Spark 4.1): 535 tests, 0 failures.