Add AutoOptimizedShuffle regression coverage [databricks]

[gerashegalov]

Fixes #13930.

Description

Databricks AQE's AutoOptimizedShuffle visits shuffle exchanges to determine and update their target partitioning. The DBR 17.3 GPU shuffle exchange now implements that contract, but there was no focused regression test ensuring the Databricks optimizer can inspect a GPU exchange without throwing UnsupportedOperationException.

This change adds test_databricks_auto_optimized_shuffle to the AQE integration suite. The test runs only on DBR 17.3 or later, explicitly enables AutoOptimizedShuffle, forces a grouped shuffle with multiple input partitions, compares CPU and GPU results, and requires GpuShuffleExchangeExec in the executed GPU plan so CPU fallback cannot hide the regression.

The test also disables ordinary AQE post-shuffle coalescing and inspects the final adaptive CPU and GPU plans. It uses the CPU AutoOptimizedShuffle partition count as the runtime oracle, then verifies that the GPU exchange's target, advertised output, GPU partitioning, and shuffle dependency all agree with that count.

Validation performed locally:

• git diff --check
• Python AST parsing of aqe_test.py and asserts.py

The strengthened Spark-backed assertions were not run locally because AutoOptimizedShuffle is proprietary Databricks behavior and requires a DBR 17.3 GPU cluster with a DBR-built RAPIDS plugin JAR.

AI assistance: The test change and this pull request description were prepared with OpenAI Codex and reviewed by the contributor before opening the pull request.

Checklists

Documentation

• Updated for new or modified user-facing features or behaviors
• No user-facing change

Testing

• Added or modified tests to cover new code paths
• Covered by existing tests
  (Please provide the names of the existing tests in the PR description.)
• Not required

Performance

• Tests ran and results are added in the PR description
• Issue filed with a link in the PR description
• Not required

[gerashegalov]

build

[gerashegalov]

build

[greptile-apps]

Greptile Summary

This PR adds focused DBR 17.3+ regression coverage for AutoOptimizedShuffle and extends the CPU/GPU collection helper with an optional CPU executed-plan assertion.

• Exercises a grouped shuffle with ordinary AQE partition coalescing disabled.
• Requires the GPU shuffle exchange and compares CPU/GPU results.
• Verifies agreement among the optimized exchange’s target, output, GPU, and dependency partition counts.
• Captures the CPU optimized partition count as the runtime oracle.

Confidence Score: 5/5

The PR appears safe to merge with no actionable correctness, security, or test-harness issue identified.

The new callback executes after the CPU query has been collected and before GPU-session configuration is applied, while the regression test explicitly verifies both GPU execution and consistent optimized partitioning.

Important Files Changed

Filename	Overview
integration_tests/src/main/python/aqe_test.py	Adds a DBR 17.3-gated AutoOptimizedShuffle regression test with CPU/GPU result and final-plan partitioning checks; no actionable defect identified.
integration_tests/src/main/python/asserts.py	Adds an optional CPU plan callback after CPU collection, mirroring the existing GPU callback without breaking existing callers.

Sequence Diagram

sequenceDiagram
  participant Test
  participant CPU as CPU Spark session
  participant GPU as GPU Spark session
  Test->>CPU: Run grouped shuffle and collect
  CPU-->>Test: Results and final adaptive plan
  Test->>Test: Capture optimized CPU partition count
  Test->>GPU: Run same grouped shuffle and collect
  GPU-->>Test: Results and final adaptive plan
  Test->>Test: Require GpuShuffleExchangeExec
  Test->>Test: Compare partitioning counts and results

Loading

_{Reviews (1): Last reviewed commit: "Strengthen AutoOptimizedShuffle regressi..." | Re-trigger Greptile}