Introduction
When we move from notebooks to production, Spark code that worked fine locally often crumbles under the weight of real data volumes, multiple environments, and team collaboration. Structuring PySpark projects: what developers should know is not about boilerplate—it's about making your team productive, your code testable, and your deployments reliable. After shipping Spark jobs across e-commerce platforms, real-time analytics systems, and data pipelines processing petabytes, we've learned that structure matters more than clever transformations.
Separate transforms from I/O
The most common mistake in PySpark production code is mixing data loading, transformation logic, and writing in a single function. This coupling makes testing impossible, forces you to spin up actual Spark sessions during development, and locks you into specific data sources.
Instead, organize your code into three layers: I/O layer (read/write), transform layer (business logic), and orchestration layer (where they meet). Your transform functions should take DataFrames as arguments and return DataFrames. The I/O layer handles connections to S3, databases, Kafka topics, or data lakes. The orchestration layer coordinates the flow.
Consider a real e-commerce scenario where we process order events. Without separation, you might write:
def process_orders():
orders_df = spark.read.parquet("s3://raw/orders")
customers_df = spark.read.parquet("s3://raw/customers")
joined = orders_df.join(customers_df, "customer_id")
enriched = joined.withColumn("order_total_usd", col("amount") * col("exchange_rate"))
enriched.write.mode("overwrite").parquet("s3://processed/orders")
This function is hard to test without touching S3. Refactor it:
# transforms.py
def enrich_orders(orders_df, customers_df):
"""Pure transformation logic, no I/O."""
joined = orders_df.join(customers_df, "customer_id")
enriched = joined.withColumn("order_total_usd", col("amount") * col("exchange_rate"))
return enriched
# io.py
class OrderDataStore:
def __init__(self, spark, bucket):
self.spark = spark
self.bucket = bucket
def read_orders(self):
return self.spark.read.parquet(f"s3://{self.bucket}/raw/orders")
def read_customers(self):
return self.spark.read.parquet(f"s3://{self.bucket}/raw/customers")
def write_orders(self, df):
df.write.mode("overwrite").parquet(f"s3://{self.bucket}/processed/orders")
# main.py (orchestration)
def main(spark, config):
store = OrderDataStore(spark, config.bucket)
orders = store.read_orders()
customers = store.read_customers()
enriched = enrich_orders(orders, customers)
store.write_orders(enriched)
Now you can unit test enrich_orders with small Spark DataFrames, mock the I/O layer, and change data sources without touching transformation logic.
Functions that take and return DataFrames
A production Spark codebase lives or dies by its function signatures. Functions accepting DataFrames as parameters are composable, testable, and debuggable. Avoid functions that reference global Spark sessions or hard-code column names.
Good pattern:
def apply_business_rules(df, min_order_value, currency_col="currency"):
"""
Filter and flag high-value orders.
Args:
df: Input DataFrame with columns: order_id, amount, currency
min_order_value: Threshold in cents
currency_col: Name of currency column (default: "currency")
Returns:
DataFrame with added column 'high_value' (boolean)
"""
return df.filter(col("amount") >= min_order_value) \
.withColumn("high_value", lit(True))
Bad pattern:
def apply_business_rules(min_order_value):
"""Implicitly uses global spark session and hard-coded paths."""
global spark
df = spark.read.parquet("s3://my-bucket/orders")
return df.filter(col("amount") >= min_order_value)
The good pattern lets you chain transformations, test with sample data, and integrate with airflow or other orchestrators. It also forces you to document what data each function expects.
When building complex pipelines, create a transformation registry or builder pattern:
class OrderPipeline:
def __init__(self, df):
self.df = df
def validate_amounts(self, min_val=0):
self.df = self.df.filter(col("amount") >= min_val)
return self
def enrich_customer_tier(self, tier_mapping):
self.df = self.df.join(tier_mapping, "customer_id", "left")
return self
def compute_final_price(self, tax_rate=0.1):
self.df = self.df.withColumn(
"final_price",
col("amount") * (1 + lit(tax_rate))
)
return self
def result(self):
return self.df
# Usage
orders = spark.read.parquet("s3://raw/orders")
processed = OrderPipeline(orders) \
.validate_amounts(min_val=100) \
.enrich_customer_tier(tier_df) \
.compute_final_price(tax_rate=0.12) \
.result()
Config per environment
Hardcoding S3 buckets, database credentials, or Kafka brokers into your code is a quick path to disaster. Production Spark projects need environment-specific configuration that changes between dev, staging, and production without touching code.
Use a config structure with environment-based overrides:
# config.py
from dataclasses import dataclass
from typing import Optional
import os
import json
@dataclass
class SparkConfig:
environment: str
source_bucket: str
dest_bucket: str
kafka_brokers: str
checkpoint_path: str
app_name: str
log_level: str = "INFO"
def load_config(env: Optional[str] = None) -> SparkConfig:
env = env or os.getenv("ENVIRONMENT", "dev")
# Load base config
with open(f"config/{env}.json") as f:
cfg = json.load(f)
# Override from environment variables
cfg["source_bucket"] = os.getenv("SOURCE_BUCKET", cfg.get("source_bucket"))
cfg["kafka_brokers"] = os.getenv("KAFKA_BROKERS", cfg.get("kafka_brokers"))
return SparkConfig(**cfg)
Your config files stay in version control without secrets:
# config/dev.json
{
"environment": "dev",
"source_bucket": "my-company-dev-raw",
"dest_bucket": "my-company-dev-processed",
"kafka_brokers": "localhost:9092",
"checkpoint_path": "/tmp/spark-checkpoints",
"app_name": "order_processor_dev",
"log_level": "DEBUG"
}
# config/prod.json
{
"environment": "prod",
"source_bucket": "my-company-prod-raw",
"dest_bucket": "my-company-prod-processed",
"kafka_brokers": "kafka-1.prod:9092,kafka-2.prod:9092",
"checkpoint_path": "s3://my-company-prod-checkpoints/order_processor",
"app_name": "order_processor_prod",
"log_level": "WARN"
}
For secrets (database passwords, API keys), use AWS Secrets Manager, HashiCorp Vault, or your orchestrator's secret management. Never check them into git or pass them as literal strings in config.
Package as wheels
Submitting loose Python files to Spark (via --py-files) works for prototyping but breaks quickly in production. Proper packaging via wheels ensures consistent dependencies, versioning, and distribution across your cluster.
Structure your project like this:
order_processor/
├── pyproject.toml
├── setup.py
├── setup.cfg
├── src/
│ └── order_processor/
│ ├── __init__.py
│ ├── transforms.py
│ ├── io.py
│ ├── config.py
│ └── cli.py
├── tests/
│ ├── conftest.py
│ ├── test_transforms.py
│ └── test_io.py
├── config/
│ ├── dev.json
│ ├── staging.json
│ └── prod.json
└── README.md
Your setup.py or pyproject.toml declares dependencies:
# setup.py
from setuptools import setup, find_packages
setup(
name="order-processor",
version="0.2.1",
description="Order processing pipeline",
package_dir={"": "src"},
packages=find_packages(where="src"),
python_requires=">=3.9",
install_requires=[
"pyspark>=3.4.0,<4",
"pyarrow>=10.0.0",
"python-dotenv>=0.21",
],
extras_require={
"dev": ["pytest>=7.0", "pytest-cov>=4.0", "black>=22.0"],
},
entry_points={
"console_scripts": [
"order-processor=order_processor.cli:main",
],
},
)
Build and distribute the wheel:
pip install build
python -m build
# Installs in your cluster
spark-submit --py-files dist/order_processor-0.2.1-py3-none-any.whl \
--conf spark.pyspark.python=/path/to/venv/bin/python \
jobs/run_order_processor.py
Or use wheels with a cluster-wide package repository (Artifactory, private PyPI) for automatic distribution. This approach also lets you version your job independently from Spark versions, and it integrates cleanly with CI/CD pipelines.
Common mistakes
- Mixing transforms and I/O: Your transformation logic depends on specific storage systems, making testing painful and refactoring dangerous.
- Global Spark session references: Functions that implicitly use
sparkcan't be composed or tested in isolation. - Hardcoded configuration: Credentials and paths in source code create security vulnerabilities and force code changes per environment.
- Loose Python files: Version skew, missing dependencies, and inconsistent environments plague teams that avoid packaging.
- No function documentation: Without docstrings, other team members (or you, six months later) won't know what columns a function expects or guarantees.
- Insufficient logging: Production Spark jobs fail silently or with cryptic errors if you don't log stage progress and data row counts.
Key takeaways
- Separate I/O (read/write), transforms (business logic), and orchestration (flow control) into distinct modules. This makes testing, refactoring, and source-switching trivial.
- Write transform functions that accept and return DataFrames, never reference global state, and document expected columns and row semantics.
- Store configuration in environment-specific files (dev.json, prod.json) and secrets in a vault. Load config at runtime, never embed it in code.
- Package your Spark jobs as wheels with proper dependency pinning. Versioning code separately from the cluster pays off when you need to roll back or upgrade.
- Test transforms with small in-memory DataFrames before deploying to production. Mock I/O layers so tests run fast and don't depend on cloud infrastructure.
- Log transformation stages, row counts, and failure modes. Future debugging depends on it.
Comments
Loading comments…