Be in charge of Query Execution in Spark SQL | by David Vrba | Towards Data Science

Querying data in Spark has become a luxury since Spark 2.x because of SQL and declarative DataFrame API. Using just few lines of high level code allows to express quite complex logic and carry out complicated transformations. The big benefit of the API is that users don’t need to think about the execution and can let the optimizer figure out the most efficient way to execute the query. And efficient query execution is often a requirement not only because the resources may become costly, but also it makes the work of the end user more comfortable by reducing the time he/she has to wait for the result of the computation.
The Spark SQL optimizer is indeed quite mature, especially now with the upcoming version 3.0 which will introduce some new internal optimizations such as dynamic partition pruning and adaptive query execution. The optimizer is internally working with a query plan and is usually able to simplify it and optimize by various rules. For example it can change the order of some transformations or drop them completely if they are not necessary for the final output. Despite all the clever optimizations there are however still situations in which a human brain can do better. In this article we will take a look at one of these cases and see how, using a simple trick, we can lead Spark towards a more efficient execution.
The code is tested in the current version of Spark which is 2.4.5 (written in June 2020) and it is checked against Spark 3.0.0-preview2 to see possible changes in the upcoming Spark 3.0.
Let me now first introduce a simple example for which we will try to achieve efficient execution. Imagine that we have data in json format with the following structure:
{"id": 1, "user_id": 100, "price": 50}{"id": 2, "user_id": 100, "price": 200}{"id": 3, "user_id": 101, "price": 120}{"id": 4, "price": 120}
Each record is like a transaction so the user_id column may contain lots of duplicated values (possibly including nulls) and besides these three columns there can be many other fields describing the transaction. Now our query will be based on a union of two similar aggregations where each of these aggregations differs by some conditions. In the first aggregation we want to take users for which the sum of the price is less than 50 and in the second aggregation we take users for which the sum of the price is more than 100. Moreover in the second aggregation we want to consider only records where user_id is not null. This model example is just a simplified version of a more complex situation that can occur in practice and for the sake of simplicity we will use it throughout the article. Here is a basic way how to express such a query using DataFrame API of PySpark (very similarly we could write it also using the Scala API):
df = spark.read.json(data_path)df_small = (  df  .groupBy("user_id")  .agg(sum("price").alias("price"))  .filter(col("price") < 50))df_big = (  df  .filter(col("user_id").isNotNull())  .groupBy("user_id")  .agg(sum("price").alias("price"))  .filter(col("price") > 100)  )result = df_small.union(df_big)
The key to achieve a good performance for your query is the ability to understand and interpret the query plan. The plan itself can be displayed by calling explain function on the Spark DataFrame or if the query is already running (or has finished) we can also go to the Spark UI and find the plan in the SQL tab.
The SQL tab has lists of completed and running queries on the cluster so by selecting our query we will see the graphical representation of the physical plan (here i removed the metrics information to make the plot smaller):
The plan has a tree structure where each node represents some operator that caries some information about the execution. We can see that in our example there are two branches with the root at the bottom and the leafs at the top where the execution starts. The leafs Scan json represent reading the data from the source, then there is a pair of HashAggregate operators which are responsible for the aggregation, and in between them there is Exchange which represents the shuffle. The Filter operators carry the information about the filtering conditions.
The plan has a typical shape for union operations, there is a new branch for each DataFrame in the union, and since in our example both DataFrames are based on the same datasource, it means that the datasource will be scanned twice. Now we can see that there is a space for improvement. Having the datasource scanned only once can lead to a nice optimization, especially in situations where I/O is expensive.
Conceptually, what we want to achieve here, is reusing some computation — scanning the data and computing the aggregation, because these are the operations that are the same in both DataFrames and in principle it should be sufficient to compute them only once.
One typical approach how to reuse a computation in Spark is using caching. There is a function cache that can be called on a DataFrame:
df.cache()
It is a lazy transformation which means that the data will be put to the caching layer after we call some action. Caching is a very common technique used in Spark however it has its limitations, especially if the cached data is large and the resources on the cluster are limited. Also one needs to be aware that storing the data in the caching layer (memory or disk) will bring some additional overhead and the operation itself is not for free. Calling cache on the whole DataFrame df is not optimal also from the reason that it will try to put all the columns to memory which may be unnecessary. More careful way is to select a superset of all columns that will be used in the following queries and then call the cache function after this select.
Apart from caching, there is another technique which is not that well described in the literature and this technique is based on reusing the Exchange. The Exchange operator represents shuffle which is a physical data movement on the cluster. This happens when the data must be reorganized (repartitioned) which is usually required for aggregations, joins and some other transformations. The important thing about shuffle is that when the data is repartitioned, Spark will always save it on the disk as the shuffle write (this is an internal behavior and it is not under control of the end user). And because it is saved on the disk, it can be reused later on if it is required. Spark will indeed reuse the data if it finds an opportunity for that. This happens each time Spark detects that the same branch from the leaf node up to an Exchange is repeating somewhere in the plan. If there is such a situation it means that these repeated branches represent identical computation and thus it is sufficient to compute it only once and then reuse it. We can recognize from the plan whether Spark found such a case, because those branches would be merged together like this:
In our example, Spark didn’t reuse the Exchange, but with a simple trick, we can push him to do so. The reason why the Exchange is not reused in our query is the Filter in the right branch that corresponds to the filtering condition user_id is not null. The filter is indeed the only difference in our two DataFrames that are in the union, so if we can eliminate this difference and make the two branches the same, Spark will take care of the rest and will reuse the Exchange.
How can we make the branches the same? Well, if the only difference is the filter, we can certainly switch the order of transformations and call the filter after the aggregation because that will make no impact on the correctness of the result that will be produced. However there is a catch! If we move the filter like this:
df_big = (  df.groupBy("user_id")  .agg(sum("price").alias("price"))  .filter(col("price") > 100)  .filter(col("price").isNotNull()))
and check the final query plan, we will see that the plan has not changed at all! The explanation is simple — the filter was moved back by the optimizer.
Conceptually it is good to understand that there are two major types of the query plan — logical plan and physical plan. And the logical plan undergoes an optimization phase before it is turned into the physical plan which is the final plan that will be executed. When we change some transformations, it is reflected in the logical plan, but then we loose control over the next steps. The optimizer will apply a set of optimization rules which are mostly based on some heuristics. The rule related to our example is called PushDownPredicate and this rule makes sure that the filters are applied as soon as possible and are pushed closer to the source. It is based on the idea that it is more efficient to first filter the data and then do the computation on the reduced dataset. This rule is indeed very useful in most of the situations, however in this very case it is fighting against us.
To achieve custom position of the Filter in the plan, we have to limit the optimizer. This is possible since Spark 2.4 because there is a configuration setting which allows us to list all the optimization rules that we want to exclude from the optimizer:
spark.conf.set("spark.sql.optimizer.excludedRules",     "org.apache.spark.sql.catalyst.optimizer.PushDownPredicate")
After setting this configuration and running the query again, we will see that now the filter stays positioned as we need. The two branches become really the same and Spark will now reuse the Exchange! The dataset will now be scanned only once and the same goes for computing the aggregation.
In Spark 3.0 the situation is changed a little bit, the optimization rule now has a different name — PushDownPredicates and there is an additional rule that is also responsible for pushing a filter PushPredicateThroughNonJoin, so we actually need to exclude both of them to achieve the desired goal.
We can see that through this technique Spark developers gave us the power to control the optimizer. But with power comes also responsibility. Let’s list a couple of points that is good to keep in mind when using this technique:
When we stop the PushDownPredicate, we will take the responsibility for all the filters in the query, not just the one that we want to reposition. There might be other filters which are important to take place as soon as possible, for example partition filters, so we need to make sure they are positioned correctly.
Limiting the optimizer and taking care of the filters is some extra work on the side of the user, so it better be worth it. In our model example the potential for speeding up the query will be in cases where I/O is expensive because we will achieve that the data will be scanned only once. This might be the case for instance with file formats which are not columnar, like json or csv, if the dataset has lots of columns.
Also if the dataset is small, it might not be worth to go the extra mile to control the optimizer because simple caching will do the job. However when the dataset is large, the overhead with storing the data in the caching layer will become apparent. On the other hand the reused Exchange would bring no additional overhead because the computed shuffle will be stored on disk anyway.
This technique is based on Spark’s internal behavior which has no official documentation and if something changes in this functionality it might be more difficult to find it out. In our example we could see that there is actually a change in Spark 3.0 where one rule is renamed and another rule is added.
We have seen that being able to achieve the optimal performance may require understanding the query plan. Spark optimizer does a very good job by optimizing our query using a set of heuristic rules. There are, however, situations in which these rules miss the most optimal configuration. Sometimes rewriting the query is good enough, but sometimes it is not, because by rewriting the query we will achieve a different logical plan but we do not have a direct control over the physical plan that will be executed. Since Spark 2.4 we can use a configuration setting excludedRules that allows us to limit the optimizer and thus navigate Spark to a more custom physical plan.
In many cases relying on the optimizer will lead to a solid plan with a quite efficient execution, however, there are cases mostly in performance critical workloads, where it might be worth to check the final plan and see whether we can improve it by taking the control over the optimizer.