Creating Custom UDFs, UDAFs, and UDTFs in PySpark

Notes and Constraints

• You have prepared the Python environment by referring to Using Conda to Package the Python Environment and Third-Party Libraries in the PySpark Program.
• In PySpark, UDFs, UDAFs, and UDTFs created using Python can be registered only as temporary UDFs, UDAFs, and UDTFs and are valid only in the current SparkSession.

Custom UDF Usage Guide

pyspark.sql.functions.udf(f: Union[Callable[[…], Any], DataTypeOrString, None] = None, returnType: DataTypeOrString = StringType()) → Union[UserDefinedFunctionLike, Callable[[Callable[[…], Any]], UserDefinedFunctionLike]]

• The user-defined functions are considered deterministic by default. Due to optimization, duplicate invocations may be eliminated or the function may even be invoked more times than it is present in the query. If your function is not deterministic, call asNondeterministic on the user-defined function. For example:

  from pyspark.sql.types import IntegerType import random 
  random_udf = udf(lambda: int(random.random() * 100), IntegerType()).asNondeterministic()

• The user-defined functions do not support conditional expressions or short circuiting in boolean expressions and end up with being executed internally. If the functions fail on special rows, the workaround is to incorporate the condition into the functions.
• The user-defined functions do not take keyword arguments on the calling side.

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType
sparkSession = SparkSession.builder.getOrCreate()
slen = udf(lambda s: len(s), IntegerType())
@udf
def to_upper(s):
    if s is not None:
        return s.upper()
@udf(returnType=IntegerType())
def add_one(x):
    if x is not None:
        return x + 1

df = sparkSession.createDataFrame([(1, "John Doe", 21)], ("id", "name", "age"))
df.select(slen("name").alias("slen(name)"), to_upper("name"), add_one("age")).show()
sparkSession.stop()

+----------+--------------+------------+                                        
|slen(name)|to_upper(name)|add_one(age)|
+----------+--------------+------------+
|         8|      JOHN DOE|          22|
+----------+--------------+------------+

Custom UDAF Usage Guide

In the current version, UDAFs cannot be directly created in PySpark. You can use pyspark.sql.functions.pandas_udf to create UDAFs. Details of the pyspark.sql.functions.pandas_udf function are as follows:

pyspark.sql.functions.pandas_udf(f=None, returnType=None, functionType=None)

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import DoubleType
import pandas as pd

sparkSession = SparkSession.builder.appName("PandasUDAFExample").getOrCreate()

@pandas_udf(DoubleType())
def pandas_mean_udaf(v: pd.Series) -> float:
    return v.mean()

data = [
    (1, 2.0),
    (1, 4.0),
    (2, 5.0),
    (2, 6.0)
]
df = sparkSession.createDataFrame(data, ["id", "value"])

result_df = df.groupBy("id").agg(pandas_mean_udaf(df["value"]).alias("mean_value"))
result_df.show()
sparkSession.stop()

+---+----------+                                                                
| id|mean_value|
+---+----------+
|  1|       3.0|
|  2|       5.5|
+---+----------+

Custom UDTF Usage Guide

For example, perform the following operation to convert one row of data into multiple rows of data:

1. Use UDF to return an array and then use explode to split the array.

   Use Python to write a UDF that converts the input data (or a field in a row) into a list (or array). Then, use the explode function built in DataFrame to split the array into multiple rows, which achieves the UDTF effect.

   Example:

   from pyspark.sql import SparkSession
   from pyspark.sql.functions import udf, explode
   from pyspark.sql.types import ArrayType, StringType
   
   sparkSession = SparkSession.builder.appName("SimulateUDTF").getOrCreate()
   
   data = [("a,b,c",), ("d,e",)]
   df = spark.createDataFrame(data, ["csv_col"])
   
   def split_string(s):
       if s:
           return s.split(",")
       return []
   
   split_udf = udf(split_string, ArrayType(StringType()))
   
   df_with_array = df.withColumn("split_col", split_udf("csv_col"))
   df_exploded = df_with_array.withColumn("word", explode("split_col"))
   
   df_exploded.show()
   sparkSession.stop()

   Execution Result:

   +-------+---------+----+
   |csv_col|split_col|word|
   +-------+---------+----+
   |  a,b,c|[a, b, c]|   a|
   |  a,b,c|[a, b, c]|   b|
   |  a,b,c|[a, b, c]|   c|
   |    d,e|   [d, e]|   d|
   |    d,e|   [d, e]|   e|
   +-------+---------+----+

2. Use flatMap of RDD to implement conversion.

   For more flexible scenarios, you can directly use the RDD API of DataFrame and perform flatMap operations to implement conversion.

   Example (splitting an input row into multiple rows):

   from pyspark.sql import SparkSession
   
   sparkSession = SparkSession.builder.appName("SimulateUDTFWithFlatMap").getOrCreate()
   
   data = [("a,b,c",), ("d,e",)]
   df = spark.createDataFrame(data, ["csv_col"])
   
   rdd = df.rdd.flatMap(lambda row: [(row.csv_col, token) for token in row.csv_col.split(",")])
   new_df = rdd.toDF(["csv_col", "word"])
   new_df.show()
   
   sparkSession.stop()

   Execution Result:

   +-------+----+
   |csv_col|word|
   +-------+----+
   |  a,b,c|   a|
   |  a,b,c|   b|
   |  a,b,c|   c|
   |    d,e|   d|
   |    d,e|   e|
   +-------+----+