I don't think you can avoid having an explode, as at a point in your algorithm you need to transform an array of string to a string.
However, you can groupBy over the two first columns, using an User-Defined Aggregate Function to compute a map representing the different values in col3 and their col4 sum, and then explode this map into col3 and sum of col4.
By doing so, memory usage should not explode, as at any point of the algorithm you don't duplicate values in col1, col2 and col4.
Here is the complete code. First you define your custom aggregator
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.Encoder
case class Record(col3: Seq[String], col4: Int)
case class Output(col3: String, col4: Int)
object ValuesMap extends Aggregator[Record, mutable.Map[String, Int], Seq[Output]] {
def zero: mutable.Map[String, Int] = mutable.Map()
def reduce(values: mutable.Map[String, Int], currentRecord: Record): mutable.Map[String, Int] = {
currentRecord.col3.foreach(
element => values.put(element, values.getOrElse(element, 0) + currentRecord.col4)
)
values
}
def merge(values1: mutable.Map[String, Int], values2: mutable.Map[String, Int]): mutable.Map[String, Int] = {
values1.foreach(
element => values2.put(element._1, values2.getOrElse(element._1, 0) + element._2)
)
values2
}
def finish(reduction: mutable.Map[String, Int]): Seq[Output] = {
reduction.map(element => Output(element._1, element._2)).toSeq
}
def bufferEncoder: Encoder[mutable.Map[String, Int]] = ExpressionEncoder[mutable.Map[String, Int]]
def outputEncoder: Encoder[Seq[Output]] = ExpressionEncoder[Seq[Output]]
}
And then you call your aggregator as an user-defined aggregate function with udaf method:
import org.apache.spark.sql.functions.{col, explode, udaf}
val values_map = udaf(ValuesMap)
val result = df.groupBy("col1", "col2")
.agg(values_map(col("col3"), col("col4")).alias("values"))
.withColumn("value", explode(col("values")))
.select(
"col1",
"col2",
"value.col3",
"value.col4"
)
If you have the following input dataframe:
+----+----+---------+----+
|col1|col2|col3 |col4|
+----+----+---------+----+
|1 |a |[x, y] |1 |
|1 |b |[y] |10 |
|1 |a |[y, z] |100 |
|2 |b |[x, y, z]|1000|
+----+----+---------+----+
you will get the following result dataframe:
+----+----+----+----+
|col1|col2|col3|col4|
+----+----+----+----+
|1 |a |z |100 |
|1 |a |y |101 |
|1 |a |x |1 |
|1 |b |y |10 |
|2 |b |z |1000|
|2 |b |y |1000|
|2 |b |x |1000|
+----+----+----+----+
