apache spark - Sum of array elements depending on value condition pyspark

Question

I have a pyspark dataframe:

id   |   column
------------------------------
1    |  [0.2, 2, 3, 4, 3, 0.5]
------------------------------
2    |  [7, 0.3, 0.3, 8, 2,]
------------------------------

I would like to create a 3 columns:

• Column 1: contain the sum of the elements < 2
• Column 2: contain the sum of the elements > 2
• Column 3: contain the sum of the elements = 2 (some times I have duplicate values so I do their sum) In case if I don't have a values I put null.

Expect result:

id   |   column               |  column<2 |  column>2   | column=2 
------------------------------|--------------------------------------------  
1    |  [0.2, 2, 3, 4, 3, 0.5]|  [0.7]    |  [12]       |  null
---------------------------------------------------------------------------
2    |  [7, 0.3, 0.3, 8, 2,]  | [0.6]     |  [15]       |  [2]
---------------------------------------------------------------------------

Can you help me please ? Thank you

Answer 1

For Spark 2.4+, you can use aggregate and filter higher-order functions like this:

df.withColumn("column<2", expr("aggregate(filter(column, x -> x < 2), 0D, (x, acc) -> acc + x)")) 
  .withColumn("column>2", expr("aggregate(filter(column, x -> x > 2), 0D, (x, acc) -> acc + x)")) 
  .withColumn("column=2", expr("aggregate(filter(column, x -> x == 2), 0D, (x, acc) -> acc + x)")) 
  .show(truncate=False)

Gives:

+---+------------------------------+--------+--------+--------+
|id |column                        |column<2|column>2|column=2|
+---+------------------------------+--------+--------+--------+
|1  |[0.2, 2.0, 3.0, 4.0, 3.0, 0.5]|0.7     |10.0    |2.0     |
|2  |[7.0, 0.3, 0.3, 8.0, 2.0]     |0.6     |15.0    |2.0     |
+---+------------------------------+--------+--------+--------+