We all love simple data structures where you can apply any SQL function to a single-value column. However, datasets also contain complex data types such as structures or arrays. Some time ago I was wondering what is the best approach for filtering them.
What would it take for you to trust your Databricks pipelines in production?
A 3-day bug hunt on a 3-person team costs up to €7,200 in lost engineering time. This workshop teaches you to prevent that — unit tests, data tests, and integration tests for PySpark and Databricks Lakeflow, including Spark Declarative Pipelines.
Konieczny
My naive guess was based on the high-order functions and supposed that the filtering of the array directly would be faster than exploding each column and filtering it just afterwards. Why? The explode+filter would create unnecessary rows, i.e. rows that could have been filtered prior to explode with a high-order filtering function, leading to a stronger memory pressure and more rows to work on. Turns out, I was wrong and happy at the same time to remember how much codegen imprisoned me a few years ago!
Code
To better understand the problem, let's suppose an orders dataset with the following structure:
root |-- order_id: long (nullable = true) |-- customer_id: integer (nullable = true) |-- order_date: date (nullable = true) |-- items: array (nullable = true) | |-- element: struct (containsNull = true) | | |-- sku: string (nullable = true) | | |-- price: double (nullable = true) | | |-- status: string (nullable = true)
I highlighted the important part which is the items array. Our job should filter the not cancelled items out. Here is how to do this with the explode+filter approach:
exploded = orders.select("order_id", "customer_id", "order_date", F.explode("items").alias("item"))
cancelled_items = exploded.filter(F.col("item.status") == "cancelled")
The code is simple, readable, but you might be thinking it would first create all items and later drop only the non cancelled ones. So to mitigate this eventual problem you rewrite the query to:
cancelled_items = orders.select(
"order_id", "customer_id", "order_date",
F.explode(F.filter("items", lambda item: item.status == "cancelled")).alias("item"),
)
It's just a typical example of premature optimization [is the root of all evil, Donald Knuth, "Structured Programming with go to Statements", 1974].
Code-GENie
Why does the high-order function bring no improvement compared to the explode+filter pair? Because of the underlying code gen! To see this in action, print the generated code for both snippets by calling this:
cancelled_items.explain(mode='codegen')
The explode+filter version does the following:
As you can see, everything - so filtering and exploding - happens in a single loop. Even though you expressed them in the code as two separate operations, Catalyst is smart enough to combine them into a single and more efficient execution unit. Knowing that, what happens when we use the high-order filtering function and explode?
To put it short, the high-order filtering happens after converting columnar Apache Parquet files to row format expected by Apache Spark. When you look at the code generated for the explode-filter solution, the conversion, filtering, explosion, and filtering are combined to a single stage. And because of that reason it should always be faster.
Data Engineering Design Patterns
Looking for a book that defines and solves most common data engineering problems? I wrote
one on that topic! You can read it online
on the O'Reilly platform,
or get a print copy on Amazon.
I also help solve your data engineering problems contact@waitingforcode.com đź“©