python - How to access a collection of items stored deep inside an array in PySpark Dataframe?

Question

I have a PySpark dataframe-

root
 |-- itemId: string (nullable = true)
 |-- contents: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- contentId: string (nullable = true)
 |    |    |-- position: integer (nullable = true)

How do I get all the contentId for every itemId? It want to have a new column (array type) containing all the contentId for every given itemId.

How do I achieve this?

Answer 1

You can explode and then group by the itemId, and use collect_list to get the array of contentId:

import pyspark.sql.functions as F

df.show(truncate=False)
+------------------------------+------+
|contents                      |itemId|
+------------------------------+------+
|[[content1, 1], [content2, 2]]|item1 |
|[[content3, 3], [content4, 4]]|item2 |
+------------------------------+------+

result = df.select('itemId', F.explode('contents').alias('contents')) 
           .groupBy('itemId') 
           .agg(F.collect_list('contents.contentId').alias('contents'))

result.show()
+------+--------------------+
|itemId|            contents|
+------+--------------------+
| item2|[content3, content4]|
| item1|[content1, content2]|
+------+--------------------+

Alternatively, you can use transform if you have Spark 3.0 or above:

import pyspark.sql.functions as F

result = df.select('itemId', F.expr("transform(contents, x -> x.contentId)").alias('contents'))

result.show()
+------+--------------------+
|itemId|            contents|
+------+--------------------+
| item1|[content1, content2]|
| item2|[content3, content4]|
+------+--------------------+