Skip to content

Non-AQE DPP subquery doesn't reuse the Comet broadcast when the dimension is a UNION #6815

Description

@andygrove

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions