Write Clean and SOLID Scala Spark Jobs | by Ronald Ángel | Towards Data Science

Creating data pipelines by writing spark jobs is nowadays easier due to the growth of new tools and data platforms that allow multiple data parties (analysts, engineers, scientists, etc.) to focus on understanding data and writing logic to get insights. Nevertheless, new tools like notebooks that allow easy scripting, sometimes are not well used and could cause a new problem: extensive data pipelines are written as simple SQL queries or scripts, neglecting important development concepts as writing clean and testable code. Thus, design principles that ensure that code is maintainable and extensible might be broken, leading to further problems in an environment where our products should be dynamic.
We will expose a process that contains a set of steps and patterns that will help you in creating better spark pipelines using a basic spark pipeline as an example.
The first and most important step (even more than code cleanliness) in every pipeline is the definition of its structure. Thus, the pipeline structure should be defined after a process of data exploration that provides the phases need to produce the expected outputs from the inputs.
Let’s work on a basic example and define a pipeline structure. Therefore, we have three datasets:
usersSourceExampleDS that contains users’ information.
genderSourceAExampleDS reads from the source ‘exampleA’ that contains genders by a specif name in a country.
genderSourceBExampleDS reads from the source ‘exampleB’ and contains another list of genders by a name. However, in this case, it does not discriminate by the country and adds a computed probability.
Then, the pipeline aim is to produce a dataset where the column gender is added to the usersSourceExampleDS as follow:
When the suffix has an explicit gender, for example, Mr or Ms add the gender right away in the column gender.
If the suffix does not have the gender then search the name on the genderSourceAExampleDS and add a new column source_a_gender. Then, search the name on the genderSourceBExampleDS and add the column source_b_gender.
Finally, when source_a_gender is not null set this value to gender, otherwise use the source_b_gender only if the probability is greater than 0.5.
Also, some metrics like male percentage and female percentage are produced into the metrics storage system.
The following phases are defined:
Data Reading: reads from the data sources. In this case, data is stored in S3.
Data Pre-processing: As we can see in the data, there is not unique ID to join or search data, then texts within the columns name, country and suffix are used. However, these columns have invalid data to be removed (NaN), multiple letter cases, acronyms, and special characters that are pre-processed in this phase.
Data Enrichment: After data is clean and ready to be used, it is transformed and parsed within this phase. Thus, new columns are added after some business validations.
Data Metrics: This phase contains tasks related to aggregations and calculations over the transformed data.
Data Writing: Finally, this phase handles writes of the genderized results to S3 and the metrics to an external tool.
Is it quite simple? Yes, We know what you are thinking. This pipeline might be written easily using spark. You open your spark context, read the datasets, parse the datasets, finally join to get the final the results and write them to the output datastores. This looks like this.
This script works fairly well; it uses some extensive spark SQL and repeats the same logic for 3 different sources though. Besides, let’s imagine the following scenario:
This is only the first version of the pipeline and multiple new genderized sources are going to be added in the future.
Multiple new shared pre-processing and transforming tasks are going to be added.
Some cleaning functions have to be tested in isolation. As an example, the function that removes acronyms. Also, the output of each intermediate step should be tested.
New extensions and tests will be automatically configured with CI/CD.
As a result, the previous source code becomes useful for some cases where an effort for extensibility is needless; else becomes non-extensible, untestable and unmaintainable. Consequently, the following steps are recommended.
Define a project structure that suits your pipelien phases, this will make your pipeline a data product.
A) Define a project structure for phases defined previously, this looks like this:
B) Define your project dependencies avoiding circular/unused dependencies and define clearly the dependency scope between test, provided and compile.
C) Start adding a general helper within each package where small and general functions are going to be defined. Define the rest of the functions in a companion object of the main class for now.
D) Define which version of the spark API you will use: RDDs, datasets, or dataframes depending on your requirements.
E) When you have multiple jobs using this same logic, think about creating a general spark-utils library that is shared for all your company jobs.
Creating jobs that behave as expected is not enough. Currently, companies cope to reduce technical debt and creating self-documented pipelines that are easy to read and maintain for all the data parties.
There are myths to be discussed here: “I write code as a data scientist” or “data engineers are too strict to read code”. We should start writing code for humans and ensure a minimum of cleanliness that guarantees our data teams’ productivity. Thus, important rules to take into account are listed below.
A) Use Meaningful Names
Using meaningful names across your pipeline is crucial. Let’s focus on Spark variables, functions, and columns.
Spark-based Variables: Use always unmutable variables and stop calling your data frames with names like “df1”. Instead, use meaningful names and suffixes that represent the API version. Examples:
val usersExampleDS = ...val usersExampleDF = ...val usersExampleRDD = ...
Do not use var, instead use the transform method as we are going to show later.
Functions: Along with meaningful names, for your functions create a whitelist of prefixes with verbs like clean, with, filter, transform/enrich, explode that are the only possible words used in your pipelines. Examples:
def cleanAcronymsFromTextdef enrichNameWithPrefix
Columns: the names of the columns should be clear and uniform. In case you use the dataset API prefer camelcase columns names to match a case class and underscore when you prefer a database notation. As an example, the following name shows clearly that the gender probability is greater than 50.
gender_probability_gt_50
In cases when you want to use a case class to define your dataset schema, define a columns mapper on the companion object and modify the column names before to set the schema. Example:
case class UserExample(name: String,                       email: String,                       country: String)object UserExample {  def columnstranslatorMap: Map[String, String] = Map(    "First Name" -> "name",    "Email" -> "email",    "Country Code" -> "country")}
Then you can read in this way:
val usersSourceExampleDS = spark.read  .option("header", "true")  .csv(PATH)  .select(UserExample.columnstranslatorMap.keys.toList.map(c => col(c).as(UserExample.columnstranslatorMap.getOrElse(c, c))): _*)  .as[UserExample]
In certain cases, the names of columns are used across multiple phases and functions. Hence, define a trait with the names of the columns and use it from the main spark job class.
trait GeneralColumnConstants {   val NameSourceColumn = "firstName"}object NamesGenderizerJob with GeneralColumnConstants {}
Finally, only in a few cases, columns should be mutable, rather add a new column to dataframes using the withColumn function to keep track of your original data.
B) Avoid Side Effects: Always prefer side-effect-free operations. As an example prefer always if/else as a one-line ternary operator.
if (true) statement1 else statement2
C) Decide wisely between null or Option: We got told that to write clean code in scala match-pattern and Optional are mandatory. However, in spark Optional could become a bottleneck since this evaluation adds an extra cost. So, the recommendation here is to use Optional at the beginning and only for those cases when the performance might be affected change values to null.
D) Avoid overusing scala implicit: Prefer always the transform function to use implicit to transform data frames leading to monkey patching.
E) Avoid UDFs whenever it is possible: UDFs are black box most of the time, then use them only when what you are doing is not possible with custom spark functions that are optimized by Spark. However, when you write UDFs please make sure you handle all the null cases.
F) Avoid overusing Accumulators: Accumulators are not meant to store your data, then use them only when some small counts are needed, for example counting the number of errors in a pipeline.
So far we can write a spark job with a proper structure and following some clean code principles. This is awesome, but it does not mean our code is already extensible and testable. In order, to reach this extensibility, the SOLID principles could play a key role. Let’s start going through these principles using our example.
A class should have one, and only one, reason to change.
All the pipeline modules, classes, or functions should have responsibility for a single part of the functionality, separating the responsibilities in simple code units. As we can see the phases defined in Step1, each phase is a responsibility unit. Besides, all the internal functions in these steps should be also responsible for a single thing, being testable and chainable in bigger functionalities. Let’s focus on a general data cleaner for the example dataset.
Above, we have two functions responsible for a single thing removeUselessChars and removeAbbreviations. Then, a chaining function called clean normalizes texts. Hence, the NamesNormalizer is only responsible for normalizing names texts. Let’s write tests in isolation for each of these functions.
Open/Closed Principle
You should be able to extend a class’s behavior, without modifying it.This principle is the foundation for building code that is maintainable and reusable.Robert C. Martin
Let’s see how to write a cleaner module for our example.
The key points to comply with this principle here are:
The abstraction CleanerService define the general cleaning functions template. This abstraction is closed for modification but opened for extensions.
Multiple implementations of cleaners, for instance: SourceACleaner, SourceBCleaner, SourceUsersCleaner extend the abstraction behavior and add concrete implementations.
Then, wherever we have to use a cleaner, Dependency Injection is preferred, we would inject a CleanerService instance instead of an instance of a lower-level class.
The cleanerServiceHelper chains multiple functions using the transform function.
Finally, see again the normalizeTexts functions, it separates the loop concern from the function that normalizes text, simply calling the previously defined TextNormalizer. Besides, using foldLeft makes functions easily reusable and maintains a DRY codebase.
Let’s see how testable it is :
Derived classes must be substitutable for their base classes.
To explain this principle easier using our example, some cleaning functions are shared by all the implementations. In this case, we can define the implementation in a general structure. For example, we can define a general behavior for deduplication in the cleanerServiceHelper and only use the concrete implementations when differs from others.
Clients should not be forced to implement interfaces they do not use.
Robert C. Martin
Do not write big monolithic traits or helpers where you had all the pipeline functions. Do you remember our first step where the phases were defined?. Those are useful here, we should be able to clean only, or for example clean and enrich without calling the rest of the phases. Thus, the recommendation is to create an uncoupled abstraction by phase. Let’s see an example of the enriching functions.
After this, the main class for a job should look like this.
Depend on abstractions, not on concretions.
In the previous example, we see how we inject the CleanerService abstraction. This is a good example of this principle in Spark. However, we could go a bit further and see how to apply dependency injection using our data frames. Let’s imagine we want to create a function that joins the usersSourceExampleDS with the genderized sources and test this in isolation. Let’s write a function for this.
then see how easy we could test.
When productionizing your spark jobs and making them maintainable, testable and extensible is a requirement; data teams should start focusing on creating pipelines that follow clean code principles and make testing tasks easier.
Some product development concepts apply at the beginning of a pipeline creation, then analyzing the phases and splitting the concerns is a key activity.
Clean code and SOLID principle suits spark, and by using them we can guarantee team productivity and code quality assurance.