4.2 Performing Batch Inference with pandas and Spark UDFs

Databricks Certified Machine Learning Associate Study Guide Exam

Track your progress as you read

Finished sections get a check mark.

1 Chapter 1: Databricks Machine Learning (38% of the Exam)

1.1 MLOps Best Practices on Databricks 1.2 Advantages of Databricks ML Runtimes and Compute Selection 1.3 Promoting Code vs. Promoting Models 1.4 How AutoML Facilitates Model and Feature Selection 1.5 AutoML Advantages in the Model Development Process 1.6 Feature Tables in Unity Catalog vs. the Workspace Feature Store 1.7 Creating a Feature Table and Writing Data to It 1.8 Training a Model with Features from a Feature Table 1.9 Scoring a Model with Features from a Feature Table 1.10 Online vs. Offline Feature Tables 1.11 Manually Logging Metrics, Artifacts, and Models in an MLflow Run 1.12 Identifying the Best Run with the MLflow Client API 1.13 Information Available in the MLflow UI 1.14 Registering a Model in the Unity Catalog Registry 1.15 Model Tags, Aliases, and Champion/Challenger Promotion

2 Chapter 2: Data Processing and ML Workflows (19% of the Exam)

2.1 Computing Summary Statistics on a Spark DataFrame 2.2 Creating Visualizations for Categorical and Continuous Features 2.3 Comparing Two Categorical or Two Continuous Features 2.4 Removing Outliers with Standard Deviation and IQR 2.5 When Log-Scale Transformation Is Appropriate 2.6 Comparing and Applying Mean, Median, and Mode Imputation 2.7 One-Hot Encoding and When It Is or Is Not Appropriate

3 Chapter 3: Model Development (31% of the Exam)

3.1 Selecting the Appropriate Algorithm for a Scenario 3.2 Methods to Mitigate Data Imbalance in Training Data 3.3 Comparing Estimators and Transformers 3.4 Developing a Training Pipeline 3.5 Tuning Hyperparameters with Hyperopt's fmin 3.6 Grid, Random, and Bayesian Search Strategies 3.7 Parallelizing Single-Node Models for Hyperparameter Tuning 3.8 Cross-Validation vs. a Train-Validation Split 3.9 Counting the Models Trained by Grid Search and Cross-Validation 3.10 Using Common Classification Metrics 3.11 Using Common Regression Metrics 3.12 Choosing the Most Appropriate Metric for the Objective 3.13 Exponentiating Log-Transformed Predictions Before Evaluation 3.14 Model Complexity and the Bias-Variance Tradeoff

4 Chapter 4: Model Deployment (12% of the Exam)

4.1 Differences and Advantages of Batch, Streaming, and Real-Time Serving 4.2 Performing Batch Inference with pandas and Spark UDFs 4.3 Streaming Inference with Delta Live Tables 4.4 Deploying a Custom Model to a Model Endpoint 4.5 Deploying, Querying, and Splitting Traffic Across Real-Time Endpoints

4.2 Performing Batch Inference with pandas and Spark UDFs

Key Takeaways Explain This Topic AI

mlflow.pyfunc.load_model(uri) loads a model on the driver and its .predict(pandas_df) scores a pandas DataFrame directly — correct for data that fits in driver memory.

mlflow.pyfunc.spark_udf(spark, uri, result_type) turns the model into a Spark UDF so partitions are scored in parallel across executors.

Vectorized @pandas_udf functions move columnar batches between the JVM and Python via Apache Arrow, so single-node models score far faster than with row-by-row Python UDFs.

result_type must match the model's output: "double" for probabilities or regression, "string" for labels, ArrayType(DoubleType()) for probability vectors.

Write predictions to Delta with append for immutable logs, overwrite for full snapshots, and MERGE when the table must hold the latest row per entity.

Last updated: August 2026

4.2 Performing Batch Inference with pandas and Spark UDFs

Distributed inference on the Databricks Lakehouse leverages Apache Spark to scale model evaluation across multi-node clusters. By integrating MLflow Model Registry with Spark SQL User-Defined Functions (UDFs) and Delta Lake , machine learning engineers can score terabyte-scale datasets in batch or process high-velocity data streams with robust ACID guarantees.

Mechanics of Distributed Model Loading & Execution

When applying a machine learning model to a distributed Spark DataFrame, understanding the execution mechanics across the driver and worker nodes is essential for performance optimization. +---------------------------------------------------------------------------------------------------+

| DISTRIBUTED MLFLOW SPARK UDF INFERENCE ARCHITECTURE |
| |
| SPARK DRIVER |
| +-------------------------------------------------------------------------------------------+ |
| | 1. Read Model Metadata from MLflow / Unity Catalog ('models:/catalog.schema.model@prod') | |
| | 2. Instantiate mlflow.pyfunc.spark_udf(spark, model_uri, result_type="double") | |
| | 3. Construct Spark Execution DAG & Distribute Tasks to Workers | |
| +-------------------------------------------------------------------------------------------+ |
| | Broadcast Model Artifacts |
| +----------------------+----------------------+ |
| | | |
| SPARK EXECUTOR 1 v SPARK EXECUTOR 2 v |
| +---------------------------------------+ +---------------------------------------+ |
| | - Load Model in Python Worker (Arrow) | | - Load Model in Python Worker (Arrow) | |
| | - Process Partition 1 (Vectorized) | | - Process Partition 2 (Vectorized) | |
| | - Compute batch predictions | | - Compute batch predictions | |
| +---------------------------------------+ +---------------------------------------+ |
| | | |
| +----------------------+----------------------+ |
| v |
| DELTA LAKE GOLD TABLE (ACID COMMIT) |
+---------------------------------------------------------------------------------------------------+
mlflow.pyfunc.load_model vs. mlflow.pyfunc.spark_udf

mlflow.pyfunc.load_model(model_uri) : Loads the serialized model into the local memory of the Spark driver as a standard Python object. Used for single-node evaluation, local testing, or custom single-node driver processing. Attempting to iterate row-by-row over a distributed DataFrame with load_model() forces all data to collect to the driver, resulting in driver Out-Of-Memory (OOM) crashes.

mlflow.pyfunc.spark_udf(spark, model_uri, result_type='double') : Creates a distributed Spark SQL UDF. The Spark engine serializes and broadcasts the model environment and weights to all executor worker nodes . Each worker loads the model in parallel and applies inference against its local DataFrame partitions.

Result Type Configuration

The result_type parameter defines the Spark SQL schema of the prediction output:

Scalar regression or binary probability: result_type="double" or result_type="float"
Categorical class label: result_type="string"
Multiclass probability distribution: result_type=ArrayType(DoubleType())
Complex outputs (e.g., class + probabilities + embeddings): result_type=StructType([...])
Scoring a pandas DataFrame Directly
When the data to score is small enough to hold in the driver's memory — a daily slice,
a test harness, an ad-hoc analysis — the simplest correct pattern needs no Spark at
all: import mlflow.pyfunc

model = mlflow.pyfunc.load_model("models:/prod_ml.finance.fraud_detector@champion")

# Bring a bounded slice into pandas, then score it in one vectorized call
pdf = (spark.table("prod_ml.finance.gold_transaction_features")
.filter("transaction_date = current_date()")
.toPandas())

pdf["fraud_probability"] = model.predict(pdf[feature_columns])

mlflow.pyfunc gives every logged model the same predict(pandas_df) -> array-like interface regardless of the framework underneath, which is why a scikit-learn model, an

XGBoost model, and a custom Python model are all called identically.

The boundary : toPandas() collects to the driver. It is the right call for
thousands or low millions of rows and the wrong call for hundreds of millions — the

same size question that decides single-node versus distributed training in Section 3.1.

Above that boundary, keep the pandas interface and distribute the execution with a

UDF.

High-Throughput Vectorized Inference with Pandas UDFs

Standard Python UDFs in PySpark operate row-by-row, incurring massive serialization and deserialization overhead between the Java Virtual Machine (JVM) and Python runtime. For single-node model libraries (Scikit-learn, XGBoost, LightGBM, PyTorch), Vectorized Pandas UDFs ( pyspark.sql.functions.pandas_udf ) utilize Apache Arrow to stream columnar batches directly between the JVM and Python worker processes. import pandas as pd

import mlflow
from pyspark.sql.functions import pandas_udf, col
from pyspark.sql.types import DoubleType

# 1. Define model URI in Unity Catalog
model_uri = "models:/prod_ml.customer_churn.churn_classifier@champion"

# 2. Vectorized Pandas UDF using Iterator of Series/DataFrames for zero-copy efficiency
@pandas_udf(DoubleType())

def predict_churn_vectorized(*cols: pd.Series) -> pd.Series:

# Load model once per worker process lifecycle
model = mlflow.pyfunc.load_model(model_uri)

# Combine incoming column Series into a single Pandas DataFrame
pdf = pd.concat(cols, axis=1)
pdf.columns = ["age", "tenure", "monthly_charges", "support_tickets", "total_spend"]

# Vectorized batch prediction (returns numpy array / Series)
predictions = model.predict(pdf)

# Return probabilities for class 1
return pd.Series(predictions[:, 1] if predictions.ndim > 1 else predictions)
Writing Scored Predictions into Delta Lake

Writing predictions into Delta Lake ensures transactional consistency, schema enforcement, time-travel auditing, and high-performance querying for downstream consumers. +---------------------------------------------------------------------------------------------------+

| DELTA LAKE INFERENCE WRITE STRATEGIES |
| |
| APPEND STRATEGY OVERWRITE STRATEGY MERGE / UPSERT |
| +---------------------------+ +---------------------------+ +-------------------+ |
| | .mode("append") | | .mode("overwrite") | | DeltaTable.merge()| |
| | Ideal for: | | Ideal for: | | Ideal for: | |
| | - Immutable audit logs | | - Daily active snapshots | | - Entity updates | |
| | - Historical time-series | | - Replacing partitions | | - Latest state | |
| | - Continuous event streams| | - Full catalog refresh | | - Deduplication | |
| +---------------------------+ +---------------------------+ +-------------------+ |
+---------------------------------------------------------------------------------------------------+
Complete Batch Scoring Pipeline with Delta Upsert (Merge) from pyspark.sql import SparkSession
from pyspark.sql.functions import struct, current_timestamp
import mlflow.pyfunc
from delta.tables import DeltaTable

spark = SparkSession.builder.getOrCreate()

# 1. Define Model URI and Register Spark UDF
model_uri = "models:/enterprise_ml.finance.fraud_detector@champion"
predict_udf = mlflow.pyfunc.spark_udf(spark, model_uri=model_uri, result_type="double")

# 2. Read feature dataset from Delta Gold table
features_df = spark.table("enterprise_ml.finance.gold_transaction_features") \
.filter("transaction_date = current_date()")

feature_columns = ["account_age", "transaction_amount", "velocity_24h", "cross_border_flag"]

# 3. Apply distributed scoring across worker partitions
scored_df = features_df.withColumn(
"fraud_probability",
predict_udf(struct(*feature_columns))
).withColumn("scored_timestamp", current_timestamp())

# 4. Upsert predictions into Gold Predictions Delta Table using ACID Merge
target_table = "enterprise_ml.finance.daily_fraud_predictions"

if not spark.catalog.tableExists(target_table):

scored_df.write.format("delta") \
.partitionBy("transaction_date") \
.saveAsTable(target_table)

else:

delta_target = DeltaTable.forName(spark, target_table)
delta_target.alias("target").merge(
source=scored_df.alias("source"),
condition="target.transaction_id = source.transaction_id"
).whenMatchedUpdate(set={
"fraud_probability": "source.fraud_probability",
"scored_timestamp": "source.scored_timestamp"
}).whenNotMatchedInsertAll().execute()