Enhancing Optimized PySpark Queries
By: Ed Turner
Tags: Optimization, Programming, Data Science, Apache Spark

As we continue increasing the volume of data we are processing and storing, and as the velocity of technological advances transforms from linear to logarithmic and from logarithmic to horizontally asymptotic, innovative approaches to improving the run-time of our software and analysis are necessary.

These necessitating innovative approaches include utilizing two very popular frameworks: Apache Spark and Apache Arrow. These two frameworks enable users to process large volumes of data in a distributive fashion. These two frameworks, also, enables users to process larger volumes of data more quickly by using vectorized approaches. These two frameworks can easily facilitate big-data analysis. However, despite these two frameworks and their ability to empower users, there is still room for improvement, specifically within the python-ecosystem. Why can we confidently identify pockets of improvement in utilizing these frameworks within python? Let’s examine some features python has.

The Features of Python

As a programming language, python enables pure flexibility. A developer does not need to specify the type of a variable before instantiating it or defining it. A developer does not need to specify a return type of a function. Python interprets the type of each object at runtime, which allows for these restrictions (or guardrails) to be removed. These features python offers lower development times with perceived productivity increases. However, it is known, these same features negatively impact program runtime. Since python interprets the type of each object at runtime, it takes a long time for some software to run, especially those that require excessive looping. Even within programs that utilize vectorized operations, given the complexity of some programmatic logic, there still can be some impact on runtime performance. Is there anything we can do to mitigate these performance-related impacts? Let us turn our attention to my solution of choice: Numba

Numba To The Rescue

Numba performs Just-In-Time compilations on python functions, very similar to how C/C++ and Java compilations are performed. Python functions that only contain standard builtin functions, or a set of NumPy functions, can be improved using Numba. Here is an example:

from time import time

from numba import jit

import numpy as np @jit(nopython=True, fastmath=True)

def numba_sum(x):

return np.sum(x) # this returns the median time of execution

def profileFunct(funct, arraySize, nTimes): _times = []

for _ in range(nTimes):

start = time()

funct(np.random.random((arraySize,)))

end = time()

_times.append(end - start) return np.median(_times) # this is the numba time

numba_times = [profileFunct( numba_sum, i, 1000 ) for i in range(100, 1001, 100)] # this is the standard numpy timing

numpy_times = [profileFunct( np.sum, i, 1000 ) for i in range(100, 1001, 100)] speed_up_lst = list(map(lambda x: x[1] / x[0], zip(numba_times, numpy_times)))

In the above example, we are only calculating the sum of a numpy.array, and then comparing the performances to each other. In this simple example, we see a moderate performance boost. Depending on your local machine, you can see a performance boost of 10% to 150% percent. Generally, you will see similar performance boosts, or even more if you are iterating in your function.

If we can speed up a simple example of addition, let’s see an example where we are using Numba + Apache Arrow + Apache Spark.

The Three Musketeers

The following is just a brief example, demonstrating the possibility of creating Just-In-Time compiled functions, and then using them with a pandas_udf.

from numba import jit

import numpy as np

import pandas as pd import pyspark.sql.functions as F

from pyspark.sql import SparkSession

from pyspark.sql.types import DoubleType spark = SparkSession.builder.appName("test").getOrCreate() spark.conf.set("spark.sql.execution.arrow.enabled", "true") df = pd.DataFrame(data=np.random.random((100,)), columns=["c1"]) sdf = spark.createDataFrame(df) # JIT compiled function

@jit(nopython=True, fastmath=True)

def numba_add_one(x):

return x + np.ones(x.shape) # this is needed to use apache arrow

@F.pandas_udf(DoubleType())

def add_one(x):

return pd.Series(numba_add_one(x.values)) sdf = sdf.withColumn("c1_add_one", add_one(F.col("c1"))) sdf.toPandas()

As stated earlier, this is a simple example. However, for more complicated applications, this is extremely valuable and will make it speed up, even the most fundamental, Apache Spark SQL queries.