Add AutoOptimizedShuffle regression coverage [databricks]

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,87 @@
 
 _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():
+    initial_shuffle_partitions = 32
+    conf = copy_and_update(_adaptive_conf, {
+        "spark.databricks.adaptive.autoOptimizeShuffle.enabled": "true",
+        # 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):
+        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")
+
+    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,
+        gpu_plan_assertion=assert_auto_optimized_shuffle,
+        cpu_plan_assertion=assert_cpu_auto_optimized_shuffle)
+
+
 def create_skew_df(spark, length):
     root = spark.range(0, length)
     mid = length / 2

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)