How to concatenate multiple rows in limited specified size array in spark-sql

user_9388998 0
Viewed 61

How to aggregate multiple rows into limited specified size array in spark-sql.

I have tried using aggregate functions but couldn't help on specifying the size of the array.

Input:
|id    |  value  |
|------|---------|
|435725|{abc,def}|
|435725|{ghi,jkl}|
|435725|{mno,pqr}|
|435725|{stu,vwx}|
|536345|{abc,def}|
|536345|{ghi,jkl}|
|536345|{mno,pqr}|
|536345|{stu,vwx}|

Output:
|   id   |       value         |
|--------|---------------------|
|435725_1|[{abc,def},{ghi,jkl}]|
|435725_2|[{mno,pqr},{stu,vwx}]|
|536345_1|[{abc,def},{ghi,jkl}]|
|536345_2|[{mno,pqr},{stu,vwx}]|
apache-spark-sql apache-spark
1 Answers

You can first split your ids so that when grouped will create arrays of same size. To do so, you can use a function based on row_number() over a Window. Then you generate your array grouping by those new ids using collect_list function.

Scala code as follow:

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{col, collect_list, concat_ws, floor, row_number}

val arraySize = 2

val window = Window.partitionBy("id").orderBy("value")

data
 .withColumn(
   "id", 
   concat_ws(
     "_", 
     col("id"), 
     floor((row_number().over(window) - 1) / arraySize) + 1
   )
 )
 .groupBy("id")
 .agg(collect_list(col("value")).as("value"))
 .orderBy("id") // optional, to order ids

By doing so, you get the following output:

+--------+----------------------+
|id      |collect_list(value)   |
+--------+----------------------+
|435725_1|[{abc,def}, {ghi,jkl}]|
|435725_2|[{mno,pqr}, {stu,vwx}]|
|536345_1|[{abc,def}, {ghi,jkl}]|
|536345_2|[{mno,pqr}, {stu,vwx}]|
+--------+----------------------+
user_68077690
Related