Materialising DataFrames in Spark

We recently dived into Spark and saw how DataFrames use lazy evaluation to defer execution until needed. This is useful for defining segments of logic to be reused, but if this causes large, complex calculations to execute multiple times, there may be a better way.

Lazy evaluation

Let’s take an example with some sales data:

sales = spark.read.format("csv") \
    .option("header", "false") \
    .schema(SalesSchema) \
    .load("Files/Sales/*.csv")

orders = spark.table("sales_orders") \
    .select("SalesOrderNumber", "OrderQuantity", "OrderSubtotal")

salesDetails = sales \
    .withColumn("SoldYear", year(sales.OrderDate)) \
    .withColumn("SoldMonth", month(sales.OrderDate)) \
    .withColumn("TotalPriceIncTax", (sales.Quantity * (sales.UnitPrice + sales.Tax))) \
    .join(orders, "SalesOrderNumber") \
    .withColumn("OrderValueBand",
        when(col("OrderSubtotal") >= 500, "Large")
        .when(col("OrderSubtotal") >= 100, "Medium")
        .otherwise("Small"))

monthlySales = salesDetails \
    .groupBy("SoldYear", "SoldMonth", "OrderValueBand") \
    .agg(sum("TotalPriceIncTax").alias("TotalSales"))

display(salesDetails)
display(monthlySales)

The display triggers execution of the DataFrames. Because DataFrames use lazy evaluation and store the plan, the monthlySales repeats the same retrieval, join, and calculations that salesDetails would have already done.

By using the Fabric Toolbox we can visualise the plan. Here’s the salesDetails showing the join:

The same logic is repeated and added to for the monthlySales data:

This is a trivial example, but with large datasets and complex calculations, we don’t always want to see steps repeated.

Let’s suppose we’re in SQL. If the salesDetails DataFrame was a temp table, we’d only read the CSV files and apply first-stage transformations once, and the monthlySales would only need to apply its own transformations on top of these.

Proactive reuse

A way to achieve this pattern is through caching. All we need is one extra method on our salesDetails DataFrame:

salesDetails = sales \
    .withColumn("SoldYear", year(sales.OrderDate)) \
    .withColumn("SoldMonth", month(sales.OrderDate)) \
    .withColumn("TotalPriceIncTax", (sales.Quantity * (sales.UnitPrice + sales.Tax))) \
    .join(orders, "SalesOrderNumber") \
    .withColumn("OrderValueBand",
        when(col("OrderSubtotal") >= 500, "Large")
        .when(col("OrderSubtotal") >= 100, "Medium")
        .otherwise("Small")) \
    .cache()

Now when the salesDetails DataFrame is actioned, it’ll be cached by Spark, and the monthlySales will calculate based on the cache without going back to the source files or performing the join.

If you want more control over where the cache lives, you can use .persist() to specify the cache location. .cache() is a shorthand that uses the default location.

The initial portion of salesDetails execution populates the cache with the same plan as above. This is followed by returning results directly from the cache:

The monthlySales now doesn’t need to execute the same work again. It only applies transformations to the cached data:

For larger datasets with complex calculations or aggregations, this optimisation can save excessive resource consumption and processing time.

With that said, it’s not a silver bullet and shouldn’t be used whenever a DataFrame definition is reused:

• Caching typically stores the results in memory, so will compete with other cached data or ongoing processing, leading to memory pressure
• Cache will persist until the end of the Spark session or until evicted. It’s recommended to remove once the work is complete, for example salesDetails.unpersist()
• If upstream work is cheap or results are only used once, the overhead of caching may cost more than it saves

Wrap up

In this post we’ve looked at materialising intermediate results. Typically reusing a DataFrame reuses the execution plan. Caching provides a boundary to reuse the data without repeating the processing.

By marking the intermediate result for caching, the first action materialises it for subsequent reuse, so it acts similarly to temp tables from the SQL world. This is particularly helpful for reusing complex logic without repeating operations over and over.

With that said, it isn’t always the answer, and should be weighed against the benefits and impacts on memory. If you have the right problem, it’s a great solution.