From e8f62f71e5621366e43f20afcf8cdf7040435ab4 Mon Sep 17 00:00:00 2001 From: Gera Shegalov Date: Fri, 28 Aug 2026 15:58:28 -0700 Subject: [PATCH 1/4] Add AutoOptimizedShuffle regression coverage Signed-off-by: Gera Shegalov --- integration_tests/src/main/python/aqe_test.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/integration_tests/src/main/python/aqe_test.py b/integration_tests/src/main/python/aqe_test.py index c2dd51536d4..828ce1574de 100755 --- a/integration_tests/src/main/python/aqe_test.py +++ b/integration_tests/src/main/python/aqe_test.py @@ -18,7 +18,7 @@ from asserts import assert_gpu_and_cpu_are_equal_collect, assert_cpu_and_gpu_are_equal_collect_with_capture from conftest import is_databricks_runtime, is_not_utc from data_gen import * -from spark_session import is_spark_400_or_later +from spark_session import is_databricks173_or_later, is_spark_400_or_later from marks import ignore_order, allow_non_gpu from spark_session import with_cpu_session, is_databricks113_or_later, is_databricks_version, is_databricks_version_or_later @@ -27,6 +27,30 @@ _adaptive_conf = { "spark.sql.adaptive.enabled": "true" } + +@pytest.mark.skipif( + not is_databricks173_or_later(), reason="Databricks 17.3+ AutoOptimizedShuffle") +@ignore_order(local=True) +def test_databricks_auto_optimized_shuffle(): + conf = copy_and_update(_adaptive_conf, { + "spark.databricks.adaptive.autoOptimizeShuffle.enabled": "true", + "spark.sql.shuffle.partitions": "32", + }) + + def do_groupby(spark): + assert spark.conf.get( + "spark.databricks.adaptive.autoOptimizeShuffle.enabled") == "true" + return spark.range(0, 4096, 1, 32) \ + .selectExpr("id % 8 AS key", "id AS value") \ + .groupBy("key").sum("value") + + assert_cpu_and_gpu_are_equal_collect_with_capture( + do_groupby, + exist_classes="GpuShuffleExchangeExec", + conf=conf, + require_non_empty=True) + + def create_skew_df(spark, length): root = spark.range(0, length) mid = length / 2 From ee89029d6e6f512f6688c929d922f6cd3ba4413e Mon Sep 17 00:00:00 2001 From: Gera Shegalov Date: Mon, 31 Aug 2026 11:15:53 -0700 Subject: [PATCH 2/4] Strengthen AutoOptimizedShuffle regression coverage Signed-off-by: Gera Shegalov --- integration_tests/src/main/python/aqe_test.py | 61 ++++++++++++++++++- integration_tests/src/main/python/asserts.py | 7 ++- 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/integration_tests/src/main/python/aqe_test.py b/integration_tests/src/main/python/aqe_test.py index 828ce1574de..8a3d5a399b7 100755 --- a/integration_tests/src/main/python/aqe_test.py +++ b/integration_tests/src/main/python/aqe_test.py @@ -32,9 +32,12 @@ not is_databricks173_or_later(), reason="Databricks 17.3+ AutoOptimizedShuffle") @ignore_order(local=True) def test_databricks_auto_optimized_shuffle(): + initial_shuffle_partitions = 32 conf = copy_and_update(_adaptive_conf, { "spark.databricks.adaptive.autoOptimizeShuffle.enabled": "true", - "spark.sql.shuffle.partitions": "32", + # Isolate pre-shuffle AOS resizing from AQE post-shuffle coalescing. + "spark.sql.adaptive.coalescePartitions.enabled": "false", + "spark.sql.shuffle.partitions": str(initial_shuffle_partitions), }) def do_groupby(spark): @@ -44,11 +47,65 @@ def do_groupby(spark): .selectExpr("id % 8 AS key", "id AS value") \ .groupBy("key").sum("value") + def collect_plan_nodes(plan): + nodes = [plan] + class_name = plan.getClass().getSimpleName() + if class_name == "AdaptiveSparkPlanExec": + nodes.extend(collect_plan_nodes(plan.executedPlan())) + elif class_name.endswith("QueryStageExec"): + nodes.extend(collect_plan_nodes(plan.plan())) + elif class_name in ("ReusedExchangeExec", "ReusedSubqueryExec"): + nodes.extend(collect_plan_nodes(plan.child())) + else: + children = plan.children().iterator() + while children.hasNext(): + nodes.extend(collect_plan_nodes(children.next())) + return nodes + + optimized_partition_counts = {} + + def assert_cpu_auto_optimized_shuffle(plan): + exchanges = [ + node for node in collect_plan_nodes(plan) + if node.getClass().getSimpleName() == "ShuffleExchangeExec" + ] + assert len(exchanges) == 1, \ + f"Expected one CPU shuffle exchange, found {len(exchanges)}:\n{plan}" + optimized_partition_counts["cpu"] = \ + exchanges[0].outputPartitioning().numPartitions() + assert optimized_partition_counts["cpu"] > 0, \ + f"Expected a positive CPU shuffle partition count:\n{plan}" + + def assert_auto_optimized_shuffle(plan): + exchanges = [ + node for node in collect_plan_nodes(plan) + if node.getClass().getSimpleName() == "GpuShuffleExchangeExec" + ] + assert len(exchanges) == 1, \ + f"Expected one GPU shuffle exchange, found {len(exchanges)}:\n{plan}" + + exchange = exchanges[0] + partition_counts = { + "target": exchange.targetOutputPartitioning().numPartitions(), + "output": exchange.outputPartitioning().numPartitions(), + "gpu": exchange.gpuOutputPartitioning().numPartitions(), + "dependency": exchange.shuffleDependencyColumnar().partitioner().numPartitions(), + } + assert len(set(partition_counts.values())) == 1, \ + f"Inconsistent optimized shuffle partition counts: {partition_counts}\n{plan}" + assert "cpu" in optimized_partition_counts, \ + "CPU AutoOptimizedShuffle partition count was not captured" + assert partition_counts["target"] == optimized_partition_counts["cpu"], \ + f"CPU and GPU AutoOptimizedShuffle partition counts differ: " \ + f"CPU={optimized_partition_counts['cpu']}, GPU={partition_counts['target']}\n{plan}" + assert_cpu_and_gpu_are_equal_collect_with_capture( do_groupby, exist_classes="GpuShuffleExchangeExec", conf=conf, - require_non_empty=True) + require_non_empty=True, + gpu_plan_assertion=assert_auto_optimized_shuffle, + cpu_plan_assertion=assert_cpu_auto_optimized_shuffle) def create_skew_df(spark, length): diff --git a/integration_tests/src/main/python/asserts.py b/integration_tests/src/main/python/asserts.py index 761686d1ec2..f50081bfee4 100644 --- a/integration_tests/src/main/python/asserts.py +++ b/integration_tests/src/main/python/asserts.py @@ -523,7 +523,8 @@ def assert_cpu_and_gpu_are_equal_collect_with_capture(func, non_exist_classes='', conf={}, require_non_empty=False, - gpu_plan_assertion=None): + gpu_plan_assertion=None, + cpu_plan_assertion=None): """Compare collected CPU/GPU results and validate the executed GPU plan. :param func: Function that creates the dataframe to collect in each Spark session. @@ -533,6 +534,8 @@ def assert_cpu_and_gpu_are_equal_collect_with_capture(func, :param require_non_empty: Require the collected CPU result to contain at least one row. :param gpu_plan_assertion: Optional callback invoked after GPU collection with the dataframe's executed JVM plan. + :param cpu_plan_assertion: Optional callback invoked after CPU collection with the + dataframe's executed JVM plan. """ (bring_back, collect_type) = _prep_func_for_compare(func, 'COLLECT_WITH_DATAFRAME') @@ -544,6 +547,8 @@ def assert_cpu_and_gpu_are_equal_collect_with_capture(func, cpu_end = time.time() if require_non_empty: assert len(from_cpu) > 0, "Expected non-empty result" + if cpu_plan_assertion: + cpu_plan_assertion(cpu_df._jdf.queryExecution().executedPlan()) print('### GPU RUN ###') gpu_start = time.time() from_gpu, gpu_df = with_gpu_session(bring_back, conf=conf) From be8a435ab188f1e36058ef90392267a25697e868 Mon Sep 17 00:00:00 2001 From: Gera Shegalov Date: Wed, 9 Sep 2026 11:05:56 -0700 Subject: [PATCH 3/4] Address AutoOptimizedShuffle test review Signed-off-by: Gera Shegalov --- integration_tests/src/main/python/aqe_test.py | 46 +++++++++---------- integration_tests/src/main/python/asserts.py | 17 ++++--- .../src/main/python/delta_lake_test.py | 2 +- .../src/main/python/group_partitions_test.py | 4 +- .../src/main/python/iceberg/iceberg_test.py | 4 +- 5 files changed, 34 insertions(+), 39 deletions(-) diff --git a/integration_tests/src/main/python/aqe_test.py b/integration_tests/src/main/python/aqe_test.py index 8a3d5a399b7..2f7fd7dbaa1 100755 --- a/integration_tests/src/main/python/aqe_test.py +++ b/integration_tests/src/main/python/aqe_test.py @@ -62,29 +62,28 @@ def collect_plan_nodes(plan): nodes.extend(collect_plan_nodes(children.next())) return nodes - optimized_partition_counts = {} - - def assert_cpu_auto_optimized_shuffle(plan): - exchanges = [ - node for node in collect_plan_nodes(plan) + def assert_auto_optimized_shuffle(cpu_plan, gpu_plan): + cpu_exchanges = [ + node for node in collect_plan_nodes(cpu_plan) if node.getClass().getSimpleName() == "ShuffleExchangeExec" ] - assert len(exchanges) == 1, \ - f"Expected one CPU shuffle exchange, found {len(exchanges)}:\n{plan}" - optimized_partition_counts["cpu"] = \ - exchanges[0].outputPartitioning().numPartitions() - assert optimized_partition_counts["cpu"] > 0, \ - f"Expected a positive CPU shuffle partition count:\n{plan}" - - def assert_auto_optimized_shuffle(plan): - exchanges = [ - node for node in collect_plan_nodes(plan) + assert len(cpu_exchanges) == 1, \ + f"Expected one CPU shuffle exchange, found {len(cpu_exchanges)}:\n{cpu_plan}" + cpu_partition_count = cpu_exchanges[0].outputPartitioning().numPartitions() + assert cpu_partition_count > 0, \ + f"Expected a positive CPU shuffle partition count:\n{cpu_plan}" + assert cpu_partition_count != initial_shuffle_partitions, \ + f"AutoOptimizedShuffle did not resize the CPU shuffle from " \ + f"{initial_shuffle_partitions} partitions:\n{cpu_plan}" + + gpu_exchanges = [ + node for node in collect_plan_nodes(gpu_plan) if node.getClass().getSimpleName() == "GpuShuffleExchangeExec" ] - assert len(exchanges) == 1, \ - f"Expected one GPU shuffle exchange, found {len(exchanges)}:\n{plan}" + assert len(gpu_exchanges) == 1, \ + f"Expected one GPU shuffle exchange, found {len(gpu_exchanges)}:\n{gpu_plan}" - exchange = exchanges[0] + exchange = gpu_exchanges[0] partition_counts = { "target": exchange.targetOutputPartitioning().numPartitions(), "output": exchange.outputPartitioning().numPartitions(), @@ -92,20 +91,17 @@ def assert_auto_optimized_shuffle(plan): "dependency": exchange.shuffleDependencyColumnar().partitioner().numPartitions(), } assert len(set(partition_counts.values())) == 1, \ - f"Inconsistent optimized shuffle partition counts: {partition_counts}\n{plan}" - assert "cpu" in optimized_partition_counts, \ - "CPU AutoOptimizedShuffle partition count was not captured" - assert partition_counts["target"] == optimized_partition_counts["cpu"], \ + f"Inconsistent optimized shuffle partition counts: {partition_counts}\n{gpu_plan}" + assert partition_counts["target"] == cpu_partition_count, \ f"CPU and GPU AutoOptimizedShuffle partition counts differ: " \ - f"CPU={optimized_partition_counts['cpu']}, GPU={partition_counts['target']}\n{plan}" + f"CPU={cpu_partition_count}, GPU={partition_counts['target']}\n{gpu_plan}" assert_cpu_and_gpu_are_equal_collect_with_capture( do_groupby, exist_classes="GpuShuffleExchangeExec", conf=conf, require_non_empty=True, - gpu_plan_assertion=assert_auto_optimized_shuffle, - cpu_plan_assertion=assert_cpu_auto_optimized_shuffle) + gpu_plan_assertion=assert_auto_optimized_shuffle) def create_skew_df(spark, length): diff --git a/integration_tests/src/main/python/asserts.py b/integration_tests/src/main/python/asserts.py index f50081bfee4..124b8c1cea9 100644 --- a/integration_tests/src/main/python/asserts.py +++ b/integration_tests/src/main/python/asserts.py @@ -523,9 +523,8 @@ def assert_cpu_and_gpu_are_equal_collect_with_capture(func, non_exist_classes='', conf={}, require_non_empty=False, - gpu_plan_assertion=None, - cpu_plan_assertion=None): - """Compare collected CPU/GPU results and validate the executed GPU plan. + gpu_plan_assertion=None): + """Compare collected CPU/GPU results and optionally validate both executed plans. :param func: Function that creates the dataframe to collect in each Spark session. :param exist_classes: Comma-separated class names required in the GPU plan. @@ -533,9 +532,7 @@ def assert_cpu_and_gpu_are_equal_collect_with_capture(func, :param conf: Spark configuration used for both executions. :param require_non_empty: Require the collected CPU result to contain at least one row. :param gpu_plan_assertion: Optional callback invoked after GPU collection with the - dataframe's executed JVM plan. - :param cpu_plan_assertion: Optional callback invoked after CPU collection with the - dataframe's executed JVM plan. + CPU and GPU dataframes' executed JVM plans. """ (bring_back, collect_type) = _prep_func_for_compare(func, 'COLLECT_WITH_DATAFRAME') @@ -547,14 +544,16 @@ def assert_cpu_and_gpu_are_equal_collect_with_capture(func, cpu_end = time.time() if require_non_empty: assert len(from_cpu) > 0, "Expected non-empty result" - if cpu_plan_assertion: - cpu_plan_assertion(cpu_df._jdf.queryExecution().executedPlan()) + cpu_plan = None + if gpu_plan_assertion: + cpu_plan = cpu_df._jdf.queryExecution().executedPlan() print('### GPU RUN ###') gpu_start = time.time() from_gpu, gpu_df = with_gpu_session(bring_back, conf=conf) gpu_end = time.time() if gpu_plan_assertion: - gpu_plan_assertion(gpu_df._jdf.queryExecution().executedPlan()) + gpu_plan = gpu_df._jdf.queryExecution().executedPlan() + gpu_plan_assertion(cpu_plan, gpu_plan) jvm = spark_jvm() if exist_classes: for clz in exist_classes.split(','): diff --git a/integration_tests/src/main/python/delta_lake_test.py b/integration_tests/src/main/python/delta_lake_test.py index 7c67d57cf2c..667240e56b6 100644 --- a/integration_tests/src/main/python/delta_lake_test.py +++ b/integration_tests/src/main/python/delta_lake_test.py @@ -1374,7 +1374,7 @@ def read_table(spark): """) return df - def assert_dv_pushdown_plan(plan): + def assert_dv_pushdown_plan(_cpu_plan, plan): from conftest import spark_jvm # Inspect the plan after collection so an adaptive plan has been finalized. diff --git a/integration_tests/src/main/python/group_partitions_test.py b/integration_tests/src/main/python/group_partitions_test.py index 359f3cbc7d0..3200d995906 100644 --- a/integration_tests/src/main/python/group_partitions_test.py +++ b/integration_tests/src/main/python/group_partitions_test.py @@ -43,7 +43,7 @@ def _collect_plan_nodes(plan): return nodes -def _assert_partial_clustering_spj_plan(plan): +def _assert_partial_clustering_spj_plan(_cpu_plan, plan): nodes = _collect_plan_nodes(plan) def nodes_of_class(class_name): @@ -83,7 +83,7 @@ def nodes_of_class(class_name): f"GroupPartitionsExec is not expected before Spark 4.2:\n{plan}" -def _assert_sorted_merge_spj_plan(plan): +def _assert_sorted_merge_spj_plan(_cpu_plan, plan): nodes = _collect_plan_nodes(plan) gpu_groups = [ node for node in nodes diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index 1d14d1c88d3..2b28a2a2be3 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -99,7 +99,7 @@ def _assert_spj_join_shape(plan, expect_spj): return scans -def _assert_partial_clustering_spj_plan(plan): +def _assert_partial_clustering_spj_plan(_cpu_plan, plan): scans = _assert_spj_join_shape(plan, expect_spj=True) exchanges = _nodes_of_class(plan, "GpuShuffleExchangeExec") @@ -247,7 +247,7 @@ def join_on_reducible_transforms(spark): join_on_reducible_transforms, conf=conf, require_non_empty=True, - gpu_plan_assertion=lambda plan: _assert_spj_join_shape( + gpu_plan_assertion=lambda _cpu_plan, plan: _assert_spj_join_shape( plan, allow_compatible_transforms)) From 7f102d55318768ec6bf589d214aa6c6a686c1ed5 Mon Sep 17 00:00:00 2001 From: Gera Shegalov Date: Wed, 9 Sep 2026 14:23:34 -0700 Subject: [PATCH 4/4] Adapt new plan assertion callbacks Signed-off-by: Gera Shegalov --- integration_tests/src/main/python/delta_lake_test.py | 2 +- integration_tests/src/main/python/iceberg/iceberg_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/integration_tests/src/main/python/delta_lake_test.py b/integration_tests/src/main/python/delta_lake_test.py index c3d66a37b63..5ba59504adf 100644 --- a/integration_tests/src/main/python/delta_lake_test.py +++ b/integration_tests/src/main/python/delta_lake_test.py @@ -1045,7 +1045,7 @@ def filtered_read(spark): lambda spark: filtered_read(spark).orderBy("id").collect(), conf=conf) assert cpu_rows == expected - def assert_gpu_pushdown(plan): + def assert_gpu_pushdown(_cpu_plan, plan): from conftest import spark_jvm callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback diff --git a/integration_tests/src/main/python/iceberg/iceberg_test.py b/integration_tests/src/main/python/iceberg/iceberg_test.py index 944fa58a1f1..9c5ad714302 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -239,7 +239,7 @@ def join_after_spj(spark): # gate this test is deliberately avoiding. expected_partitions = (2 if partition_filter else 4) + (1 if partially_clustered else 0) - def assert_plan(plan): + def assert_plan(_cpu_plan, plan): scans = _assert_spj_join_shape(plan, expect_spj=True) counts = [scan.outputPartitioning().numPartitions() for scan in scans] assert counts == [expected_partitions] * len(scans), \