diff --git a/integration_tests/src/main/python/aqe_test.py b/integration_tests/src/main/python/aqe_test.py index c2dd51536d4..2f7fd7dbaa1 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,83 @@ _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 + + 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(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(gpu_exchanges) == 1, \ + f"Expected one GPU shuffle exchange, found {len(gpu_exchanges)}:\n{gpu_plan}" + + exchange = gpu_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{gpu_plan}" + assert partition_counts["target"] == cpu_partition_count, \ + f"CPU and GPU AutoOptimizedShuffle partition counts differ: " \ + 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) + + def create_skew_df(spark, length): root = spark.range(0, length) mid = length / 2 diff --git a/integration_tests/src/main/python/asserts.py b/integration_tests/src/main/python/asserts.py index 761686d1ec2..124b8c1cea9 100644 --- a/integration_tests/src/main/python/asserts.py +++ b/integration_tests/src/main/python/asserts.py @@ -524,7 +524,7 @@ def assert_cpu_and_gpu_are_equal_collect_with_capture(func, conf={}, require_non_empty=False, gpu_plan_assertion=None): - """Compare collected CPU/GPU results and validate the executed GPU plan. + """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. @@ -532,7 +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. + CPU and GPU dataframes' executed JVM plans. """ (bring_back, collect_type) = _prep_func_for_compare(func, 'COLLECT_WITH_DATAFRAME') @@ -544,12 +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" + 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 b0f4abf4be2..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 @@ -1537,7 +1537,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 3af47cc737e..9c5ad714302 100644 --- a/integration_tests/src/main/python/iceberg/iceberg_test.py +++ b/integration_tests/src/main/python/iceberg/iceberg_test.py @@ -100,7 +100,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") @@ -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), \ @@ -330,7 +330,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))