[CARMEL-5986] Make Coalesce/Rebucketing effective when bucketing disabled by planner (#961)

* [CARMEL-5986] Make Coalesce/Rebucketing effective when bucketing disabled by planner

* fix ut

Author: xingchaozh

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

@@ -329,12 +329,12 @@ object QueryExecution {
       PlanSubqueries(sparkSession),
       EliminateHintPlaceHolder,
       EnsureRequirements,
-      DisableUnnecessaryBucketedScan,
       // `RemoveRedundantSorts` needs to be added after `EnsureRequirements` to guarantee the same
       // number of partitions when instantiating PartitioningCollection.
       RemoveRedundantSorts,
       EnsureRepartitionForWriting,
       EliminateShuffleExec,
+      DisableUnnecessaryBucketedScan,
       ApplyColumnarRulesAndInsertTransitions(sparkSession.sessionState.columnarRules),
       CollapseCodegenStages(),
       ReuseExchange,

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala

@@ -91,10 +91,10 @@ case class AdaptiveSparkPlanExec(
   private def queryStagePreparationRules: Seq[Rule[SparkPlan]] = Seq(
     eliminateHintPlaceHolder,
     ensureRequirements,
-    DisableUnnecessaryBucketedScan,
     removeRedundantSorts,
     EnsureRepartitionForWriting,
-    EliminateShuffleExec
+    EliminateShuffleExec,
+    DisableUnnecessaryBucketedScan
   ) ++ context.session.sessionState.queryStagePrepRules
 
   @transient private val initialPlan = context.session.withActive {

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

@@ -19,10 +19,9 @@ 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.{FileSourceScanExec, FilterExec, ProjectExec, SortExec, SparkPlan}
+import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
 import org.apache.spark.sql.execution.exchange.Exchange
-import org.apache.spark.sql.internal.SQLConf
 
 /**
  * Disable unnecessary bucketed table scan based on actual physical query plan.
@@ -120,7 +119,12 @@ object DisableUnnecessaryBucketedScan extends Rule[SparkPlan] {
   }
 
   private def hasInterestingPartition(plan: SparkPlan): Boolean = {
-    plan.requiredChildDistribution.exists {
+    val isReplacedShuffle = plan match {
+      case _: RebucketingExec | _: CoalesceExec => true
+      case _ => false
+    }
+
+    isReplacedShuffle || plan.requiredChildDistribution.exists {
       case _: ClusteredDistribution | _: HashClusteredDistribution | AllTuples => true
       case _ => false
     }

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

@@ -43,20 +43,25 @@ abstract class DynamicBloomFilterPruningSuiteBase
   private val tableFormat: String = "parquet"
   private var originalAutoBroadcastJoinThreshold: Long = _
   private var originalParquetCompressionCodec: String = _
+  private var originalAutoBucketedScan: Boolean = _
 
 
   override def beforeAll(): Unit = {
     super.beforeAll()
     originalAutoBroadcastJoinThreshold = conf.autoBroadcastJoinThreshold
     originalParquetCompressionCodec = conf.parquetCompressionCodec
+    originalAutoBucketedScan = conf.autoBucketedScanEnabled
     conf.setConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD, -1L)
+    conf.setConf(SQLConf.AUTO_BUCKETED_SCAN_ENABLED, true)
   }
 
   override def afterAll(): Unit = {
     spark.sessionState.conf.unsetConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED)
     spark.sessionState.conf.unsetConf(SQLConf.ADAPTIVE_EXECUTION_FORCE_APPLY)
+    spark.sessionState.conf.unsetConf(SQLConf.AUTO_BUCKETED_SCAN_ENABLED)
     conf.setConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD, originalAutoBroadcastJoinThreshold)
     conf.setConf(SQLConf.PARQUET_COMPRESSION, originalParquetCompressionCodec)
+    conf.setConf(SQLConf.AUTO_BUCKETED_SCAN_ENABLED, originalAutoBucketedScan)
     super.afterAll()
   }
 

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

@@ -1381,7 +1381,8 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
           SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
           SQLConf.ENABLE_COALESCE.key -> "true",
           SQLConf.ENABLE_REBUCKETING.key -> "false",
-          SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen) {
+          SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen,
+          SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "true") {
           Seq("inner").foreach { joinType =>
             val df = spark.sql(s"select a.* from " +
               s"(select * from $tblName distribute by col2) a " +

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

@@ -174,7 +174,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, 2, 3), // TODO 2->3 if ENABLE_REBUCKETING=false
+         """.stripMargin, 3, 3), // TODO 3->2 if ENABLE_REBUCKETING=false
           (
             """
          SELECT /*+ merge(t1) broadcast(t3)*/ * FROM t1 JOIN t2 JOIN t3
@@ -184,7 +184,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, 2, 3), // TODO 2->3 if ENABLE_REBUCKETING=false
+         """.stripMargin, 3, 3), // TODO 3->2 if ENABLE_REBUCKETING=false
           // Multiple joins on non-bucketed columns
           (
             """
@@ -195,7 +195,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, 1, 3), // TODO 1->2 if ENABLE_REBUCKETING=false
+         """.stripMargin, 2, 3), // TODO 2->1 if ENABLE_REBUCKETING=false
           (
             """
          SELECT /*+ merge(t1, t3)*/ * FROM t1 JOIN t2 JOIN t3

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

@@ -177,7 +177,8 @@ class MultipleBucketJoinSuite extends SharedSparkSession with AdaptiveSparkPlanH
   test("Join two bucketed tables with double bucket number of another") {
     val res = 0 until 2000 map(i => Row(i, i.toString, i, i.toString))
     withSQLConf(SQLConf.ENABLE_REBUCKETING.key -> "false",
-      SQLConf.ENABLE_COALESCE.key -> "true") {
+      SQLConf.ENABLE_COALESCE.key -> "true",
+      SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "true") {
       testBucketing(s"SELECT * FROM $table1 a JOIN $table2 b ON a.key = b.key",
         BucketedTableTestSpec(bucketSpec = bucketSpec1, expectSort = true),
         BucketedTableTestSpec(bucketSpec = bucketSpec2),
@@ -305,7 +306,8 @@ class MultipleBucketJoinSuite extends SharedSparkSession with AdaptiveSparkPlanH
         """.stripMargin
     val res = 0 until 2000 map(i => Row(i, i.toString, "part0", i, i.toString, "part1"))
     withSQLConf(SQLConf.ENABLE_REBUCKETING.key -> "false",
-      SQLConf.ENABLE_COALESCE.key -> "true") {
+      SQLConf.ENABLE_COALESCE.key -> "true",
+      SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "true") {
       testBucketing(query,
         BucketedTableTestSpec(bucketSpec = bucketSpec1, expectSort = true),
         BucketedTableTestSpec(bucketSpec = bucketSpec2),
@@ -338,7 +340,8 @@ class MultipleBucketJoinSuite extends SharedSparkSession with AdaptiveSparkPlanH
     val res1 = 0 until 2000 map(i => Row(i, i.toString, "part0", i, i.toString, "part0"))
     val res2 = 0 until 2000 map(i => Row(i, i.toString, "part1", i, i.toString, "part0"))
     withSQLConf(SQLConf.ENABLE_REBUCKETING.key -> "false",
-      SQLConf.ENABLE_COALESCE.key -> "true") {
+      SQLConf.ENABLE_COALESCE.key -> "true",
+      SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "true") {
       testBucketing(
         query,
         BucketedTableTestSpec(bucketSpec = bucketSpec1, expectSort = true),
@@ -377,7 +380,8 @@ class MultipleBucketJoinSuite extends SharedSparkSession with AdaptiveSparkPlanH
     val res1 = 0 until 2000 map(i => Row(i, i.toString, "part0", i, i.toString, "part0"))
     val res2 = 0 until 2000 map(i => Row(i, i.toString, "part0", i, i.toString, "part1"))
     withSQLConf(SQLConf.ENABLE_REBUCKETING.key -> "false",
-      SQLConf.ENABLE_COALESCE.key -> "true") {
+      SQLConf.ENABLE_COALESCE.key -> "true",
+      SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "true") {
       testBucketing(
         query,
         BucketedTableTestSpec(bucketSpec = bucketSpec1, expectSort = true),
@@ -415,7 +419,8 @@ class MultipleBucketJoinSuite extends SharedSparkSession with AdaptiveSparkPlanH
     val res4 = 0 until 2000 map(i => Row(i, i.toString, "part1", i, i.toString, "part1"))
 
     withSQLConf(SQLConf.ENABLE_REBUCKETING.key -> "false",
-      SQLConf.ENABLE_COALESCE.key -> "true") {
+      SQLConf.ENABLE_COALESCE.key -> "true",
+      SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "true") {
       testBucketing(
         query,
         BucketedTableTestSpec(bucketSpec = bucketSpec1, expectSort = true),