Data Vault 2.0 на Spark: Hubs, Links, Satellites и параллельный инжект

from pyspark.sql import functions as F

# Правило 1: Нормализация кейса и пробелов
customer_hk = F.sha2(
    F.trim(F.lower(F.col("customer_id"))),
    256
)
# 'CUST-99', ' cust-99 ', 'Cust-99' - все дают одинаковый хэш

# Правило 2: Составной бизнес-ключ - конкатенация с фиксированным разделителем
# '||' как разделитель избегает коллизий: concat('AB','CD') vs concat('A','BCD')
order_product_lnk_hk = F.sha2(
    F.concat_ws("||",
        F.trim(F.lower(F.col("order_id"))),
        F.trim(F.lower(F.col("product_id")))
    ),
    256
)

# Правило 3: NULL в бизнес-ключе - НИКОГДА не хэшировать
# coalesce нельзя: sha2('') даст один хэш для всех NULL-клиентов → Data Skew!
# Правильно: фильтровать битые записи ДО хэширования
valid_df = bronze_df.filter(F.col("customer_id").isNotNull())
dead_letter_df = bronze_df.filter(F.col("customer_id").isNull())

# Считаем hash_diff как хэш от всех атрибутов сателлита
hash_diff = F.sha2(
    F.concat_ws("||",
        F.coalesce(F.col("full_name"),  F.lit("NULL")),
        F.coalesce(F.col("city"),       F.lit("NULL")),
        F.coalesce(F.col("segment"),    F.lit("NULL")),
        F.coalesce(F.col("email"),      F.lit("NULL")),
    ),
    256
)

# Новые данные сравниваем с последней версией в Satellite
# Вставляем только те строки, где hash_diff изменился
new_records = incoming_df.join(
    last_version_df.select("hub_customer_hk", "hash_diff"),
    on="hub_customer_hk",
    how="left_anti"  # строки, которых нет в last_version
).union(
    incoming_df.join(
        last_version_df.select("hub_customer_hk", "hash_diff"),
        on=["hub_customer_hk"],
        how="inner"
    ).filter(F.col("incoming_hash_diff") != F.col("existing_hash_diff"))
)

# Эффективный APPEND с дедупликацией для Hub
existing_keys = spark.table("lakehouse.silver.hub_customer") \
    .select("business_key")

new_hub_rows = staged_df \
    .select("hk_customer", "customer_id", "load_ts", "record_source") \
    .distinct() \
    .join(existing_keys,
          staged_df["customer_id"] == existing_keys["business_key"],
          how="left_anti")  # только те, которых ещё нет

new_hub_rows.writeTo("lakehouse.silver.hub_customer").append()

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder \
    .appName("DataVault-2.0-Pipeline") \
    .config("spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions") \
    .config("spark.sql.catalog.lakehouse",
            "org.apache.iceberg.spark.SparkCatalog") \
    .config("spark.sql.catalog.lakehouse.type", "hadoop") \
    .config("spark.sql.catalog.lakehouse.warehouse", "/tmp/warehouse") \
    .getOrCreate()

RECORD_SOURCE = "ecom_orders_api"

# ─── 0. Создаём DV-таблицы ────────────────────────────────────────────────────
spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.silver.hub_customer (
        hub_hk         STRING     COMMENT 'SHA-256 от business_key',
        business_key   STRING     COMMENT 'Оригинальный customer_id из источника',
        load_ts        TIMESTAMP,
        record_source  STRING
    ) USING iceberg
    PARTITIONED BY (days(load_ts))
""")

spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.silver.hub_order (
        hub_hk         STRING,
        business_key   STRING,
        load_ts        TIMESTAMP,
        record_source  STRING
    ) USING iceberg
    PARTITIONED BY (days(load_ts))
""")

spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.silver.hub_product (
        hub_hk         STRING,
        business_key   STRING,
        load_ts        TIMESTAMP,
        record_source  STRING
    ) USING iceberg
    PARTITIONED BY (days(load_ts))
""")

spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.silver.lnk_order_customer (
        lnk_hk         STRING     COMMENT 'SHA-256 от order_id||customer_id',
        hk_order       STRING     COMMENT 'FK → hub_order',
        hk_customer    STRING     COMMENT 'FK → hub_customer',
        load_ts        TIMESTAMP,
        record_source  STRING
    ) USING iceberg
    PARTITIONED BY (days(load_ts))
""")

spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.silver.lnk_order_product (
        lnk_hk         STRING,
        hk_order       STRING,
        hk_product     STRING,
        load_ts        TIMESTAMP,
        record_source  STRING
    ) USING iceberg
    PARTITIONED BY (days(load_ts))
""")

spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.silver.sat_customer_details (
        hk_customer    STRING     COMMENT 'FK → hub_customer',
        load_ts        TIMESTAMP  COMMENT 'Начало действия этой версии',
        load_end_ts    TIMESTAMP  COMMENT 'NULL = текущая версия',
        hash_diff      STRING     COMMENT 'SHA-256 от всех атрибутов',
        record_source  STRING,
        full_name      STRING,
        city           STRING,
        segment        STRING,
        email          STRING
    ) USING iceberg
    PARTITIONED BY (days(load_ts))
""")

spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.silver.sat_order_context (
        hk_order       STRING,
        load_ts        TIMESTAMP,
        load_end_ts    TIMESTAMP,
        hash_diff      STRING,
        record_source  STRING,
        product_name   STRING,
        amount         DECIMAL(10,2),
        status         STRING
    ) USING iceberg
    PARTITIONED BY (days(load_ts))
""")

# ─── 1. Сырые данные из Bronze ────────────────────────────────────────────────
bronze_data = [
    # order_id, customer_id, product_id, customer_name, city, segment, email, product_name, amount, status
    ("ORD-100", "CUST-99", "PROD-001", "Ivan Ivanov",  "Kazan",  "STANDARD", "ivan@m.ru",    "Laptop",  1200.0, "CREATED"),
    ("ORD-100", "CUST-99", "PROD-002", "Ivan Ivanov",  "Kazan",  "STANDARD", "ivan@m.ru",    "Mouse",     50.0, "CREATED"),
    ("ORD-101", "CUST-42", "PROD-001", "Alice Smith",  "London", "PREMIUM",  "alice@m.com",  "Laptop",   990.0, "PAID"),
    # Битая запись: customer_id = NULL → должна попасть в DLQ
    ("ORD-102", None,      "PROD-003", None,           None,     None,       None,           "Book",      15.0, "CREATED"),
]
cols = ["order_id", "customer_id", "product_id", "customer_name", "city",
        "segment", "email", "product_name", "amount", "status"]
bronze_df = spark.createDataFrame(bronze_data, cols)

# ─── 2. Валидация: фильтрация битых записей в DLQ ─────────────────────────────
# Ключевое правило: НИКОГДА не хэшировать NULL бизнес-ключи
valid_df      = bronze_df.filter(F.col("customer_id").isNotNull())
dead_letter_df = bronze_df.filter(F.col("customer_id").isNull())

if dead_letter_df.count() > 0:
    print(f"⚠️  DLQ: {dead_letter_df.count()} записей с NULL business_key отфильтрованы")
    # В production: сохранить dead_letter_df в отдельную таблицу для разбора

# ─── 3. Staging: вычисляем все хэши локально ─────────────────────────────────
load_ts = F.current_timestamp()

staged_df = valid_df \
    .withColumn("hk_customer",
        F.sha2(F.trim(F.lower(F.col("customer_id"))), 256)) \
    .withColumn("hk_order",
        F.sha2(F.trim(F.lower(F.col("order_id"))), 256)) \
    .withColumn("hk_product",
        F.sha2(F.trim(F.lower(F.col("product_id"))), 256)) \
    .withColumn("lnk_order_customer_hk",
        F.sha2(F.concat_ws("||",
            F.trim(F.lower(F.col("order_id"))),
            F.trim(F.lower(F.col("customer_id")))
        ), 256)) \
    .withColumn("lnk_order_product_hk",
        F.sha2(F.concat_ws("||",
            F.trim(F.lower(F.col("order_id"))),
            F.trim(F.lower(F.col("product_id")))
        ), 256)) \
    .withColumn("hash_diff_customer",
        F.sha2(F.concat_ws("||",
            F.coalesce(F.col("customer_name"), F.lit("NULL")),
            F.coalesce(F.col("city"),          F.lit("NULL")),
            F.coalesce(F.col("segment"),       F.lit("NULL")),
            F.coalesce(F.col("email"),         F.lit("NULL")),
        ), 256)) \
    .withColumn("hash_diff_order",
        F.sha2(F.concat_ws("||",
            F.coalesce(F.col("product_name"), F.lit("NULL")),
            F.coalesce(F.col("amount").cast("string"), F.lit("NULL")),
            F.coalesce(F.col("status"),       F.lit("NULL")),
        ), 256)) \
    .withColumn("load_ts", load_ts) \
    .withColumn("record_source", F.lit(RECORD_SOURCE))

staged_df.cache()  # кэшируем - от него пойдут параллельные ветки

# ─── 4. Load Pattern: Hub Customer ───────────────────────────────────────────
existing_cust_keys = spark.table("lakehouse.silver.hub_customer").select("business_key")

hub_customer_new = staged_df \
    .select("hk_customer", "customer_id", "load_ts", "record_source") \
    .withColumnRenamed("hk_customer", "hub_hk") \
    .withColumnRenamed("customer_id", "business_key") \
    .distinct() \
    .join(existing_cust_keys, on="business_key", how="left_anti")

hub_customer_new.writeTo("lakehouse.silver.hub_customer").append()
print(f"hub_customer: +{hub_customer_new.count()} новых записей")

# ─── 5. Load Pattern: Hub Order ───────────────────────────────────────────────
existing_order_keys = spark.table("lakehouse.silver.hub_order").select("business_key")

hub_order_new = staged_df \
    .select("hk_order", "order_id", "load_ts", "record_source") \
    .withColumnRenamed("hk_order", "hub_hk") \
    .withColumnRenamed("order_id", "business_key") \
    .distinct() \
    .join(existing_order_keys, on="business_key", how="left_anti")

hub_order_new.writeTo("lakehouse.silver.hub_order").append()

# ─── 6. Load Pattern: Hub Product ─────────────────────────────────────────────
existing_prod_keys = spark.table("lakehouse.silver.hub_product").select("business_key")

hub_product_new = staged_df \
    .select("hk_product", "product_id", "load_ts", "record_source") \
    .withColumnRenamed("hk_product", "hub_hk") \
    .withColumnRenamed("product_id", "business_key") \
    .distinct() \
    .join(existing_prod_keys, on="business_key", how="left_anti")

hub_product_new.writeTo("lakehouse.silver.hub_product").append()

# ─── 7. Load Pattern: Link Order-Customer ────────────────────────────────────
existing_loc_links = spark.table("lakehouse.silver.lnk_order_customer").select("lnk_hk")

lnk_order_customer_new = staged_df \
    .select("lnk_order_customer_hk", "hk_order", "hk_customer", "load_ts", "record_source") \
    .withColumnRenamed("lnk_order_customer_hk", "lnk_hk") \
    .distinct() \
    .join(existing_loc_links, on="lnk_hk", how="left_anti")

lnk_order_customer_new.writeTo("lakehouse.silver.lnk_order_customer").append()

# ─── 8. Load Pattern: Link Order-Product ─────────────────────────────────────
existing_lop_links = spark.table("lakehouse.silver.lnk_order_product").select("lnk_hk")

lnk_order_product_new = staged_df \
    .select("lnk_order_product_hk", "hk_order", "hk_product", "load_ts", "record_source") \
    .withColumnRenamed("lnk_order_product_hk", "lnk_hk") \
    .distinct() \
    .join(existing_lop_links, on="lnk_hk", how="left_anti")

lnk_order_product_new.writeTo("lakehouse.silver.lnk_order_product").append()

# ─── 9. Load Pattern: Satellite Customer Details ──────────────────────────────
# Для Satellite: вставляем только строки с изменённым hash_diff
existing_sat_cust = spark.table("lakehouse.silver.sat_customer_details") \
    .filter(F.col("load_end_ts").isNull()) \
    .select("hk_customer", "hash_diff")

sat_customer_new = staged_df \
    .select("hk_customer", "load_ts", "hash_diff_customer",
            "record_source", "customer_name", "city", "segment", "email") \
    .withColumn("load_end_ts", F.lit(None).cast("timestamp")) \
    .withColumnRenamed("hash_diff_customer", "hash_diff") \
    .withColumnRenamed("customer_name", "full_name") \
    .distinct() \
    .join(existing_sat_cust,
          on=["hk_customer", "hash_diff"],
          how="left_anti")  # только строки с изменившимися атрибутами

sat_customer_new.writeTo("lakehouse.silver.sat_customer_details").append()

# ─── 10. Load Pattern: Satellite Order Context ────────────────────────────────
existing_sat_order = spark.table("lakehouse.silver.sat_order_context") \
    .filter(F.col("load_end_ts").isNull()) \
    .select("hk_order", "hash_diff")

sat_order_new = staged_df \
    .select("hk_order", "load_ts", "hash_diff_order",
            "record_source", "product_name", "amount", "status") \
    .withColumn("load_end_ts", F.lit(None).cast("timestamp")) \
    .withColumnRenamed("hash_diff_order", "hash_diff") \
    .withColumn("amount", F.col("amount").cast("decimal(10,2)")) \
    .distinct() \
    .join(existing_sat_order,
          on=["hk_order", "hash_diff"],
          how="left_anti")

sat_order_new.writeTo("lakehouse.silver.sat_order_context").append()

staged_df.unpersist()

print("\n✅ Data Vault 2.0 загрузка завершена!")

# ─── 11. Пример реконструкции сущности ───────────────────────────────────────
print("\n=== Реконструкция: текущее состояние клиентов ===")
spark.sql("""
    SELECT
        h.business_key AS customer_id,
        s.full_name,
        s.city,
        s.segment,
        s.load_ts
    FROM lakehouse.silver.hub_customer h
    JOIN lakehouse.silver.sat_customer_details s
        ON h.hub_hk = s.hk_customer
    WHERE s.load_end_ts IS NULL
    ORDER BY h.business_key
""").show(truncate=False)

-- PIT-таблица для hub_customer
SELECT
    hc.hub_hk      AS hk_customer,
    snap.snap_date,
    -- Для каждого satellite: ближайший load_ts не позже snap_date
    MAX(sc.load_ts) AS sat_cust_details_load_ts,
    MAX(se.load_ts) AS sat_cust_email_load_ts
FROM hub_customer hc
CROSS JOIN snapshot_dates snap
LEFT JOIN sat_customer_details sc
    ON hc.hub_hk = sc.hk_customer
    AND sc.load_ts <= snap.snap_date
    AND (sc.load_end_ts > snap.snap_date OR sc.load_end_ts IS NULL)
LEFT JOIN sat_customer_email se
    ON hc.hub_hk = se.hk_customer
    AND se.load_ts <= snap.snap_date
    AND (se.load_end_ts > snap.snap_date OR se.load_end_ts IS NULL)
GROUP BY hc.hub_hk, snap.snap_date