PySpark collect_list vs collect_set, and explode vs explode_outer

Arrays turn up in PySpark all the time: a list of tags on a post, the items in a basket, the pages in a visit. Two functions pack rows into an array, and two functions unpack an array into rows. They look alike, and each pair differs in one detail that decides whether your numbers are right.

What you will learn

• What collect_list and collect_set return, and which one keeps duplicates.
• What happens to NULL values and to the order of the items.
• What explode does with a NULL or empty array, and why explode_outer exists.
• How to check your row counts after either one.
• How to answer the interview question that asks for the difference.

1The data

A shop has 6 order rows. Asha bought a pen twice and a book once. Ravi bought a book, and one of his rows has no item recorded. Meena bought a pen.

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.master("local[1]").getOrCreate()

orders = spark.createDataFrame(
    [(1, "Asha", "pen"), (2, "Asha", "pen"), (3, "Asha", "book"),
     (4, "Ravi", "book"), (5, "Ravi", None), (6, "Meena", "pen")],
    "order_id int, customer string, item string",
)
orders.show()

+--------+--------+----+
|order_id|customer|item|
+--------+--------+----+
|       1|    Asha| pen|
|       2|    Asha| pen|
|       3|    Asha|book|
|       4|    Ravi|book|
|       5|    Ravi|NULL|
|       6|   Meena| pen|
+--------+--------+----+

The goal is 1 row per customer, with all their items in one array column.

2collect_list and collect_set

Both are aggregate functions, so you use them after groupBy, the same way you use sum or count. The difference is only what they do with repeated values.

(orders.groupBy("customer")
    .agg(F.sort_array(F.collect_list("item")).alias("items_list"),
         F.sort_array(F.collect_set("item")).alias("items_set"))
    .orderBy("customer").show(truncate=False))

+--------+----------------+-----------+
|customer|items_list      |items_set  |
+--------+----------------+-----------+
|Asha    |[book, pen, pen]|[book, pen]|
|Meena   |[pen]           |[pen]      |
|Ravi    |[book]          |[book]     |
+--------+----------------+-----------+

Asha's list has 3 items and her set has two, because the pen is counted once in the set. Ravi has only one item in both, and his NULL did not appear in either array.

Counting what you collected

Because NULL is skipped, the size of a list can be smaller than the number of rows in the group.

(orders.groupBy("customer")
    .agg(F.count("*").alias("order_rows"),
         F.count("item").alias("items_not_null"),
         F.size(F.collect_list("item")).alias("list_size"),
         F.size(F.collect_set("item")).alias("set_size"))
    .orderBy("customer").show())

+--------+----------+--------------+---------+--------+
|customer|order_rows|items_not_null|list_size|set_size|
+--------+----------+--------------+---------+--------+
|    Asha|         3|             3|        3|       2|
|   Meena|         1|             1|        1|       1|
|    Ravi|         2|             1|        1|       1|
+--------+----------+--------------+---------+--------+

Ravi has 2 order rows but a list of size 1. The list size matches count("item"), which also skips NULL. If you wanted the number of orders, count the rows, not the array.

The same functions work in Spark SQL. This query returns the same sets:

orders.createOrReplaceTempView("orders")
spark.sql("SELECT customer, sort_array(collect_set(item)) AS items FROM orders GROUP BY customer ORDER BY customer").show(truncate=False)

+--------+-----------+
|customer|items      |
+--------+-----------+
|Asha    |[book, pen]|
|Meena   |[pen]      |
|Ravi    |[book]     |
+--------+-----------+

3explode and explode_outer

Now the other direction. Keep the customer baskets as a table with one array per customer, and add 2 customers who have nothing: Zoya's array is NULL, and Kiran's array exists but is empty.

baskets = (orders.groupBy("customer")
    .agg(F.sort_array(F.collect_list("item")).alias("items")))
extra = spark.createDataFrame([("Zoya", None), ("Kiran", [])], "customer string, items array<string>")
all_baskets = baskets.union(extra)
all_baskets.orderBy("customer").show(truncate=False)

+--------+----------------+
|customer|items           |
+--------+----------------+
|Asha    |[book, pen, pen]|
|Kiran   |[]              |
|Meena   |[pen]           |
|Ravi    |[book]          |
|Zoya    |NULL            |
+--------+----------------+

explode gives one output row per array element.

all_baskets.select("customer", F.explode("items").alias("item")).orderBy("customer", "item").show()

+--------+----+
|customer|item|
+--------+----+
|    Asha|book|
|    Asha| pen|
|    Asha| pen|
|   Meena| pen|
|    Ravi|book|
+--------+----+

Kiran and Zoya are gone. An empty array has no elements to turn into rows, and a NULL array has none either, so the row disappears and Spark gives no warning. explode_outer keeps them.

all_baskets.select("customer", F.explode_outer("items").alias("item")).orderBy("customer", "item").show()

+--------+----+
|customer|item|
+--------+----+
|    Asha|book|
|    Asha| pen|
|    Asha| pen|
|   Kiran|NULL|
|   Meena| pen|
|    Ravi|book|
|    Zoya|NULL|
+--------+----+

The row counts show the difference: explode returns 5 rows and explode_outer returns 7. If you need to know which items sit at which position, posexplode and posexplode_outer add a position column, which starts at 0.

(all_baskets.where("customer in ('Asha', 'Zoya')")
    .select("customer", F.posexplode_outer("items").alias("pos", "item"))
    .orderBy("customer", "pos").show())

+--------+----+----+
|customer| pos|item|
+--------+----+----+
|    Asha|   0|book|
|    Asha|   1| pen|
|    Asha|   2| pen|
|    Zoya|NULL|NULL|
+--------+----+----+

4Going around the loop

If you collect rows and explode them again, do you get the same rows back? It depends on which function you used.

(orders.groupBy("customer").agg(F.sort_array(F.collect_set("item")).alias("items"))
    .select("customer", F.explode("items").alias("item"))
    .orderBy("customer", "item").show())

+--------+----+
|customer|item|
+--------+----+
|    Asha|book|
|    Asha| pen|
|   Meena| pen|
|    Ravi|book|
+--------+----+

Asha's second pen is gone, because collect_set removed it before explode ran. collect_list followed by explode returns the original items. So use collect_set only when you really want distinct values.

5Mistakes to avoid

• Trusting the order. Sort the array, or collect a struct and sort that, if the order carries meaning.
• Collecting a very large group. One array lives inside 1 row on one machine. If a customer has millions of events, collect_list can run a task out of memory. Aggregate to counts or top items instead.
• Expecting NULL to be kept. Neither collect_list nor collect_set keeps NULL, so the array can be shorter than the group.
• Losing rows without noticing. After a plain explode, compare the number of distinct keys before and after.

6Check yourself

• Asha has the items pen, pen and book. What does collect_set return for her, and how many items does collect_list return?
• A customer has an empty array. Which of explode and explode_outer returns a row for them?
• You explode a column and then sum an order amount. Why can the total be too high?

Practise it: open Explode Array Column and Items per user as an array in the editor.