Why PySpark for ML: when your data outgrows pandas and a single machine

LEARN · DISTRIBUTED ML WITH PYSPARK

If you build machine learning systems in Python, pandas is often the first tool you reach for. It is simple, expressive, and excellent for exploration. For many datasets, it is exactly the right choice.

The problem starts when your data preparation workload grows beyond one machine.

A notebook that once felt instant can become a struggle:

Loading data consumes most of your RAM.
Feature engineering becomes slower than model training.
Large joins and aggregations become difficult to run.
Repeating experiments requires rebuilding expensive datasets.

This is where PySpark becomes useful.

PySpark is the Python interface for Apache Spark, a distributed data processing engine designed to process large datasets by splitting work across partitions. It is not “pandas but faster.” It is a different execution model built for workloads where a single process is no longer enough.

This guide covers:

What PySpark does in modern machine learning workflows.
How Spark DataFrames, lazy evaluation, and DAG execution work.
How to run Spark locally.
When Spark is the wrong tool.
A real security lesson from a Spark vulnerability.
How PySpark fits into a practical DistributedML pipeline.

What is PySpark?

PySpark lets Python developers use Spark’s distributed computing capabilities, including:

DataFrame transformations.
SQL-style analytics.
Large-scale joins and aggregations.
Feature engineering pipelines.
Data preparation for machine learning.

A PySpark application can run:

On a laptop in local mode.
On a single powerful server.
On a distributed cluster.
In managed cloud data platforms.

Modern Spark development primarily uses the DataFrame API because Spark can analyze and optimize these operations before execution.

Apache Spark 4.0 was released in 2025 as a major version with significant improvements across SQL, streaming, Python support, and performance. Many production environments continue to use Spark 3.x while planning upgrades, so teams should check compatibility before moving versions.

The important idea is not the version number. The important idea is choosing an execution engine that matches the scale of your workload.

Why ML pipelines outgrow pandas

Machine learning projects often spend more time preparing data than training models.

A typical workflow looks like this:

Raw data
   |
   v
Cleaning
   |
   v
Feature engineering
   |
   v
Aggregations
   |
   v
Training dataset
   |
   v
Model training

Spark is particularly useful for the data preparation stages:

Combining millions or billions of records.
Creating historical features.
Building repeatable training datasets.
Processing event streams.
Running scheduled feature pipelines.

A rough comparison:

Dataset size	pandas	PySpark
Hundreds of MB	Usually excellent	Often unnecessary
A few GB	Depends on hardware	Can become useful
Tens of GB	May require careful optimization	Often a good fit
Hundreds of GB or more	Increasingly difficult	Designed for this scale

There is no universal point where Spark becomes necessary.

The better question is:

Can this workload be processed comfortably on one machine, repeatedly and reliably?

Running PySpark locally

You do not need a cluster to learn Spark.

Install PySpark:

pip install pyspark

A minimal local Spark application:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = (
    SparkSession.builder
    .appName("PySparkIntro")
    .master("local[*]")
    .getOrCreate()
)

data = [
    ("Alice", 34),
    ("Bob", 27),
    ("Carol", 42),
]

df = spark.createDataFrame(data, ["name", "age"])

result = (
    df
    .filter(F.col("age") >= 30)
    .withColumn("age_next_year", F.col("age") + 1)
)

result.show()

spark.stop()

The local[*] setting tells Spark to use available CPU cores on the current machine.

This is useful for learning and testing, but local mode does not make a laptop equivalent to a cluster. Memory, CPU, and storage limits still apply.

Spark DataFrames: familiar syntax, different execution

If you know pandas, Spark DataFrames feel familiar.

You can:

Select columns.
Filter rows.
Create calculated fields.
Group records.
Join datasets.

Example:

from pyspark.sql import functions as F

customer_summary = (
    df
    .groupBy("country")
    .agg(
        F.count("*").alias("customers"),
        F.avg("age").alias("average_age")
    )
)

customer_summary.show()

The major difference is what happens internally.

A pandas operation usually executes immediately against data stored in your Python process.

Spark works differently.

When you write transformations, you are describing a computation that Spark can optimize before running it.

Lazy evaluation: Spark waits on purpose

Spark uses lazy evaluation.

Transformations build a plan, but Spark does not execute the computation immediately.

For example:

from pyspark.sql import functions as F

processed = (
    df
    .filter(F.col("age") >= 30)
    .withColumn("age_squared", F.col("age") * F.col("age"))
)

At this point, Spark has a plan.

The actual work begins when you call an action:

processed.show()

or:

count = processed.count()

print(count)

This gives Spark an opportunity to optimize the complete workload.

Instead of blindly executing every operation one by one, Spark can examine the entire pipeline and choose a more efficient strategy.

The DAG: Spark’s execution blueprint

Spark represents computations as a Directed Acyclic Graph, commonly called a DAG.

The DAG describes:

Where data enters the system.
Which transformations are required.
Which operations can run in parallel.
Where data must move between workers.

Spark uses this information to optimize execution by:

Removing unnecessary operations.
Combining compatible steps.
Improving execution order.
Reducing expensive data movement.

You can inspect a planned execution:

processed.explain()

Learning to read execution plans is one of the most valuable Spark skills.

A slow Spark job is not always fixed by adding more machines. The issue may be:

An inefficient join.
Too much data movement.
Reading unnecessary columns.
Poor partition choices.

A workload where Spark starts to make sense

Spark is useful when data processing becomes too large for a single in-memory workflow.

For example:

from pyspark.sql import functions as F

large_df = (
    spark.range(100_000_000)
    .withColumn("group_id", F.col("id") % 100)
)

summary = (
    large_df
    .groupBy("group_id")
    .count()
)

summary.show()

This creates a distributed DataFrame with 100 million rows.

Whether it completes quickly depends on available resources. The point is the execution model: Spark can divide this type of computation across partitions and workers when running in an environment designed for distributed processing.

Trying to create the same scale of data in a single pandas process can quickly become a memory bottleneck.

Why storage format matters

Large ML systems rarely keep everything as CSV files.

Columnar formats such as Apache Parquet are common because they support:

Efficient compression.
Faster analytical reads.
Reading only required columns.
Better integration with distributed processing.

Example:

features = spark.read.parquet("data/features")

training_data = features.select(
    "user_id",
    "purchase_amount",
    "session_length"
)

training_data.show()

For large pipelines, data layout can matter as much as the compute engine.

A poorly organized dataset can make even powerful infrastructure inefficient.

PySpark in real machine learning systems

Spark is often used to prepare data rather than replace specialized model-training frameworks.

A common architecture looks like this:

Use Spark for cleaning and feature engineering.
Create a training dataset.
Train a model with a framework suited to the model type.
Deploy the model separately.

Spark commonly works alongside tools such as:

scikit-learn for traditional machine learning.
XGBoost for gradient-boosted models.
PyTorch or TensorFlow for deep learning workloads.

The exact combination depends on the problem, data size, and deployment requirements.

Spark’s value is often making the data ready.

When Spark is worth using

PySpark is a strong choice when:

Your datasets no longer fit comfortably into memory.
Feature engineering dominates your pipeline runtime.
Multiple teams need repeatable data processing.
You repeatedly optimize code to avoid memory failures.
Distributed execution reduces operational cost or runtime.

The biggest benefit comes from solving a real scaling problem.

When NOT to use Spark

Spark is powerful, but it is not the default answer for every dataset.

Avoid Spark when:

Your data is small.
pandas completes the job quickly.
You are exploring a dataset interactively.
You need extremely low-latency online inference.
Cluster overhead costs more than the computation.

Spark introduces additional complexity:

Distributed debugging.
Execution-plan analysis.
Cluster configuration.
Data movement between workers.
More operational responsibility.

For a 100 MB dataset, Spark may be slower and harder to maintain than a simple pandas script.

Good engineering is not choosing the most powerful tool. It is choosing the simplest tool that reliably solves the problem.

Cherry on the cake: a Spark vulnerability that changed security conversations

Performance is not the only concern in distributed systems.

A real example is CVE-2022-33891, a vulnerability affecting Apache Spark’s command-line interface under certain configurations. The issue involved insufficient restrictions around command argument handling and could allow command execution with the privileges of the Spark process when exploited in affected setups.

The lesson was operational rather than theoretical:

Distributed compute platforms need regular patching.
Security configuration matters as much as code.
A powerful data platform is also a production system that needs the same security discipline as an application.

Teams running Spark should keep versions supported, review authentication and authorization settings, and avoid treating internal data infrastructure as automatically trusted.

A practical PySpark learning path

A good progression is:

Install PySpark locally.
Learn DataFrame transformations.
Understand transformations versus actions.
Practice reading and writing Parquet.
Use explain() to understand execution plans.
Move to a cluster only when your workload requires it.

The biggest beginner mistake is thinking Spark is simply a faster pandas replacement.

Spark is a distributed execution engine. It becomes valuable when your data preparation challenges become distributed computing challenges.

Start with small experiments, measure where your pipeline slows down, and build a PySpark workflow when it genuinely improves your machine learning process. Your next step: run the examples above on your own dataset, compare the bottlenecks, and decide where Spark can make your ML pipeline more scalable.