PySpark for Machine Learning

Table of Contents

Introduction

Apache Spark is the de facto standard for large-scale distributed data processing. Its machine learning library, MLlib, provides scalable implementations of common ML algorithms that run seamlessly on clusters of hundreds of machines.

PySpark, the Python API for Spark, allows data scientists to use familiar Python syntax while leveraging the full power of Spark's distributed computing engine. This tutorial covers the complete PySpark ML workflow: from creating a SparkSession and manipulating DataFrames, through building classification, regression, and clustering models, to deploying models with Spark Structured Streaming.

Whether you are processing gigabytes or terabytes of data, the patterns in this tutorial apply equally.

Prerequisites

• Python 3.9+
• Java 11 or 17 (required by Spark)
• Basic understanding of machine learning concepts

Install PySpark:

pip install pyspark numpy pandas

For a cluster deployment, ensure Spark is installed on all nodes. For local development, PySpark includes a built-in standalone Spark instance.

Spark Basics and SparkSession

Every PySpark application begins with a SparkSession, which is the unified entry point for all Spark functionality.

Creating a SparkSession

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("PySpark ML Tutorial") \
    .master("local[]") \

    .config("spark.driver.memory", "4g") \
    .config("spark.sql.shuffle.partitions", "8") \
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \
    .getOrCreate()

Verify the session
print(f"Spark version: {spark.version}")
print(f"App name: {spark.sparkContext.appName}")
print(f"Master: {spark.sparkContext.master}")
print(f"Default parallelism: {spark.sparkContext.defaultParallelism}")

from pyspark.sql.types import StructType, StructField, FloatType, IntegerType, StringType
import numpy as np

Create from Python data
data = [
    (1, "Alice", 28, 55000.0, "Engineering"),
    (2, "Bob", 35, 72000.0, "Marketing"),
    (3, "Carol", 42, 88000.0, "Engineering"),
    (4, "Dave", 31, 61000.0, "Sales"),
    (5, "Eve", 26, 48000.0, "Marketing"),
]

schema = StructType([
    StructField("id", IntegerType(), False),
    StructField("name", StringType(), False),
    StructField("age", IntegerType(), False),
    StructField("salary", FloatType(), False),
    StructField("department", StringType(), False),
])

df = spark.createDataFrame(data, schema=schema)
df.show()
df.printSchema()

Read from CSV
dfcsv = spark.read.csv("data.csv", header=True, inferSchema=True)


Read from Parquet (optimized columnar format)
dfparquet = spark.read.parquet("data.parquet")

Read from JSON
dfjson = spark.read.json("data.json")

from pyspark.sql import functions as F

Select and filter
df.select("name", "age", "salary").filter(F.col("age") > 30).show()

Add computed columns
df = df.withColumn("salarymonthly", F.col("salary") / 12)
df = df.withColumn("agegroup", F.when(F.col("age") < 30, "Young")

                                  .when(F.col("age") < 40, "Mid")
                                  .otherwise("Senior"))

Aggregations
df.groupBy("department").agg(
    F.avg("salary").alias("avgsalary"),
    F.count("").alias("count"),
    F.max("age").alias("maxage"),

).show()

Handle missing values
dfclean = df.na.fill({"salary": 0.0, "department": "Unknown"})
dfclean = df.na.drop(subset=["age", "salary"])


Sort
df.orderBy(F.col("salary").desc()).show()

Statistics
df.describe("age", "salary").show()
print(f"Correlation (age, salary): {df.stat.corr('age', 'salary'):.4f}")

from pyspark.ml.feature import StringIndexer, OneHotEncoder

Generate a larger synthetic dataset
np.random.seed(42)
departments = ["Engineering", "Marketing", "Sales", "HR", "Finance"]
rows = [(int(i), str(np.random.choice(departments)),
         int(np.random.randint(22, 60)),
         float(np.random.uniform(30000, 120000)),
         int(np.random.choice([0, 1])))
        for i in range(10000)]

schema = StructType([
    StructField("id", IntegerType()), StructField("department", StringType()),
    StructField("age", IntegerType()), StructField("salary", FloatType()),
    StructField("label", IntegerType()),
])
dfml = spark.createDataFrame(rows, schema)

String indexing
deptindexer = StringIndexer(inputCol="department", outputCol="deptindex")
dfml = deptindexer.fit(dfml).transform(dfml)

One-hot encoding
deptencoder = OneHotEncoder(inputCol="deptindex", outputCol="deptonehot")

dfml = deptencoder.fit(dfml).transform(dfml)


dfml.select("department", "deptindex", "deptonehot").show(5, truncate=False)

Scaling and Normalizing Numeric Features

from pyspark.ml.feature import VectorAssembler, StandardScaler, MinMaxScaler

Assemble numeric features into a single vector
assembler = VectorAssembler(
    inputCols=["age", "salary"],
    outputCol="numericfeatures"

)
dfml = assembler.transform(dfml)


Standard scaling (zero mean, unit variance)
stdscaler = StandardScaler(
    inputCol="numericfeatures",

    outputCol="scaledfeatures",
    withMean=True,
    withStd=True,
)
scalermodel = stdscaler.fit(dfml)

dfml = scalermodel.transform(dfml)

MinMax scaling
minmaxscaler = MinMaxScaler(inputCol="numericfeatures", outputCol="minmaxfeatures")

dfml = minmaxscaler.fit(dfml).transform(dfml)

from pyspark.ml.feature import Bucketizer, PCA, PolynomialExpansion

Bucketizing continuous features
bucketizer = Bucketizer(
    splits=[0, 25, 35, 45, 100],
    inputCol="age",
    outputCol="agebucket"
)
dfml = bucketizer.transform(dfml)

PCA for dimensionality reduction
pca = PCA(k=2, inputCol="scaledfeatures", outputCol="pcafeatures")
pcamodel = pca.fit(dfml)
dfml = pcamodel.transform(dfml)

print(f"PCA explained variance: {pcamodel.explainedVariance}")

Polynomial expansion
poly = PolynomialExpansion(degree=2, inputCol="numericfeatures", outputCol="polyfeatures")
dfml = poly.transform(dfml)

Classification with Spark MLlib

Logistic Regression

from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.evaluation import BinaryClassificationEvaluator, MulticlassClassificationEvaluator

Prepare features
featurecols = ["age", "salary", "deptindex"]
assembler = VectorAssembler(inputCols=featurecols, outputCol="features")

dfprepared = assembler.transform(dfml)


Split data
traindf, testdf = dfprepared.randomSplit([0.8, 0.2], seed=42)

Train logistic regression
lr = LogisticRegression(
    featuresCol="features",
    labelCol="label",
    maxIter=100,
    regParam=0.01,
    elasticNetParam=0.8,
)
lrmodel = lr.fit(traindf)

Predictions
predictions = lrmodel.transform(testdf)
predictions.select("label", "prediction", "probability").show(10, truncate=False)

Evaluation
binaryeval = BinaryClassificationEvaluator(labelCol="label")

auc = binaryeval.evaluate(predictions, {binaryeval.metricName: "areaUnderROC"})


multieval = MulticlassClassificationEvaluator(labelCol="label")
accuracy = multieval.evaluate(predictions, {multieval.metricName: "accuracy"})
f1 = multieval.evaluate(predictions, {multieval.metricName: "f1"})

print(f"AUC-ROC:  {auc:.4f}")
print(f"Accuracy: {accuracy:.4f}")
print(f"F1 Score: {f1:.4f}")

Random Forest Classifier

from pyspark.ml.classification import RandomForestClassifier

rf = RandomForestClassifier(
    featuresCol="features",
    labelCol="label",
    numTrees=100,
    maxDepth=10,
    featureSubsetStrategy="sqrt",
    seed=42,
)
rfmodel = rf.fit(traindf)

rfpredictions = rfmodel.transform(testdf)

rfaccuracy = multieval.evaluate(rfpredictions, {multieval.metricName: "accuracy"})

print(f"Random Forest Accuracy: {rfaccuracy:.4f}")

Feature importance
importance = rfmodel.featureImportances.toArray()

for col, imp in zip(featurecols, importance):
    print(f"  {col}: {imp:.4f}")

Gradient-Boosted Trees

from pyspark.ml.classification import GBTClassifier

gbt = GBTClassifier(
    featuresCol="features",
    labelCol="label",
    maxIter=50,
    maxDepth=5,
    stepSize=0.1,
    seed=42,
)
gbtmodel = gbt.fit(traindf)

gbtpredictions = gbtmodel.transform(testdf)

gbtaccuracy = multieval.evaluate(gbtpredictions, {multieval.metricName: "accuracy"})

print(f"GBT Accuracy: {gbtaccuracy:.4f}")

Regression with Spark MLlib

Linear Regression

from pyspark.ml.regression import LinearRegression
from pyspark.ml.evaluation import RegressionEvaluator

Prepare regression data (predict salary from age and department)
regassembler = VectorAssembler(inputCols=["age", "deptindex"], outputCol="features")
dfreg = regassembler.transform(dfml)


trainreg, testreg = dfreg.randomSplit([0.8, 0.2], seed=42)

linreg = LinearRegression(

    featuresCol="features",
    labelCol="salary",
    maxIter=100,
    regParam=0.01,
    elasticNetParam=0.5,
)
linmodel = linreg.fit(trainreg)

print(f"Coefficients: {linmodel.coefficients}")

print(f"Intercept: {linmodel.intercept:.2f}")

regpredictions = linmodel.transform(testreg)


regeval = RegressionEvaluator(labelCol="salary")
rmse = regeval.evaluate(regpredictions, {regeval.metricName: "rmse"})

mae = regeval.evaluate(regpredictions, {regeval.metricName: "mae"})
r2 = regeval.evaluate(regpredictions, {regeval.metricName: "r2"})


print(f"RMSE: {rmse:.2f}")
print(f"MAE:  {mae:.2f}")
print(f"R2:   {r2:.4f}")

from pyspark.ml.regression import RandomForestRegressor

rfreg = RandomForestRegressor(
    featuresCol="features",
    labelCol="salary",
    numTrees=100,
    maxDepth=10,
    seed=42,
)
rfregmodel = rfreg.fit(trainreg)
rfregpreds = rfregmodel.transform(testreg)


rfrmse = regeval.evaluate(rfregpreds, {regeval.metricName: "rmse"})
print(f"RF Regression RMSE: {rfrmse:.2f}")

from pyspark.ml.clustering import KMeans
from pyspark.ml.evaluation import ClusteringEvaluator

clusterassembler = VectorAssembler(inputCols=["age", "salary"], outputCol="features")
dfcluster = clusterassembler.transform(dfml)


Standard scale features for clustering
from pyspark.ml.feature import StandardScaler
clusterscaler = StandardScaler(inputCol="features", outputCol="scaledfeatures", withMean=True, withStd=True)

dfcluster = clusterscaler.fit(dfcluster).transform(dfcluster)


Find optimal k using silhouette score
silhouettescores = []
for k in range(2, 8):
    kmeans = KMeans(featuresCol="scaledfeatures", k=k, seed=42, maxIter=50)

    kmmodel = kmeans.fit(dfcluster)

    predictions = kmmodel.transform(dfcluster)

    evaluator = ClusteringEvaluator(featuresCol="scaledfeatures")
    score = evaluator.evaluate(predictions)
    silhouettescores.append((k, score))

    print(f"k={k}: Silhouette Score = {score:.4f}")

Train with optimal k
optimalk = max(silhouettescores, key=lambda x: x[1])[0]

print(f"\nOptimal k: {optimalk}")

kmeansfinal = KMeans(featuresCol="scaledfeatures", k=optimalk, seed=42)

kmfinalmodel = kmeansfinal.fit(dfcluster)


Cluster centers
for i, center in enumerate(kmfinalmodel.clusterCenters()):

    print(f"Cluster {i}: {center}")

from pyspark.ml.clustering import GaussianMixture

gmm = GaussianMixture(featuresCol="scaledfeatures", k=3, seed=42, maxIter=100)
gmmmodel = gmm.fit(dfcluster)

gmmpredictions = gmmmodel.transform(dfcluster)

gmmpredictions.select("age", "salary", "prediction", "probability").show(10, truncate=False)

Pipeline API

Spark ML Pipelines chain multiple transformers and estimators into a single reproducible workflow.

Building a Complete Pipeline

from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler, StandardScaler
from pyspark.ml.classification import RandomForestClassifier

Define pipeline stages
deptindexer = StringIndexer(inputCol="department", outputCol="deptindex", handleInvalid="keep")
deptencoder = OneHotEncoder(inputCol="deptindex", outputCol="deptencoded")

assembler = VectorAssembler(inputCols=["age", "salary", "deptencoded"], outputCol="rawfeatures")

scaler = StandardScaler(inputCol="rawfeatures", outputCol="features", withMean=False, withStd=True)
classifier = RandomForestClassifier(
    featuresCol="features", labelCol="label", numTrees=100, maxDepth=10, seed=42
)

Create and fit the pipeline
pipeline = Pipeline(stages=[deptindexer, deptencoder, assembler, scaler, classifier])

trainraw, testraw = dfml.randomSplit([0.8, 0.2], seed=42)

pipelinemodel = pipeline.fit(trainraw)


Predict using the full pipeline
pipelinepredictions = pipelinemodel.transform(testraw)
pipelinepredictions.select("department", "age", "salary", "label", "prediction").show(10)


accuracy = MulticlassClassificationEvaluator(labelCol="label").evaluate(
    pipelinepredictions, {MulticlassClassificationEvaluator.metricName: "accuracy"}
)
print(f"Pipeline Accuracy: {accuracy:.4f}")

Cross-Validation and Hyperparameter Tuning

Grid Search with Cross-Validation

from pyspark.ml.tuning import CrossValidator, ParamGridBuilder

Build parameter grid
paramgrid = ParamGridBuilder() \

    .addGrid(classifier.numTrees, [50, 100, 200]) \
    .addGrid(classifier.maxDepth, [5, 10, 15]) \
    .addGrid(classifier.featureSubsetStrategy, ["sqrt", "log2"]) \
    .build()

print(f"Total parameter combinations: {len(paramgrid)}")

Create cross-validator
evaluator = BinaryClassificationEvaluator(labelCol="label", metricName="areaUnderROC")

cv = CrossValidator(
    estimator=pipeline,
    estimatorParamMaps=paramgrid,

    evaluator=evaluator,
    numFolds=5,
    parallelism=4,       # evaluate 4 models in parallel
    seed=42,
)

Fit cross-validator
cvmodel = cv.fit(trainraw)


Best model results
bestmodel = cvmodel.bestModel

bestrf = bestmodel.stages[-1]  # last stage is the classifier

print(f"Best numTrees: {bestrf.getNumTrees}")
print(f"Best maxDepth: {bestrf.getOrDefault('maxDepth')}")


cvpredictions = cvmodel.transform(testraw)
cvauc = evaluator.evaluate(cvpredictions)
print(f"Cross-Validated AUC: {cvauc:.4f}")


All fold metrics
print("\nAverage metrics per parameter combination:")
for i, metric in enumerate(cvmodel.avgMetrics):
    print(f"  Combination {i}: AUC = {metric:.4f}")

Train-Validation Split (Faster Alternative)

from pyspark.ml.tuning import TrainValidationSplit

tvs = TrainValidationSplit(
    estimator=pipeline,
    estimatorParamMaps=paramgrid,

    evaluator=evaluator,
    trainRatio=0.8,
    parallelism=4,
    seed=42,
)

tvsmodel = tvs.fit(trainraw)

tvspredictions = tvsmodel.transform(testraw)
tvsauc = evaluator.evaluate(tvspredictions)
print(f"TVS AUC: {tvsauc:.4f}")

# Save the entire pipeline model
pipelinemodel.write().overwrite().save("models/rfpipeline")


Save the cross-validated best model
cvmodel.bestModel.write().overwrite().save("models/bestcvmodel")

Load models
from pyspark.ml import PipelineModel

loadedpipeline = PipelineModel.load("models/rfpipeline")
loadedpredictions = loadedpipeline.transform(testraw)

loadedaccuracy = MulticlassClassificationEvaluator(labelCol="label").evaluate(
    loadedpredictions, {MulticlassClassificationEvaluator.metricName: "accuracy"}

)
print(f"Loaded model accuracy: {loadedaccuracy:.4f}")

Exporting for Non-Spark Environments

# Convert Spark ML model to PMML or ONNX for deployment outside Spark
Option 1: Extract model parameters manually
rfmodel = pipelinemodel.stages[-1]
print(f"Number of trees: {rfmodel.getNumTrees}")

print(f"Feature importances: {rfmodel.featureImportances}")

Option 2: Use MLeap for portable serialization
pip install mleap
from mleap.pyspark.sparksupport import SimpleSparkSerializer

pipelinemodel.serializeToBundle("jar:file:/path/to/model.zip", testraw)

from pyspark.sql.types import StructType, StructField, StringType, IntegerType, FloatType

Define schema for incoming stream data
streamschema = StructType([
    StructField("id", IntegerType()),
    StructField("department", StringType()),
    StructField("age", IntegerType()),
    StructField("salary", FloatType()),
])

Read from a streaming source (e.g., files arriving in a directory)
streamingdf = spark.readStream \

    .schema(streamschema) \
    .option("maxFilesPerTrigger", 1) \
    .json("data/streaminput/")


Apply the trained pipeline model to streaming data
streamingpredictions = loadedpipeline.transform(streamingdf)

Write predictions to output
query = streamingpredictions \

    .select("id", "department", "age", "salary", "prediction") \
    .writeStream \
    .outputMode("append") \
    .format("json") \
    .option("path", "data/streamoutput/") \
    .option("checkpointLocation", "data/checkpoints/") \
    .trigger(processingTime="10 seconds") \
    .start()

Monitor the stream
print(f"Stream is active: {query.isActive}")
print(f"Stream status: {query.status}")

query.awaitTermination()  # uncomment to run indefinitely
query.stop()  # stop the stream

Kafka Integration

# Read from Kafka
kafkadf = spark.readStream \

    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "ml-input") \
    .load()

Parse the Kafka value (JSON string) into structured data
from pyspark.sql.functions import fromjson, col

parseddf = kafkadf.select(
    fromjson(col("value").cast("string"), streamschema).alias("data")
).select("data.")


Apply ML model
predictions = loadedpipeline.transform(parseddf)

Write predictions back to Kafka
query = predictions.select(
    col("id").cast("string").alias("key"),
    col("prediction").cast("string").alias("value"),
).writeStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("topic", "ml-output") \
    .option("checkpointLocation", "data/kafkacheckpoints/") \

    .start()

traindf.cache()
print(f"Cached partitions: {traindf.rdd.getNumPartitions()}")

from pyspark.sql.functions import broadcast
result = largedf.join(broadcast(small_df), on="key")

Conclusion

This tutorial covered the complete PySpark machine learning workflow:

• SparkSession and DataFrames provide the foundation for distributed data processing with a SQL-like API.
• Feature engineering at scale is handled by MLlib transformers including StringIndexer, OneHotEncoder, VectorAssembler, and StandardScaler.
• Classification, regression, and clustering algorithms in MLlib scale horizontally across clusters while maintaining familiar scikit-learn-like interfaces.
• Pipelines chain all preprocessing and modeling steps into a single, reproducible, and serializable object.
• Cross-validation with ParamGridBuilder enables systematic hyperparameter tuning with parallel evaluation.
• Model persistence allows saving and loading complete pipeline models for production deployment.
• Structured Streaming enables real-time inference by applying trained models to live data streams from files or Kafka.

PySpark is the right tool when your data exceeds the capacity of a single machine or when you need the operational benefits of Spark's ecosystem (fault tolerance, streaming, and integration with Hadoop, Hive, and Kafka). For smaller datasets, scikit-learn remains simpler and faster. The key is choosing the right tool for the scale of your problem.