[CARMEL-5997] Support more sql patterns for deciding bucketed scan dynamically (#975)

* [CARMEL-5997] Support more sql patterns for deciding bucketed scan dynamically

* Fix ut

Author: xingchaozh

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

@@ -19,9 +19,11 @@ package org.apache.spark.sql.execution.bucketing
 
 import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, ClusteredDistribution, HashClusteredDistribution}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.execution._
+import org.apache.spark.sql.execution.{CoalesceExec, FileSourceScanExec, FilterExec, LocalLimitExec, ProjectExec, RebucketingExec, SortExec, SparkPlan, UnionExec, WindowSortLimitExec}
 import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
 import org.apache.spark.sql.execution.exchange.Exchange
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, BroadcastRangeJoinExec}
+import org.apache.spark.sql.internal.SQLConf
 
 /**
  * Disable unnecessary bucketed table scan based on actual physical query plan.
@@ -81,7 +83,7 @@ object DisableUnnecessaryBucketedScan extends Rule[SparkPlan] {
    *
    * @param withInterestingPartition The traversed plan has operator with interesting partition.
    * @param withExchange The traversed plan has [[Exchange]] operator.
-   * @param withAllowedNode The traversed plan has only [[isAllowedUnaryExecNode]] operators.
+   * @param withAllowedNode The traversed plan has only [[isAllowedExecNode]] operators.
    */
   private def disableBucketWithInterestingPartition(
                                                      plan: SparkPlan,
@@ -114,7 +116,7 @@ object DisableUnnecessaryBucketedScan extends Rule[SparkPlan] {
           _,
           withInterestingPartition,
           withExchange,
-          withAllowedNode && isAllowedUnaryExecNode(o)))
+          withAllowedNode && isAllowedExecNode(o)))
     }
   }
 
@@ -131,14 +133,20 @@ object DisableUnnecessaryBucketedScan extends Rule[SparkPlan] {
   }
 
   /**
-   * Check if the operator is allowed single-child operator.
+   * Check if the operator is allowed operator.
    * We may revisit this method later as we probably can
    * remove this restriction to allow arbitrary operator between
    * bucketed table scan and operator with interesting partition.
    */
-  private def isAllowedUnaryExecNode(plan: SparkPlan): Boolean = {
+  private def isAllowedExecNode(plan: SparkPlan): Boolean = {
     plan match {
-      case _: SortExec | _: ProjectExec | _: FilterExec => true
+      case _: SortExec | _: ProjectExec | _: FilterExec |
+           _: BroadcastHashJoinExec |
+           _: BroadcastNestedLoopJoinExec |
+           _: BroadcastRangeJoinExec |
+           _: UnionExec |
+           _: WindowSortLimitExec |
+           _: LocalLimitExec => true
       case partialAgg: BaseAggregateExec =>
         partialAgg.requiredChildDistributionExpressions.isEmpty
       case _ => false

sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala

@@ -72,7 +72,10 @@ abstract class DisableUnnecessaryBucketedScanSuite extends QueryTest
 
     def checkNumBucketedScan(query: String, expectedNumBucketedScan: Int,
                              logEnabled: Boolean = false): Unit = {
-      val plan = sql(query).queryExecution.executedPlan
+      val df = sql(query)
+      df.collect()
+
+      val plan = df.queryExecution.executedPlan
 
       if (logEnabled) {
         // scalastyle:off println
@@ -157,7 +160,8 @@ abstract class DisableUnnecessaryBucketedScanSuite extends QueryTest
 //  }
 
   test("SPARK-32859: disable unnecessary bucketed table scan - multiple joins test") {
-    withSQLConf(SQLConf.ENABLE_REBUCKETING.key -> "true") {
+    withSQLConf(SQLConf.ENABLE_REBUCKETING.key -> "false",
+      SQLConf.ENABLE_COALESCE.key -> "true") {
       withTable("t1", "t2", "t3") {
         df1.write.format("parquet").bucketBy(8, "i").saveAsTable("t1")
         df2.write.format("parquet").bucketBy(8, "i").saveAsTable("t2")
@@ -174,7 +178,7 @@ abstract class DisableUnnecessaryBucketedScanSuite extends QueryTest
             """
          SELECT /*+ broadcast(t1) merge(t3)*/ * FROM t1 JOIN t2 JOIN t3
          ON t1.i = t2.i AND t2.i = t3.i
-         """.stripMargin, 3, 3), // TODO 3->2 if ENABLE_REBUCKETING=false
+         """.stripMargin, 2, 3),
           (
             """
          SELECT /*+ merge(t1) broadcast(t3)*/ * FROM t1 JOIN t2 JOIN t3
@@ -184,7 +188,7 @@ abstract class DisableUnnecessaryBucketedScanSuite extends QueryTest
             """
          SELECT /*+ merge(t1, t3)*/ * FROM t1 JOIN t2 JOIN t3
          ON t1.i = t2.i AND t2.i = t3.i
-         """.stripMargin, 3, 3), // TODO 3->2 if ENABLE_REBUCKETING=false
+         """.stripMargin, 3, 3),
           // Multiple joins on non-bucketed columns
           (
             """
@@ -195,7 +199,7 @@ abstract class DisableUnnecessaryBucketedScanSuite extends QueryTest
             """
          SELECT /*+ merge(t1, t3)*/ * FROM t1 JOIN t2 JOIN t3
          ON t1.i = t2.j AND t2.j = t3.i
-         """.stripMargin, 2, 3), // TODO 2->1 if ENABLE_REBUCKETING=false
+         """.stripMargin, 2, 3),
           (
             """
          SELECT /*+ merge(t1, t3)*/ * FROM t1 JOIN t2 JOIN t3
@@ -280,12 +284,26 @@ abstract class DisableUnnecessaryBucketedScanSuite extends QueryTest
          (SELECT t2.i FROM t2 GROUP BY t2.i)
          """.stripMargin, 1, 2),
         // Non-allowed operator in sub-plan
+        (
+          """
+         SELECT j, COUNT(*)
+         FROM (SELECT * FROM (SELECT i, j FROM t1 DISTRIBUTE BY i, j) ORDER BY i, j LIMIT 10)
+         GROUP BY j
+         """.stripMargin, 1, 1),
+        // Union all in sub-plan
         (
           """
          SELECT COUNT(*)
          FROM (SELECT t1.i FROM t1 UNION ALL SELECT t2.i FROM t2)
          GROUP BY i
-         """.stripMargin, 2, 2),
+         """.stripMargin, 0, 2),
+        // Union in sub-plan
+        (
+          """
+         SELECT COUNT(*)
+         FROM (SELECT t1.i FROM t1 UNION SELECT t2.i FROM t2)
+         GROUP BY i
+         """.stripMargin, 0, 2),
         // Multiple [[Exchange]] in sub-plan
         (
           """
@@ -299,14 +317,81 @@ abstract class DisableUnnecessaryBucketedScanSuite extends QueryTest
          GROUP BY j
          """.stripMargin, 0, 1),
         // No bucketed table scan in plan
+        (
+          """
+         SELECT i, j, COUNT(*)
+         FROM (SELECT t1.j, t3.i FROM t1 JOIN t3 ON t1.j = t3.j)
+         GROUP BY i, j
+         """.stripMargin, 0, 0),
+
+        // Broadcast hash join in plan
+        (
+          """
+         SELECT i, j, COUNT(*)
+         FROM (SELECT t1.j, t3.i FROM t1 JOIN t3 ON t1.i = t3.j)
+         GROUP BY i, j
+         """.stripMargin, 0, 1),
+
+        // Broadcast nested loop join in plan
+        (
+          """
+         SELECT i, j, COUNT(*)
+         FROM (SELECT t1.j, t3.i FROM t1 JOIN t3 ON t1.i != t3.j)
+         GROUP BY i, j
+         """.stripMargin, 0, 1),
+
+        // Broadcast range join in plan
+        (
+          """
+         SELECT i, j, COUNT(*)
+         FROM (SELECT t1.j, t3.i FROM t1 JOIN t3 ON t1.i between t3.i and t3.j)
+         GROUP BY i, j
+         """.stripMargin, 0, 1),
+
+        // Local limit
         (
           """
          SELECT j, COUNT(*)
-         FROM (SELECT t1.j FROM t1 JOIN t3 ON t1.j = t3.j)
+         FROM (SELECT t1.j FROM t1 JOIN t3 ON t1.i = t3.j limit 200)
          GROUP BY j
-         """.stripMargin, 0, 0)
+         """.stripMargin, 0, 1),
+
+        // WindowSortLimitExec
+        (
+          """
+         SELECT * FROM
+         (SELECT j, row_number() OVER (partition by j order by i) rn, i
+         FROM t1) t
+         WHERE rn <= 3 order by j, rn, i
+         """.stripMargin, 0, 1),
+
+        (
+          """
+         SELECT * FROM
+         (SELECT j, row_number() OVER (partition by j order by i) rn, i
+         FROM t1) t
+         WHERE rn <= 3
+         """.stripMargin, 0, 1),
+
+        (
+          """
+         SELECT * FROM
+         (SELECT j, row_number() OVER (partition by j order by i) rn, i
+         FROM t1 order by j) t
+         WHERE rn <= 3
+         """.stripMargin, 0, 1),
+
+        (
+          """
+         SELECT * FROM
+         (SELECT j, row_number() OVER (partition by j order by i) rn, i
+         FROM t1 order by j) t
+         WHERE rn <= 3 order by j, rn, i
+         """.stripMargin, 0, 1)
       ).foreach { case (query, numScanWithAutoScanEnabled, numScanWithAutoScanDisabled) =>
-        checkDisableBucketedScan(query, numScanWithAutoScanEnabled, numScanWithAutoScanDisabled)
+        withSQLConf(SQLConf.RANGE_JOIN_ENABLED.key -> "true") {
+          checkDisableBucketedScan(query, numScanWithAutoScanEnabled, numScanWithAutoScanDisabled)
+        }
       }
     }
   }