Describe the bug
With AQE off, a dynamic partition pruning subquery doesn't reuse the join's CometBroadcastExchange when the dimension is a UNION ALL. The dimension is computed and broadcast a second time for the DPP subquery.
CometExecRule.rewriteInSubqueryPlan only replaces the SubqueryBroadcastExec with a CometSubqueryBroadcastExec when the exchange's child, after stripping the columnar-to-row transition, is a CometNativeExec. CometUnionExec extends CometExec but not CometNativeExec, so the check fails and the subquery keeps a Spark BroadcastExchange over CometColumnarToRow(CometUnion(...)). ReuseExchangeAndSubquery can't match that to the join's CometBroadcastExchange. CometCoalesceExec, CometTakeOrderedAndProjectExec and a bare CometSparkToColumnarExec would hit the same check.
The same query with AQE on reuses the broadcast through CometPlanAdaptiveDynamicPruningFilters. #6786 adds the same CometNativeExec check to that rule, which would make the AQE path behave like this one too.
Steps to reproduce
Reproduced on main at f9dd86d (Spark 4.1) with Comet's test session:
spark.range(0, 1000).selectExpr("CAST(id % 10 AS STRING) AS fk", "id AS v")
.write.partitionBy("fk").parquet(fact)
spark.range(0, 5).selectExpr("CAST(id AS STRING) AS value", "id AS w").write.parquet(d1)
spark.range(5, 10).selectExpr("CAST(id AS STRING) AS value", "id AS w").write.parquet(d2)
// spark.sql.adaptive.enabled=false, spark.sql.optimizer.dynamicPartitionPruning.enabled=true
sql("""SELECT /*+ BROADCAST(d) */ f.v FROM u_fact f JOIN
(SELECT value FROM u_d1 WHERE w < 2 UNION ALL SELECT value FROM u_d2 WHERE w > 8) d
ON f.fk = d.value""")
The executed plan:
CometBroadcastHashJoin [cast(fk#34 as bigint)], [cast(value#35 as bigint)], Inner, BuildRight
:- CometNativeScan parquet [v#33L,fk#34] ...
: +- SubqueryBroadcast dynamicpruning#39, [0], [cast(value#35 as bigint)]
: +- BroadcastExchange HashedRelationBroadcastMode(...)
: +- CometColumnarToRow
: +- CometUnion Union, [value#35]
: :- CometProject / CometFilter / CometNativeScan parquet (d1)
: +- CometProject / CometFilter / CometNativeScan parquet (d2)
+- CometBroadcastExchange [value#35]
+- CometUnion Union, [value#35]
:- ... (same d1 branch)
+- ... (same d2 branch)
There is no ReusedExchange in the plan. With AQE on, the DPP subquery is a CometSubqueryBroadcast over ReusedExchange [value#12], CometBroadcastExchange [value#12].
Expected behavior
The DPP subquery reuses the join's CometBroadcastExchange, as it does for a dimension whose top operator is a CometNativeExec, and as the AQE path does today. Results are correct either way. The cost is a second evaluation of the dimension and a second broadcast job.
Additional context
Found while reviewing #6786. One possible fix is to accept any Comet columnar plan in that check, for example cometChild.isInstanceOf[CometPlan] && cometChild.supportsColumnar. That change, applied to the AQE rule in #6786, restored reuse there and kept the existing DPP tests in CometExecSuite passing.
Describe the bug
With AQE off, a dynamic partition pruning subquery doesn't reuse the join's
CometBroadcastExchangewhen the dimension is aUNION ALL. The dimension is computed and broadcast a second time for the DPP subquery.CometExecRule.rewriteInSubqueryPlanonly replaces theSubqueryBroadcastExecwith aCometSubqueryBroadcastExecwhen the exchange's child, after stripping the columnar-to-row transition, is aCometNativeExec.CometUnionExecextendsCometExecbut notCometNativeExec, so the check fails and the subquery keeps a SparkBroadcastExchangeoverCometColumnarToRow(CometUnion(...)).ReuseExchangeAndSubquerycan't match that to the join'sCometBroadcastExchange.CometCoalesceExec,CometTakeOrderedAndProjectExecand a bareCometSparkToColumnarExecwould hit the same check.The same query with AQE on reuses the broadcast through
CometPlanAdaptiveDynamicPruningFilters. #6786 adds the sameCometNativeExeccheck to that rule, which would make the AQE path behave like this one too.Steps to reproduce
Reproduced on
mainat f9dd86d (Spark 4.1) with Comet's test session:The executed plan:
There is no
ReusedExchangein the plan. With AQE on, the DPP subquery is aCometSubqueryBroadcastoverReusedExchange [value#12], CometBroadcastExchange [value#12].Expected behavior
The DPP subquery reuses the join's
CometBroadcastExchange, as it does for a dimension whose top operator is aCometNativeExec, and as the AQE path does today. Results are correct either way. The cost is a second evaluation of the dimension and a second broadcast job.Additional context
Found while reviewing #6786. One possible fix is to accept any Comet columnar plan in that check, for example
cometChild.isInstanceOf[CometPlan] && cometChild.supportsColumnar. That change, applied to the AQE rule in #6786, restored reuse there and kept the existing DPP tests inCometExecSuitepassing.