address comments

mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala

@@ -45,7 +45,6 @@ import org.apache.spark.mllib.linalg.CholeskyDecomposition
 import org.apache.spark.mllib.optimization.NNLS
 import org.apache.spark.rdd.{DeterministicLevel, RDD}
 import org.apache.spark.sql._
-import org.apache.spark.sql.catalyst.expressions.aggregate.CollectOrdered
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
 import org.apache.spark.storage.StorageLevel
@@ -501,17 +500,14 @@ class ALSModel private[ml] (
         }
       }.toDF(srcOutputColumn, dstOutputColumn, ratingColumn)
 
-    val aggFunc = CollectOrdered(struct(ratingColumn, dstOutputColumn).expr, num, true)
-      .toAggregateExpression(false)
-
     val arrayType = ArrayType(
       new StructType()
         .add(dstOutputColumn, IntegerType)
         .add(ratingColumn, FloatType)
     )
 
     ratings.groupBy(srcOutputColumn)
-      .agg(new Column(aggFunc))
+      .agg(collect_top_k(struct(ratingColumn, dstOutputColumn), num, false))
       .as[(Int, Seq[(Float, Int)])]
       .map(t => (t._1, t._2.map(p => (p._2, p._1))))
       .toDF(srcOutputColumn, recommendColumn)

mllib/src/test/scala/org/apache/spark/ml/recommendation/CollectOrderedSuite.scala → mllib/src/test/scala/org/apache/spark/ml/recommendation/CollectTopKSuite.scala

@@ -18,11 +18,10 @@
 package org.apache.spark.ml.recommendation
 
 import org.apache.spark.ml.util.MLTest
-import org.apache.spark.sql.{Column, DataFrame}
-import org.apache.spark.sql.catalyst.expressions.aggregate.CollectOrdered
-import org.apache.spark.sql.functions.{col, struct}
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.functions.{col, collect_top_k, struct}
 
-class CollectOrderedSuite extends MLTest {
+class CollectTopKSuite extends MLTest {
 
   import testImplicits._
 
@@ -44,16 +43,10 @@ class CollectOrderedSuite extends MLTest {
     ).toDF("user", "item", "score")
   }
 
-  private def collect_ordered(e: Column, num: Int, reverse: Boolean): Column = {
-    new Column(CollectOrdered(e.expr, num, reverse)
-      .toAggregateExpression(false)
-    )
-  }
-
   test("k smallest with k < #items") {
     val k = 2
     val topK = dataFrame.groupBy("user")
-      .agg(collect_ordered(col("score"), k, false))
+      .agg(collect_top_k(col("score"), k, true))
       .as[(Int, Seq[Float])]
       .collect()
 
@@ -69,7 +62,7 @@ class CollectOrderedSuite extends MLTest {
   test("k smallest with k > #items") {
     val k = 5
     val topK = dataFrame.groupBy("user")
-      .agg(collect_ordered(col("score"), k, false))
+      .agg(collect_top_k(col("score"), k, true))
       .as[(Int, Seq[Float])]
       .collect()
 
@@ -85,7 +78,7 @@ class CollectOrderedSuite extends MLTest {
   test("k largest with k < #items") {
     val k = 2
     val topK = dataFrame.groupBy("user")
-      .agg(collect_ordered(struct("score", "item"), k, true))
+      .agg(collect_top_k(struct("score", "item"), k, false))
       .as[(Int, Seq[(Float, Int)])]
       .map(t => (t._1, t._2.map(p => (p._2, p._1))))
       .collect()
@@ -102,7 +95,7 @@ class CollectOrderedSuite extends MLTest {
   test("k largest with k > #items") {
     val k = 5
     val topK = dataFrame.groupBy("user")
-      .agg(collect_ordered(struct("score", "item"), k, true))
+      .agg(collect_top_k(struct("score", "item"), k, false))
       .as[(Int, Seq[(Float, Int)])]
       .map(t => (t._1, t._2.map(p => (p._2, p._1))))
       .collect()

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala

@@ -196,15 +196,16 @@ case class CollectSet(
 }
 
 /**
- * Collect the top-k elements. This expression is dedicated only for MLLIB.
+ * Collect the top-k elements. This expression is dedicated only for Spark-ML.
+ * @param reverse when true, returns the smallest k elements.
  */
-case class CollectOrdered(
+case class CollectTopK(
     child: Expression,
     num: Int,
     reverse: Boolean = false,
     mutableAggBufferOffset: Int = 0,
     inputAggBufferOffset: Int = 0) extends Collect[BoundedPriorityQueue[Any]] {
-  require(num > 0)
+  assert(num > 0)
 
   def this(child: Expression, num: Int) = this(child, num, false, 0, 0)
   def this(child: Expression, num: Int, reverse: Boolean) = this(child, num, reverse, 0, 0)
@@ -219,19 +220,19 @@ case class CollectOrdered(
   }
 
   override def createAggregationBuffer(): BoundedPriorityQueue[Any] =
-    new BoundedPriorityQueue[Any](num)(ordering.reverse)
+    new BoundedPriorityQueue[Any](num)(ordering)
 
   override def eval(buffer: BoundedPriorityQueue[Any]): Any =
-    new GenericArrayData(buffer.toArray.sorted(ordering))
+    new GenericArrayData(buffer.toArray.sorted(ordering.reverse))
 
-  override def prettyName: String = "collect_ordered"
+  override def prettyName: String = "collect_top_k"
 
-  override protected def withNewChildInternal(newChild: Expression): CollectOrdered =
+  override protected def withNewChildInternal(newChild: Expression): CollectTopK =
     copy(child = newChild)
 
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): CollectOrdered =
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): CollectTopK =
     copy(mutableAggBufferOffset = newMutableAggBufferOffset)
 
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): CollectOrdered =
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): CollectTopK =
     copy(inputAggBufferOffset = newInputAggBufferOffset)
 }

sql/core/src/main/scala/org/apache/spark/sql/functions.scala

@@ -367,6 +367,9 @@ object functions {
    */
   def collect_set(columnName: String): Column = collect_set(Column(columnName))
 
+  private[spark] def collect_top_k(e: Column, num: Int, reverse: Boolean): Column =
+    withAggregateFunction { CollectTopK(e.expr, num, reverse) }
+
   /**
    * Aggregate function: returns the Pearson Correlation Coefficient for two columns.
    *