How to limit functions.collect_set in Spark SQL? - apache-spark

Question

I'm dealing with a column of numbers in a large spark DataFrame, and I would like to create a new column that stores an aggregated list of unique numbers that appear in that column.
Basically exactly what functions.collect_set does. However, i only need up to 1000 elements in the aggregated list. Is there any way to pass that parameter somehow to functions.collect_set(), or any other way to get only up to 1000 elements in the aggregated list, without using a UDAF?
Since the column is so large, I'd like to avoid collecting all elements and trimming the list afterwards.
Thanks!

Answer 1

Spark 2.4
As pointed out in a comment, Spark 2.4.0 comes with slice standard function which can do this sort of thing.
val usage = sql("describe function slice").as[String].collect()(2)
scala> println(usage)
Usage: slice(x, start, length) - Subsets array x starting from index start (array indices start at 1, or starting from the end if start is negative) with the specified length.
That gives the following query:
val q = input
.groupBy('key)
.agg(collect_set('id) as "collect")
.withColumn("three_only", slice('collect, 1, 3))
scala> q.show(truncate = false)
+---+--------------------------------------+------------+
|key|collect |three_only |
+---+--------------------------------------+------------+
|0 |[0, 15, 30, 45, 5, 20, 35, 10, 25, 40]|[0, 15, 30] |
|1 |[1, 16, 31, 46, 6, 21, 36, 11, 26, 41]|[1, 16, 31] |
|3 |[33, 48, 13, 38, 3, 18, 28, 43, 8, 23]|[33, 48, 13]|
|2 |[12, 27, 37, 2, 17, 32, 42, 7, 22, 47]|[12, 27, 37]|
|4 |[9, 19, 34, 49, 24, 39, 4, 14, 29, 44]|[9, 19, 34] |
+---+--------------------------------------+------------+
Before Spark 2.4
I'd use a UDF that would do what you want after collect_set (or collect_list) or a much harder UDAF.
Given more experience with UDFs, I'd go with that first. Even though UDFs are not optimized, for this use case it's fine.
val limitUDF = udf { (nums: Seq[Long], limit: Int) => nums.take(limit) }
val sample = spark.range(50).withColumn("key", $"id" % 5)
scala> sample.groupBy("key").agg(collect_set("id") as "all").show(false)
+---+--------------------------------------+
|key|all |
+---+--------------------------------------+
|0 |[0, 15, 30, 45, 5, 20, 35, 10, 25, 40]|
|1 |[1, 16, 31, 46, 6, 21, 36, 11, 26, 41]|
|3 |[33, 48, 13, 38, 3, 18, 28, 43, 8, 23]|
|2 |[12, 27, 37, 2, 17, 32, 42, 7, 22, 47]|
|4 |[9, 19, 34, 49, 24, 39, 4, 14, 29, 44]|
+---+--------------------------------------+
scala> sample.
groupBy("key").
agg(collect_set("id") as "all").
withColumn("limit(3)", limitUDF($"all", lit(3))).
show(false)
+---+--------------------------------------+------------+
|key|all |limit(3) |
+---+--------------------------------------+------------+
|0 |[0, 15, 30, 45, 5, 20, 35, 10, 25, 40]|[0, 15, 30] |
|1 |[1, 16, 31, 46, 6, 21, 36, 11, 26, 41]|[1, 16, 31] |
|3 |[33, 48, 13, 38, 3, 18, 28, 43, 8, 23]|[33, 48, 13]|
|2 |[12, 27, 37, 2, 17, 32, 42, 7, 22, 47]|[12, 27, 37]|
|4 |[9, 19, 34, 49, 24, 39, 4, 14, 29, 44]|[9, 19, 34] |
+---+--------------------------------------+------------+
See functions object (for udf function's docs).

Answer 2

I'm using a modified copy of the collect_set and collect_list functions; because of code scopes, the modified copies must be in the same package path as the originals. The linked code works for Spark 2.1.0; if you are using a prior version, method signatures may be different.
Throw this file (https://gist.github.com/lokkju/06323e88746c85b2ce4de3ea9cdef9bc) into your project as src/main/org/apache/spark/sql/catalyst/expression/collect_limit.scala
use it as:
import org.apache.spark.sql.catalyst.expression.collect_limit._
df.groupBy('set_col).agg(collect_set_limit('set_col,1000)

Answer 3

scala> df.show
+---+-----+----+--------+
| C0| C1| C2| C3|
+---+-----+----+--------+
| 10| Name|2016| Country|
| 11|Name1|2016|country1|
| 10| Name|2016| Country|
| 10| Name|2016| Country|
| 12|Name2|2017|Country2|
+---+-----+----+--------+
scala> df.groupBy("C1").agg(sum("C0"))
res36: org.apache.spark.sql.DataFrame = [C1: string, sum(C0): bigint]
scala> res36.show
+-----+-------+
| C1|sum(C0)|
+-----+-------+
|Name1| 11|
|Name2| 12|
| Name| 30|
+-----+-------+
scala> df.limit(2).groupBy("C1").agg(sum("C0"))
res33: org.apache.spark.sql.DataFrame = [C1: string, sum(C0): bigint]
scala> res33.show
+-----+-------+
| C1|sum(C0)|
+-----+-------+
| Name| 10|
|Name1| 11|
+-----+-------+
scala> df.groupBy("C1").agg(sum("C0")).limit(2)
res2: org.apache.spark.sql.DataFrame = [C1: string, sum(C0): bigint]
scala> res2.show
+-----+-------+
| C1|sum(C0)|
+-----+-------+
|Name1| 11|
|Name2| 12|
+-----+-------+
scala> df.distinct
res8: org.apache.spark.sql.DataFrame = [C0: int, C1: string, C2: int, C3: string]
scala> res8.show
+---+-----+----+--------+
| C0| C1| C2| C3|
+---+-----+----+--------+
| 11|Name1|2016|country1|
| 10| Name|2016| Country|
| 12|Name2|2017|Country2|
+---+-----+----+--------+
scala> df.dropDuplicates(Array("c1"))
res11: org.apache.spark.sql.DataFrame = [C0: int, C1: string, C2: int, C3: string]
scala> res11.show
+---+-----+----+--------+
| C0| C1| C2| C3|
+---+-----+----+--------+
| 11|Name1|2016|country1|
| 12|Name2|2017|Country2|
| 10| Name|2016| Country|
+---+-----+----+--------+

Answer 4

As other answers have mentioned, the performant way of doing this would be to write a UDAF. Unfortunately the UDAF API is actually not as extensible as the aggregate functions that ship with spark. However you can use their internal APIs to build on the internal functions to do what you need.
Here is an implementation for collect_set_limit that is mostly a copy past of Spark's internal CollectSet AggregateFunction. I would just extend it but its a case class. Really all that's needed is to override update and merge methods to respect a passed in limit:
case class CollectSetLimit(
child: Expression,
limitExp: Expression,
mutableAggBufferOffset: Int = 0,
inputAggBufferOffset: Int = 0) extends Collect[mutable.HashSet[Any]] {
val limit = limitExp.eval( null ).asInstanceOf[Int]
def this(child: Expression, limit: Expression) = this(child, limit, 0, 0)
override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
copy(mutableAggBufferOffset = newMutableAggBufferOffset)
override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
copy(inputAggBufferOffset = newInputAggBufferOffset)
override def createAggregationBuffer(): mutable.HashSet[Any] = mutable.HashSet.empty
override def update(buffer: mutable.HashSet[Any], input: InternalRow): mutable.HashSet[Any] = {
if( buffer.size < limit ) super.update(buffer, input)
else buffer
}
override def merge(buffer: mutable.HashSet[Any], other: mutable.HashSet[Any]): mutable.HashSet[Any] = {
if( buffer.size >= limit ) buffer
else buffer ++= other.take( limit - buffer.size )
}
override def prettyName: String = "collect_set_limit"
}
And to actually register it, we can do it through Spark's internal FunctionRegistry which takes in the name and the builder which is effectively a function that creates a CollectSetLimit using the provided expressions:
val collectSetBuilder = (args: Seq[Expression]) => CollectSetLimit( args( 0 ), args( 1 ) )
FunctionRegistry.builtin.registerFunction( "collect_set_limit", collectSetBuilder )
Edit:
Turns out adding it to the builtin only works if you haven't created the SparkContext yet as it makes an immutable clone on startup. If you have an existing context then this should work to add it with reflection:
val field = classOf[SessionCatalog].getFields.find( _.getName.endsWith( "functionRegistry" ) ).get
field.setAccessible( true )
val inUseRegistry = field.get( SparkSession.builder.getOrCreate.sessionState.catalog ).asInstanceOf[FunctionRegistry]
inUseRegistry.registerFunction( "collect_set_limit", collectSetBuilder )

Answer 5

use take
val firstThousand = rdd.take(1000)
Will return the first 1000.
Collect also has a filter function that can be provided. That would allow you to be more specific as to what is returned.