Performance in Apache Spark: Benchmark 9 Different Techniques
By: David Vrba
Tags: Benchmark, Spark Sql, Apache Spark, Data Science, Data Engineering

In Apache Spark, it is quite common that the same transformation can be achieved in different ways. This is also a consequence of the continuous development of Spark because new techniques and functions are available in new releases.

In this article, we will compare 9 different techniques that lead to the same result, but each of them has a very different performance. We will show it on one particular example related to arrays processing, but the results can be generalized beyond.

Benchmarks in Spark

When measuring performance for different transformations in Spark one needs to be careful to avoid some pitfalls. First of all, we need to materialize the transformations by calling an action to run a Spark job (to understand the difference between transformations and actions in Spark, see my recent article). One action that can come to your mind is the function count(). However with the count, there is this problem, that to evaluate how many rows are represented by the resulting DataFrame, Spark doesn’t always need to execute all transformations. The Spark optimizer can simplify the query plan in such a way that the actual transformation that you need to measure will be skipped because it is simply not needed for finding out the final count. So to make sure that really all transformations are executed, it is better to call write and save the result somewhere. But saving the output has another drawback because the execution time will be affected by the writing process.

In Spark 3.0 the situation for benchmarking simplified and doing performance benchmarks became much more convenient thanks to the noop write format, which is a new feature in Spark 3.0. We can simply specify it as the write format and it will materialize the query and execute all the transformations but it will not write the result anywhere.

(

final_df

.write

.mode("overwrite")

.format("noop")

.save()

)

The use case for which we are going to benchmark the queries is related to the following example. Imagine a dataset where we have an array of words on each row. We want to add a new column that is also of array type and the values of this new array correspond to the length of the words in the original array as you can see in the image below:

Image by author

Technique I — Higher-Order Functions (HOFs)

Conceptually thinking, the transformation that we want to do here is that from one array we want to create another array — and we want to do it for each row in the dataset. In functional programming languages, there is usually a map function that is called on the array (or another collection) and it takes another function as an argument, this function is then applied on each element of the array as you can see in the image below

Image by author

Since Spark 2.4, this concept is also supported in Spark SQL and this map function is called transform (note that besides transform there are also other HOFs available in Spark, such as filter, exists, and other). The support was first only in the SQL API, so if you want to use it with the DataFrames DSL (in 2.4) you have to call it inside expr (or selectExpr). Since Spark 3.0, however, the HOFs are supported also in the Scala API and since Spark 3.1.1 the support was also added in Python API. Using HOFs, the transformation can be implemented as follows:

final_df = (

df

.withColumn("tag_lengths",expr("TRANSFORM(tags, x -> length(x))"))

)

In Scala API — since Spark 3.0:

val final_df = df

.withColumn("tag_lengths", transform($"tags", x => length(x)))

In Python API — since Spark 3.1.1

final_df = (

df

.withColumn("tag_lengths", transform("tags", lambda x: length(x)))

)

Technique II — Python UDF (vanilla)

Since there is no direct native function that would carry out the transformation that we want to do, one might be tempted to implement a user-defined function (UDF) for it, and indeed before the support for HOFs in Spark 2.4, UDF was a very common technique to solve problems with arrays in Spark. In PySpark there are two major types of UDFs, the first one is an ordinary UDF — we call it here a vanilla UDF, the second type is a Pandas UDF and we will measure their performance separately. The transformation with the vanilla UDF can be written as follows:

@udf("array<int>")

def pythonUDF(tags):

return [len(tag) for tag in tags] final_df = df.withColumn("tag_lengths", pythonUDF("tags"))

Technique III — Python UDF (Pandas)

Because Spark is not able to translate the Python code from the UDF to JVM instructions, the Python UDF has to be executed on the Python worker unlike the rest of the Spark job which is executed in JVM. In order to do that, Spark has to transfer the data from JVM to the Python worker. With the vanilla Python UDF, this transfer is realized by converting the data to pickled bytes and sending it to Python which is not very efficient and has quite a big memory footprint. This was improved in Spark 2.3 by integrating Spark with Apache Arrow and supporting so-called Pandas UDFs. The idea is that if the UDF is implemented using Pandas API and the function is registered as Pandas UDF, the data will be transferred using the Arrow format, which is an in-memory data format that was specifically designed to exchange the data between different systems efficiently. Moreover, the Pandas UDF can leverage vectorized execution on the Python side which will be very efficient especially if the UDF will do some numerical computations and will use vectorized operations that are supported by NumPy or SciPy. Using the Pandas UDF our array transformation can be implemented as follows:

@pandas_udf("array<int>")

def pandasUDF(tags):

return tags.map(lambda x: [len(tag) for tag in x]) final_df = df.withColumn("tag_lengths", pandasUDF("tags"))

Technique IV — Scala UDF

Using this technique, we will test how the performance changes if we use Scala instead of Python. If you use only native transformations supported by the DataFrame API and don’t use any UDFs, the performance will be the same, no matter if you use Python or Scala because the native transformations create a query plan under the hood and the plan will be the same regardless of the used language (see my other article to learn more about query plans). The execution of the plan and execution of the Spark job will then happen inside JVM. The situation is however different if you use a UDF because in this case the transformations inside the UDF can not be converted to the query plan and as we described above, in the case of Python, the execution will take place on the Python worker which is related to the overhead with transferring the data. If the UDF is implemented in Scala, we can avoid the data transfer to Python and the execution will stay in JVM. There is however still some overhead related to converting the data from an internal data format to java objects. For the native transformations, Spark is using an internal data format and all the native transformations generate a very efficient Java code at the runtime. With the UDF Spark doesn’t know how to generate the code and has to convert the data to Java objects, then it executes your UDF on it, and afterward, it converts the data back to the internal format. The code with the Scala UDF can be written as follows:

def tag_lengths(tags: Seq[String]): Seq[Int] = {

tags.map(tag => tag.length)

} val scalaUDF = udf(tag_lengths(_:Seq[String]): Seq[Int]) val final_df = df.withColumn("tag_lengths", scalaUDF($"tags"))

Technique V — Scala UDF called from Python

Above we described the difference between using Scala and Python as the language for Spark — if we use only native transformations, there is no impact on performance, however, if we use UDFs, there is going to be a performance hit. For Python developers, this may become unfortunate if the performance is critical and yet the situation requires using UDF. Interestingly, there is a mixed approach, in which we can write the Spark application in Python, and in Scala, we implement only the UDF to keep the execution inside JVM.

// in Scala:

// this is inside src/main/scala/vrbad/sandbox/CustomFunctions.scala package vrbad.sandbox import org.apache.spark.sql.SparkSession object CustomFunctions { def sizeElements(elements: Seq[String]): Seq[Int] = {

elements.map(el => el.length)

} def registerUDF(_spark: SparkSession) = {

val f = sizeElements(_)

_spark.udf.register("sizeElements", f)

}

}

# in Python: java_instance = spark._jvm.vrbad.sandbox.CustomFunctions

java_instance.registerUDF(spark._jsparkSession) final_df = (

df.withColumn("tag_lengths", expr("sizeElements(tags)"))

)

The Scala file is compiled using sbt into a jar file which is uploaded to the cluster and it can be used inside a PySpark application, as you can see in the code above.

Technique VI — exploding the array

Exploding the arrays allows for accessing each element of the array separately. The explode function will simply generate a new row for each element in the array, so the DataFrame becomes bigger in terms of how many rows it represents. We can compute the length of each element and after that, we can group these results into arrays and thus shrink the DataFrame back to its original size:

from pyspark.sql.functions import explode, length, collect_list final_df = (

df.withColumn("tag", explode("tags"))

.withColumn("tag_size", length("tag"))

.groupBy("id")

.agg(

collect_list("tag_size").alias("tag_lengths")

)

)

It is essential to understand that this technique has five important caveats, where one of them is directly related to performance, two are related to the correctness of the result, one to the shape of the final output, and one to the null handling:

groupBy is a transformation that requires specific partitioning, namely the data has to be repartitioned by the grouping key. And if this requirement is not met, Spark will have to shuffle the data which will have a direct impact on performance because shuffle is in general expensive. The collect_list function is not deterministic, so the elements in the final arrays can end up in different order as are the elements in the original array. The approach assumes that the grouping key contains unique values before calling the explode, because if there are duplicities, they will be removed after the groupBy, so in such a case we would end up with fewer rows than we started with — the duplicities would be grouped together. The output DataFrame after the groupBy will contain only the grouping key and the result of aggregation and all other original fields are lost, so if you need them in the final output, you need to join the result with the original DataFrame, which will of course bring additional slow down. The explode function will not generate any row if the array is null, so here again, after the grouping, we can end up with fewer rows than we started with. This can be controlled with the function explode_outer, which will generate a row (with null value) even if the array is null.

Technique VII — exploding the array with bucketed source

The shuffle mentioned in the previous point can be avoided if the data has the correct partitioning (see my other article where I describe this in detail). This can be achieved if the datasource is bucketed by the grouping key. Here we will test how it speeds up the performance, the query is the same as in Technique VI, but the source will be bucketed by the grouping key — id.

Technique VIII — RDD API in Python

The RDD API is a low-level API that was used especially before Spark 2.0. Since Spark 2.x, it is recommended to avoid it in favor of DataFrames or SQL. The RDDs are still used in Spark internally, the query plan is always compiled into RDD DAG (directed acyclic graph) during execution. The point is that Spark does it in a very efficient way, first it optimizes the query plan, and then it also generates highly efficient Java code that operates on internal data representation — a so-called Tungsten format. However, if you write the RDD code by yourself, no optimization happens under the hood and the data will be represented as Java objects with much bigger memory consumption as compared to Tungsten.

For the sake of completeness, we will test here also the RDD approach and in Python, the code may look like this

from pyspark.sql.types import * r = (

df.rdd

.map(lambda row: (row[0], [len(tag) for tag in row[1]]))

) data_schema = StructType(

[

StructField("id", StringType()),

StructField("tag_lengths", ArrayType(LongType()))

]

) final_df = spark.createDataFrame(r, data_schema)

Technique IX — RDD API in Scala

The last technique tested in this benchmark is the same as the previous one — the RDD approach, but now in Scala:

val final_df = df.rdd

.map(row =>

(row(0).asInstanceOf[String],

row(1).asInstanceOf[Seq[String]].map(tag => tag.length))

).toDF("id", "tag_lengths")

Benchmark details

All queries are executed using Spark 3.1.1 (the most recent release at the time of writing — March 2021) on the Databricks platform with the runtime 8.0. The used cluster has 3 workers m5d.2xlarge (24 cores altogether) and the input dataset is stored on S3 (AWS distributed object store) in Apache Parquet file format and it has 1 047 385 835 rows (slightly over one billion rows).

Benchmark results

In the plot below you can see the execution time for all 9 techniques:

Image by author

The fastest is technique I — the native approach with higher-order functions which took 66 seconds. Slightly slower were techniques V and IV — the Scala UDF called from Python application (89 seconds) and Scala UDF called from Scala application (94s). The next was the query written using RDD API in Scala, surprisingly it took only 104 seconds. Much slower was the approach using explode with groupBy and if we use bucketed source and avoid the shuffle it runs 244 seconds, on the other hand, if we don’t use bucketing the execution is slower and takes 302 seconds. Very comparable is the Pandas UDF — 299 seconds. The vanilla Python UDF took 386 seconds and finally the slowest was the RDD API used from Python (1020 seconds). Let’s see some highlights:

The native approach with HOFs is the most efficient — it is not surprising, it can leverage all the internal features such as Spark optimizer, code generation, or internal Tungsten data format. The Scala UDF is much faster than Python UDF — we expected that because Scala UDF is executed inside JVM and avoids the data transfer between JVM and Python workers. Scala UDF called from Python was slightly faster than Scala UDF called from Scala. Here I assumed these two techniques to be equivalent in terms of performance and I don’t really see any reason why it should be faster when called from PySpark application, however, the difference is pretty small — only 5 seconds, so it might be also something unrelated happening on the cluster. The difference between the techniques VI and VII (explode + groupBy without bucketing and with bucketing) was not very huge. It means that the shuffle was not a big bottleneck in the first case. This is because the shuffled data was not very big, and from the SparkUI, I could see that the Shuffle Write was only 26 GB. The dataset is big in terms of rows (over 1 billion) but it has only two columns (id and tags) that are being shuffled, and on a cluster with 24 cores, such shuffle doesn’t bring here a big overhead. The Pandas UDF was faster than the vanilla UDF but it was not as fast as one might expect. There are actually cases where the performance of Pandas UDF can become close to the performance of Scala UDF. But those cases are usually situations in which your UDF is doing some numerical computation and is written in terms of SciPy or NumPy and so Pandas can leverage vectorized execution with BLAS or LAPACK libraries under the hood. This is not our case — we are not doing any numerical computation but we are still getting faster execution as compared to vanilla UDF because Pandas UDF is using Apache Arrow format for the data exchange. The RDD API in Scala was in our case quite fast. Well, in the case of the DataFrame API, Spark is using an optimizer under the hood, which will not be used for RDDs. But here we are doing really only one transformation on the array and so the optimizer will not bring much benefit. The RDD API will not leverage generated code and Tungsten data format and its performance becomes comparable with the Scala UDF. The RDD API in Python is really very slow and we can see that all other techniques are much faster.

Conclusion

In this article, we tested the performance of 9 techniques for a particular use case in Apache Spark — processing arrays. We have seen that best performance was achieved with higher-order functions which are supported since Spark 2.4 in SQL, since 3.0 in Scala API and since 3.1.1 in Python API. We also compared different approaches for user-defined functions, such as vanilla Python UDF, Pandas UDF, Scala UDF, and a mixed approach in which we call Scala UDF from a PySpark application. We have also seen how it compares with exploding the arrays and pointed out some caveats related to this approach. Finally, we also included the legacy RDD API in our tests.