Filtering based on value and creating list in spark dataframe

Question

I am new to spark and I am trying to do the following, using Pyspark:

I have a dataframe with 3 columns, "id", "number1", "number2".

For each value of "id" I have multiple rows and what I want to do is create a list of tuples with all the rows that correspond to each id.

Eg, for the following dataframe

id | number1 | number2 |
a  |       1 |       1 |
a  |       2 |       2 |
b  |       3 |       3 |
b  |       4 |       4 |

the desired outcome would be 2 lists as such:

[(1, 1), (2, 2)] 

and

[(3, 3), (4, 4)]

I'm not sure how to approach this, since I'm a newbie. I have managed to get a list of the distinct ids doing the following

distinct_ids = [x for x in df.select('id').distinct().collect()]

In pandas that I'm more familiar with, now I would loop through the dataframe for each distinct id and gather all the rows for it, but I'm sure this is far from optimal.

Can you give me any ideas? Groupby comes to mind but I'm not sure how to approach

question from:https://stackoverflow.com/questions/65847380/filtering-based-on-value-and-creating-list-in-spark-dataframe

Answer 1

You can use groupby and aggregate using collect_list and array:

import pyspark.sql.functions as F

df2 = df.groupBy('id').agg(F.collect_list(F.array('number1', 'number2')).alias('number'))

df2.show()
+---+----------------+
| id|          number|
+---+----------------+
|  b|[[3, 3], [4, 4]]|
|  a|[[1, 1], [2, 2]]|
+---+----------------+

And if you want to get back a list of tuples,

result = [[tuple(j) for j in i] for i in [r[0] for r in df2.select('number').orderBy('number').collect()]]

which gives result as [[(1, 1), (2, 2)], [(3, 3), (4, 4)]]

If you want a numpy array, you can do

result = np.array([r[0] for r in df2.select('number').collect()])

which gives

array([[[3, 3],
        [4, 4]],

       [[1, 1],
        [2, 2]]])