fix: CometBroadcastExchangeExec should respect AQE shuffle partition coalescing

[andygrove]

Summary

• Fix CometBroadcastExchangeExec to properly use AQE's coalesced shuffle partition specs
• When AQEShuffleReadExec wraps a CometShuffleExchangeExec, use getShuffleRDD(partitionSpecs) to get the coalesced RDD instead of bypassing AQE

Problem

When AQE (Adaptive Query Execution) coalesces shuffle partitions based on runtime statistics, CometBroadcastExchangeExec was bypassing this optimization by reading directly from the inner shuffle plan. This caused broadcast exchanges to spawn many tasks (e.g., 200) even when AQE determined that only 1 task was needed due to small data volume after filtering.

For example, in TPC-H Q18, after the filter sum(l_quantity) > 313, there are only ~900 rows to broadcast. AQE correctly identifies this and coalesces 200 shuffle partitions into 1. However, Comet's broadcast exchange was ignoring this and still reading from all 200 original partitions.

Solution

Added proper handling for AQEShuffleReadExec wrapping CometShuffleExchangeExec:

1. Extract partition specs from AQEShuffleReadExec
2. Use CometShuffleExchangeExec.getShuffleRDD(partitionSpecs) to get the coalesced RDD
3. Serialize batches from the coalesced RDD for broadcasting

Test plan

• Existing unit tests pass (CometExecSuite, CometShuffleSuite)
• Verified with TPC-H Q18: broadcast stages now use 1 task instead of 200 when AQE coalesces partitions

🤖 Generated with Claude Code

[codecov-commenter]

Codecov Report

❌ Patch coverage is 46.15385% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 60.01%. Comparing base (f09f8af) to head (ed7b8e0).
⚠️ Report is 864 commits behind head on main.

Files with missing lines	Patch %	Lines
...e/spark/sql/comet/CometBroadcastExchangeExec.scala	46.15%	4 Missing and 3 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #3235      +/-   ##
============================================
+ Coverage     56.12%   60.01%   +3.89%     
- Complexity      976     1426     +450     
============================================
  Files           119      170      +51     
  Lines         11743    15787    +4044     
  Branches       2251     2611     +360     
============================================
+ Hits           6591     9475    +2884     
- Misses         4012     4989     +977     
- Partials       1140     1323     +183     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mbutrovich]

Should we run TPC-H benchmarks on this one?

[wForget]

In what scenarios will the CoalesceShufflePartitions rule be applied to BroadcastExchange? Could you share the spark plan?