[CARMEL-6608] Increase bucket table scan partitions (#1269)

apache · Mar 10, 2023 · c3f17a5 · c3f17a5
1 parent 434b16e
commit c3f17a5
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala
@@ -955,8 +955,11 @@ case class FileSourceScanExec(
       var partitionNumber = FilePartition.
         minPartitionNumberBySpecifiedSize(fsRelation.sparkSession, selectedPartitions, originSize)
 
-      if (partitionNumber > conf.autoBucketedScanMaxPartitions) {
-        partitionNumber = conf.autoBucketedScanMaxPartitions
+      val maxBucketedScanMaxPartitions =
+        Math.max(conf.autoBucketedScanMaxPartitions, relation.bucketSpec.get.numBuckets)
+
+      if (partitionNumber > maxBucketedScanMaxPartitions) {
+        partitionNumber = maxBucketedScanMaxPartitions
 
         FilePartition.maxSplitBytesBySpecifiedNumber(
           fsRelation.sparkSession, selectedPartitions, partitionNumber)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -6694,11 +6694,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
             assert(scan.get.optionalBucketSet.isEmpty)
             assert(scan.get.disableBucketedScan)
 
-            if (maxPartitions > 1) {
-              assert(scan.get.inputRDD.getNumPartitions > 1)
-            } else {
-              assert(scan.get.inputRDD.getNumPartitions == 1)
-            }
+            assert(scan.get.inputRDD.getNumPartitions == 2)
           }
         }
       }