PySpark Big Data Processing: Production Patterns That Scale

# io.thecodeforge — Python tutorial

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, TimestampType, LongType

# Production SparkSession — never use defaults for anything beyond local testing.
# spark.sql.shuffle.partitions defaults to 200, which is catastrophically wrong
# for both small datasets (200 near-empty tasks) and massive ones (OOM per task).
# Rule of thumb: target 100-200MB of data per partition after a shuffle.
spark = (
    SparkSession.builder
    .appName("transaction-risk-aggregator")          # Shows up in Spark UI — make it meaningful
    .config("spark.sql.shuffle.partitions", "400")   # Tune to (total_data_GB * 1024 / 200) roughly
    .config("spark.sql.adaptive.enabled", "true")    # AQE: lets Spark re-optimize mid-execution (Spark 3.x)
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true")  # AQE merges small post-shuffle partitions
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")  # Faster than Java serializer
    .config("spark.sql.autoBroadcastJoinThreshold", "50mb")  # Broadcast joins under 50MB — avoids shuffle
    .getOrCreate()
)

# Silence INFO log spam in local development — you want WARN or ERROR in dev
spark.sparkContext.setLogLevel("WARN")

# Define schema explicitly — NEVER rely on inferSchema in production.
# inferSchema scans the entire dataset twice and makes wrong type guesses on messy data.
transaction_schema = StructType([
    StructField("transaction_id", StringType(), nullable=False),
    StructField("account_id",     StringType(), nullable=False),
    StructField("merchant_id",    StringType(), nullable=False),
    StructField("amount_usd",     DoubleType(), nullable=False),
    StructField("transaction_ts", TimestampType(), nullable=False),
    StructField("risk_score",     DoubleType(), nullable=True),   # nullable=True for fields that may be missing
])

# Read parquet — the production format. CSV is for demos and data exports only.
# partitionBy on ingestion (year/month/day) means Spark prunes irrelevant files at read time.
transaction_df = (
    spark.read
    .schema(transaction_schema)          # Enforce schema — fail fast on bad data
    .option("mergeSchema", "false")       # Don't silently accept schema drift
    .parquet("/data/transactions/year=2024/month=01/")  # Partition-pruned read path
)

# This is a lazy transformation — zero data has moved. Zero files have been opened.
# Spark just recorded the intent in the query plan.
high_risk_transactions = (
    transaction_df
    .filter(F.col("risk_score") > 0.85)          # Predicate pushdown: Spark pushes this into the reader
    .filter(F.col("amount_usd") > 100.0)          # Stack filters — Catalyst combines them into one pass
    .select("transaction_id", "account_id", "merchant_id", "amount_usd", "risk_score")  # Column pruning
    .withColumn("risk_tier",                       # Derived column — evaluated lazily with the rest
        F.when(F.col("risk_score") > 0.95, F.lit("CRITICAL"))
         .when(F.col("risk_score") > 0.85, F.lit("HIGH"))
         .otherwise(F.lit("MEDIUM"))
    )
)

# Calling explain() shows the physical plan WITHOUT executing the query.
# Use this to verify predicate pushdown is happening and catch accidental cross joins.
high_risk_transactions.explain(mode="formatted")

# .count() is an action — THIS is when Spark actually reads the files and runs the DAG.
record_count = high_risk_transactions.count()
print(f"High-risk transactions in January 2024: {record_count:,}")

# Cache only if you're reusing the DataFrame multiple times in the same job.
# Caching a DataFrame you only use once wastes memory and slows you down.
high_risk_transactions.cache()  # Stores in memory after first action triggers it

# Second action — hits cache instead of re-reading parquet files
high_risk_transactions.show(5, truncate=False)

# io.thecodeforge — Python tutorial

from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, IntegerType

spark = (
    SparkSession.builder
    .appName("merchant-risk-join")
    .config("spark.sql.shuffle.partitions", "400")
    .config("spark.sql.adaptive.enabled", "true")
    .config("spark.sql.adaptive.skewJoin.enabled", "true")  # AQE detects and splits skewed partitions
    .config("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256mb")  # Flag partitions over 256MB as skewed
    .getOrCreate()
)

# --- SCENARIO: Join 2B transactions against 50K merchant records ---
# merchant_lookup is small (~5MB). transactions_df is massive (~800GB).

merchant_schema = StructType([
    StructField("merchant_id",       StringType(), nullable=False),
    StructField("merchant_name",     StringType(), nullable=True),
    StructField("merchant_category", StringType(), nullable=True),
    StructField("country_code",      StringType(), nullable=True),
])

transaction_schema = StructType([
    StructField("transaction_id", StringType(), nullable=False),
    StructField("account_id",     StringType(), nullable=False),
    StructField("merchant_id",    StringType(), nullable=False),
    StructField("amount_usd",     DoubleType(), nullable=False),
])

merchant_df = spark.read.schema(merchant_schema).parquet("/data/merchants/")
transaction_df = spark.read.schema(transaction_schema).parquet("/data/transactions/")

# ---------------------------------------------------------------
# APPROACH 1: Broadcast join — the right move when one side is small.
# Spark ships the entire small DataFrame to every executor.
# Zero shuffle on the large side. This is the fastest possible join.
# Use when the smaller DataFrame fits in executor memory (typically under 100MB).
# ---------------------------------------------------------------
merchant_enriched_df = transaction_df.join(
    F.broadcast(merchant_df),  # Explicitly broadcast the small side — don't rely on autoBroadcastJoinThreshold alone
    on="merchant_id",
    how="left"   # left join: keep all transactions even if merchant metadata is missing
)

# ---------------------------------------------------------------
# APPROACH 2: Salting — the manual fix for a skewed shuffle join
# when BOTH sides are too large to broadcast.
#
# Skew example: merchant MER-0001 has 40% of all transactions.
# Without salting, one executor handles 40% of the data. With salting,
# we artificially split that hot key across N buckets.
# ---------------------------------------------------------------

SALT_BUCKETS = 20  # Spread the hot key across 20 partitions instead of 1

# Add a random salt to the large (skewed) side
transaction_salted_df = transaction_df.withColumn(
    "salt",
    (F.rand() * SALT_BUCKETS).cast(IntegerType())  # Random int 0–19 per row
).withColumn(
    "salted_merchant_id",
    F.concat(F.col("merchant_id"), F.lit("_"), F.col("salt").cast(StringType()))  # "MER-0001_7"
)

# Explode the small side to match all salt values
# This increases merchant_df by SALT_BUCKETS times — acceptable when it's small
merchant_exploded_df = merchant_df.withColumn(
    "salt",
    F.explode(F.array([F.lit(i) for i in range(SALT_BUCKETS)]))  # One row per bucket: 0,1,2...19
).withColumn(
    "salted_merchant_id",
    F.concat(F.col("merchant_id"), F.lit("_"), F.col("salt").cast(StringType()))  # "MER-0001_7"
)

# Now join on the salted key — hot key is distributed across 20 partitions
skew_fixed_df = (
    transaction_salted_df.join(
        merchant_exploded_df,
        on="salted_merchant_id",
        how="left"
    )
    .drop("salt", "salted_merchant_id")  # Clean up the salt columns post-join
)

# ---------------------------------------------------------------
# Verify the join result — check for nulls that indicate join misses
# ---------------------------------------------------------------
join_miss_count = skew_fixed_df.filter(F.col("merchant_name").isNull()).count()
total_count = skew_fixed_df.count()

print(f"Total transactions enriched: {total_count:,}")
print(f"Transactions with missing merchant data: {join_miss_count:,}")
print(f"Enrichment rate: {((total_count - join_miss_count) / total_count) * 100:.2f}%")

# Check for data skew in the final partition distribution
# High stddev relative to mean = skew still present
skew_fixed_df.groupBy(F.spark_partition_id()).count().describe().show()

# io.thecodeforge — Python tutorial

from pyspark.sql import SparkSession, functions as F, Window
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, TimestampType, IntegerType
import pandas as pd

spark = (
    SparkSession.builder
    .appName("account-behaviour-features")
    .config("spark.sql.shuffle.partitions", "200")
    .config("spark.sql.adaptive.enabled", "true")
    .getOrCreate()
)

transaction_schema = StructType([
    StructField("transaction_id", StringType(),   nullable=False),
    StructField("account_id",     StringType(),   nullable=False),
    StructField("merchant_id",    StringType(),   nullable=False),
    StructField("amount_usd",     DoubleType(),   nullable=False),
    StructField("transaction_ts", TimestampType(), nullable=False),
])

transaction_df = spark.read.schema(transaction_schema).parquet("/data/transactions/")

# ---------------------------------------------------------------
# NATIVE AGGREGATIONS — no UDFs needed for standard feature engineering
# ---------------------------------------------------------------

# Window spec: per account, ordered by time, unbounded lookback
# This is the pattern for running totals, rolling averages, lag features
account_time_window = (
    Window
    .partitionBy("account_id")         # Each account gets its own independent window
    .orderBy("transaction_ts")         # Ordered chronologically within the partition
    .rowsBetween(Window.unboundedPreceding, Window.currentRow)  # All rows up to current
)

# 30-day rolling window — for detecting velocity anomalies
account_30day_window = (
    Window
    .partitionBy("account_id")
    .orderBy(F.col("transaction_ts").cast("long"))  # Cast to epoch seconds for range arithmetic
    .rangeBetween(-30 * 24 * 3600, 0)  # 30 days in seconds — rangeBetween works on the ORDER BY column value
)

feature_df = (
    transaction_df
    .withColumn(
        "account_lifetime_spend",
        F.sum("amount_usd").over(account_time_window)  # Running total spend per account
    )
    .withColumn(
        "account_txn_rank",
        F.rank().over(account_time_window)  # Transaction sequence number within account
    )
    .withColumn(
        "rolling_30d_spend",
        F.sum("amount_usd").over(account_30day_window)  # 30-day rolling spend
    )
    .withColumn(
        "rolling_30d_txn_count",
        F.count("transaction_id").over(account_30day_window)  # Transaction velocity in 30 days
    )
    .withColumn(
        "prev_transaction_amount",
        F.lag("amount_usd", 1).over(
            Window.partitionBy("account_id").orderBy("transaction_ts")  # Previous txn amount for delta calculation
        )
    )
    .withColumn(
        "amount_delta_vs_prev",
        F.col("amount_usd") - F.coalesce(F.col("prev_transaction_amount"), F.lit(0.0))  # coalesce handles first txn (null prev)
    )
    .withColumn(
        "is_large_transaction",
        F.when(F.col("amount_usd") > F.col("rolling_30d_spend") * 0.5, F.lit(True))  # >50% of 30d spend in one txn
         .otherwise(F.lit(False))
    )
)

# ---------------------------------------------------------------
# PANDAS UDF — use ONLY when native functions genuinely can't do it.
# Example: calling a custom ML scorer or regex library.
# Vectorized: receives a pd.Series, returns a pd.Series. Arrow handles serialization.
# DO NOT use row-level Python UDFs (udf decorator) — they kill throughput.
# ---------------------------------------------------------------
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import DoubleType

@pandas_udf(DoubleType())  # Vectorized UDF — Arrow-serialized batches, not row-by-row
def compute_velocity_score(rolling_count: pd.Series, rolling_spend: pd.Series) -> pd.Series:
    """
    Custom velocity scoring that combines transaction count and spend rate.
    This is genuinely custom logic — no Spark native equivalent.
    Receives full column batches as pandas Series — highly efficient.
    """
    # Normalize: high count + high spend = high velocity risk
    count_score = (rolling_count / rolling_count.max()).fillna(0)   # Relative count percentile
    spend_score = (rolling_spend / rolling_spend.max()).fillna(0)   # Relative spend percentile
    return (count_score * 0.4 + spend_score * 0.6).round(4)        # Weighted composite score

scored_df = feature_df.withColumn(
    "velocity_risk_score",
    compute_velocity_score(
        F.col("rolling_30d_txn_count").cast(DoubleType()),
        F.col("rolling_30d_spend")
    )
)

# Write partitioned output — partitionBy creates a directory hierarchy for efficient reads
(
    scored_df
    .repartition(200, "account_id")  # Co-locate same account on same partition before write
    .write
    .mode("overwrite")
    .partitionBy("account_id")       # Downstream reads can filter on account_id with zero shuffle
    .parquet("/data/features/account_behaviour/")
)

print("Feature generation complete. Verifying output...")
output_df = spark.read.parquet("/data/features/account_behaviour/")
output_df.select(
    "transaction_id", "account_id", "amount_usd",
    "rolling_30d_spend", "rolling_30d_txn_count",
    "is_large_transaction", "velocity_risk_score"
).orderBy("account_id", "transaction_ts").show(8, truncate=False)

# io.thecodeforge — Python tutorial

from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, TimestampType

spark = (
    SparkSession.builder
    .appName("daily-risk-report-writer")
    .config("spark.sql.shuffle.partitions", "400")
    .config("spark.sql.adaptive.enabled", "true")
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
    .config("spark.sql.adaptive.coalescePartitions.minPartitionSize", "128mb")  # AQE won't merge below this
    .getOrCreate()
)

transaction_schema = StructType([
    StructField("transaction_id", StringType(),    nullable=False),
    StructField("account_id",     StringType(),    nullable=False),
    StructField("merchant_id",    StringType(),    nullable=False),
    StructField("amount_usd",     DoubleType(),    nullable=False),
    StructField("transaction_ts", TimestampType(), nullable=False),
    StructField("risk_tier",      StringType(),    nullable=True),
    StructField("country_code",   StringType(),    nullable=True),
])

processed_df = spark.read.schema(transaction_schema).parquet("/data/enriched_transactions/")

# Add partition columns before writing — these become directory names in the output path
dated_df = processed_df.withColumn(
    "report_date",
    F.to_date(F.col("transaction_ts"))   # Extract date part for daily partitioning
).withColumn(
    "report_hour",
    F.hour(F.col("transaction_ts"))       # Hour-level partitioning for intraday queries
)

# ---------------------------------------------------------------
# APPROACH 1: Large dataset — repartition to control output file size.
# repartition() DOES cause a full shuffle. Only use when partition count
# and distribution matter more than the shuffle cost.
# Target: ~200MB per output file for S3/HDFS performance
# ---------------------------------------------------------------

# Check current partition count before writing
current_partitions = dated_df.rdd.getNumPartitions()
print(f"Current partition count before write: {current_partitions}")

# Repartition by the SAME columns you're partitioning the output by.
# This co-locates rows that will land in the same output directory,
# which means each task writes to exactly one output file — no file merging needed.
dated_df_repartitioned = dated_df.repartition(
    200,                  # 200 output files per (report_date, risk_tier) combination
    "report_date",        # Primary distribution key matches partitionBy key
    "risk_tier"           # Secondary distribution key
)

# Write with dynamic partition overwrite — ONLY overwrites the specific partitions present
# in this DataFrame, not the entire output table.
# Without this setting, mode('overwrite') nukes the ENTIRE output path.
spark.conf.set("spark.sql.sources.partitionOverwriteMode", "dynamic")

(
    dated_df_repartitioned
    .write
    .mode("overwrite")                   # Dynamic overwrite: only affects partitions in this batch
    .partitionBy("report_date", "risk_tier")  # Creates /report_date=2024-01-15/risk_tier=HIGH/ directories
    .option("compression", "snappy")     # Snappy: good balance of compression ratio vs CPU cost
    .parquet("/data/risk_reports/daily/")
)

# ---------------------------------------------------------------
# APPROACH 2: Small-to-medium result set — coalesce to avoid small files.
# coalesce() reduces partition count WITHOUT a full shuffle (narrow transformation).
# Use when you want fewer, larger output files and shuffle cost matters.
# ---------------------------------------------------------------
daily_summary_df = (
    processed_df
    .groupBy("country_code", "risk_tier", F.to_date("transaction_ts").alias("report_date"))
    .agg(
        F.count("transaction_id").alias("transaction_count"),
        F.sum("amount_usd").alias("total_amount_usd"),
        F.avg("amount_usd").alias("avg_amount_usd"),
        F.countDistinct("account_id").alias("unique_accounts")
    )
)

# After groupBy, Spark has 400 shuffle partitions — but this summary is small (~100K rows).
# coalesce to a reasonable number of files without paying for a re-shuffle.
daily_summary_df.coalesce(10)  # 10 files for a small summary dataset is fine

(
    daily_summary_df.coalesce(10)
    .write
    .mode("overwrite")
    .option("compression", "snappy")
    .parquet("/data/risk_reports/daily_summary/")
)

# ---------------------------------------------------------------
# Validate the output: check file count and size distribution
# In production, wire this into your data quality framework
# ---------------------------------------------------------------
output_df = spark.read.parquet("/data/risk_reports/daily_summary/")
output_df.printSchema()
output_df.orderBy("report_date", "country_code", "risk_tier").show(10)

print(f"Output partition count: {output_df.rdd.getNumPartitions()}")
print(f"Total summary records: {output_df.count():,}")