Scaling single-node ML on Spark: pandas API on Spark, applyInPandas, and grouped training

LEARN · DISTRIBUTED ML WITH PYSPARK

Many machine learning workflows start on one machine:

Load data into pandas.
Create features.
Split data by an entity such as a store, product, customer, or device.
Train a model for each entity.
Generate predictions.

This approach is productive, but it eventually reaches limits:

The dataset no longer fits in memory.
Feature engineering becomes too slow.
Thousands of small model-training jobs run sequentially.
Retraining becomes difficult to operationalize.

Apache Spark helps, but the key question is not “How do I add Spark?”

The better question is:

Do I need to distribute the data, or do I need to distribute the models?

Two PySpark patterns solve different problems:

pandas API on Spark (PandasOnSpark): keep a pandas-style API while Spark executes operations across a cluster.
grouped training with applyInPandas(): distribute independent pandas and scikit-learn workloads across Spark executors.

This distinction is the foundation of scalable single-node ML migration.

The first decision: distribute data or distribute models?

When the data is the problem

Use Spark when the dataset itself is too large for a single machine.

Examples:

Billions of transactions.
Multi-year event logs.
Large feature tables.
IoT streams.

The solution is usually:

Store data in distributed formats.
Transform data with Spark DataFrames or pandas API on Spark.
Reduce data before expensive model training when possible.

Here, Spark distributes the data processing.

When the models are the problem

Many ML systems do not need one huge model. They need many independent models.

Examples:

One forecast model per store.
One anomaly detector per machine.
One regression model per product.
One customer-segment model per region.

A local pandas workflow might look like:

for store_id, store_data in sales.groupby("store_id"):
    model = train_model(store_data)
    predictions = model.predict(store_data)

The data for each store may be manageable, but thousands of training loops become slow.

The Spark solution is to distribute the model training tasks.

Choosing the right pattern

Problem	Recommended approach
One large table with distributed transformations	Spark DataFrame APIs
Existing pandas-style transformations	pandas API on Spark
Thousands of independent small models	`groupBy().applyInPandas()`
Large neural network training	Specialized distributed training systems

Spark is not automatically faster for every workload. It works best when computation can be split into independent tasks or distributed operations.

Building a modern PySpark environment

Compatibility matters in ML pipelines because Spark, pandas, NumPy, and scikit-learn interact closely.

For current production systems, validate versions together. A typical 2025-era stack may include:

Apache Spark 4.x with PySpark 4.x.
pandas 2.x.
scikit-learn 1.5+ or newer compatible releases.

The exact versions should follow your platform’s tested dependency matrix rather than an arbitrary installation command.

For example, a development environment can start with:

pip install pyspark pandas scikit-learn

In production, pin versions through your deployment tooling, test upgrades in staging, and verify Python worker compatibility with your Spark runtime.

Moving pandas workflows to pandas API on Spark

The pandas API on Spark package (pyspark.pandas) lets many pandas-style operations execute using Spark.

A local pandas workflow:

import pandas as pd

sales = pd.DataFrame(
    {
        "store_id": [1, 1, 2],
        "price": [10.0, 20.0, 15.0],
        "quantity": [2, 3, 4],
    }
)

sales["revenue"] = sales["price"] * sales["quantity"]

summary = (
    sales.groupby("store_id")["revenue"]
    .sum()
    .reset_index()
)

print(summary)

The pandas API on Spark version:

import pyspark.pandas as ps

sales = ps.DataFrame(
    {
        "store_id": [1, 1, 2],
        "price": [10.0, 20.0, 15.0],
        "quantity": [2, 3, 4],
    }
)

sales["revenue"] = sales["price"] * sales["quantity"]

summary = (
    sales.groupby("store_id")["revenue"]
    .sum()
    .reset_index()
)

summary.to_spark().show()

The syntax is familiar, but execution is different:

pandas runs inside one Python process.
pandas API on Spark creates Spark execution plans.
Spark distributes partitions across executors.
Data stays distributed until explicitly collected.

Avoid breaking distributed execution

The most common scaling mistake is accidentally moving data back to the driver.

This can destroy the benefit of Spark:

local_df = spark_df.toPandas()

toPandas() collects all rows into driver memory. It is useful for intentionally small results, debugging, or visualization, but it should not be part of a large pipeline.

A safer pattern is:

small_result = (
    spark_df
    .groupBy("category")
    .count()
    .toPandas()
)

The expensive aggregation happens on Spark. Only the small final result is collected.

For performance-sensitive systems, Spark-native DataFrame operations are often clearer than forcing every transformation through a pandas-style API.

Grouped training with applyInPandas()

The most powerful pattern for many ML teams is distributed fan-out training.

Imagine a retailer with:

Store identifiers.
Product sales history.
Promotion features.
Weather signals.
Demand targets.

A local script might train one model per store. With Spark, each store can become an independent training task.

applyInPandas() works by:

Grouping Spark rows by a key.
Sending each group as a pandas DataFrame to a Python worker.
Running custom pandas/scikit-learn code.
Combining the returned pandas DataFrames.

A grouped training function:

import pandas as pd
from sklearn.linear_model import LinearRegression


def train_store_model(pdf: pd.DataFrame) -> pd.DataFrame:
    features = [
        "feature_1",
        "feature_2",
    ]

    model = LinearRegression()

    model.fit(
        pdf[features],
        pdf["target"],
    )

    result = pdf[
        [
            "store_id",
            "target",
        ]
    ].copy()

    result["prediction"] = model.predict(
        pdf[features]
    )

    return result

Spark requires an output schema:

from pyspark.sql.types import (
    DoubleType,
    LongType,
    StructField,
    StructType,
)


schema = StructType(
    [
        StructField("store_id", LongType()),
        StructField("target", DoubleType()),
        StructField("prediction", DoubleType()),
    ]
)

Apply the training function:

predictions = (
    spark_df
    .groupBy("store_id")
    .applyInPandas(
        train_store_model,
        schema=schema,
    )
)

predictions.show()

The ML code remains close to familiar pandas and scikit-learn code, while Spark manages parallel execution.

Production rules for grouped training

Keep groups reasonably sized

applyInPandas() distributes groups, not individual rows inside a group.

Good candidates:

Thousands of stores with manageable histories.
Many products with similar data volumes.
Device telemetry split by machine.

Risky candidates:

One group containing billions of rows.
Extremely uneven group sizes.

Inspect the distribution first:

group_sizes = (
    spark_df
    .groupBy("store_id")
    .count()
    .orderBy("count", ascending=False)
)

group_sizes.show(10)

Keep functions self-contained

Distributed tasks can execute anywhere and in any order.

Avoid:

Global mutable state.
Hidden local files.
Assumptions about execution order.

Prefer functions that receive data, train, and return results.

Measure Spark overhead

Spark introduces costs:

Task scheduling.
JVM-to-Python communication.
Data serialization.
Python worker startup.

If every group contains only a few rows, Spark overhead may exceed the actual model training time.

Grouped training works best when each task performs meaningful computation.

The cherry on the cake: forecasting thousands of store-SKU combinations

A surprising real-world pattern is that many companies do not need one giant forecasting model. They need a large number of smaller models.

A retailer might have:

20,000 stores.
100,000 products.
Daily sales history.
Local promotions.
Seasonal effects.

A forecasting system may create models for:

Store A + Product X.
Store A + Product Y.
Store B + Product X.

This is a natural fit for grouped training.

However, Spark only solves the parallel execution layer. A production ML platform still needs answers for:

How models are versioned.
How failed groups are retried.
How new products are handled.
How predictions are monitored.
Where artifacts are stored.

Security lesson: distributed systems still need hardening

Spark clusters often process sensitive business data, so security must be part of the architecture.

A historical example is CVE-2022-33891, a command injection vulnerability affecting Apache Spark versions before patched releases. The issue demonstrated how unsafe handling of user-controlled input could become a serious risk in distributed data platforms.

The lesson remains relevant:

Keep Spark versions updated.
Monitor dependency vulnerabilities.
Restrict cluster network access.
Use authentication and authorization controls.
Separate development and production environments.

Scaling computation does not remove the need for secure engineering.

When to use each approach

Use pandas API on Spark when:

Your team already knows pandas.
You need distributed tabular transformations quickly.
The workload is mostly data preparation.

Use applyInPandas() when:

Groups are independent.
Each group fits in executor memory.
Existing pandas/scikit-learn training code is valuable.
Parallel execution offsets Spark overhead.

Avoid forcing these patterns onto:

Shared-state ML algorithms.
Extremely unbalanced groups.
Tiny tasks dominated by scheduling costs.
Large neural network training workloads.

Final checklist

Before scaling a pandas ML workflow with Spark:

Decide whether you need to distribute data or models.
Use pandas API on Spark for pandas-style distributed transformations.
Avoid unnecessary toPandas() calls.
Measure group sizes before using applyInPandas().
Keep grouped training functions deterministic.
Monitor memory, serialization, and execution time.
Test Spark, pandas, and scikit-learn versions together.
Include security reviews in cluster operations.

Your next exercise: take one real per-entity ML pipeline, measure its current runtime, migrate transformations with pandas API on Spark, then convert one sequential training loop into grouped training with applyInPandas().

Benchmark the result, learn where the bottlenecks moved, and build your Spark ML architecture from measurements rather than assumptions.