PySpark User Defined Functions (UDF)

Code description

User defined functions (UDF) in PySpark can be used to extend built-in function library to provide extra functionality, for example, creating a function to extract values from XML, etc. 

This code snippet shows you how to implement an UDF in PySpark. It shows two slightly different approaches - one use udf decorator and another without. 

Output:

+--------------+-------+--------+--------+
|       str_col|int_col|str_len1|str_len2|
+--------------+-------+--------+--------+
|Hello Kontext!|    100|      14|      14|
|Hello Context!|    100|      14|      14|
+--------------+-------+--------+--------+

Code snippet

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType
from pyspark.sql.functions import udf

app_name = "PySpark  Exmaple"
master = "local"

spark = SparkSession.builder \
    .appName(app_name) \
    .master(master) \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")

data = [['Hello Kontext!', 100], ['Hello Context!', 100]]

# Define the schema for the input data
schema = StructType([StructField('str_col', StringType(), nullable=True),
                     StructField('int_col', IntegerType(), nullable=True)])

# Create a DataFrame with the schema provided
df = spark.createDataFrame(data=data, schema=schema)


@udf(IntegerType())
def custom_udf1(str):
    return len(str)


def custom_func2(str):
    return len(str)


custom_udf2 = udf(custom_func2, returnType="int")

df = df.withColumn('str_len1', custom_udf1(df.str_col)).withColumn(
    'str_len2', custom_udf2(df.str_col))

df.show()