Module 7: Data Ingestion

land_orders(1, 1, 1000)
display(dbutils.fs.ls(f"{landing}/batch_1"))

%sql
CREATE TABLE IF NOT EXISTS orders_copy

def copy_into_orders():
    return spark.sql(f"""
      COPY INTO orders_copy
      FROM '{landing}'
      FILEFORMAT = JSON
      PATTERN = 'batch_*/*.json'
      COPY_OPTIONS ('mergeSchema' = 'true')
    """)


copy_into_orders().show()      # 1000 rows loaded

land_orders(2, 1001, 2000)
copy_into_orders().show()                      # 1000 new rows
print(spark.table("orders_copy").count())      # 2000

def run_autoloader():
    stream = (
        spark.readStream.format("cloudFiles")
        .option("cloudFiles.format", "json")
        .option("cloudFiles.inferColumnTypes", "true")
        .option("cloudFiles.schemaLocation", f"{checkpoints}/orders_bronze_schema")
        .option("pathGlobFilter", "*.json")
        .load(landing)
    )
    query = (
        stream.withColumn("_source_file", F.col("_metadata.file_path"))
        .writeStream
        .option("checkpointLocation", f"{checkpoints}/orders_bronze")
        .trigger(availableNow=True)
        .toTable("orders_bronze")
    )
    query.awaitTermination()


run_autoloader()
print(spark.table("orders_bronze").count())      # 2000: both batches

land_orders(3, 1, 10, extra=("coupon", F.lit("SPRING")))

try:
    run_autoloader()
except Exception as e:
    print("stopped:", str(e).splitlines()[0])    # new column found

run_autoloader()
spark.table("orders_bronze").printSchema()     # has "coupon"

spark.sql("DROP TABLE IF EXISTS orders_bronze")
dbutils.fs.rm(f"{checkpoints}/orders_bronze", True)
dbutils.fs.rm(f"{checkpoints}/orders_bronze_schema", True)
run_autoloader()
print(spark.table("orders_bronze").count())     # 2010 if you ran topic 5: every file again

# classic compute pattern: do NOT hard-code passwords
df = (
    spark.read.format("jdbc")
    .option("url", "jdbc:postgresql://db.example.com:5432/shop")
    .option("dbtable", "public.orders")
    .option("user", dbutils.secrets.get("shop", "user"))
    .option("password", dbutils.secrets.get("shop", "password"))
    .option("partitionColumn", "order_id")     # numeric, date or timestamp
    .option("lowerBound", 1)
    .option("upperBound", 2_000_000)
    .option("numPartitions", 8)                # 8 parallel queries
    .load()
)

import json
import time
from datetime import datetime, timezone


def fetch_page(page, page_size=5):
    """Stand-in for requests.get(url, params={'page': page}).json()"""
    start = (page - 1) * page_size
    items = [{"id": i, "name": f"item-{i}"} for i in range(start + 1, start + page_size + 1) if i <= 12]
    return {"items": items, "next_page": page + 1 if items and start + page_size < 12 else None}


def fetch_with_retry(page, attempts=3):
    for n in range(1, attempts + 1):
        try:
            return fetch_page(page)
        except Exception:
            if n == attempts:
                raise
            time.sleep(2 ** n)          # backoff: 2s, 4s


run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
api_dir = f"{base}/landing/api_items/run_{run_id}"
dbutils.fs.mkdirs(api_dir)

page, written = 1, 0
while page:
    body = fetch_with_retry(page)
    for item in body["items"]:
        item["_ingested_at"] = datetime.now(timezone.utc).isoformat()
        item["_source"] = "demo-api"
    with open(f"{api_dir}/page_{page:04d}.json", "w") as f:
        f.write("\n".join(json.dumps(i) for i in body["items"]))
    written += len(body["items"])
    page = body["next_page"]

print(written)      # 12 items over 3 pages

# target: current customers
spark.sql("CREATE OR REPLACE TABLE customers_cdc AS SELECT * FROM customers")

# change feed: insert 201, update 1 and 2, delete 3; two updates for customer 1 (keep the latest)
changes = spark.createDataFrame(
    [
        (1, "Customer 1", "FR", "2024-02-01", "U", 1),
        (1, "Customer 1", "DE", "2024-02-01", "U", 2),
        (2, "Customer 2", "ES", "2024-03-01", "U", 1),
        (3, None, None, None, "D", 1),
        (201, "Customer 201", "NL", "2025-01-01", "I", 1),
    ],
    "customer_id INT, name STRING, country STRING, signup_date STRING, op STRING, seq INT",
).withColumn("signup_date", F.to_date("signup_date"))

# keep only the latest change per key
from pyspark.sql.window import Window

w = Window.partitionBy("customer_id").orderBy(F.desc("seq"))
latest = changes.withColumn("rn", F.row_number().over(w)).filter("rn = 1").drop("rn")
latest.createOrReplaceTempView("changes_latest")

%sql
MERGE INTO customers_cdc AS t
USING changes_latest AS s
ON t.customer_id = s.customer_id
WHEN MATCHED AND s.op = 'D' THEN DELETE
WHEN MATCHED AND s.op = 'U' THEN UPDATE SET
  name = s.name, country = s.country, signup_date = s.signup_date
WHEN NOT MATCHED AND s.op <> 'D' THEN INSERT
  (customer_id, name, country, signup_date) VALUES (s.customer_id, s.name, s.country, s.signup_date)

# initial load: orders up to id 1000
spark.sql("DROP TABLE IF EXISTS orders_incr")
spark.table("orders").filter("order_id <= 1000").write.saveAsTable("orders_incr")


def incremental_load():
    mark = spark.table("orders_incr").agg(F.max("order_id")).first()[0] or 0
    new = spark.table("orders").filter(F.col("order_id") > mark)
    new.write.mode("append").saveAsTable("orders_incr")
    return new.count()


print(incremental_load())     # 1000
print(incremental_load())     # 0: nothing new

batch = spark.table("orders").filter("order_id BETWEEN 1 AND 50")
batch.createOrReplaceTempView("batch_v")

# running this twice gives the same result
spark.sql("""
  MERGE INTO orders_incr t USING batch_v s
  ON t.order_id = s.order_id
  WHEN MATCHED THEN UPDATE SET *
  WHEN NOT MATCHED THEN INSERT *
""")
print(spark.table("orders_incr").count())     # 2000 both times

(batch.write.mode("overwrite")
      .option("replaceWhere", "order_id BETWEEN 1 AND 50")
      .saveAsTable("orders_incr"))

bad_dir = f"{base}/landing/messy"
dbutils.fs.mkdirs(bad_dir)
with open(f"{bad_dir}/messy.json", "w") as f:
    f.write('{"id": 1, "amount": 10.5}\n')
    f.write('{"id": 2, "amount": 7.0}\n')
    f.write('this line is not json\n')
    f.write('{"id": 4, "amount": 3.25}\n')

schema = "id INT, amount DOUBLE, _corrupt_record STRING"
(spark.read.schema(schema).option("mode", "PERMISSIVE").json(bad_dir)
      .write.mode("overwrite").saveAsTable("messy_parsed"))

parsed = spark.table("messy_parsed")
good = parsed.filter(F.col("_corrupt_record").isNull()).drop("_corrupt_record")
bad = (parsed.filter(F.col("_corrupt_record").isNotNull())
             .withColumn("_quarantined_at", F.current_timestamp()))

print(good.count(), bad.count())       # 3 good, 1 bad
bad.write.mode("overwrite").saveAsTable("quarantine_messy")