Make reducer take generic arguments

Author: szehon-ho

sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/ReducibleFunction.java

@@ -60,46 +60,47 @@
 @Evolving
 public interface ReducibleFunction<I, O> {
 
-    /**
-     * This method is for bucket functions.
-     *
-     * If this bucket function is 'reducible' on another bucket function,
-     * return the {@link Reducer} function.
-     * <p>
-     * Example to return reducer for reducing f_source = bucket(4, x) on f_target = bucket(2, x)
-     * <ul>
-     *     <li>thisFunction = bucket</li>
-     *     <li>otherFunction = bucket</li>
-     *     <li>thisNumBuckets = Int(4)</li>
-     *     <li>otherNumBuckets = Int(2)</li>
-     * </ul>
-     *
-     * @param otherFunction the other bucket function
-     * @param thisNumBuckets number of buckets for this bucket function
-     * @param otherNumBuckets number of buckets for the other bucket function
-     * @return a reduction function if it is reducible, null if not
-     */
-    default Reducer<I, O> reducer(ReducibleFunction<?, ?> otherFunction,
-      int thisNumBuckets,
-      int otherNumBuckets) {
-      return reducer(otherFunction);
-    }
+  /**
+   * This method is for parameterized functions.
+   *
+   * If this parameterized function is 'reducible' on another bucket function,
+   * return the {@link Reducer} function.
+   * <p>
+   * Example to return reducer for reducing f_source = bucket(4, x) on f_target = bucket(2, x)
+   * <ul>
+   *     <li>thisFunction = bucket</li>
+   *     <li>thisParam = Int(4)</li>
+   *     <li>otherFunction = bucket</li>
+   *     <li>otherParam = Int(2)</li>
+   * </ul>
+   *
+   * @param thisParam parameter for this function
+   * @param otherFunction the other parameterized function
+   * @param otherParam parameter for the other function
+   * @return a reduction function if it is reducible, null if not
+   */
+  default Reducer<I, O> reducer(
+    Object thisParam,
+    ReducibleFunction<?, ?> otherFunction,
+    Object otherParam) {
+    throw new UnsupportedOperationException();
+  }
 
-    /**
-     * This method is for all other functions.
-     *
-     * If this function is 'reducible' on another function, return the {@link Reducer} function.
-     * <p>
-     * Example of reducing f_source = days(x) on f_target = hours(x)
-     * <ul>
-     *     <li>thisFunction = days</li>
-     *     <li>otherFunction = hours</li>
-     * </ul>
-     *
-     * @param otherFunction the other function
-     * @return a reduction function if it is reducible, null if not.
-     */
-    default Reducer<I, O> reducer(ReducibleFunction<?, ?> otherFunction) {
-      return reducer(otherFunction, 0, 0);
-    }
+  /**
+   * This method is for all other functions.
+   *
+   * If this function is 'reducible' on another function, return the {@link Reducer} function.
+   * <p>
+   * Example of reducing f_source = days(x) on f_target = hours(x)
+   * <ul>
+   *     <li>thisFunction = days</li>
+   *     <li>otherFunction = hours</li>
+   * </ul>
+   *
+   * @param otherFunction the other function
+   * @return a reduction function if it is reducible, null if not.
+   */
+  default Reducer<I, O> reducer(ReducibleFunction<?, ?> otherFunction) {
+    throw new UnsupportedOperationException();
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala

@@ -70,10 +70,9 @@ case class TransformExpression(
     } else {
       (function, other.function) match {
         case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
-          val reducer = f.reducer(o, numBucketsOpt.getOrElse(0), other.numBucketsOpt.getOrElse(0))
-          val otherReducer =
-            o.reducer(f, other.numBucketsOpt.getOrElse(0), numBucketsOpt.getOrElse(0))
-          reducer != null || otherReducer != null
+          val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
+          val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
+          thisReducer.isDefined || otherReducer.isDefined
         case _ => false
       }
     }
@@ -91,14 +90,24 @@ case class TransformExpression(
   def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
     (function, other.function) match {
       case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
-        val reducer = e1.reducer(e2,
-          numBucketsOpt.getOrElse(0),
-          other.numBucketsOpt.getOrElse(0))
-        Option(reducer)
+        reducer(e1, numBucketsOpt, e2, other.numBucketsOpt)
       case _ => None
     }
   }
 
+  // Return a Reducer for a reducible function on another reducible function
+  private def reducer(thisFunction: ReducibleFunction[_, _],
+    thisNumBucketsOpt: Option[Int],
+    otherFunction: ReducibleFunction[_, _],
+    otherNumBucketsOpt: Option[Int]): Option[Reducer[_, _]] = {
+    val res = (thisNumBucketsOpt, otherNumBucketsOpt) match {
+      case (Some(numBuckets), Some(otherNumBuckets)) =>
+        thisFunction.reducer(numBuckets, otherFunction, otherNumBuckets)
+      case _ => thisFunction.reducer(otherFunction)
+    }
+    Option(res)
+  }
+
   override def dataType: DataType = function.resultType()
 
   override protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression =

sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala

@@ -86,11 +86,15 @@ object BucketFunction extends ScalarFunction[Int] with ReducibleFunction[Int, In
     (input.getLong(1) % input.getInt(0)).toInt
   }
 
-  override def reducer(func: ReducibleFunction[_, _],
-    thisNumBuckets: Int,
-    otherNumBuckets: Int): Reducer[Int, Int] = {
+  override def reducer(
+    thisNumBucketsArg: Object,
+    otherFunc: ReducibleFunction[_, _],
+    otherNumBucketsArg: Object): Reducer[Int, Int] = {
 
-    if (func == BucketFunction) {
+    val thisNumBuckets = thisNumBucketsArg.asInstanceOf[Int]
+    val otherNumBuckets = otherNumBucketsArg.asInstanceOf[Int]
+
+    if (otherFunc == BucketFunction) {
       if ((thisNumBuckets > otherNumBuckets)
         && (thisNumBuckets % otherNumBuckets == 0)) {
         BucketReducer(thisNumBuckets, otherNumBuckets)