4.2 Performing Batch Inference with pandas and Spark UDFs

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()