Simplify bucket reducer logic

szehon-ho · Apr 3, 2024 · 0356f9e · 0356f9e
1 parent ad61b0c
commit 0356f9e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 13 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
@@ -568,7 +568,8 @@ case class EnsureRequirements(
         child, values, joinKeyPositions, reducers, applyPartialClustering, replicatePartitions))
   }
 
-  private def reduceCommonPartValues(commonPartValues: Seq[(InternalRow, Int)],
+  private def reduceCommonPartValues(
+      commonPartValues: Seq[(InternalRow, Int)],
       expressions: Seq[Expression],
       reducers: Option[Seq[Option[Reducer[_, _]]]]) = {
     reducers match {

diff --git a/.../src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala b/.../src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala
@@ -92,20 +92,12 @@ object BucketFunction extends ScalarFunction[Int] with ReducibleFunction[Int, In
       otherNumBuckets: Int): Reducer[Int, Int] = {
 
     if (otherFunc == BucketFunction) {
-      if ((thisNumBuckets > otherNumBuckets)
-        && (thisNumBuckets % otherNumBuckets == 0)) {
-        BucketReducer(thisNumBuckets, otherNumBuckets)
-      } else {
-        val gcd = this.gcd(thisNumBuckets, otherNumBuckets)
-        if (gcd != thisNumBuckets) {
-          BucketReducer(thisNumBuckets, gcd)
-        } else {
-          null
-        }
+      val gcd = this.gcd(thisNumBuckets, otherNumBuckets)
+      if (gcd != thisNumBuckets) {
+        return BucketReducer(thisNumBuckets, gcd)
       }
-    } else {
-      null
     }
+    null
   }
 
   private def gcd(a: Int, b: Int): Int = BigInt(a).gcd(BigInt(b)).toInt