S049 · Stream Processing — Watermarks, Windows, Exactly-Once

# stream_processing.py — event-time windowed aggregation with watermark + idempotent upsert.
# Run: spark-submit --packages org.apache.spark:spark-sql-kafka-0-10_2.12:3.5.0,\
#                            org.postgresql:postgresql:42.7.0 stream_processing.py
 
from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, StringType, TimestampType
 
spark = (
    SparkSession.builder
    .appName("stream_processing")
    .config("spark.sql.streaming.checkpointLocation", "/tmp/ck/click_agg")
    .config("spark.sql.shuffle.partitions", "8")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("WARN")
 
# ------------------------------------------------------------------
# Source: JSON events on a Kafka topic
# ------------------------------------------------------------------
schema = (StructType()
    .add("user_id", StringType())
    .add("event", StringType())
    .add("ts", TimestampType()))     # event time inside the payload
 
raw = (
    spark.readStream
    .format("kafka")
    .option("kafka.bootstrap.servers", "localhost:9092")
    .option("subscribe", "clicks")
    .option("startingOffsets", "latest")
    .load()
)
 
parsed = (
    raw.selectExpr("CAST(value AS STRING) AS json")
       .select(F.from_json("json", schema).alias("e"))
       .select("e.*")
)
 
# ------------------------------------------------------------------
# Watermark: ‘I promise no more events \< now - 2 minutes’
# Windowing: 1-minute tumbling by event time
# ------------------------------------------------------------------
windowed = (
    parsed
    .withWatermark("ts", "2 minutes")
    .groupBy(
        F.window("ts", "1 minute"),
        F.col("user_id")
    )
    .agg(F.count("*").alias("n"))
    .select(
        F.col("window.start").alias("window_start"),
        F.col("window.end").alias("window_end"),
        "user_id",
        "n"
    )
)
 
# ------------------------------------------------------------------
# Sink: foreachBatch → Postgres UPSERT (idempotent = effectively-once)
# ------------------------------------------------------------------
JDBC_URL   = "jdbc:postgresql://localhost:5432/streamdemo"
JDBC_PROPS = {"user": "postgres", "password": "demo", "driver": "org.postgresql.Driver"}
 
def upsert_to_pg(batch_df, batch_id: int) -> None:
    if batch_df.isEmpty(): return
    tmp = f"click_agg_stg_{batch_id}"
    (batch_df.write
        .mode("overwrite")
        .jdbc(url=JDBC_URL, table=tmp, properties=JDBC_PROPS))
 
    import psycopg2
    conn = psycopg2.connect("host=localhost dbname=streamdemo user=postgres password=demo")
    with conn, conn.cursor() as cur:
        cur.execute(f"""
            INSERT INTO click_agg (window_start, window_end, user_id, n)
            SELECT window_start, window_end, user_id, n FROM {tmp}
            ON CONFLICT (window_start, user_id)
              DO UPDATE SET n = EXCLUDED.n, window_end = EXCLUDED.window_end;
            DROP TABLE {tmp};
        """)
    conn.close()
 
query = (
    windowed.writeStream
    .outputMode("update")           # only changed rows per batch
    .foreachBatch(upsert_to_pg)
    .trigger(processingTime="30 seconds")
    .start()
)
query.awaitTermination()