How to perform arithmetic operations in a DataFrame groupBy aggregation in Spark? [duplicate] - scala

Question

This question already has answers here:
Cumulative product in Spark
(4 answers)
Closed 4 years ago.
I have a dataframe as follows:
val df = Seq(("x", "y", 1),("x", "z", 2),("x", "a", 4), ("x", "a", 5), ("t", "y", 1), ("t", "y2", 6), ("t", "y3", 3), ("t", "y4", 5)).toDF("F1", "F2", "F3")
+---+---+---+
| F1| F2| F3|
+---+---+---+
| x| y| 1|
| x| z| 2|
| x| a| 4|
| x| a| 5|
| t| y| 1|
| t| y2| 6|
| t| y3| 3|
| t| y4| 5|
+---+---+---+
How do I groupBy on column "F1", and multiply on "F3"?
For sum, I can do as follows, but not sure which function to use for multiplication.
df.groupBy("F1").agg(sum("F3")).show
+---+-------+
| F1|sum(F3)|
+---+-------+
| x| 12|
| t| 15|
+---+-------+

Answer 1

val df = Seq(("x", "y", 1),("x", "z", 2),("x", "a", 4), ("x", "a", 5), ("t", "y", 1), ("t", "y2", 6), ("t", "y3", 3), ("t", "y4", 5)).toDF("F1", "F2", "F3")
import org.apache.spark.sql.Row
val x=df.select($"F1",$"F3").groupByKey{case r=>r.getString(0)}.reduceGroups{ ((r),(r2)) =>Row(r.getString(0),r.getInt(1)*r2.getInt(1)) }
x.show()
+-----+------------------------------------------+
|value|ReduceAggregator(org.apache.spark.sql.Row)|
+-----+------------------------------------------+
| x| [x, 40]|
| t| [t, 90]|
+-----+------------------------------------------+

Answer 2

Define a custom aggregation function as follows :
class Product extends UserDefinedAggregateFunction {
// This is the input fields for your aggregate function.
override def inputSchema: org.apache.spark.sql.types.StructType =
StructType(StructField("value", LongType) :: Nil)
// This is the internal fields you keep for computing your aggregate.
override def bufferSchema: StructType = StructType(
StructField("product", LongType) :: Nil
)
// This is the output type of your aggregatation function.
override def dataType: DataType = LongType
override def deterministic: Boolean = true
// This is the initial value for your buffer schema.
override def initialize(buffer: MutableAggregationBuffer): Unit = {
buffer(0) = 1L
}
// This is how to update your buffer schema given an input.
override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
buffer(0) = buffer.getAs[Long](0) * input.getAs[Long](0)
}
// This is how to merge two objects with the bufferSchema type.
override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
buffer1(0) = buffer1.getAs[Long](0) * buffer2.getAs[Long](0)
}
// This is where you output the final value, given the final value of your bufferSchema.
override def evaluate(buffer: Row): Any = {
buffer.getLong(0)
}
}
Then use it in aggregation as follows :
val product = new Product
val df = Seq(("x", "y", 1),("x", "z", 2),("x", "a", 4), ("x", "a", 5), ("t", "y", 1), ("t", "y2", 6), ("t", "y3", 3), ("t", "y4", 5)).toDF("F1", "F2", "F3")
df.groupBy("F1").agg(product(col("F3"))).show
Here's the output :
+---+-----------+
| F1|product(F3)|
+---+-----------+
| x| 40|
| t| 90|
+---+-----------+