Writing Spark jobs for reuse | Spark & Data Eng

Writing Spark jobs for reuse: what developers should know

When you're running Spark jobs in production across multiple teams and projects, the difference between a one-off script and a maintainable, reusable job becomes painfully clear. We've learned that writing Spark jobs for reuse requires discipline around parameters, configuration, testing, and observability—the same principles that apply to microservices, but with data pipelines. This article walks through the patterns that let you build jobs once and deploy them confidently across dozens of use cases.

Pass dates and paths as arguments

The first step toward reusability is removing hardcoded values from your code. In production Spark jobs, dates and file paths change per run, per environment, and per business unit. Instead of baking them into your Python or Scala file, accept them as command-line arguments or through a job launcher framework.

Here's a typical pattern we use:

python


import argparse
from datetime import datetime
from pyspark.sql import SparkSession

def main():
    parser = argparse.ArgumentParser(description="Process user activity logs")
    parser.add_argument("--input-path", required=True, help="S3 path to input Parquet")
    parser.add_argument("--output-path", required=True, help="S3 path for output Parquet")
    parser.add_argument("--date", required=True, help="Processing date (YYYY-MM-DD)")
    parser.add_argument("--environment", default="dev", help="Environment: dev, staging, prod")
    
    args = parser.parse_args()
    
    spark = SparkSession.builder.appName("UserActivityProcessor").getOrCreate()
    
    try:
        df = spark.read.parquet(args.input_path)
        
        # Filter to specific date
        df_filtered = df.filter(df.event_date == args.date)
        
        # Process...
        result_df = df_filtered.groupBy("user_id").count()
        
        # Write to environment-specific output
        result_df.write.mode("overwrite").parquet(args.output_path)
        
        print(f"Job completed. Wrote {result_df.count()} records to {args.output_path}")
        
    finally:
        spark.stop()

if __name__ == "__main__":
    main()

This approach lets you invoke the job like this:

python


spark-submit \
  --master yarn \
  job.py \
  --input-path s3://raw-data/events/ \
  --output-path s3://processed-data/user-activity/ \
  --date 2024-01-15 \
  --environment prod

The same job runs on different dates, environments, and datasets without modification. Your orchestration layer (Airflow, Step Functions, Databricks Jobs) controls the arguments, not your code.

Separate config from code

As jobs grow more complex, you'll have Spark configurations, database credentials, schema definitions, and business logic thresholds. Don't scatter these throughout your codebase. Create a dedicated configuration module or file.

For smaller teams, a YAML or JSON config file works well:

python


# config.yaml
environments:
  dev:
    spark_config:
      spark.sql.shuffle.partitions: 50
      spark.sql.adaptive.enabled: true
    database_url: jdbc:postgresql://dev-db:5432/analytics
    min_events_per_user: 1
  
  prod:
    spark_config:
      spark.sql.shuffle.partitions: 500
      spark.sql.adaptive.enabled: true
    database_url: jdbc:postgresql://prod-db:5432/analytics
    min_events_per_user: 10

schemas:
  user_events:
    - name: user_id
      type: string
      nullable: false
    - name: event_type
      type: string
    - name: timestamp
      type: timestamp

Load and apply this in your job startup:

python


import yaml
from pyspark.sql.types import StructType, StructField, StringType, TimestampType

def load_config(config_file: str, environment: str) -> dict:
    with open(config_file, 'r') as f:
        full_config = yaml.safe_load(f)
    return full_config[environment]

def apply_spark_config(spark, config: dict):
    for key, value in config['spark_config'].items():
        spark.conf.set(key, value)

def get_schema(config: dict, schema_name: str):
    fields = []
    for field_def in config['schemas'][schema_name]:
        field_type = StringType() if field_def['type'] == 'string' else TimestampType()
        fields.append(StructField(field_def['name'], field_type, field_def['nullable']))
    return StructType(fields)

# In main():
config = load_config('config.yaml', args.environment)
apply_spark_config(spark, config)
min_events = config['min_events_per_user']
schema = get_schema(config, 'user_events')

For larger organizations, consider using a secrets manager (AWS Secrets Manager, HashiCorp Vault) for sensitive values and environment variables for environment-specific URLs. This keeps your config version-controllable and your secrets secure.

Return DataFrames from functions

One of the most effective reusability patterns is breaking your Spark pipeline into small, composable functions that accept and return DataFrames. This makes your code testable, readable, and easy to recombine in different jobs.

Instead of writing a monolithic processing function, structure it like this:

python


from pyspark.sql import DataFrame
from pyspark.sql.functions import col, when, to_timestamp

def load_events(spark: SparkSession, path: str, schema: StructType) -> DataFrame:
    """Load events from Parquet, apply schema and basic validation."""
    df = spark.read.schema(schema).parquet(path)
    return df.filter(col("user_id").isNotNull())

def enrich_events(df: DataFrame, lookup_path: str) -> DataFrame:
    """Join events with user metadata."""
    lookup = df.sparkSession.read.parquet(lookup_path)
    return df.join(lookup, on="user_id", how="left")

def compute_aggregations(df: DataFrame, min_events: int) -> DataFrame:
    """Group events, compute metrics, filter by minimum threshold."""
    agg_df = df.groupBy("user_id").agg(
        count("*").alias("event_count"),
        max("timestamp").alias("last_event_time"),
        avg("event_value").alias("avg_value")
    )
    return agg_df.filter(col("event_count") >= min_events)

def apply_business_rules(df: DataFrame) -> DataFrame:
    """Apply segment classification, flags, etc."""
    return df.withColumn(
        "segment",
        when(col("event_count") > 100, "power_user")
        .when(col("event_count") > 10, "active_user")
        .otherwise("casual_user")
    )

# In main():
events_df = load_events(spark, args.input_path, schema)
enriched_df = enrich_events(events_df, args.lookup_path)
agg_df = compute_aggregations(enriched_df, config['min_events_per_user'])
final_df = apply_business_rules(agg_df)
final_df.write.mode("overwrite").parquet(args.output_path)

Now, if another team needs a different pipeline—say, enriched events without aggregation—they import and reuse parts of it. You can also test each function in isolation with small test DataFrames, which brings us to the next point.

Log job metrics

Production observability is non-negotiable. When a Spark job runs on a 10 TB dataset, you need to know whether it processed 100M or 0 records, whether it spilled to disk, and where time was spent. Log metrics at every step.

Here's a practical approach:

python


import logging
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def log_step(name: str, df: DataFrame, duration_seconds: float = None):
    """Log record count, schema, and optionally duration after a processing step."""
    record_count = df.count()
    logger.info(f"Step: {name} | Records: {record_count} | Columns: {len(df.columns)}")
    if duration_seconds:
        logger.info(f"  Duration: {duration_seconds:.2f}s | Records/sec: {record_count/duration_seconds:.0f}")

def main():
    config = load_config('config.yaml', args.environment)
    spark = SparkSession.builder.appName("UserActivityProcessor").getOrCreate()
    
    job_start = datetime.now()
    logger.info(f"Job started at {job_start} with args: {vars(args)}")
    
    try:
        start = datetime.now()
        events_df = load_events(spark, args.input_path, schema)
        log_step("load_events", events_df, (datetime.now() - start).total_seconds())
        
        start = datetime.now()
        enriched_df = enrich_events(events_df, args.lookup_path)
        log_step("enrich_events", enriched_df, (datetime.now() - start).total_seconds())
        
        start = datetime.now()
        agg_df = compute_aggregations(enriched_df, config['min_events_per_user'])
        log_step("compute_aggregations", agg_df, (datetime.now() - start).total_seconds())
        
        start = datetime.now()
        final_df = apply_business_rules(agg_df)
        log_step("apply_business_rules", final_df, (datetime.now() - start).total_seconds())
        
        # Log output metrics
        partition_count = final_df.rdd.getNumPartitions()
        logger.info(f"Output partitions: {partition_count}")
        
        start = datetime.now()
        final_df.write.mode("overwrite").parquet(args.output_path)
        write_duration = (datetime.now() - start).total_seconds()
        logger.info(f"Write completed in {write_duration:.2f}s to {args.output_path}")
        
        job_duration = (datetime.now() - job_start).total_seconds()
        logger.info(f"Job completed successfully in {job_duration:.2f}s")
        
    except Exception as e:
        logger.error(f"Job failed after {(datetime.now() - job_start).total_seconds():.2f}s", exc_info=True)
        raise
    finally:
        spark.stop()

Send these logs to your centralized logging system (ELK, CloudWatch, Datadog). Over time, you'll see patterns: which transformations are slow, which jobs spike CPU or memory, and which business logic thresholds need adjustment.

Common mistakes

Hardcoding paths and dates in code: This forces developers to edit source before every run and makes it impossible to version your job as a single artifact. Always accept these as arguments.
No validation of input arguments: Misspelled paths silently create empty DataFrames. Validate arguments early and fail fast with clear error messages.
Calling DataFrame.count() for logging without caching: If you count a large DataFrame just to log it, you trigger a full re-compute. Either cache the DataFrame or use explain() to understand execution, not count().
Storing secrets in config files: A config file in Git is convenient but dangerous. Use a secrets manager or environment variables for credentials.
Writing monolithic pipelines that can't be tested in isolation: If your entire job is one function, you can't unit test it. Break it into smaller, DataFrame-returning functions.
Not handling schema drift: When your upstream data schema changes, your job fails silently or produces wrong results. Define and validate schemas explicitly.

Key takeaways

Pass dates, paths, and environment-specific settings as command-line arguments, never hardcode them in source.
Separate configuration (thresholds, Spark settings, schemas) from code using YAML, JSON, or environment variables; use secrets managers for credentials.
Structure your Spark logic as small functions that accept and return DataFrames; this makes your code composable, testable, and reusable across jobs.
Log record counts, durations, and processing metrics at every pipeline step so you can debug production issues and optimize performance.
Validate inputs early, handle schema drift explicitly, and use exception handlers to log failures with context for downstream alerting.

Writing Spark jobs for reuse: what developers should know

Writing Spark jobs for reuse: what developers should know

Pass dates and paths as arguments

Separate config from code

Return DataFrames from functions

Log job metrics

Common mistakes

Key takeaways

Comments

Leave a comment

Related articles

Validation with FluentValidation: a practical .NET guide

Use records for DTOs: a practical guide for developers

EF Core from zero to production: lessons from production

Configuring EF Core relationships in .NET & ASP.NET Core