Concept: User-defined functions

User Defined Functions let you use your own arbitrary Python in PySpark. For example, you could use a UDF to parse information from a complicated text format in each row of your dataset.

After declaration, a UDF works similarly to built in PySpark functions such as concat, date_diff, trim, etc.

Motivation

Unintuitively, under normal circumstances data is never actually brought into your Python code. When you manipulate DataFrames using PySpark, you are describing the steps that the Spark cluster should take in a distributed, parallel fashion to get your final DataFrame. This allows Spark and Foundry to scale almost ad infinitum, but introduces the minor setup of UDFs for injecting code to run within the cluster on actual data. PySpark sends your UDF code to each server running your query.

Consider alternatives to UDFs

Avoid UDFs when PySpark's built-in functions can express the same logic because UDFs are significantly slower. If you must use a UDF, see the PySpark style guide for best practices.

Example

Copied!
1
"Weather report: rain 55-62"

Suppose we want to get the low temperature from the following weather format, in this case 55. We can write the following ordinary Python function,

Copied!
1
2
def extract_low_temperature(weather_report):
    return int(weather_report.split(' ')[-1].split('-')[0])

Create a UDF around the extract_low_temperature function to integrate it into your PySpark query. Creating a UDF involves providing the function and its expected return type in PySpark's type system.

Copied!
1
2
3
4
5
# Import the necessary type
from pyspark.sql.types import IntegerType

# Wrap our function as a UDF
low_temp_udf = F.udf(extract_low_temperature, IntegerType())

Now the UDF can be used on a DataFrame, taking a whole column as an argument.

Copied!
1
df = df.withColumn('low', low_temp_udf(F.col('weather_report')))

Reading from Multiple Columns

A UDF can take arbitrary column arguments. The column arguments correspond to the function arguments.

Copied!
1
2
3
4
5
6
7
8
9
10
11
from pyspark.sql.types import StringType

def weather_quality(temperature, windy):
    if temperature > 70 and windy == False:
        return "good"
    else:
        return "bad"

weather_udf = F.udf(weather_quality, StringType())

df = df.withColumn('quality', weather_udf(F.col('temp'), F.col('wind')))

Debugging UDFs

Spark does not include logs written inside UDFs in the top-level driver logs. To inspect values inside a UDF, create a temporary UDF that returns debug information as a column. For more techniques, see Logging.