Apache Spark for Large-Scale Processing

Introduction

Apache Spark has become the de-facto engine for large-scale data processing and distributed machine learning. Originally developed at UC Berkeley's AMPLab in 2009 and open-sourced through the Apache Software Foundation, Spark replaced MapReduce as the primary big-data compute layer by offering in-memory processing that can be up to 100× faster than disk-based Hadoop jobs.

For ML engineers, Spark is indispensable. It powers feature engineering pipelines that process terabytes of raw data, runs distributed model training through MLlib, and integrates natively with data lakes, feature stores, and orchestration tools like Airflow. Understanding Spark internals — not just its API — is essential for building performant, cost-efficient ML systems.

This post covers Spark from the ground up: its cluster architecture, the evolution from RDDs to DataFrames and Datasets, Spark SQL for declarative analytics, MLlib for distributed training, feature engineering patterns at scale, Kubernetes-native deployment, and battle-tested performance tuning strategies.

Spark Architecture

Spark follows a master–worker architecture. Every Spark application consists of a driver process and a set of executor processes coordinated by a cluster manager. Understanding these three components is crucial for debugging, tuning, and capacity planning.

Driver

The driver is the JVM process that runs your main() function. It is responsible for three things: (1) maintaining the SparkSession / SparkContext, (2) converting your high-level API calls into a physical execution plan (a DAG of stages and tasks), and (3) scheduling tasks on executors and collecting results. The driver holds all metadata — lineage graphs, broadcast variables, and accumulator values. If the driver OOMs, your entire application fails.

Executors

Executors are JVM processes launched on worker nodes. Each executor is allocated a fixed amount of memory and CPU cores. Executors run the actual computation (tasks), store intermediate data in memory or on disk, and report status back to the driver. A single executor can run multiple tasks concurrently — one per core slot.

Cluster Manager

The cluster manager allocates resources across applications. Spark supports four cluster managers: Standalone (built-in), YARN (Hadoop), Mesos (deprecated), and Kubernetes (now the recommended cloud-native option). The cluster manager is external to Spark itself — it only manages container/process lifecycle.

Execution Model: DAG → Stages → Tasks

When you call an action (e.g., .count(), .collect(), .write), Spark's Catalyst optimizer compiles the logical plan into an optimized physical plan expressed as a DAG (Directed Acyclic Graph). The DAG is split at shuffle boundaries into stages. Each stage is divided into tasks — one task per partition. Tasks are the smallest unit of work shipped to executors.

# Conceptual execution flow
User Code  →  Logical Plan  →  Catalyst Optimizer  →  Physical Plan
                                                          ↓
                                                   DAG of Stages
                                                          ↓
                                              Stage 1: [Task₁, Task₂, …, Taskₙ]
                                              Stage 2: [Task₁, Task₂, …, Taskₘ]
                                                          ↓
                                              Scheduled on Executors

RDDs, DataFrames & Datasets

Spark's API has evolved through three generations. Each layer adds more structure, enabling progressively better optimization while preserving fault tolerance through lineage-based recomputation.

RDDs (Resilient Distributed Datasets)

RDDs are Spark's original abstraction — an immutable, partitioned collection of objects distributed across the cluster. They support two types of operations: transformations (lazy, produce a new RDD) and actions (trigger execution, return a result). RDDs give you fine-grained control over partitioning and data placement, but they are opaque to Spark's optimizer — the engine cannot inspect or rearrange your lambda functions.

# RDD example — word count
from pyspark import SparkContext

sc = SparkContext("local[*]", "WordCount")
lines = sc.textFile("s3://data/books/*.txt")

counts = (
    lines.flatMap(lambda line: line.split(" "))
         .map(lambda word: (word.lower(), 1))
         .reduceByKey(lambda a, b: a + b)
)

counts.saveAsTextFile("s3://output/word-counts")

DataFrames

DataFrames (introduced in Spark 1.3) are distributed tables with named columns and known types. Because the schema is explicit, the Catalyst optimizer can apply predicate pushdown, column pruning, constant folding, and join reordering — optimizations impossible with raw RDDs. DataFrames are the primary API for PySpark users.

# DataFrame example — aggregation
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder \
    .appName("SalesAgg") \
    .getOrCreate()

df = spark.read.parquet("s3://lake/sales/")

summary = (
    df.filter(F.col("amount") > 0)
      .groupBy("region", "product_category")
      .agg(
          F.sum("amount").alias("total_sales"),
          F.count("order_id").alias("num_orders"),
          F.avg("amount").alias("avg_order_value")
      )
      .orderBy(F.desc("total_sales"))
)

summary.write.mode("overwrite").parquet("s3://lake/sales_summary/")

Datasets (Scala/Java only)

Datasets combine the strong typing of RDDs with the optimization benefits of DataFrames. They are available only in Scala and Java — PySpark exposes DataFrames exclusively. In Scala, DataFrame is simply an alias for Dataset[Row].

Spark SQL

Spark SQL lets you run ANSI SQL queries directly against DataFrames, Hive tables, Parquet/ORC files, and JDBC sources. Under the hood, the same Catalyst optimizer compiles SQL text into the identical physical plan that the DataFrame API would produce. This means there is zero performance difference between writing SQL and writing DataFrame chains — choose whichever is more readable.

Registering Views and Querying

# Register a DataFrame as a temporary SQL view
df.createOrReplaceTempView("sales")

result = spark.sql("""
    SELECT region,
           product_category,
           SUM(amount)   AS total_sales,
           COUNT(*)      AS num_orders,
           AVG(amount)   AS avg_order_value
    FROM   sales
    WHERE  amount > 0
    GROUP  BY region, product_category
    ORDER  BY total_sales DESC
""")

result.show(20, truncate=False)

Catalog and Metastore Integration

Spark SQL integrates with the Hive Metastore, AWS Glue Data Catalog, and Unity Catalog (Databricks) to provide a persistent table namespace. This enables schema evolution, partition discovery, and cross-engine interoperability (Presto, Trino, Athena can read the same tables).

Adaptive Query Execution (AQE)

Spark 3.0+ introduced AQE, which re-optimizes the query plan at runtime based on actual shuffle statistics. AQE can dynamically coalesce small post-shuffle partitions, switch join strategies (sort-merge → broadcast) when one side is small, and optimize skewed joins by splitting hot partitions. Enable it with:

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

MLlib for Distributed ML

MLlib is Spark's built-in machine learning library. It provides distributed implementations of common algorithms — classification, regression, clustering, collaborative filtering, and dimensionality reduction — all operating on DataFrames through a Pipeline API inspired by scikit-learn.

The Pipeline API

An MLlib pipeline chains together Transformers (e.g., StringIndexer, VectorAssembler) and Estimators (e.g., LogisticRegression, GBTClassifier) into a single reproducible workflow. Calling .fit() on the pipeline produces a PipelineModel that can be serialized and deployed.

# End-to-end MLlib classification pipeline
from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, VectorAssembler, StandardScaler
from pyspark.ml.classification import GBTClassifier
from pyspark.ml.evaluation import BinaryClassificationEvaluator

# Load data
df = spark.read.parquet("s3://lake/churn_features/")
train, test = df.randomSplit([0.8, 0.2], seed=42)

# Feature pipeline
indexer   = StringIndexer(inputCol="plan_type", outputCol="plan_idx")
assembler = VectorAssembler(
    inputCols=["plan_idx", "tenure_months", "monthly_charges",
               "total_charges", "support_tickets"],
    outputCol="raw_features"
)
scaler = StandardScaler(inputCol="raw_features", outputCol="features")
gbt    = GBTClassifier(labelCol="churned", featuresCol="features", maxIter=100)

pipeline = Pipeline(stages=[indexer, assembler, scaler, gbt])

# Train
model = pipeline.fit(train)

# Evaluate
predictions = model.transform(test)
evaluator   = BinaryClassificationEvaluator(
    labelCol="churned", metricName="areaUnderROC"
)
auc = evaluator.evaluate(predictions)
print(f"Test AUC-ROC: {auc:.4f}")

# Save model
model.write().overwrite().save("s3://models/churn_gbt_v1/")

Hyperparameter Tuning with CrossValidator

MLlib's CrossValidator and TrainValidationSplit perform distributed hyperparameter search. Each fold/split is trained as a separate Spark job, utilizing the full cluster for parallelism. Use ParamGridBuilder to define the search space.

from pyspark.ml.tuning import CrossValidator, ParamGridBuilder

param_grid = ParamGridBuilder() \
    .addGrid(gbt.maxDepth, [5, 8, 12]) \
    .addGrid(gbt.stepSize, [0.05, 0.1, 0.2]) \
    .addGrid(gbt.subsamplingRate, [0.7, 0.9]) \
    .build()

cv = CrossValidator(
    estimator=pipeline,
    estimatorParamMaps=param_grid,
    evaluator=evaluator,
    numFolds=3,
    parallelism=4   # train 4 models concurrently
)

cv_model = cv.fit(train)
best_auc = evaluator.evaluate(cv_model.transform(test))
print(f"Best CV AUC-ROC: {best_auc:.4f}")

Supported Algorithms

Feature Engineering at Scale

In ML systems, feature engineering often accounts for 60–80% of total development time. Spark excels here because it can process billions of rows with complex joins, window functions, and aggregations that would be impossible in-memory on a single machine.

Common Feature Patterns

• Temporal aggregations — rolling averages, exponential decay, time-since-event
• Cross-entity joins — user × product interaction matrices
• Window functions — rank, lag, lead, running sums per group
• Text featurization — TF-IDF, word2vec, n-gram hashing via MLlib transformers
• Geospatial binning — H3/S2 cell assignment for location features

# Window-based feature engineering
from pyspark.sql import Window
from pyspark.sql import functions as F

# Define a 30-day rolling window per user
w30 = Window.partitionBy("user_id") \
            .orderBy(F.col("event_ts").cast("long")) \
            .rangeBetween(-30 * 86400, 0)

features = (
    events
    .withColumn("purchases_30d",
        F.count(F.when(F.col("event_type") == "purchase", 1)).over(w30))
    .withColumn("spend_30d",
        F.sum(F.when(F.col("event_type") == "purchase", F.col("amount"))).over(w30))
    .withColumn("avg_session_dur",
        F.avg("session_duration").over(w30))
    .withColumn("days_since_last_purchase",
        F.datediff(F.col("event_ts"), F.max(
            F.when(F.col("event_type") == "purchase", F.col("event_ts"))
        ).over(w30)))
)

# Persist to feature store (Delta Lake)
features.write.format("delta") \
    .mode("overwrite") \
    .partitionBy("ds") \
    .save("s3://feature-store/user_engagement/")

Avoiding Skew in Feature Joins

When joining a large fact table against a dimension with high-cardinality keys, data skew can cause a single task to process disproportionately more data. Mitigation strategies include: salting (appending a random suffix to the skewed key), broadcast joins (if the small side fits in memory), and leveraging AQE's automatic skew handling.

# Salting technique to handle skewed joins
import random

SALT_BUCKETS = 10

# Salt the large (skewed) side
large_df = large_df.withColumn(
    "salt", (F.rand() * SALT_BUCKETS).cast("int")
).withColumn(
    "salted_key", F.concat(F.col("join_key"), F.lit("_"), F.col("salt"))
)

# Explode the small side to match all salt values
small_df = small_df.crossJoin(
    spark.range(SALT_BUCKETS).withColumnRenamed("id", "salt")
).withColumn(
    "salted_key", F.concat(F.col("join_key"), F.lit("_"), F.col("salt"))
)

joined = large_df.join(small_df, "salted_key", "inner").drop("salt", "salted_key")

Spark on Kubernetes

Since Spark 2.3 (GA in 3.1), Kubernetes has been a first-class cluster manager. Running Spark on K8s brings significant operational advantages: no dedicated Hadoop/YARN cluster, unified infrastructure with other microservices, autoscaling via the Horizontal Pod Autoscaler, GPU scheduling for DL workloads, and multi-tenant isolation through namespaces and resource quotas.

How It Works

When you spark-submit with --master k8s://, Spark creates a driver pod. The driver pod then requests executor pods from the Kubernetes API server. Each executor pod runs a single JVM. When the job finishes, executor pods are deleted automatically. The driver pod persists until you retrieve logs.

# spark-submit on Kubernetes
spark-submit \
  --master k8s://https://k8s-api.internal:6443 \
  --deploy-mode cluster \
  --name feature-pipeline \
  --conf spark.kubernetes.namespace=ml-pipelines \
  --conf spark.kubernetes.container.image=registry.io/spark:3.5.1 \
  --conf spark.kubernetes.authenticate.driver.serviceAccountName=spark-sa \
  --conf spark.executor.instances=20 \
  --conf spark.executor.memory=8g \
  --conf spark.executor.cores=4 \
  --conf spark.driver.memory=4g \
  --conf spark.kubernetes.executor.volumes.persistentVolumeClaim.spark-data.mount.path=/data \
  --conf spark.kubernetes.executor.volumes.persistentVolumeClaim.spark-data.options.claimName=spark-pvc \
  s3a://code/feature_pipeline.py

Dynamic Allocation on K8s

Spark supports dynamic allocation on Kubernetes, meaning executor pods are created and terminated based on workload demand. Idle executors are removed after a configurable timeout, saving cluster resources. This is essential for cost efficiency in cloud environments where you pay per vCPU-second.

# Enable dynamic allocation on K8s
spark.conf.set("spark.dynamicAllocation.enabled", "true")
spark.conf.set("spark.dynamicAllocation.shuffleTracking.enabled", "true")
spark.conf.set("spark.dynamicAllocation.minExecutors", "2")
spark.conf.set("spark.dynamicAllocation.maxExecutors", "50")
spark.conf.set("spark.dynamicAllocation.executorIdleTimeout", "60s")

Spark Operator for K8s

The spark-on-k8s-operator (GoogleCloudPlatform/spark-on-k8s-operator) is a Kubernetes operator that manages Spark applications as custom resources (SparkApplication CRDs). It handles driver/executor lifecycle, automatic retries, scheduled jobs (cron), and integrates with Prometheus for metrics. This is the recommended approach for production Spark-on-K8s deployments.

Performance Tuning

Tuning Spark is part science, part art. The key is understanding the relationship between partitions, memory, serialization, and shuffle. Below are the most impactful levers.

1. Partition Count

The number of partitions determines parallelism. Too few partitions underutilize the cluster; too many create scheduling overhead and small-file problems. A good heuristic: 2–4 partitions per executor core, with each partition holding 128–256 MB of data. Use .repartition(n) (full shuffle) or .coalesce(n) (narrow dependency, reduce-only) to adjust.

2. Memory Configuration

Executor memory is divided into: execution memory (shuffles, sorts, joins), storage memory (cached DataFrames), and overhead (off-heap, PySpark interop). The unified memory manager (default since Spark 1.6) dynamically balances execution and storage. Key settings:

# Memory tuning checklist
spark.executor.memory          = 8g      # JVM heap
spark.executor.memoryOverhead  = 2g      # off-heap (PySpark needs more)
spark.memory.fraction          = 0.6     # fraction of heap for execution+storage
spark.memory.storageFraction   = 0.5     # initial split (dynamic by default)
spark.sql.shuffle.partitions   = 200     # default shuffle partitions (tune!)

3. Serialization

Spark uses Java serialization by default, which is slow and verbose. Switching to Kryo serialization can reduce shuffle sizes by 2–5× and improve task throughput. Register your custom classes for best results.

spark.conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
spark.conf.set("spark.kryo.registrationRequired", "false")

4. Broadcast Joins

When one side of a join is small enough to fit in executor memory (typically < 200 MB), Spark can broadcast it to all executors, eliminating the shuffle entirely. This is one of the most effective optimizations.

# Force a broadcast join
from pyspark.sql.functions import broadcast

result = large_df.join(
    broadcast(small_lookup_df),
    "key_column",
    "left"
)

# Or set the auto-broadcast threshold
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "200m")

5. Caching Strategy

Use .cache() (alias for .persist(MEMORY_AND_DISK)) when a DataFrame is reused multiple times in the same job. But cache judiciously — cached data competes with execution memory. Always .unpersist() when done.

6. Avoiding Shuffles

Shuffles are the most expensive operation in Spark. They serialize data, write it to local disk, transfer it over the network, and deserialize it on the receiving end. Minimize shuffles by: co-partitioning data before joins, using broadcast joins, pre-aggregating before a groupBy, and avoiding .repartition() unless necessary.

Tuning Quick Reference

Summary

Apache Spark remains the backbone of large-scale data processing for ML. Its architecture — a driver coordinating executors via a cluster manager — enables linear horizontal scaling. The evolution from RDDs to DataFrames unlocked the Catalyst optimizer and Tungsten execution engine, delivering order-of-magnitude performance gains with simpler code.

MLlib provides production-ready distributed ML algorithms and a pipeline API that mirrors scikit-learn's design. For feature engineering, Spark's window functions, join optimizations, and integration with Delta Lake / feature stores make it indispensable. Running on Kubernetes brings cloud-native autoscaling, GPU support, and operational simplicity. And performance tuning — partition sizing, memory config, broadcast joins, AQE — is the difference between a job that takes 3 hours and one that takes 15 minutes.