Apache Spark & Scala

How to Make Joins in Spark Dataset API More Type-Safe

Published on December 10, 2022 by Viliam Simko · Updated September 12, 2024

I'm a senior freelance developer specializing in Scala and Apache Spark. I primarily work with companies in the fintech domain and help them tackle their big data challenges.

In this blog post, I want to address a common issue that we encounter when writing data transformations in Spark using the Dataset API. As developers, we often face a trade-off between optimizing for type-safety or performance. In my experience, this decision is particularly challenging in the fintech domain where the data can be complex and requires careful handling.

Over the years, I have found that prioritizing type-safety can make a significant difference in the reliability and accuracy of data transformations. In this article, I will share some tips and tricks that I have learned along the way to help you navigate this dilemma and make informed decisions when working with Spark's Dataset API.

The Problem: Type-Safe Outer Joins

Let's explore the limitations of Spark's Dataset API when it comes to type-safe outer join operations.

The join operation is one of the most common operations in Spark. The Dataset API also provides joinWith, which preserves the types of the two joined datasets. However, its return type does not express the optionality introduced by an outer join.

For example, when performing a left outer join, the right-hand side may be missing. In ordinary Scala code, we would naturally represent this as Option[R]. Spark's joinWith, however, exposes the result as Dataset[(L, R)], so the possibility of a missing right-hand value is not reflected in the Scala type.

This means that when using a left outer join, we need to explicitly convert the R part to an Option[R] before mapping over the result.

leftDS
  .joinWith(rightDS, joinCondition, "left_outer")
  .as[L, Option[R]]
  .map { case (left: L, right: Option[R]) =>
    // now we can safely map
  }

This works, but having to repeat the same conversion throughout a codebase quickly becomes cumbersome. It also makes the intended semantics of the join less visible at the call site.

Extending Outer Joins in Spark

After encountering the same pattern repeatedly, I decided to refactor it into a separate abstraction. Using an implicit class, I created several useful methods that extend the joining functionality of a standard Dataset.

The original implementation was built with Spark 3.3.1 and Scala 2.13. The idea is to provide a small, domain-specific layer over Spark's existing Dataset API, so that the Scala types better reflect the semantics of the join being performed.

Implicit classes are a feature of Scala that allows us to extend the functionality of an existing API without accessing its source code. This can significantly improve the developer experience when used carefully.

In our case, we define extension methods for Dataset that handle the conversion of the right-hand side of a left outer join to Option[R].

import org.apache.spark.sql.{Column, Dataset, Encoders}
import scala.reflect.runtime.universe.TypeTag

implicit class DatasetTypesafeJoins[L <: Product : TypeTag](left: Dataset[L]) {

  type JoinResult = (L, Option[R])

  def joinLeftOuterWith[R <: Product : TypeTag](
      right: Dataset[R],
      condition: Column
  ): Dataset[JoinResult] = {
    left
      .joinWith(right, condition, "left_outer")
      .as(Encoders.product[JoinResult])
  }

  // ... more extension methods
}

The JoinResult type alias makes the intended result easier to understand: a successful match contains an R, while a missing right-hand record is represented by None.

We also need to import TypeTag from scala.reflect.runtime.universe. It is required by the encoder machinery used by Encoders.product for deriving the encoder for the result type.

import scala.reflect.runtime.universe.TypeTag

Other Useful Extension Methods

Once the pattern is extracted into an extension API, other join operations can be exposed in the same style. For example:

joinLeftOuterWith(right: Dataset[R], condition: Column)
joinLeftOuterWith(right: Dataset[R], usingColumn: String)
joinLeftAntiWith(right: Dataset[R], condition: Column)
joinLeftAntiWith(right: Dataset[R], usingColumn: String)
crossJoinWith(right: Dataset[R])
joinLeftInnerWith(right: Dataset[R], usingColumn: String)

The exact set of helpers depends on the conventions and join patterns used in a particular project. The important part is that frequently repeated, potentially error-prone conversions can be hidden behind a small and consistently typed API.

Scala 3 Compatibility

Unfortunately, the solution presented here relies on Spark's Scala 2.x encoder machinery and is therefore not available in the officially supported Scala 3 setup. Spark's Encoders.product relies on Scala 2 reflection facilities such as TypeTag, which prevents this approach from working directly with Scala 3.

There are community-driven solutions that address this limitation, but they are not yet mature or stable enough to be considered a standard approach. For the implementation presented in this article, Scala 2.13 remains the intended environment.

Challenges I Faced

The biggest challenge was finding a convenient way to express the nullability introduced by an outer join without giving up the benefits of Spark's typed Dataset API.

Spark's join operation itself is flexible, but the type returned by joinWith does not encode the fact that one side of an outer join may be missing. The extension methods provide a lightweight way to bridge that gap.

Key Learnings

Refactoring our code using this new API simplified many of our existing transformations. Making joins more type-safe helps catch potential mistakes at compile time, reducing debugging effort and avoiding some classes of runtime errors.

The broader lesson is that a small abstraction around a general-purpose API can make domain semantics much clearer. In this case, Option[R] makes the possibility of a missing right-hand record explicit in the Scala type system.

Tips and Advice

I wish such functionality were available directly in the standard Spark Dataset API. Until then, a small set of project-specific extension methods can make repeated join patterns considerably easier to use.

If you need help with refactoring Spark transformations or implementing similar type-safe abstractions, feel free to get in touch.

Final Thoughts and Next Steps

There are many other useful extension methods that can be built around the Spark Dataset API. In future articles, we can also look at patterns involving Spark and Kafka that I have gathered from real-world projects.

In summary, making joins more type-safe in Spark's Dataset API can make data transformations easier to reason about and safer to refactor. By making the optionality of outer-join results explicit in Scala types, we can write transformations with clearer contracts and fewer assumptions about missing data.

← Back to all articles