Module 21: Data Engineering Design Patterns

# full load: replace the target with the whole source
spark.table("orders").write.mode("overwrite").saveAsTable("pat_orders_full")

spark.sql("CREATE TABLE IF NOT EXISTS pat_state (source STRING, high_water TIMESTAMP)")
spark.sql("CREATE TABLE IF NOT EXISTS pat_orders_incr AS SELECT * FROM orders WHERE 1 = 0")

def load_increment(source_table="orders", target="pat_orders_incr", col="order_ts"):
    row = spark.table("pat_state").filter(F.col("source") == source_table).first()
    last = row["high_water"] if row else None

    src = spark.table(source_table)
    new = src.filter(F.col(col) > F.lit(last)) if last else src
    new = new.orderBy(col, "order_id").limit(1000)        # a small, repeatable batch for the demo

    new.createOrReplaceTempView("new_rows")
    spark.sql(f"""
      MERGE INTO {target} t USING new_rows s ON t.order_id = s.order_id
      WHEN MATCHED THEN UPDATE SET * WHEN NOT MATCHED THEN INSERT *""")

    new_max = new.agg(F.max(col)).first()[0]
    if new_max:                                            # move the watermark only after success
        spark.sql("DELETE FROM pat_state WHERE source = '%s'" % source_table)
        spark.createDataFrame([(source_table, new_max)], "source STRING, high_water TIMESTAMP") \
             .write.mode("append").saveAsTable("pat_state")
    return new.count()

print(load_increment(), spark.table("pat_orders_incr").count())     # 1000 1000
print(load_increment(), spark.table("pat_orders_incr").count())     # 1000 about 2000 (equal timestamps at the boundary can shift it)

# a small change feed: operation + sequence number
feed = spark.createDataFrame(
    [(1, "Customer 1", "NL", "UPDATE", 3), (1, "Customer 1", "DE", "UPDATE", 5),
     (2, "Customer 2", "FR", "DELETE", 4), (999, "New", "US", "INSERT", 6)],
    "customer_id INT, name STRING, country STRING, op STRING, seq INT")

spark.sql("CREATE OR REPLACE TABLE pat_customers AS SELECT * FROM customers")

latest = (feed.withColumn("rn", F.row_number().over(Window.partitionBy("customer_id").orderBy(F.desc("seq"))))
              .filter("rn = 1").drop("rn"))              # keep only the last event per key
latest.createOrReplaceTempView("latest")
spark.sql("""
  MERGE INTO pat_customers t USING latest s ON t.customer_id = s.customer_id
  WHEN MATCHED AND s.op = 'DELETE' THEN DELETE
  WHEN MATCHED THEN UPDATE SET name = s.name, country = s.country
  WHEN NOT MATCHED AND s.op != 'DELETE' THEN INSERT (customer_id, name, country, signup_date)
       VALUES (s.customer_id, s.name, s.country, NULL)""")
print(spark.table("pat_customers").count())     # 200: one deleted, one inserted

day = "2025-03-01"
(spark.table("orders").filter(F.to_date("order_ts") == day)
      .withColumn("order_date", F.to_date("order_ts"))
      .write.mode("overwrite").option("replaceWhere", f"order_date = '{day}'")
      .saveAsTable("pat_orders_by_day"))                 # run it twice: same table

%sql
CREATE OR REPLACE TABLE pat_products AS SELECT * FROM products;

MERGE INTO pat_products t
USING (SELECT 1 AS product_id, 'Renamed product' AS product_name, 'Tools' AS category, 9.99 AS unit_price) s
ON t.product_id = s.product_id
WHEN MATCHED THEN UPDATE SET *
WHEN NOT MATCHED THEN INSERT *

%sql
-- the point-in-time join: facts get the dimension row valid at the event time
SELECT f.order_id, d.country
FROM orders f
JOIN dim_customer_scd2 d
  ON f.customer_id = d.customer_id
 AND CAST(f.order_ts AS DATE) >= d.valid_from
 AND (d.valid_to IS NULL OR CAST(f.order_ts AS DATE) < d.valid_to)
LIMIT 5

(spark.table("orders").withColumn("event_type", F.lit("order_placed"))
      .select("order_id", "event_type", F.col("order_ts").alias("event_ts"),
              F.to_json(F.struct("customer_id", "product_id", "quantity", "status")).alias("payload"))
      .write.mode("overwrite").saveAsTable("pat_order_events"))

%sql
-- derive current state from the log
SELECT * FROM pat_order_events
QUALIFY row_number() OVER (PARTITION BY order_id ORDER BY event_ts DESC) = 1
LIMIT 5

df = spark.table("orders_raw")

rules = (F.col("customer_id").isNotNull() & F.col("quantity").isNotNull() & (F.col("quantity") > 0))

good = df.filter(rules)
bad = (df.filter(~rules | rules.isNull())
         .withColumn("_reason", F.when(F.col("customer_id").isNull(), "missing customer")
                                 .when(F.col("quantity").isNull(), "missing quantity")
                                 .otherwise("invalid quantity"))
         .withColumn("_quarantined_at", F.current_timestamp()))

good.write.mode("overwrite").saveAsTable("pat_orders_good")
bad.write.mode("overwrite").saveAsTable("pat_orders_quarantine")
print(good.count(), bad.count())

import random, time

def retry(fn, attempts=4, base=1.0, retry_on=(ConnectionError, TimeoutError)):
    for i in range(attempts):
        try:
            return fn()
        except retry_on as e:
            if i == attempts - 1:
                raise                                    # give up: the job fails loudly
            wait = base * (2 ** i) + random.uniform(0, base)   # 1s, 2s, 4s ... plus jitter
            print(f"attempt {i + 1} failed ({e}); retrying in {wait:.1f}s")
            time.sleep(wait)

calls = {"n": 0}
def flaky():
    calls["n"] += 1
    if calls["n"] < 3:
        raise ConnectionError("temporary")
    return "ok"

print(retry(flaky))       # succeeds on the third attempt

spark.sql("CREATE TABLE IF NOT EXISTS pat_progress (chunk STRING, done_at TIMESTAMP)")

done = {r.chunk for r in spark.table("pat_progress").collect()}
for day in ["2025-03-01", "2025-03-02", "2025-03-03"]:
    if day in done:
        continue                                          # resume: skip finished chunks
    # ... process the day (idempotent write) ...
    spark.createDataFrame([(day,)], "chunk STRING").withColumn("done_at", F.current_timestamp()) \
         .write.mode("append").saveAsTable("pat_progress")

def run_range(start, end):
    days = [r.d for r in spark.sql(
        f"SELECT explode(sequence(DATE'{start}', DATE'{end}', INTERVAL 1 DAY)) AS d").collect()]
    for d in days:
        build_day(str(d))                    # the same function the daily job calls

def build_day(day):
    (spark.table("orders").filter(F.to_date("order_ts") == day)
          .groupBy(F.to_date("order_ts").alias("order_date"))
          .agg(F.count("*").alias("orders"))
          .write.mode("overwrite").option("replaceWhere", f"order_date = '{day}'")
          .saveAsTable("pat_daily_orders"))

run_range("2025-03-01", "2025-03-03")

spark.table("orders").limit(100).write.mode("overwrite").saveAsTable("pat_orders_evo")

# new batch with an extra column
new = spark.table("orders").limit(5).withColumn("coupon", F.lit("SPRING10"))

# default: Delta rejects the unknown column
try:
    new.write.mode("append").saveAsTable("pat_orders_evo")
except Exception as e:
    print("rejected:", str(e)[:80])

# explicit opt-in: add the column to the table
new.write.mode("append").option("mergeSchema", "true").saveAsTable("pat_orders_evo")
print(spark.table("pat_orders_evo").columns)        # now includes coupon

contract = {
    "table": "pat_orders_full",      # the copy of orders from topic 1
    "owner": "sales-data",
    "columns": {"order_id": "int", "customer_id": "int", "quantity": "int", "status": "string"},
    "not_null": ["order_id"],
    "unique": ["order_id"],
    "freshness_hours": 24,
}

def check_contract(spark, c):
    df = spark.table(c["table"])
    types = dict(df.dtypes)
    problems = []
    for col, t in c["columns"].items():
        if types.get(col) != t:
            problems.append(f"{col}: expected {t}, got {types.get(col)}")
    for col in c["not_null"]:
        if df.filter(F.col(col).isNull()).count():
            problems.append(f"{col} has nulls")
    for col in c["unique"]:
        if df.count() != df.select(col).distinct().count():
            problems.append(f"{col} is not unique")
    return problems

print(check_contract(spark, contract))      # [] = holds. A type message just means: adjust the contract to the real types

config = {
    "tables": [
        {"name": "pat_customers_cfg", "source": "customers", "key": "customer_id", "filter": "country IS NOT NULL"},
        {"name": "pat_products_cfg",  "source": "products",  "key": "product_id",  "filter": "unit_price > 0"},
    ]
}

for t in config["tables"]:
    df = spark.table(t["source"]).filter(t["filter"]).dropDuplicates([t["key"]])
    df.write.mode("overwrite").saveAsTable(t["name"])
    print(t["name"], df.count())

spark.sql("""
  CREATE OR REPLACE TABLE pat_meta AS
  SELECT * FROM VALUES
    ('customers', 'pat_customers_meta', 'customer_id', 'country IS NOT NULL', true),
    ('products',  'pat_products_meta',  'product_id',  'unit_price > 0',      true),
    ('orders',    'pat_orders_meta',    'order_id',    'quantity > 0',        false)
  AS t(source, target, key, filter_expr, enabled)
""")

for r in spark.table("pat_meta").filter("enabled").collect():
    df = spark.table(r.source).filter(r.filter_expr).dropDuplicates([r.key])
    df.write.mode("overwrite").saveAsTable(r.target)
    print(r.target, df.count())

from typing import Callable
from pyspark.sql import DataFrame

Transform = Callable[[DataFrame], DataFrame]

def trim_strings(df: DataFrame) -> DataFrame:
    for name, t in df.dtypes:
        if t == "string":
            df = df.withColumn(name, F.trim(F.col(name)))
    return df

def drop_duplicates_on(*keys: str) -> Transform:
    return lambda df: df.dropDuplicates(list(keys))

def add_ingest_metadata(source: str) -> Transform:
    return lambda df: df.withColumn("_ingested_at", F.current_timestamp()).withColumn("_source", F.lit(source))

def apply(df: DataFrame, *steps: Transform) -> DataFrame:
    for step in steps:
        df = step(df)
    return df

result = apply(spark.table("customers"),
               trim_strings, drop_duplicates_on("customer_id"), add_ingest_metadata("customers"))
print(result.count())