1616
1717from private_optimizer_common import (
1818 assert_rule_fires ,
19+ assert_rule_skipped ,
1920 private_optimizer_conf ,
2021)
2122from spark_session import is_databricks_runtime
2223
2324
24- @pytest .mark .private_optimizer
25- @pytest .mark .skipif (
26- is_databricks_runtime (),
27- reason = "Databricks executor-broadcast AQE can put the materialized shuffle on the "
28- "BHJ build side; this marker test covers streamed-side skew split. "
29- "See https://github.com/NVIDIA/cudf-spark/issues/15136" )
30- def test_optimize_skewed_bhj_join (spark_tmp_path ):
31- """OptimizeSkewedBHJJoinRule splits a skewed partition on the streamed side
32- of an AQE broadcast hash join. Needs a runtime broadcast (static
33- autoBroadcastJoinThreshold=-1, adaptive.autoBroadcastJoinThreshold=10m) so
34- the streamed side is a materialized shuffle stage, plus small skew
35- thresholds. Marker: the shuffle reader is 'coalesced and skewed'.
25+ SKEWED_BHJ_MARKER = "coalesced and skewed"
26+ # The DB GpuBroadcastHashJoinExec plan ends with isNullAwareAntiJoin and
27+ # executorBroadcast. This marker proves the shim guard's exact precondition was
28+ # reached rather than accepting any broadcast hash join as a skipped-rule path.
29+ DB_EXECUTOR_BROADCAST_MARKER = "GpuBuildRight, false, true"
30+ # The keyed exchange identifies the explicitly repartitioned streamed input.
31+ # GpuShuffleCoalesce alone is not specific because aggregate stages can add it.
32+ DB_STREAMED_SHUFFLE_MARKER = "GpuColumnarExchange gpuhashpartitioning(key1"
3633
37- The rule additionally short-circuits in OptimizeSkewedBHJJoinRule.apply when
38- AQEUtils.isOptimizeSkewBHJSupported is false, so if a future Spark/runtime
39- drops support the rule becomes a no-op and the marker assertion below fails
40- loudly rather than passing silently.
4134
42- Validated with a small GLOBAL aggregate over the materialized skewed join;
43- a GROUP BY on the skew key is intentionally avoided here."""
44- conf_extra = {
35+ def _skewed_bhj_conf_extra ():
36+ return {
4537 "spark.sql.adaptive.enabled" : "true" ,
4638 "spark.sql.adaptive.skewJoin.enabled" : "true" ,
4739 "spark.sql.autoBroadcastJoinThreshold" : "-1" ,
@@ -53,19 +45,72 @@ def test_optimize_skewed_bhj_join(spark_tmp_path):
5345 "spark.sql.adaptive.localShuffleReader.enabled" : "false" ,
5446 }
5547
56- def fn (spark ):
57- spark .range (0 , 2000 , 1 , 10 ).selectExpr (
58- "CASE WHEN id < 1000 THEN 249 ELSE id END AS key2" , "id AS value2"
59- ).createOrReplaceTempView ("skewData2" )
60- spark .range (0 , 1000 , 1 , 10 ).selectExpr (
61- "CASE WHEN id < 250 THEN 249 WHEN id >= 750 THEN 1000 ELSE id END AS key1" , "id AS value1"
62- ).createOrReplaceTempView ("skewData1" )
63- return spark .sql (
64- "SELECT count(*) AS cnt, min(value2) AS mn, max(value2) AS mx, sum(value1) AS sm "
65- "FROM skewData1 JOIN skewData2 ON key1 = key2" )
6648
49+ def _skewed_bhj_confs ():
50+ conf_extra = _skewed_bhj_conf_extra ()
6751 on = private_optimizer_conf (
68- {"spark.rapids.sql.adaptive.skewJoin.broadcast.enabled" : "true" }, extra_conf = conf_extra )
52+ {"spark.rapids.sql.adaptive.skewJoin.broadcast.enabled" : "true" },
53+ extra_conf = conf_extra )
6954 off = private_optimizer_conf (
70- {"spark.rapids.sql.adaptive.skewJoin.broadcast.enabled" : "false" }, extra_conf = conf_extra )
71- assert_rule_fires (fn , on , off , marker = "coalesced and skewed" , physical = True )
55+ {"spark.rapids.sql.adaptive.skewJoin.broadcast.enabled" : "false" },
56+ extra_conf = conf_extra )
57+ return on , off
58+
59+
60+ def _skewed_bhj_global_agg (spark ):
61+ spark .range (0 , 2000 , 1 , 10 ).selectExpr (
62+ "CASE WHEN id < 1000 THEN 249 ELSE id END AS key2" , "id AS value2"
63+ ).createOrReplaceTempView ("skewData2" )
64+ spark .range (0 , 1000 , 1 , 10 ).selectExpr (
65+ "CASE WHEN id < 250 THEN 249 WHEN id >= 750 THEN 1000 ELSE id END AS key1" ,
66+ "id AS value1"
67+ ).repartition (100 , "key1" ).createOrReplaceTempView ("skewData1" )
68+ return spark .sql (
69+ "SELECT /*+ BROADCAST(skewData2) */ "
70+ "count(*) AS cnt, min(value2) AS mn, max(value2) AS mx, sum(value1) AS sm "
71+ "FROM skewData1 JOIN skewData2 ON key1 = key2" )
72+
73+
74+ @pytest .mark .private_optimizer
75+ @pytest .mark .skipif (
76+ is_databricks_runtime (),
77+ reason = "The positive rule-fire assertion is Apache-only; the Databricks "
78+ "executor-broadcast guarded path is covered by the skipped-path test below. "
79+ "See https://github.com/NVIDIA/cudf-spark/issues/15136" )
80+ def test_optimize_skewed_bhj_join (spark_tmp_path ):
81+ """OptimizeSkewedBHJJoinRule splits a skewed partition on the streamed side
82+ of an AQE broadcast hash join. The broadcast hint fixes the build side while
83+ the explicit key repartition makes the streamed side a materialized shuffle
84+ stage; small skew thresholds make the key 249 partition eligible to split.
85+ Marker: the shuffle reader is 'coalesced and skewed'.
86+
87+ The rule additionally short-circuits in OptimizeSkewedBHJJoinRule.apply when
88+ AQEUtils.isOptimizeSkewBHJSupported is false, so if a future Spark/runtime
89+ drops support the rule becomes a no-op and the marker assertion below fails
90+ loudly rather than passing silently.
91+
92+ Validated with a small GLOBAL aggregate over the materialized skewed join;
93+ a GROUP BY on the skew key is intentionally avoided here."""
94+ on , off = _skewed_bhj_confs ()
95+ assert_rule_fires (_skewed_bhj_global_agg , on , off , marker = SKEWED_BHJ_MARKER ,
96+ physical = True )
97+
98+
99+ @pytest .mark .private_optimizer
100+ @pytest .mark .skipif (
101+ not is_databricks_runtime (),
102+ reason = "Databricks-only coverage for executor-broadcast AQE fallback. "
103+ "Apache runtime coverage is in test_optimize_skewed_bhj_join." )
104+ def test_optimize_skewed_bhj_join_skips_on_databricks_executor_broadcast (spark_tmp_path ):
105+ """Databricks executor-broadcast AQE is expected to skip the streamed-side
106+ rewrite even when the streamed input is a materialized skewed shuffle stage.
107+ Without the executor-broadcast shim guard this shape is eligible for the
108+ rewrite, as verified by the positive Apache test above. The streamed-side
109+ skew marker must remain absent while CPU and GPU results still match. The
110+ required markers verify a non-null-aware, executor-broadcast build-right join
111+ whose streamed input retains the expected GPU hash-partitioning exchange."""
112+ on , off = _skewed_bhj_confs ()
113+ assert_rule_skipped (_skewed_bhj_global_agg , on , off , marker = SKEWED_BHJ_MARKER ,
114+ physical = True , required_on_markers = (
115+ DB_EXECUTOR_BROADCAST_MARKER ,
116+ DB_STREAMED_SHUFFLE_MARKER ))
0 commit comments