Spark for Large Datasets — PySpark Guide for Data Engineers

"When pandas hits a wall, Spark hands you a map, a bus, and a helmet." — Your future cluster admin

Hook: Why Spark, now that you know pandas and SQLAlchemy?

You’ve been happily chaining pandas and SQLAlchemy queries in CPU-land, and you've trained models with PyTorch on tidy datasets. Then someone says: "We have 2 TB of logs, 1000 parquet partitions, and a weekly retrain schedule." Your laptop sighs. Your pandas pipeline quits on principle. Enter Spark — the distributed, fault-tolerant engine that scales ETL, analytics, and feature engineering across a cluster.

This guide shows how Spark (PySpark) fits into the pipeline between raw data sources and your PyTorch models or production APIs — with practical patterns, code snippets, and deployment pointers.

What Spark is (quick refresher)

• Spark is a distributed compute engine optimized for big data transformations and SQL-style operations.
• PySpark is Spark's Python API — you write Python, Spark executes across nodes.
• Key ideas: lazy evaluation, transformations vs actions, shuffles, and RDD/DataFrame abstractions.

Why it matters: Spark lets you preprocess multi-hundred-GB datasets (filtering, joins, aggregations, feature extraction) before handing a manageable sample or feature store to PyTorch.

Core concepts you’ll actually use

Transformations vs Actions

• Transformations (map/withColumn/filter/join) build a DAG — nothing runs yet.
• Actions (count, collect, write) trigger execution.

Partitioning & Shuffles

• Partitions = Spark’s unit of parallel work. Good partitioning reduces expensive shuffles.
• Joins and groupBy can cause shuffles — plan to minimize them.

Caching & Persistence

• Use cache() for repeated reads on the same intermediate DataFrame.
• Use the correct storage level (MEMORY_ONLY, MEMORY_AND_DISK).

Broadcast joins

• When one table is small (< ~100MB), broadcast it to every executor for a fast join.

File formats & predicate pushdown

• Use Parquet/ORC with columnar storage and predicate pushdown. Avoid CSV for large production datasets.

Quick PySpark cheatsheet (code)

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

spark = SparkSession.builder 
    .appName('feature-pipeline') 
    .config('spark.sql.shuffle.partitions', '200') 
    .getOrCreate()

# Read from S3/HDFS/Local parquet
df = spark.read.parquet('s3a://my-bucket/raw/events/')

# Light transforms
df = (df.filter(col('event_type') == 'purchase')
        .withColumn('amount_usd', col('amount') * col('rate_usd')))

# Cache if reused
df.cache()

# Broadcast join when right-side is small
from pyspark.sql.functions import broadcast
small = spark.read.parquet('s3a://my-bucket/dim/product_lookup/')
joined = df.join(broadcast(small), on='product_id')

# Write partitioned parquet for downstream use
joined.write.mode('overwrite').partitionBy('year', 'month').parquet('s3a://my-bucket/processed/purchases/')

Real-world patterns (from messy to model-ready)

1. Ingest: read from S3 / HDFS / Kafka / JDBC (yes, SQLAlchemy taught you the JDBC idea)
2. Clean & normalize: schema enforcement, type casts, UDFs (sparingly), vectorized functions
3. Feature engineering: aggregations, window functions, joins to dimension tables
4. Persist features: write as Parquet/Delta to a feature store or S3
5. Consume for training: either export to TFRecords/Parquet and use DataLoader or sample/collect for local training

Note: avoid toPandas() on full datasets — collect only sampled partitions or features. Use Spark to produce sharded TFRecords if your PyTorch training expects that format.

Performance tips & gotchas (so you don't cry at 3am)

• Schema first: define schemas to avoid expensive type inference.
• File sizes matter: aim for ~128MB–512MB per file to balance task overhead vs locality.
• Use partition pruning: write partitioned output by date/region and push filters into reads.
• Reduce wide transformations: chain filters early to cut data before joins.
• Avoid Python UDFs when possible — use built-in SQL functions or Pandas UDFs (vectorized) for better perf.
• Monitor the Spark UI: stages, tasks, and SQL tab show hotspots.

Integrating with ML workflows (PyTorch & MLlib)

• Spark is superb at generating large-scale features and sampling large corpora for training.
• For distributed ML inside Spark, there’s MLlib (logistic regression, tree ensembles) — but for deep learning you’ll typically:
  • Use Spark to preprocess and write Parquet/TFRecord shards to S3/HDFS
  • Use a distributed training framework (Horovod/PyTorch Lightning on GPU clusters)
  • Or export features to a feature store that your training job queries

Quick example: writing TFRecords from Spark partitions for PyTorch ingestion (pattern):

• Use mapPartitions to serialize examples, write each partition to a TFRecord/flat file.

Deployment & operationalization

• Run jobs via spark-submit, Airflow, or Databricks jobs.
• Cluster managers: YARN, Kubernetes, Mesos, or managed services (EMR, Databricks).
• Packaging: keep driver logic small; bundle dependencies with --py-files or use Docker images for Kubernetes.
• CI/CD: use your GitHub workflows (remember your Git/GitHub Workflows chapter) to build artifacts, run unit tests, and trigger job deployments.

Example spark-submit:

spark-submit \
  --master yarn \
  --deploy-mode cluster \
  --conf spark.executor.memory=8g \
  --conf spark.executor.cores=4 \
  my_job.py --date 2026-03-01

For Databricks or EMR, use their job APIs to schedule and monitor runs; add alerting on job failures.

Fault tolerance & observability

• Spark retries failed tasks automatically; durable storage (HDFS/S3) persists outputs.
• Use the Spark UI + cluster metrics + application logs for debugging.
• Add test datasets and unit tests for transformations (small sample inputs) to your GitHub CI.

When NOT to use Spark

• Small datasets (< a few GB) — pandas is simpler and often faster.
• Low-latency, single-record transactions — use streaming systems or databases.

Quick checklist before launching a Spark pipeline

• Defined schema for input data
• Parquet/Delta output with partitioning
• Broadcast small joins and reduce shuffles
• File sizes balanced (128–512MB rule of thumb)
• Caching applied for repeated reuse
• CI/CD pipeline triggers spark-submit (or managed job)
• Observability: metrics, logs, alerts

Takeaways — the TL;DR you can paste into a README

• Spark scales your pandas/SQL pipelines to big datasets and prepares model-scale features for deep learning.
• Use Parquet, partitioning, and broadcast joins to keep jobs fast and cheap.
• Deploy with spark-submit, Databricks, EMR, or Kubernetes and automate with your existing GitHub workflows.

Final thought: Spark doesn’t make your data clean — but it gives you the machinery to clean mountains of it without losing your sanity.

If you want, I can: show a full end-to-end PySpark notebook that reads raw logs, builds features, writes partitioned Parquet, and produces TFRecord shards ready for PyTorch. Pick a storage backend (S3 or HDFS) and your cluster manager and I’ll scaffold it.