04.7 PySpark UDFs

Architect goal: Know when custom Python logic is justified, how UDF execution crosses the Spark/Python boundary, why native expressions are usually preferable, when Pandas UDFs or Arrow-optimized UDFs help, and how to profile UDF-heavy workloads.

1. What is a UDF?

A user-defined function lets you apply custom logic that is not conveniently expressed with Spark's built-in functions.

Example:

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

def normalize(value):
    if value is None:
        return None
    return value.strip().upper()

normalize_udf = udf(normalize, StringType())

result = df.withColumn(
    "normalized_name",
    normalize_udf("name")
)

The function becomes part of the Spark expression graph, but the function body executes in Python rather than as a native Spark SQL expression.

2. The first architecture decision

Before writing a UDF, ask:

Can Spark SQL / PySpark express this?
        │
       yes
        ↓
Use native expression
        │
       no
        ↓
Can a SQL UDF solve it?
        │
       yes
        ↓
Use SQL UDF
        │
       no
        ↓
Is Python really required?
        │
       yes
        ↓
Evaluate Pandas / Arrow / Python UDF options

Databricks recommends native Spark operations when an equivalent exists because moving data between Spark and Python introduces serialization overhead. citeturn515884search10turn515884search1

3. Native expression vs Python UDF

Suppose the requirement is uppercase.

Prefer native Spark

from pyspark.sql import functions as F

df.withColumn(
    "normalized_name",
    F.upper(F.trim("name"))
)

Avoid unnecessary UDF

@udf("string")
def normalize(value):
    return value.strip().upper() if value else None

The native expression keeps the operation inside Spark's optimized execution path.

Architect rule

A UDF should solve a real expressiveness gap, not replace a function you forgot existed.

4. Python UDF execution model

Conceptually:

Spark JVM / execution engine
        │
        │ serialize data
        ▼
Python worker
        │
        │ execute Python function
        ▼
Python result
        │
        │ serialize / transfer
        ▼
Spark execution

The exact implementation depends on the UDF type and runtime, but the important architectural boundary is Spark execution ↔ Python execution.

That boundary can become expensive at large scale.

5. Scalar Python UDF

A scalar UDF produces one result per input row.

from pyspark.sql.functions import udf

@udf("double")
def add_tax(amount):
    if amount is None:
        return None
    return amount * 1.18

result = df.withColumn("gross_amount", add_tax("amount"))

Use when the transformation genuinely requires Python logic.

6. Type discipline

Declare the return type explicitly:

@udf("double")
def add_tax(amount):
    ...

or:

from pyspark.sql.types import DoubleType

add_tax = udf(add_tax_function, DoubleType())

Treat the return type as part of the data contract.

For production pipelines, also test:

• NULL input
• Empty strings
• Invalid values
• Boundary values
• Unexpected Python return types

7. NULL behavior

Do not assume your Python function will receive only valid values.

@udf("string")
def normalize(value):
    if value is None:
        return None
    return value.strip().upper()

Define explicit business behavior for NULL, empty, malformed, and out-of-range values.

8. Python UDFs and optimization

Native Spark expressions can participate in Catalyst optimization and native execution more directly.

A Python UDF creates a Python execution boundary, limiting some optimization opportunities around the function.

That does not mean every UDF is catastrophically slow. It means the cost must be justified and measured.

9. Pandas UDFs

Pandas UDFs use Apache Arrow to transfer data between Spark and Python and operate on Pandas Series/DataFrames in batches.

Conceptually:

Spark column batches
        ↓
      Arrow
        ↓
Pandas Series / DataFrame
        ↓
vectorized Python logic
        ↓
      Arrow
        ↓
Spark

Apache Spark documents Pandas UDFs as vectorized operations using Arrow. citeturn515884search0turn515884search9

Databricks documents that Pandas UDFs can substantially outperform row-at-a-time Python UDFs because of reduced serialization overhead. citeturn515884search5

10. Series-to-Series Pandas UDF

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

@pandas_udf("double")
def normalize_amount(s: pd.Series) -> pd.Series:
    return s.fillna(0) * 1.18

result = df.withColumn(
    "gross_amount",
    normalize_amount("amount")
)

The function operates on batches rather than invoking Python once for every Spark row.

11. Pandas UDF does not mean “free performance”

Pandas UDFs still involve Python execution.

They can also introduce:

• Arrow conversion
• Pandas memory consumption
• Batch-size considerations
• Python dependency management
• Type-conversion constraints

Use them when vectorization provides a real advantage.

12. Grouped Pandas operations

Pandas UDFs can also be used for grouped calculations and Pandas Function APIs such as applyInPandas().

Example:

def normalize_group(pdf):
    pdf["scaled"] = pdf["amount"] / pdf["amount"].max()
    return pdf

result = (
    df.groupBy("customer_id")
      .applyInPandas(
          normalize_group,
          schema=df.schema.add("scaled", "double")
      )
)

Be especially careful with group size.

A very large group can require substantial Python/Pandas memory.

13. Iterator Pandas UDFs

Iterator-based Pandas UDFs can be useful when expensive initialization can be reused across batches.

Conceptually:

for batch in iterator:
    yield transform(batch)

Apache Spark documents iterator variants for Pandas UDFs and their use cases. citeturn515884search0

14. Arrow-optimized Python UDFs

Current PySpark also supports Arrow-optimized Python UDF execution.

Example:

from pyspark.sql.functions import udf

@udf("int", useArrow=True)
def string_length(value):
    return len(value) if value else 0

Spark's current documentation describes Arrow Python UDFs as row-oriented UDFs that use Arrow for efficient data transfer/serialization. citeturn515884search0turn515884search6

Important version note

Arrow-related UDF behavior and defaults have changed across recent Spark releases. Treat runtime/version documentation as authoritative rather than assuming behavior from older tutorials. citeturn515884search0turn515884search6

15. UDF categories

Approach	Execution idea	Starting recommendation
Native PySpark / SQL	Native Spark expression	Preferred
SQL UDF	SQL-defined reusable logic	Excellent when expressible in SQL
Scalar Python UDF	Python function per row	Use only when justified
Arrow-optimized Python UDF	Python UDF with Arrow transfer	Evaluate when appropriate
Pandas UDF	Vectorized Python/Pandas	Strong option for Python-specific vectorizable logic
Pandas Function API	DataFrame/group-oriented Python	Use for transformations requiring DataFrame semantics
Python UDTF	Python function returning rows	Specialized use case; check runtime support

Databricks currently documents Python scalar/non-scalar UDFs and Python UDTFs as distinct capabilities, with availability depending on compute type and runtime. citeturn515884search1turn515884search8

16. Python UDTFs

A UDTF returns multiple rows rather than a single scalar value.

Conceptually:

input arguments
      ↓
Python UDTF
      ↓
table / multiple rows

Databricks currently documents Python UDTFs as a Public Preview capability and recommends Unity Catalog registration for governed reuse. citeturn515884search8

Do not use a UDTF simply because a scalar UDF is inconvenient. Choose based on output semantics.

17. SQL UDF vs Python UDF

If business logic can be expressed in SQL, a SQL UDF can be easier to govern and execute without introducing a Python boundary.

Example concept:

CREATE FUNCTION normalize_amount(amount DECIMAL(18,2))
RETURNS DECIMAL(18,2)
RETURN amount * 1.18;

Exact syntax and governance depend on the Databricks SQL/runtime environment.

18. UDF and Photon

Do not assume:

“Photon is enabled, therefore my Python UDF is automatically native Photon code.”

Photon accelerates supported execution paths, while Python UDFs still involve Python execution.

The architect question is:

Which part of the workload is native engine work, and where does execution cross into Python?

19. UDFs and data movement

At scale, the cost model is roughly:

Spark processing
 + serialization
 + Python execution
 + serialization
 + downstream Spark processing

For millions or billions of rows, even small per-row overhead can become significant.

20. Common anti-patterns

Anti-pattern 1 — UDF for a built-in function

Bad:

@udf("string")
def upper_name(x):
    return x.upper()

Better:

F.upper("name")

Anti-pattern 2 — UDF inside a huge fact transformation

If a native expression can do the work, use it.

Anti-pattern 3 — Grouped Pandas UDF on giant groups

A single massive group can overwhelm Python/Pandas memory.

Anti-pattern 4 — Returning complex objects casually

Define stable Spark return types and test serialization.

Anti-pattern 5 — Ignoring dependencies

A production Python UDF requires reproducible Python dependencies across compute.

Anti-pattern 6 — Assuming Pandas UDF is always fastest

Measure against native expressions and other alternatives.

21. When a UDF is justified

Good candidates include:

• Specialized domain algorithms
• Python-only libraries
• Complex parsing unavailable in native functions
• Proprietary business logic
• Mathematical logic difficult to express safely in SQL
• Reusable custom transformations

But first prove that native SQL/PySpark cannot express the requirement adequately.

22. UDF design checklist

Before production:

1. Can native Spark do it?
2. Can SQL express it?
3. Is Python genuinely required?
4. Is the operation row-wise or vectorizable?
5. Would Pandas/Arrow reduce overhead?
6. What is the expected data volume?
7. What is the largest group/batch?
8. What are the NULL/error semantics?
9. What dependencies are required?
10. How will it be tested?
11. How will it be monitored?
12. What happens when the Python worker fails?

23. Profiling UDF workloads

Current PySpark provides profiling support for Python, Pandas, and Arrow UDFs.

A SparkSession-based profiler can be enabled with:

spark.conf.set(
    "spark.sql.pyspark.udf.profiler",
    "perf"
)

Spark also documents memory profiling for UDF workloads. citeturn515884search4

Use profiling to replace assumptions with evidence.

24. Hands-on lab — native vs Python UDF

Create a large DataFrame.

Implement the same transformation using:

1. Native F.upper()
2. Scalar Python UDF
3. Pandas UDF

Compare:

• Runtime
• Physical plan
• Python execution
• CPU
• Memory

The goal is not to memorize a benchmark. It is to understand why execution differs.

25. Hands-on lab — vectorization

Implement:

gross_amount = amount × tax_rate

using:

• Native expression
• Scalar UDF
• Pandas UDF

Then test with increasing data volumes.

Record the crossover behavior.

26. Hands-on lab — grouped Pandas function

Create customer groups with deliberately different sizes.

Use applyInPandas() and observe how group size affects Python memory.

This teaches an important production principle:

Distributed grouping does not make an individual Python/Pandas group infinitely small.

27. Hands-on lab — profiling

Enable UDF profiling and execute a representative workload.

Inspect:

• Function time
• Python overhead
• Memory behavior
• Hot functions

Use profiling evidence to decide whether optimization is worthwhile.

28. Production troubleshooting

When a UDF workload is slow:

Step 1 — Remove the UDF mentally

Ask whether native Spark can express the transformation.

Step 2 — Inspect the plan

Use:

result.explain("formatted")

Step 3 — Measure Python overhead

Use Spark/PySpark profiling where appropriate.

Step 4 — Check data volume

A tiny UDF on a billion rows is not tiny operationally.

Step 5 — Check group size

Especially for grouped Pandas operations.

Step 6 — Check serialization and memory

Look for executor pressure and Python-worker issues.

Step 7 — Compare alternatives

Benchmark native, SQL, Pandas/Arrow, and scalar Python approaches where applicable.

29. Architect decision matrix

Requirement	Preferred starting point
Built-in transformation exists	Native PySpark
Logic naturally expressed in SQL	SQL expression / SQL UDF
Python-only scalar logic	Scalar/Arrow Python UDF after measurement
Vectorizable Python logic	Pandas UDF candidate
Group/DataFrame Python logic	Pandas Function API candidate
Function returns multiple rows	UDTF candidate
Billion-row simple transformation	Avoid Python boundary
Unknown performance	Benchmark + profile

30. Interview questions

1. What is a PySpark UDF?
2. Why can Python UDFs be slower?
3. Native function vs Python UDF?
4. What is a Pandas UDF?
5. How does Arrow help?
6. Pandas UDF vs scalar Python UDF?
7. What is a Python UDTF?
8. When would you choose a UDF?
9. Can Photon make Python UDF execution native?
10. What happens at the Spark/Python boundary?
11. How do you profile UDF performance?
12. Why can grouped Pandas UDFs consume large memory?
13. How would you migrate a slow Python UDF to native Spark?
14. What dependency-management issues exist with UDFs?
15. How would you design a governed enterprise UDF?

31. Cheat sheet

# Native — preferred
df.withColumn("name2", F.upper(F.trim("name")))

# Scalar Python UDF
from pyspark.sql.functions import udf

@udf("double")
def add_tax(x):
    return x * 1.18 if x is not None else None

# Pandas UDF
import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def add_tax_vectorized(x: pd.Series) -> pd.Series:
    return x.fillna(0) * 1.18

# Arrow-optimized Python UDF
@udf("int", useArrow=True)
def string_length(x):
    return len(x) if x else 0

# Profiling
spark.conf.set("spark.sql.pyspark.udf.profiler", "perf")

32. Key takeaways

1. Native Spark expressions are the default choice.
2. A Python UDF creates a Python execution boundary.
3. Serialization and Python execution can become expensive at scale.
4. Pandas UDFs use Arrow and vectorized execution.
5. Pandas UDFs are not automatically faster than native Spark.
6. Grouped Python operations require careful memory analysis.
7. Arrow-optimized Python UDF behavior is runtime/version dependent.
8. UDF return types and NULL semantics are part of the data contract.
9. Profile before optimizing.
10. The best UDF is often the UDF you successfully replaced with a native expression.

Next

04.8 PySpark Pandas API — Pandas API on Spark, pandas interoperability, index semantics, distributed execution, migration patterns, limitations, and when to use it versus native PySpark.

Official references

• PySpark UDF
• PySpark Pandas UDF
• Apache Arrow in PySpark
• PySpark UDF/UDTF guide
• Databricks UDFs
• Databricks Pandas UDFs