[FEA] Support AutoOptimizedShuffle for Databricks
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.