spark-rapids icon indicating copy to clipboard operation
spark-rapids copied to clipboard

[FEA] Support AutoOptimizedShuffle for Databricks

Open jihoonson opened this issue 1 month ago • 0 comments

Is your feature request related to a problem? Please describe.

On Databricks, AQE can optimize the shuffle at query time using the auto-optimized shuffle feature. This feature can be enabled using the spark.databricks.adaptive.autoOptimizeShuffle.enabled config. When the feature is enabled, the plugin currently throws the UnsupportedOperationException as below.

java.lang.UnsupportedOperationException
	at org.apache.spark.rapids.shims.GpuShuffleExchangeExec.targetOutputPartitioning(GpuShuffleExchangeExec.scala:52)
	at com.databricks.sql.execution.adaptive.ShuffleSizeVisitor.visitExchange(AutoOptimizedShuffle.scala:975)
	at com.databricks.sql.execution.adaptive.ShuffleSizeVisitor.visitExchange(AutoOptimizedShuffle.scala:431)
	at com.databricks.sql.execution.SparkPlanVisitor.visit(SparkPlanVisitor.scala:37)
	at com.databricks.sql.execution.SparkPlanVisitor.visit$(SparkPlanVisitor.scala:34)
...

Describe the solution you'd like The plugin should support the auto-optimized shuffle. Perhaps we can add a GPU version of AutoOptimizedShuffle.

Describe alternatives you've considered

We might consider disabling the auto-optimized shuffle feature automatically when it is detected in the plugin. Instead of throwing an exception, the plugin will run successfully, but can be slower compared to when the shuffle is optimized.

Additional context

More information about the auto-optimized shuffle can be found in https://www.databricks.com/blog/2020/10/21/faster-sql-adaptive-query-execution-in-databricks.html.

jihoonson avatar Dec 03 '25 01:12 jihoonson