Using Pandas API on Spark · Topic 1 of 1

Pandas API on Spark and Pandas UDFs

Pandas API on Spark for the Databricks exam: its advantages over pandas, converting DataFrames, and creating and invoking pandas UDFs with type hints.

4 min

In a nutshell

import pyspark.pandas as ps gives pandas syntax on distributed data.
to_pandas() collects to the driver; avoid it on large data.
sdf.pandas_api() and psdf.to_spark() convert between the two Spark side types.

What it is

The pandas API on Spark (pyspark.pandas) gives you the familiar pandas DataFrame and Series API, but the data is distributed and processed by Spark. Code that works on one machine in pandas can scale to a cluster with few changes.

A pandas UDF is a different tool. It is a user defined function in a PySpark query that receives pandas objects, works on a batch of rows at once, and returns pandas objects.

How it works

Advantages of the pandas API on Spark

Advantage	Meaning
Familiar API	pandas users keep their syntax and habits
Scales out	Data larger than one machine's memory can be processed
Single codebase	The same code style serves small and large data
Spark engine	Work is planned by Spark's optimizer and runs in parallel
Interoperability	Convert to a Spark DataFrame when you need the full PySpark API

Evaluation is lazy, like Spark, so work happens when a result is needed.

Converting between types

import pyspark.pandas as ps
psdf = ps.DataFrame({"a": [1, 2, 3]})
sdf = psdf.to_spark()
back = sdf.pandas_api()
pdf = psdf.to_pandas()

ps.from_pandas(pdf) creates a distributed DataFrame from a pandas one. to_pandas() collects everything to the driver, so use it only on small data. to_spark() and pandas_api() convert to and from a Spark DataFrame.

Costs to watch

The API is not fully identical to pandas, so some code needs changes. A pandas on Spark DataFrame carries an index, and Spark creates a distributed-sequence index by default when none is given. The sequence index type computes on a single partition and is discouraged for large data. Operations like sort_values shuffle data across the cluster.

Creating a pandas UDF

Declare a pandas UDF with the pandas_udf decorator, the return type, and Python type hints that choose the form.

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

@pandas_udf("double")
def add_tax(price: pd.Series) -> pd.Series:
    return price * 1.2

The decorator needs the return type, given as a DDL string or a Spark type. Data moves between the JVM and Python in Arrow batches, which is why a pandas UDF is faster than a plain Python UDF that handles one row at a time.

Invoking a pandas UDF

Call it like any column function inside select, withColumn or agg.

df.select(add_tax("price").alias("price_tax")).show()

Forms selected by type hints

Type hint	Form	Behavior
`pd.Series` to `pd.Series`	Scalar	Returns a Series the same length as the input
`pd.Series` to a scalar such as `float`	Grouped aggregate	Reduces a column to one value per group
`Iterator[pd.Series]` to `Iterator[pd.Series]`	Iterator	Setup code runs once per task, not per batch

A grouped aggregate pandas UDF is used with groupBy().agg(...). For a whole pandas DataFrame per group, groupBy("id").applyInPandas(func, schema) applies a function to each group and returns a DataFrame with the schema you declare.

How the exam asks about this

"Apply a function to a column in batches rather than row by row." A pandas UDF, declared with pandas_udf.
"A UDF must compute the mean per group." Use a Series to scalar pandas UDF in groupBy().agg().
"Convert a pandas on Spark DataFrame to use the full PySpark API." Call to_spark().

What to remember for the exam

import pyspark.pandas as ps gives pandas syntax on distributed data.
to_pandas() collects to the driver; avoid it on large data.
sdf.pandas_api() and psdf.to_spark() convert between the two Spark side types.
A pandas UDF is vectorized, using Arrow batches; a plain @udf runs row by row.
Type hints choose the form: Series to Series is scalar, Series to scalar is grouped aggregate.
An Iterator form runs setup code once per task.

Watch out

Calling to_pandas() on a large pandas on Spark DataFrame pulls all rows into driver memory and can crash the driver. Filter or aggregate first.

Common mistakes

Assuming pandas code runs unchanged. The pandas API on Spark covers much of pandas, but not all of it. Missing or different behavior needs a workaround such as converting with to_pandas() on a small sample.

Returning a wrong length from a scalar pandas UDF. A Series to Series UDF must return a Series of the same length as its input. A different length raises an error.

Forgetting the return type. Without a declared return type in pandas_udf, Spark cannot build the schema. Declare it explicitly.

Check yourself

Why is a pandas UDF usually faster than a plain Python UDF?
It processes batches of rows as pandas objects transferred in Arrow format, instead of calling Python once per row.
Which type hint pattern makes a grouped aggregate pandas UDF?
A function from a pd.Series to a scalar value such as float, used inside groupBy().agg().
What does psdf.to_pandas() do?
It collects the whole distributed DataFrame to the driver as a single machine pandas DataFrame, which is risky for large data.

Topic 1 of 1 in Using Pandas API on Spark

Test yourself on Using Pandas API on Spark

The full Databricks Associate Developer for Apache Spark question bank has 600 questions, with timed practice exams and an explanation for every answer.

Get Full Access Try Sample Exam

This study guide is independent and is not affiliated with, endorsed by, or sponsored by Databricks. Vendor names, certification names and exam codes are trademarks of their respective owners and are used for identification only.