Introduction

You can run machine learning code inside a Spark transformation, but the right pattern depends on whether you are doing inference or training. Per-record inference can live inside map or, more efficiently, mapPartitions. Training a separate model inside each map call is usually a design mistake because it destroys batching, wastes startup time, and fights Spark’s execution model.

Use Spark for Distributed Data, Not Tiny Per-Row Model Creation

The basic anti-pattern is creating or training a model for every element.

# Bad idea
rdd.map(lambda row: train_model(row))

That causes repeated model construction, heavy serialization costs, and poor use of executors. Spark is good at moving a function over partitions of data. It is not good at turning every row into a separate machine-learning job.

The usual good pattern is one of these:

• train a model once and use Spark for distributed inference,
• load a model once per partition and score records in that partition,
• or use Spark MLlib so training itself is expressed as a Spark job.

For Lightweight Inference, Broadcast the Model

If the model is small and serializable, broadcasting it can work well.

1from pyspark.sql import SparkSession
2from sklearn.linear_model import LogisticRegression
3import numpy as np
4
5spark = SparkSession.builder.master("local[*]").appName("broadcast-model").getOrCreate()
6sc = spark.sparkContext
7
8X = np.array([[0.0], [1.0], [2.0], [3.0]])
9y = np.array([0, 0, 1, 1])
10model = LogisticRegression().fit(X, y)
11
12broadcast_model = sc.broadcast(model)
13rows = sc.parallelize([[0.5], [2.5], [1.5]])
14
15predictions = rows.map(lambda row: int(broadcast_model.value.predict([row])[0])).collect()
16print(predictions)
17
18spark.stop()

This works because the model is small enough to ship to executors and cheap enough to run per record. It is a reasonable pattern for classic tabular inference.

Prefer mapPartitions When Model Setup Is Expensive

If loading the model takes noticeable time, do not do it once per row. Do it once per partition.

1from pyspark.sql import SparkSession
2
3spark = SparkSession.builder.master("local[*]").appName("partition-model").getOrCreate()
4sc = spark.sparkContext
5
6def score_partition(rows):
7    multiplier = 10
8    for row in rows:
9        yield row * multiplier
10
11rdd = sc.parallelize([1, 2, 3, 4], 2)
12print(rdd.mapPartitions(score_partition).collect())
13
14spark.stop()

The example is intentionally simple, but the structure is what matters. In real code, score_partition would load a model once, process all rows in the iterator, and yield predictions. This is usually the right place for deep-learning inference or any model with nontrivial initialization cost.

Use Spark ML APIs for Training When Possible

If you truly want distributed training over Spark data, prefer Spark-native algorithms or frameworks that integrate explicitly with Spark. That keeps optimization, shuffling, and aggregation aligned with Spark’s execution engine.

Trying to call a standalone training routine inside map usually creates many unrelated local models instead of one global model. Unless that is your actual goal, it is the wrong abstraction.

A quick rule is:

• use map or mapPartitions for inference,
• use MLlib or another integrated training framework for training,
• and only train per partition when your use case truly wants one local model per partition.

Watch Serialization and Dependency Boundaries

A model object that works locally may fail on executors because it cannot be pickled cleanly, because a native dependency is missing, or because version differences exist between driver and workers. That is why Spark inference jobs should be tested in an environment that resembles the actual cluster, not only in local mode.

Also remember that Python UDF overhead, JVM-to-Python boundaries, and model size can dominate runtime. Sometimes the right optimization is not changing the mapping function. It is changing how and where the model is served.

Common Pitfalls

• Training a new model inside every map call.
• Broadcasting a large model that should instead be loaded once per partition.
• Assuming local tests prove cluster serialization and dependency setup are correct.
• Using per-row Python inference when vectorized or partition-based scoring would be much cheaper.
• Treating Spark as a general task queue instead of a distributed data engine.

Summary

• Running ML code inside Spark is normal for inference, but not all placements are equally efficient.
• 'map is acceptable for cheap scoring functions; mapPartitions is better when model setup is expensive.'
• Broadcast small, serializable models for simple inference workflows.
• Use Spark-native ML tooling for distributed training rather than training inside map.
• Optimize around executor setup, serialization, and partition behavior, not just the prediction function itself.