wangyum opened a new pull request, #59636: URL: https://github.com/apache/spark/pull/59636
### What changes were proposed in this pull request? `CoalesceBucketsInJoin` could read the bucket count of a broadcast join's build side and stamp the streamed side's bucketed scan with it, when `spark.sql.bucketing.coalesceBucketsInJoin.enabled` is true and a sort-merge or shuffled-hash join's child is a `BuildLeft` broadcast hash join or broadcast nested-loop join. This PR makes `ExtractJoinWithBuckets` find each join input's bucketed scan by walking only the child that actually exposes its partitioning to the parent join: through `FilterExec` and `ProjectExec`, and into the streamed child of a broadcast hash or nested-loop join. The extractor now returns that scan together with its bucket count, and the rewrite updates only that one scan, so a bucketed build scan keeps its own bucket count. ### Why are the changes needed? This is a correctness bug. `getBucketSpec` used `collectFirst`, whose pre-order walk visits the left child first. With a `BuildLeft` broadcast join, the left child is the broadcast build side, so a bucketed build scan's count was reported as the join input's bucket count, and `updateNumCoalescedBucketsInScan` then stamped every bucketed scan in that subtree, including the streamed scan whose partitioning the parent join actually uses. The coalesced scan groups bucket files by their original bucket id modulo the coalesced count while advertising `HashPartitioning` of that count. When the streamed table's bucket count is not divisible by the coalesced count (e.g. 100 coalesced to 8), `hash(key) % 100 % 8` and `hash(key) % 8` land in different partitions for some keys. The other side of the sort-merge join is a real table of the coalesced bucket count, `EnsureRequirements` inserts no shuffle, and the join silently drops matching rows (or null-pads them for outer joins). Repro: ```sql SET spark.sql.adaptive.enabled = false; SET spark.sql.autoBroadcastJoinThreshold = -1; SET spark.sql.bucketing.coalesceBucketsInJoin.enabled = true; SET spark.sql.optimizer.excludedRules = org.apache.spark.sql.catalyst.optimizer.ReorderJoin; CREATE TABLE dim16 (i INT, k INT) USING parquet CLUSTERED BY (i) INTO 16 BUCKETS; CREATE TABLE fact100 (i INT, k INT, v STRING) USING parquet CLUSTERED BY (i) INTO 100 BUCKETS; CREATE TABLE other8 (i INT, v STRING) USING parquet CLUSTERED BY (i) INTO 8 BUCKETS; INSERT INTO dim16 VALUES (0, 1); CREATE OR REPLACE TEMP VIEW picked AS SELECT CAST(id AS INT) AS i FROM range(0, 4000) WHERE pmod(pmod(hash(CAST(id AS INT)), 100), 8) <> pmod(hash(CAST(id AS INT)), 8) ORDER BY id LIMIT 1; INSERT INTO fact100 SELECT i, 1, 'fact' FROM picked; INSERT INTO other8 SELECT i, 'other' FROM picked; SELECT /*+ BROADCAST(dim16), MERGE(fact100, other8) */ fact100.i, fact100.v, other8.v FROM dim16 JOIN fact100 ON dim16.k = fact100.k JOIN other8 ON fact100.i = other8.i; -- returns 0 rows before this PR; 1 row after ``` The bug also blocked a valid coalescing: with a 16-bucket streamed scan hidden behind a 100-bucket build scan, the old code saw 100 first and declined the whole pair instead of coalescing 16 to 8. Note `spark.sql.bucketing.coalesceBucketsInJoin.enabled` defaults to false, so default plans are not affected. ### Does this PR introduce _any_ user-facing change? Yes, a bug fix: with bucket coalescing in join enabled, queries whose sort-merge or shuffled-hash join reads through a `BuildLeft` broadcast join no longer drop or null-pad rows that should match, and one additional valid coalescing case (streamed 16 buckets coalesced to 8 behind a 100-bucket build scan) now avoids a shuffle. No API, config, or default behavior change. ### How was this patch tested? Existing and new tests in `CoalesceBucketsInJoinSuite`: - `SPARK-60382: broadcast join coalescing uses the streamed side bucket count` builds the physical plans directly and checks that a 100-bucket streamed scan against an 8-bucket other side is left untouched in both `BuildLeft` and the divisible 16-to-8 case only rewrites the streamed scan, for both broadcast hash and broadcast nested-loop joins. - `SPARK-60382: broadcast hash join bucket coalescing keeps the sort-merge match` runs the end-to-end SQL above and asserts the matching row is returned while the 100-bucket scan is not stamped with 8. - All existing tests in the suite still pass: `build/sbt 'sql/testOnly org.apache.spark.sql.execution.bucketing.CoalesceBucketsInJoinSuite'`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Cursor (GLM 5.3) Made with [Cursor](https://cursor.com) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
