UDFs or native functions

Spark runs on the JVM. A Python UDF (*user-defined function*) forces every row to leave the JVM, be serialized to a Python process, come back and be deserialized. On top of that, the optimizer cannot see what the UDF does: it does not push a filter through it nor combine it with other steps.

The same rule, both ways

from pyspark.sql import functions as F
from pyspark.sql.types import StringType

@F.udf(returnType=StringType())
def band_udf(amount):
    if amount is None:
        return "no amount"
    return "high" if amount >= 500 else "low"

with_udf = orders.withColumn("band", band_udf("amount"))

native = orders.withColumn(
    "band",
    F.when(F.col("amount").isNull(), "no amount")
     .when(F.col("amount") >= 500, "high")
     .otherwise("low"),
)

Both give the same result; the native one becomes code generated on the JVM. Almost everything has a ready-made function in pyspark.sql.functions: text (regexp_extract, split), dates (to_date, date_trunc), JSON (from_json), arrays (transform, filter) and maps.

If you need Python: pandas UDF

A *pandas UDF* (vectorized) receives a whole batch as a pandas.Series, exchanged through Apache Arrow. It pays the round trip once per batch, not once per row.

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def with_tax(amount: pd.Series) -> pd.Series:
    return amount * 1.1

orders.withColumn("with_tax", with_tax("amount"))

• A Python UDF also hides the type: declare returnType, otherwise the default is string.
• Test the UDF logic as a plain Python function, outside Spark: it is the cheapest test there is.