Jupyter and Apache Zeppelin is a good place to experiment with data. Unfortunately, the specifics of notebooks do not encourage to organize the code, including its decomposition and readability. We can copy cells to Intellij IDEA and build JAR, but the effect will not be stunning. We can copy cells to Intellij IDEA and build JAR, but the effect will not be stunning. In this article you will learn how to make more readable Scala Apache Spark code in Intellij IDEA.
0. The base code
It is a simple application which:
- downloads groceries data from a file;
- filters fruits;
- normalizes names;
- calculates the quantity of each fruit.
val spark = SparkSession
.builder
.appName("MyAwesomeApp")
.master("local[*]")
.getOrCreate()
import spark.implicits._
val groceries = spark.read
.option("inferSchema", "true")
.option("header", "true")
.csv("some-data.csv")
val sumOfFruits = groceries
.filter($"type" === "fruit")
.withColumn("normalized_name", lower($"name"))
.groupBy("normalized_name")
.agg(
sum(($"quantity")).as("sum")
)
val fruits = groceries.filter($"type" === "fruit")
val normalizedFruits = fruits.withColumn("normalized_name", lower($"name"))
val sumOfFruits = normalizedFruits
.groupBy("normalized_name")
.agg(
sum(($"quantity")).as("sum")
)
sumOfFruits.show()1. Extract Methods
Let’s use the power of IDE, more precisely the Extract Method. It allows you to easily create a method from a selected piece of code. This way, let’s try to create methods corresponding to each step in the application.
It doesn’t work!?
def main(args: Array[String]) {
val spark = SparkSession
.builder
.appName("MyAwesomeApp")
.master("local[*]")
.getOrCreate()
import spark.implicits._
val groceries: DataFrame = getGroceries
val fruits: Dataset[Row] = filterFruits(groceries)
val normalizedFruits: DataFrame = withNormalizedName(fruits)
val sumOfFruits: DataFrame = sumByNormalizedName(normalizedFruits)
sumOfFruits.show()
}
private def sumByNormalizedName(normalizedFruits: DataFrame) = {
val sumOfFruits = normalizedFruits
.groupBy("normalized_name")
.agg(
sum(($"quantity")).as("sum")
)
sumOfFruits
}
private def withNormalizedName(fruits: Dataset[Row]) = {
val normalizedFruits = fruits.withColumn("normalized_name", lower($"name"))
normalizedFruits
}
private def filterFruits(groceries: DataFrame) = {
val fruits = groceries.filter($"type" === "fruit")
fruits
}
private def getGroceries: DataFrame = {
val groceries = spark.read
.option("inferSchema","true")
.option("header","true")
.csv("some-data.csv")
groceries
}The code in the main method is already more readable… but this code does not work. We want to use SparkSession and spark.implicits._ in the methods. Unfortunately these values are not within the scope of methods.
2. SparkSession overdose
We can fix this by passing on SparkSession in every method. Unfortunately, this is a pain in the ass. We also have to import spark.implicits._ every time. I’m to lazy for this solution 😁.
private def sumByNormalizedName(normalizedFruits: DataFrame, spark: SparkSession) = {
import spark.implicits._
val sumOfFruits = normalizedFruits
.groupBy("normalized_name")
.agg(
sum(($"quantity")).as("sum")
)
sumOfFruits
}
private def withNormalizedName(fruits: Dataset[Row], spark: SparkSession) = {
import spark.implicits._
val normalizedFruits = fruits.withColumn("normalized_name", lower($"name"))
normalizedFruits
}
private def filterFruits(groceries: DataFrame, spark: SparkSession) = {
import spark.implicits._
val fruits = groceries.filter($"type" === "fruit")
fruits
}
private def getGroceries(spark: SparkSession): DataFrame = {
val groceries = spark.read
.option("inferSchema","true")
.option("header","true")
.csv("some-data.csv")
groceries
}3. SparkSession at your service
We need to provide access to SparkSession in a slightly different way. The SparkJob object will help.
package pl.wiadrodanych.demo.base
import org.apache.spark.sql.SparkSession
trait SparkJob {
val spark: SparkSession = SparkSession
.builder
.appName("SomeApp")
.master("local[*]")
.getOrCreate()
}
object SparkJob extends SparkJob {}Now we can import SparkJob and spark.implicits._ in the application. The code looks better. We can reuse the methods.Now we can import SparkJob and spark.implicits._ in the application. The code looks better. We can reuse the methods.
import org.apache.spark.sql._
import org.apache.spark.sql.functions._
import pl.wiadrodanych.demo.base.SparkJob
import pl.wiadrodanych.demo.base.SparkJob.spark.implicits._
object NiceApp {
val spark = SparkJob.spark
def main(args: Array[String]) = {
val groceries: DataFrame = getGroceries
val fruits: Dataset[Row] = filterFruits(groceries)
val normalizedFruits: DataFrame = addNormalizedNameColumn(fruits)
val sumOfFruits: DataFrame = sumByNormalizedName(normalizedFruits)
sumOfFruits.show()
}
private def sumByNormalizedName(normalizedFruits: DataFrame) = {
val sumOfFruits = normalizedFruits
.groupBy("normalized_name")
.agg(
sum(($"quantity")).as("sum")
)
sumOfFruits
}
private def addNormalizedNameColumn(fruits: Dataset[Row]) = {
val normalizedFruits = fruits.withColumn("normalized_name", lower($"name"))
normalizedFruits
}
private def filterFruits(groceries: DataFrame) = {
val fruits = groceries.filter($"type" === "fruit")
fruits
}
private def getGroceries: DataFrame = {
val groceries = spark.read
.option("inferSchema", "true")
.option("header", "true")
.csv("some-data.csv")
groceries
}4. Implicit class / Extension method
I wrote a lot of C# code in my life. An interesting and useful concept is Extension Method. It allows you to “add” methods to an existing type/class without modifying it. Below is an example. Instead of writing
int numberA = 1int numberB = 2val sum = Sum(numberA, numberB)...public int Sum(int numberA, int numberB){ return numberA + numberB}We can write
int numberA = 1int numberB = 2val sum = numberA.Add(numberB)...public static int Add(this int numberA, int numberB){ return numberA + numberB}The difference in readability can be seen in the following example
Sum(A, Sum(B, Sum(C,Sum (D,...))))// VSA.Add(B).Add(C).Add(D)...In Scala we can get a similar mechanism using Implicit class. Below is the reorganized logic of the reviewed Apache Spark application.
package pl.wiadrodanych.demo.extensions
import org.apache.spark.sql._
import org.apache.spark.sql.functions._
import pl.wiadrodanych.demo.base.SparkJob.spark.implicits._
object GroceryDataFrameExtensions {
implicit class RichDataFrame(df: DataFrame) {
def sumByNormalizedName: DataFrame = {
val sumOfFruits = df
.groupBy("normalized_name")
.agg(
sum(($"quantity")).as("sum")
)
sumOfFruits
}
def addNormalizedNameColumn: DataFrame = {
val normalizedFruits = df.withColumn("normalized_name", lower($"name"))
normalizedFruits
}
def filterFruits: DataFrame = {
val fruits = df.filter($"type" === "fruit")
fruits
}
}
}Application logic has moved to another object and the code can be read like prose.
package pl.wiadrodanych.demo
import org.apache.spark.sql.DataFrame
import pl.wiadrodanych.demo.NiceApp.spark
import pl.wiadrodanych.demo.extensions.GroceryDataFrameExtensions._
object CoolApp {
def main(args: Array[String]) = {
val result = getGroceries
.filterFruits
.addNormalizedNameColumn
.sumByNormalizedName
result.show
}
private def getGroceries: DataFrame = {
val groceries = spark.read
.option("inferSchema", "true")
.option("header", "true")
.csv("some-data.csv")
groceries
}
}Let’s go back to what the application was supposed to do:
- download groceries data from a file
- filter fruits
- normalize names
- calculate the quantity of each fruit
Maybe not word for word, but you know what it is about 😁.
EDIT: Dataset transform
While the previous way is cool, it can sometimes be misleading. To separate the business code from the base class, we can use Dataset.transform. You will find details in this article from MungingData.
Repository
Please share what you think about this in comment secion. What is your way of making the code readable?