Capstone: Fraud Detection Pipeline

from pyspark.sql import functions as F

parsed = (
    raw_kafka_df
    .select(F.from_json(F.col("value").cast("string"), TRANSACTION_SCHEMA).alias("tx"))
    .select("tx.*")
    .filter(F.col("transaction_id").isNotNull())
)

# Apply a 10-minute watermark on the event-time column.
# Spark tracks max(transaction_time) seen so far.
# The watermark boundary = max_seen - 10 minutes.
# Any event whose transaction_time falls BEFORE that boundary
# is considered "too late" and is silently dropped.
with_watermark = parsed.withWatermark("transaction_time", "10 minutes")

from delta.tables import DeltaTable
from pyspark.sql.functions import col, to_date

# Static reference tables are loaded ONCE here and cached.
# Loading inside process_batch would hit the DB on every micro-batch.
customers_df = load_customers_batch(spark)
blacklist_df = load_merchant_blacklist(spark)
customers_df.cache()
blacklist_df.cache()

def process_batch(micro_batch_df, batch_id):
    if micro_batch_df.isEmpty():
        return

    # Cache the raw micro-batch so Kafka bytes are not re-read for each sink
    micro_batch_df.persist()

    # Stream-batch join: enrich each micro-batch with customer and merchant data
    enriched = enrich_with_customers(micro_batch_df, customers_df)
    enriched = enrich_with_blacklist(enriched, blacklist_df)

    # Compute 5-minute tumbling velocity counts on this batch
    velocity = transaction_velocity_window(enriched)

    # Join window counts back using time bounds so each transaction
    # matches only the window it actually falls inside.
    enriched = enriched.join(
        velocity.select(
            col("account_id").alias("_va"),
            col("window_start"),
            col("window_end"),
            col("tx_count"),
            col("total_amount").alias("window_total_amount"),
        ),
        on=(
            (col("account_id") == col("_va"))
            & (col("transaction_time") >= col("window_start"))
            & (col("transaction_time") < col("window_end"))
        ),
        how="left",
    ).drop("_va")

    # Score every row against the 6 fraud rules
    scored = apply_fraud_rules(enriched, tx_count_col="tx_count")
    scored.persist()

    # Sink 1: all enriched transactions, partitioned by date (append)
    (scored
        .withColumn("tx_date", to_date("transaction_time"))
        .write.format("delta")
        .mode("append")
        .partitionBy("tx_date")
        .option("mergeSchema", "true")
        .save(f"{DELTA_BASE_PATH}/enriched_transactions"))

    # Sink 2: only rows that triggered the fraud threshold
    (scored.filter("is_fraud_alert = true")
        .write.format("delta")
        .mode("append")
        .save(f"{DELTA_BASE_PATH}/fraud_alerts"))

    # Sink 3: window aggregates - upsert (idempotent re-runs safe)
    upsert_to_delta(f"{DELTA_BASE_PATH}/window_aggregates", velocity,
                    merge_keys=["account_id", "window_start"])
    scored.unpersist()
    micro_batch_df.unpersist()

query = (
    with_watermark
    .writeStream
    .foreachBatch(process_batch)
    .option("checkpointLocation", f"{CHECKPOINT_BASE}/fraud_detection")
    .trigger(processingTime="10 seconds")
    .outputMode("append")
    .start()
)

from delta.tables import DeltaTable
from pyspark.sql import DataFrame

def upsert_to_delta(delta_path: str, micro_batch_df: DataFrame, merge_keys: list[str]):
    """
    Idempotent merge: update existing rows, insert new ones.
    Safe to re-run if a micro-batch is replayed (exactly-once semantics).
    """
    spark = micro_batch_df.sparkSession
    if micro_batch_df.isEmpty():
        return

    if not DeltaTable.isDeltaTable(spark, delta_path):
        # First write creates the table
        micro_batch_df.write.format("delta").mode("overwrite").save(delta_path)
        return

    merge_condition = " AND ".join(f"t.{k} = s.{k}" for k in merge_keys)
    (DeltaTable.forPath(spark, delta_path)
        .alias("t")
        .merge(micro_batch_df.alias("s"), merge_condition)
        .whenMatchedUpdateAll()
        .whenNotMatchedInsertAll()
        .execute())

from pyspark.sql import functions as F

HIGH_AMOUNT_THRESHOLD = 10_000.0
VELOCITY_TX_LIMIT     = 5
NEW_ACCOUNT_AGE_DAYS  = 30
FRAUD_SCORE_THRESHOLD = 0.35

def apply_fraud_rules(df, tx_count_col=None):
    df = df.withColumn("flag_high_amount",
        F.when(F.col("amount") > HIGH_AMOUNT_THRESHOLD, F.lit(1)).otherwise(F.lit(0)))

    df = df.withColumn("flag_blacklisted_merchant",
        F.col("is_blacklisted").cast("int"))

    df = df.withColumn("flag_new_account",
        F.when(F.col("account_age_days") < NEW_ACCOUNT_AGE_DAYS, F.lit(1)).otherwise(F.lit(0)))

    df = df.withColumn("flag_flagged_account",
        F.col("is_flagged").cast("int"))

    df = df.withColumn("flag_international_mismatch",
        F.when(
            (F.col("is_international") == 1) &
            (F.col("country") != F.col("customer_country")),
            F.lit(1),
        ).otherwise(F.lit(0)))

    if tx_count_col and tx_count_col in df.columns:
        df = df.withColumn("flag_velocity_burst",
            F.when(F.col(tx_count_col) > VELOCITY_TX_LIMIT, F.lit(1)).otherwise(F.lit(0)))
    else:
        df = df.withColumn("flag_velocity_burst", F.lit(0))

    raw_score = (
        F.col("flag_high_amount")              * 0.30
        + F.col("flag_blacklisted_merchant")   * 0.40
        + F.col("flag_new_account")            * 0.10
        + F.col("flag_flagged_account")        * 0.35
        + F.col("flag_international_mismatch") * 0.20
        + F.col("flag_velocity_burst")         * 0.50
    )
    df = df.withColumn("fraud_score", F.least(raw_score.cast("double"), F.lit(1.0)))
    df = df.withColumn("is_fraud_alert", F.col("fraud_score") >= FRAUD_SCORE_THRESHOLD)
    return df

# Job 2: Account Risk Scorer (account_risk.py)
# Uses a SLIDING window (30 min / 5 min slide) + outputMode("update")

account_events = (
    spark.readStream.format("kafka")
    .option("kafka.bootstrap.servers", KAFKA_BROKERS)
    .option("subscribe", ACCOUNT_EVENTS_TOPIC)
    .load()
    .select(F.from_json(F.col("value").cast("string"), ACCOUNT_EVENT_SCHEMA).alias("d"))
    .select("d.*")
    .withWatermark("event_time", "15 minutes")   # 15-min late-data tolerance
)

# Sliding window: each FAILED_LOGIN can appear in up to 6 overlapping windows
risk_windows = failed_login_sliding_window(account_events, "30 minutes", "5 minutes")

TIER_WEIGHTS = {"STANDARD": 0.10, "PREMIUM": 0.06, "BUSINESS": 0.04}

# Load and cache once, not inside the batch function
customers_df = load_customers_batch(spark).select("account_id", "tier", "is_flagged")
customers_df.cache()

def _add_tier_weights(df):
    """Map customer tier to a risk weight multiplier column."""
    return df.withColumn(
        "tier_weight",
        F.when(F.col("tier") == "STANDARD", F.lit(0.10))
         .when(F.col("tier") == "PREMIUM",  F.lit(0.06))
         .when(F.col("tier") == "BUSINESS", F.lit(0.04))
         .otherwise(F.lit(0.10)),
    )

def score_batch(micro_batch_df, batch_id):
    if micro_batch_df.isEmpty():
        return

    micro_batch_df.persist()

    enriched = micro_batch_df.join(
        F.broadcast(customers_df), on="account_id", how="left"
    )

    scored = (
        _add_tier_weights(enriched)
        .withColumn("risk_score",
            F.least(
                (F.col("failed_login_count") * F.col("tier_weight")).cast("double"),
                F.lit(1.0),
            ))
        .withColumn("is_high_risk", F.col("risk_score") >= 0.5)
    )
    upsert_to_delta(f"{DELTA_BASE_PATH}/account_risk_scores", scored,
                    merge_keys=["account_id", "window_start"])
    micro_batch_df.unpersist()

query = (
    risk_windows
    .writeStream
    .foreachBatch(score_batch)
    .option("checkpointLocation", f"{CHECKPOINT_BASE}/account_risk")
    .trigger(processingTime="30 seconds")
    .outputMode("update")    # emit only windows that changed this micro-batch
    .start()
)