[SPARK-32859][SQL] Introduce physical rule to decide bucketing dynamically

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -2764,6 +2764,14 @@ object SQLConf {
       .booleanConf
       .createWithDefault(false)
 
+  val DYNAMIC_DECIDE_BUCKETING_ENABLED =
+    buildConf("spark.sql.sources.dynamic.decide.bucketing.enabled")
+      .doc("When true, dynamically decide whether to do bucketed scan on input tables " +
+        "based on query plan.")
+      .version("3.1.0")
+      .booleanConf
+      .createWithDefault(true)
+
   /**
    * Holds information about keys that have been deprecated.
    *
@@ -3386,6 +3394,8 @@ class SQLConf extends Serializable with Logging {
 
   def truncateTrashEnabled: Boolean = getConf(SQLConf.TRUNCATE_TRASH_ENABLED)
 
+  def dynamicDecideBucketingEnabled: Boolean = getConf(SQLConf.DYNAMIC_DECIDE_BUCKETING_ENABLED)
+
   /** ********************** SQLConf functionality methods ************ */
 
   /** Set Spark SQL configuration properties. */

sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala

@@ -155,6 +155,7 @@ case class RowDataSourceScanExec(
  * @param partitionFilters Predicates to use for partition pruning.
  * @param optionalBucketSet Bucket ids for bucket pruning.
  * @param optionalNumCoalescedBuckets Number of coalesced buckets.
+ * @param optionalDynamicDecideBucketing Disable bucketing dynamically based on query plan.
  * @param dataFilters Filters on non-partition columns.
  * @param tableIdentifier identifier for the table in the metastore.
  */
@@ -165,6 +166,7 @@ case class FileSourceScanExec(
     partitionFilters: Seq[Expression],
     optionalBucketSet: Option[BitSet],
     optionalNumCoalescedBuckets: Option[Int],
+    optionalDynamicDecideBucketing: Option[Boolean],
     dataFilters: Seq[Expression],
     tableIdentifier: Option[TableIdentifier])
   extends DataSourceScanExec {
@@ -257,7 +259,8 @@ case class FileSourceScanExec(
 
   // exposed for testing
   lazy val bucketedScan: Boolean = {
-    if (relation.sparkSession.sessionState.conf.bucketingEnabled && relation.bucketSpec.isDefined) {
+    if (relation.sparkSession.sessionState.conf.bucketingEnabled && relation.bucketSpec.isDefined
+      && optionalDynamicDecideBucketing.getOrElse(true)) {
       val spec = relation.bucketSpec.get
       val bucketColumns = spec.bucketColumnNames.flatMap(n => toAttribute(n))
       bucketColumns.size == spec.bucketColumnNames.size
@@ -623,6 +626,7 @@ case class FileSourceScanExec(
         filterUnusedDynamicPruningExpressions(partitionFilters), output),
       optionalBucketSet,
       optionalNumCoalescedBuckets,
+      optionalDynamicDecideBucketing,
       QueryPlan.normalizePredicates(dataFilters, output),
       None)
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala

@@ -35,7 +35,7 @@ import org.apache.spark.sql.catalyst.rules.{PlanChangeLogger, Rule}
 import org.apache.spark.sql.catalyst.util.StringUtils.PlanStringConcat
 import org.apache.spark.sql.catalyst.util.truncatedString
 import org.apache.spark.sql.execution.adaptive.{AdaptiveExecutionContext, InsertAdaptiveSparkPlan}
-import org.apache.spark.sql.execution.bucketing.CoalesceBucketsInJoin
+import org.apache.spark.sql.execution.bucketing.{CoalesceBucketsInJoin, PlanBucketing}
 import org.apache.spark.sql.execution.dynamicpruning.PlanDynamicPruningFilters
 import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ReuseExchange}
 import org.apache.spark.sql.execution.streaming.{IncrementalExecution, OffsetSeqMetadata}
@@ -344,6 +344,7 @@ object QueryExecution {
       PlanSubqueries(sparkSession),
       RemoveRedundantProjects(sparkSession.sessionState.conf),
       EnsureRequirements(sparkSession.sessionState.conf),
+      PlanBucketing(sparkSession.sessionState.conf),
       ApplyColumnarRulesAndInsertTransitions(sparkSession.sessionState.conf,
         sparkSession.sessionState.columnarRules),
       CollapseCodegenStages(sparkSession.sessionState.conf),

sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/PlanBucketing.scala

@@ -0,0 +1,144 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.bucketing
+
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Partial, PartialMerge}
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, HashClusteredDistribution}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{FileSourceScanExec, FilterExec, ProjectExec, SortExec, SparkPlan}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.exchange.Exchange
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Plans bucketing dynamically based on actual physical query plan.
+ * NOTE: this rule is designed to be applied right after [[EnsureRequirements]],
+ * where all [[ShuffleExchangeExec]] and [[SortExec]] have been added to plan properly.
+ *
+ * When BUCKETING_ENABLED and DYNAMIC_DECIDE_BUCKETING_ENABLED are set to true, go through
+ * query plan to check where bucketed table scan is unnecessary, and disable bucketed table
+ * scan if needed.
+ *
+ * For all operators which [[hasInterestingPartition]] (i.e., require [[ClusteredDistribution]]
+ * or [[HashClusteredDistribution]]), check if the sub-plan for operator has [[Exchange]] and
+ * bucketed table scan (and only allow certain operators in plan, see details in
+ * [[canDisableBucketedScan]]). If yes, disable the bucketed table scan in the sub-plan.
+ *
+ * Examples:
+ * (1).join:
+ *         SortMergeJoin(t1.i = t2.j)
+ *            /            \
+ *        Sort(i)        Sort(j)
+ *          /               \
+ *      Shuffle(i)       Scan(t2: i, j)
+ *        /         (bucketed on column j, enable bucketed scan)
+ *   Scan(t1: i, j)
+ * (bucketed on column j, DISABLE bucketed scan)
+ *
+ * (2).aggregate:
+ *         HashAggregate(i, ..., Final)
+ *                      |
+ *                  Shuffle(i)
+ *                      |
+ *         HashAggregate(i, ..., Partial)
+ *                      |
+ *                    Filter
+ *                      |
+ *                  Scan(t1: i, j)
+ *  (bucketed on column j, DISABLE bucketed scan)
+ *
+ * The idea of [[hasInterestingPartition]] is inspired from "interesting order" in
+ * the paper "Access Path Selection in a Relational Database Management System"
+ * (http://www.inf.ed.ac.uk/teaching/courses/adbs/AccessPath.pdf).
+ */
+case class PlanBucketing(conf: SQLConf) extends Rule[SparkPlan] {
+  private def disableBucketWithInterestingPartition(plan: SparkPlan): SparkPlan = {
+    var hasPlanWithInterestingPartition = false
+
+    val newPlan = plan.transformUp {
+      case p if hasInterestingPartition(p) =>
+        hasPlanWithInterestingPartition = true
+
+        val newChildren = p.children.map(child => {
+          if (canDisableBucketedScan(child)) {
+            disableBucketedScan(child)
+          } else {
+            child
+          }
+        })
+        p.withNewChildren(newChildren)
+      case other => other
+    }
+
+    if (hasPlanWithInterestingPartition) {
+      newPlan
+    } else {
+      // Disable all bucketed scans if there's no operator with interesting partition
+      // found in query plan.
+      disableBucketedScan(newPlan)
+    }
+  }
+
+  private def hasInterestingPartition(plan: SparkPlan): Boolean = {
+    plan.requiredChildDistribution.exists {
+      case _: ClusteredDistribution | _: HashClusteredDistribution => true
+      case _ => false
+    }
+  }
+
+  private def canDisableBucketedScan(plan: SparkPlan): Boolean = {
+    val hasExchange = plan.find(_.isInstanceOf[Exchange]).isDefined
+    val hasBucketedScan = plan.find {
+      case scan: FileSourceScanExec => isBucketedScanWithoutFilter(scan)
+      case _ => false
+    }.isDefined
+
+    def isAllowedPlan(plan: SparkPlan): Boolean = plan match {
+      case _: SortExec | _: Exchange | _: ProjectExec | _: FilterExec |
+           _: FileSourceScanExec => true
+      case partialAgg: BaseAggregateExec =>
+        val modes = partialAgg.aggregateExpressions.map(_.mode)
+        modes.nonEmpty && modes.forall(mode => mode == Partial || mode == PartialMerge)
+      case _ => false
+    }
+    val onlyHasAllowedPlans = plan.find(!isAllowedPlan(_)).isEmpty
+    hasExchange && hasBucketedScan && onlyHasAllowedPlans
+  }
+
+  private def disableBucketedScan(plan: SparkPlan): SparkPlan = {
+    plan.transformUp {
+      case scan: FileSourceScanExec if isBucketedScanWithoutFilter(scan) =>
+        scan.copy(optionalDynamicDecideBucketing = Some(false))
+    }
+  }
+
+  private def isBucketedScanWithoutFilter(scan: FileSourceScanExec): Boolean = {
+    // Do not disable bucketing for the scan if it has filter pruning,
+    // because bucketing is still useful here to save CPU/IO cost with
+    // only reading selected bucket files.
+    scan.bucketedScan && scan.optionalBucketSet.isEmpty
+  }
+
+  def apply(plan: SparkPlan): SparkPlan = {
+    if (!conf.bucketingEnabled || !conf.dynamicDecideBucketingEnabled) {
+      plan
+    } else {
+      disableBucketWithInterestingPartition(plan)
+    }
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala

@@ -220,6 +220,7 @@ object FileSourceStrategy extends Strategy with PredicateHelper with Logging {
           partitionKeyFilters.toSeq,
           bucketSet,
           None,
+          None,
           dataFilters,
           table.map(_.identifier))
 

sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala

@@ -348,7 +348,7 @@ class DataFrameJoinSuite extends QueryTest
             }
             assert(broadcastExchanges.size == 1)
             val tables = broadcastExchanges.head.collect {
-              case FileSourceScanExec(_, _, _, _, _, _, _, Some(tableIdent)) => tableIdent
+              case FileSourceScanExec(_, _, _, _, _, _, _, _, Some(tableIdent)) => tableIdent
             }
             assert(tables.size == 1)
             assert(tables.head === TableIdentifier(table1Name, Some(dbName)))

sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala

@@ -1314,7 +1314,7 @@ class SubquerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
       // need to execute the query before we can examine fs.inputRDDs()
       assert(stripAQEPlan(df.queryExecution.executedPlan) match {
         case WholeStageCodegenExec(ColumnarToRowExec(InputAdapter(
-            fs @ FileSourceScanExec(_, _, _, partitionFilters, _, _, _, _)))) =>
+            fs @ FileSourceScanExec(_, _, _, partitionFilters, _, _, _, _, _)))) =>
           partitionFilters.exists(ExecSubqueryExpression.hasSubquery) &&
             fs.inputRDDs().forall(
               _.asInstanceOf[FileScanRDD].filePartitions.forall(

sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala

@@ -72,7 +72,8 @@ class CoalesceBucketsInJoinSuite extends SQLTestUtils with SharedSparkSession {
       bucketSpec = Some(BucketSpec(setting.numBuckets, setting.cols.map(_.name), Nil)),
       fileFormat = new ParquetFileFormat(),
       options = Map.empty)(spark)
-    FileSourceScanExec(relation, setting.cols, relation.dataSchema, Nil, None, None, Nil, None)
+    FileSourceScanExec(relation, setting.cols, relation.dataSchema, Nil, None, None,
+      None, Nil, None)
   }
 
   private def run(setting: JoinSetting): Unit = {

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategySuite.scala

@@ -262,20 +262,22 @@ class FileSourceStrategySuite extends QueryTest with SharedSparkSession with Pre
           "p1=2/file7_0000" -> 1),
         buckets = 3)
 
-    // No partition pruning
-    checkScan(table) { partitions =>
-      assert(partitions.size == 3)
-      assert(partitions(0).files.size == 5)
-      assert(partitions(1).files.size == 0)
-      assert(partitions(2).files.size == 2)
-    }
-
-    // With partition pruning
-    checkScan(table.where("p1=2")) { partitions =>
-      assert(partitions.size == 3)
-      assert(partitions(0).files.size == 3)
-      assert(partitions(1).files.size == 0)
-      assert(partitions(2).files.size == 1)
+    withSQLConf(SQLConf.DYNAMIC_DECIDE_BUCKETING_ENABLED.key -> "false") {
+      // No partition pruning
+      checkScan(table) { partitions =>
+        assert(partitions.size == 3)
+        assert(partitions(0).files.size == 5)
+        assert(partitions(1).files.size == 0)
+        assert(partitions(2).files.size == 2)
+      }
+
+      // With partition pruning
+      checkScan(table.where("p1=2")) { partitions =>
+        assert(partitions.size == 3)
+        assert(partitions(0).files.size == 3)
+        assert(partitions(1).files.size == 0)
+        assert(partitions(2).files.size == 1)
+      }
     }
   }
 

sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala

@@ -432,22 +432,24 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils
         // join1 is a broadcast join where df2 is broadcasted. Note that output partitioning on the
         // streamed side (t1) is HashPartitioning (bucketed files).
         val join1 = t1.join(df2, t1("i1") === df2("i2") && t1("j1") === df2("j2"))
-        val plan1 = join1.queryExecution.executedPlan
-        assert(collect(plan1) { case e: ShuffleExchangeExec => e }.isEmpty)
-        val broadcastJoins = collect(plan1) { case b: BroadcastHashJoinExec => b }
-        assert(broadcastJoins.size == 1)
-        assert(broadcastJoins(0).outputPartitioning.isInstanceOf[PartitioningCollection])
-        val p = broadcastJoins(0).outputPartitioning.asInstanceOf[PartitioningCollection]
-        assert(p.partitionings.size == 4)
-        // Verify all the combinations of output partitioning.
-        Seq(Seq(t1("i1"), t1("j1")),
-          Seq(t1("i1"), df2("j2")),
-          Seq(df2("i2"), t1("j1")),
-          Seq(df2("i2"), df2("j2"))).foreach { expected =>
-          val expectedExpressions = expected.map(_.expr)
-          assert(p.partitionings.exists {
-            case h: HashPartitioning => expressionsEqual(h.expressions, expectedExpressions)
-          })
+        withSQLConf(SQLConf.DYNAMIC_DECIDE_BUCKETING_ENABLED.key -> "false") {
+          val plan1 = join1.queryExecution.executedPlan
+          assert(collect(plan1) { case e: ShuffleExchangeExec => e }.isEmpty)
+          val broadcastJoins = collect(plan1) { case b: BroadcastHashJoinExec => b }
+          assert(broadcastJoins.size == 1)
+          assert(broadcastJoins(0).outputPartitioning.isInstanceOf[PartitioningCollection])
+          val p = broadcastJoins(0).outputPartitioning.asInstanceOf[PartitioningCollection]
+          assert(p.partitionings.size == 4)
+          // Verify all the combinations of output partitioning.
+          Seq(Seq(t1("i1"), t1("j1")),
+            Seq(t1("i1"), df2("j2")),
+            Seq(df2("i2"), t1("j1")),
+            Seq(df2("i2"), df2("j2"))).foreach { expected =>
+            val expectedExpressions = expected.map(_.expr)
+            assert(p.partitionings.exists {
+              case h: HashPartitioning => expressionsEqual(h.expressions, expectedExpressions)
+            })
+          }
         }
 
         // Join on the column from the broadcasted side (i2, j2) and make sure output partitioning