Skip to content

pandas User-Defined Functions

pandas User-Defined Functions (Vectorized User-Defined Functions or pandas UDFs) are user-defined functions that are executed using Apache Arrow to transfer data and pandas to work with the data, which allows for vectorized operations.

Pandas UDFs are defined using @pandas_udf decorator.

A Pandas UDF behaves as a regular PySpark function API in general.

As of Spark 3.0.0 (SPARK-28264), using Python type hints in pandas UDF is encouraged (instead of specifying pandas UDF type via functionType argument).

The return type (type hint) of a user-defined function should be as follows:

  • pandas.Series (pandas) in most cases
  • pandas.DataFrame (pandas) for struct input or output

@pandas_udf Decorator


pandas_udf function is used a decorator (using @pandas_udf annotation).

Python Decorators

Learn more in PEP 318 – Decorators for Functions and Methods.

pandas_udf belongs to pyspark.sql.functions module.

from pyspark.sql.functions import pandas_udf


functionType can be one of PandasUDFTypes (but is currently discouraged in favour of type hints).

@pandas_udf(returnType = "long", functionType = PandasUDFType.GROUPED_AGG)
def my_udaf(names: pd.Series) -> 'long':
  return pd.Series(names.count())

functionType is also known as evalType.

SQL_SCALAR_PANDAS_UDF is the default scalar UDF type.


@pandas_udf decorator can optionally specify a return type (as the first positional argument or using returnType).

A return type can be one of the names of pyspark.sql.types.DataType instances or the DataType themselves.


pandas UDAFs

pandas User-Defined Aggregate Functions.


import pandas as pd
from pyspark.sql.functions import pandas_udf
def to_upper(s: pd.Series) -> pd.Series:
    return s.str.upper()
def my_concat(names: pd.Series, ages: pd.Series) -> pd.Series:
    return pd.Series([f"{n} is {a} years old" for (n, a) in zip(names, ages)])
pandas_df = pd.DataFrame({
  'name': ['jacek', 'agata', 'iweta', 'patryk', 'maksym'],
  'age': [50, 49, 29, 26, 11]
df = spark.createDataFrame(pandas_df)
|  name|age|
| jacek| 50|
| agata| 49|
| iweta| 29|
|patryk| 26|
|maksym| 11|
>>> df.printSchema()
 |-- name: string (nullable = true)
 |-- age: long (nullable = true)
|     JACEK|
|     AGATA|
|     IWETA|
|    PATRYK|
|    MAKSYM|
+----------+, df.age)).show(truncate = False)
|my_concat(name, age)  |
|jacek is 50 years old |
|agata is 49 years old |
|iweta is 29 years old |
|patryk is 26 years old|
|maksym is 11 years old|