SCD в PySpark: Type 1, Type 2 и Type 3 через Iceberg MERGE INTO

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

spark = SparkSession.builder \
    .appName("SCD-Type1") \
    .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()

# Создаём Dimension таблицу
spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.gold.dim_customer_scd1 (
        customer_id   STRING     COMMENT 'Бизнес-ключ',
        full_name     STRING,
        email         STRING,
        city          STRING,
        country       STRING,
        segment       STRING,
        _updated_at   TIMESTAMP
    ) USING iceberg
""")

# Начальная загрузка
spark.sql("""
    INSERT INTO lakehouse.gold.dim_customer_scd1 VALUES
    ('cust-001', 'Иван Иванов',  'ivan@mail.ru',  'Казань',  'RU', 'STANDARD', TIMESTAMP '2026-01-01 00:00:00'),
    ('cust-002', 'Alice Smith',  'alice@mail.com', 'London',  'GB', 'PREMIUM',  TIMESTAMP '2026-01-01 00:00:00')
""")

# Батч обновлений из Silver
updates_data = [
    # Иван исправил email (опечатка) + переехал в Москву
    ("cust-001", "Иван Иванов", "ivan.ivanov@mail.ru", "Москва", "RU", "PREMIUM"),
    # Новый клиент
    ("cust-003", "Bob Johnson",  "bob@example.com",     "Berlin", "DE", "STANDARD"),
]
updates_df = spark.createDataFrame(
    updates_data,
    ["customer_id", "full_name", "email", "city", "country", "segment"]
).withColumn("_updated_at", F.current_timestamp())

updates_df.createOrReplaceTempView("scd1_updates")

# SCD Type 1: простой UPSERT
spark.sql("""
    MERGE INTO lakehouse.gold.dim_customer_scd1 AS target
    USING scd1_updates AS source
        ON target.customer_id = source.customer_id
    WHEN MATCHED THEN
        UPDATE SET
            target.full_name   = source.full_name,
            target.email       = source.email,
            target.city        = source.city,
            target.country     = source.country,
            target.segment     = source.segment,
            target._updated_at = source._updated_at
    WHEN NOT MATCHED THEN
        INSERT *
""")

print("SCD Type 1 результат:")
spark.table("lakehouse.gold.dim_customer_scd1").show(truncate=False)
# Иван теперь в Москве. История "Казань" потеряна навсегда.

# Добавляем в WHEN MATCHED условие, что данные реально изменились
spark.sql("""
    MERGE INTO lakehouse.gold.dim_customer_scd1 AS target
    USING scd1_updates AS source
        ON target.customer_id = source.customer_id
    WHEN MATCHED AND (
        target.full_name != source.full_name OR
        target.email     != source.email     OR
        target.city      != source.city      OR
        target.segment   != source.segment
    ) THEN
        UPDATE SET ...
    WHEN NOT MATCHED THEN
        INSERT *
""")

current_segment    STRING   -- текущий тарифный план: PREMIUM
previous_segment   STRING   -- предыдущий: STANDARD
segment_changed_at DATE     -- когда произошло изменение

# Создаём таблицу со структурой Type 3
spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.gold.dim_customer_scd3 (
        customer_id        STRING,
        full_name          STRING,
        current_segment    STRING,
        previous_segment   STRING   COMMENT 'NULL если изменений не было',
        segment_changed_at DATE     COMMENT 'Дата последнего изменения сегмента',
        current_city       STRING,
        previous_city      STRING,
        city_changed_at    DATE
    ) USING iceberg
""")

# Начальная загрузка
spark.sql("""
    INSERT INTO lakehouse.gold.dim_customer_scd3 VALUES
    ('cust-001', 'Иван Иванов', 'STANDARD', NULL, NULL, 'Казань', NULL, NULL),
    ('cust-002', 'Alice Smith',  'PREMIUM',  NULL, NULL, 'London', NULL, NULL)
""")

# Батч: Иван апгрейдился до PREMIUM, переехал в Москву
scd3_updates = [
    ("cust-001", "Иван Иванов", "PREMIUM", "Москва"),
]
scd3_df = spark.createDataFrame(
    scd3_updates,
    ["customer_id", "full_name", "new_segment", "new_city"]
)
scd3_df.createOrReplaceTempView("scd3_updates")

# SCD Type 3: при MATCHED - сдвигаем current → previous, записываем new → current
spark.sql("""
    MERGE INTO lakehouse.gold.dim_customer_scd3 AS target
    USING scd3_updates AS source
        ON target.customer_id = source.customer_id
    WHEN MATCHED AND target.current_segment != source.new_segment THEN
        UPDATE SET
            target.previous_segment   = target.current_segment,
            target.current_segment    = source.new_segment,
            target.segment_changed_at = CURRENT_DATE(),
            target.previous_city      = target.current_city,
            target.current_city       = source.new_city,
            target.city_changed_at    = CURRENT_DATE()
    WHEN NOT MATCHED THEN
        INSERT (customer_id, full_name, current_segment, previous_segment,
                segment_changed_at, current_city, previous_city, city_changed_at)
        VALUES (source.customer_id, source.full_name, source.new_segment, NULL,
                CURRENT_DATE(), source.new_city, NULL, NULL)
""")

print("SCD Type 3 результат:")
spark.table("lakehouse.gold.dim_customer_scd3").show(truncate=False)

-- Этот код НЕ работает для SCD Type 2!
MERGE INTO dim_customer AS target
USING updates AS source ON target.customer_id = source.customer_id AND target.is_current = true
WHEN MATCHED AND data_changed THEN
    -- Хотим: закрыть старую строку И вставить новую
    UPDATE SET target.is_current = false, target.valid_to = today
    -- INSERT новой строки здесь невозможен!
WHEN NOT MATCHED THEN
    INSERT new row

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

spark = SparkSession.builder \
    .appName("SCD-Type2-Full-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()

# ─── 1. Создаём SCD Type 2 таблицу ───────────────────────────────────────────
spark.sql("""
    CREATE TABLE IF NOT EXISTS lakehouse.gold.dim_users_scd2 (
        user_id       STRING     COMMENT 'Бизнес-ключ',
        full_name     STRING,
        city          STRING,
        segment       STRING,
        valid_from    DATE       COMMENT 'Начало действия этой версии',
        valid_to      DATE       COMMENT '9999-12-31 для текущей версии',
        is_current    BOOLEAN    COMMENT 'true только для одной строки по user_id'
    ) USING iceberg
""")

# ─── 2. Начальная загрузка ────────────────────────────────────────────────────
initial_date = "2026-01-01"
spark.sql(f"""
    INSERT INTO lakehouse.gold.dim_users_scd2 VALUES
    ('usr-001', 'Иван Иванов',  'Казань', 'STANDARD', DATE '{initial_date}', DATE '9999-12-31', true),
    ('usr-002', 'Alice Smith',  'London', 'PREMIUM',  DATE '{initial_date}', DATE '9999-12-31', true),
    ('usr-003', 'Bob Johnson',  'Berlin', 'STANDARD', DATE '{initial_date}', DATE '9999-12-31', true)
""")

print("=== Состояние ПОСЛЕ начальной загрузки ===")
spark.table("lakehouse.gold.dim_users_scd2").orderBy("user_id", "valid_from").show(truncate=False)


# ─── 3. Функция SCD Type 2 обновления ─────────────────────────────────────────
def apply_scd2_update(updates_df, effective_date: str):
    """
    Применяет SCD Type 2 обновление к dim_users_scd2.

    Паттерн Source Expansion:
    - Первая копия обновлений (с реальным user_id) → WHEN MATCHED → UPDATE (закрыть старую строку)
    - Вторая копия (с NULL merge_key) → WHEN NOT MATCHED → INSERT (открыть новую версию)

    Условие совпадения ONLY для is_current=true гарантирует, что мы работаем
    только с актуальной версией, не трогая историческиe записи.
    """
    updates_df.createOrReplaceTempView("raw_scd2_updates")

    # Source Expansion через UNION ALL
    # Первая часть: строки для UPDATE (merge_key = реальный user_id, совпадёт в WHEN MATCHED)
    # Вторая часть: строки для INSERT (merge_key = NULL, не совпадёт, попадёт в WHEN NOT MATCHED)
    spark.sql(f"""
        CREATE OR REPLACE TEMP VIEW staged_scd2_updates AS

        -- Строки для обновления существующих записей (WHEN MATCHED → UPDATE)
        SELECT
            u.user_id    AS merge_key,
            u.user_id,
            u.full_name,
            u.city,
            u.segment,
            DATE '{effective_date}' AS effective_date
        FROM raw_scd2_updates u

        UNION ALL

        -- Строки для вставки новых версий (WHEN NOT MATCHED → INSERT).
        -- Выбираем только записи, где данные РЕАЛЬНО изменились.
        -- NULL в merge_key гарантирует, что MERGE не найдёт совпадений
        -- и отправит эти строки в ветку WHEN NOT MATCHED.
        SELECT
            NULL          AS merge_key,
            u.user_id,
            u.full_name,
            u.city,
            u.segment,
            DATE '{effective_date}' AS effective_date
        FROM raw_scd2_updates u
        JOIN lakehouse.gold.dim_users_scd2 t
            ON u.user_id = t.user_id AND t.is_current = true
        WHERE
            u.city     <> t.city     OR
            u.segment  <> t.segment  OR
            u.full_name <> t.full_name
    """)

    # Основной MERGE
    spark.sql(f"""
        MERGE INTO lakehouse.gold.dim_users_scd2 AS target
        USING staged_scd2_updates AS source
            ON target.user_id = source.merge_key AND target.is_current = true
        WHEN MATCHED AND (
            target.city      <> source.city      OR
            target.segment   <> source.segment   OR
            target.full_name <> source.full_name
        ) THEN
            -- Закрываем текущую версию строки
            UPDATE SET
                target.valid_to   = source.effective_date,
                target.is_current = false
        WHEN NOT MATCHED THEN
            -- Вставляем новую актуальную версию
            INSERT (user_id, full_name, city, segment, valid_from, valid_to, is_current)
            VALUES (
                source.user_id,
                source.full_name,
                source.city,
                source.segment,
                source.effective_date,
                DATE '9999-12-31',
                true
            )
    """)


# ─── 4. Батч 1: Иван переезжает в Москву ─────────────────────────────────────
print("\n=== Батч 1: Иван переезжает в Москву (2026-03-01) ===")
batch1_data = [
    ("usr-001", "Иван Иванов", "Москва",  "STANDARD"),  # переезд
    ("usr-002", "Alice Smith", "London",  "PREMIUM"),   # не изменился
]
batch1_df = spark.createDataFrame(batch1_data, ["user_id", "full_name", "city", "segment"])
apply_scd2_update(batch1_df, "2026-03-01")

spark.table("lakehouse.gold.dim_users_scd2").orderBy("user_id", "valid_from").show(truncate=False)


# ─── 5. Батч 2: Иван апгрейдится до PREMIUM ──────────────────────────────────
print("\n=== Батч 2: Иван апгрейдился до PREMIUM (2026-05-15) ===")
batch2_data = [
    ("usr-001", "Иван Иванов", "Москва", "PREMIUM"),   # апгрейд
    ("usr-004", "New User",    "Paris",  "TRIAL"),     # новый клиент
]
batch2_df = spark.createDataFrame(batch2_data, ["user_id", "full_name", "city", "segment"])
apply_scd2_update(batch2_df, "2026-05-15")

print("Финальное состояние таблицы:")
spark.table("lakehouse.gold.dim_users_scd2").orderBy("user_id", "valid_from").show(truncate=False)


# ─── 6. Point-in-time запросы ─────────────────────────────────────────────────
print("\n=== Point-in-time запрос: состояние клиентов на 15 февраля 2026 ===")
spark.sql("""
    SELECT user_id, full_name, city, segment, valid_from, valid_to
    FROM lakehouse.gold.dim_users_scd2
    WHERE DATE '2026-02-15' BETWEEN valid_from AND valid_to
    ORDER BY user_id
""").show(truncate=False)
# Иван должен быть в Казани, Alice в London - это состояние до его переезда

print("\n=== Текущее состояние всех клиентов ===")
spark.sql("""
    SELECT user_id, full_name, city, segment
    FROM lakehouse.gold.dim_users_scd2
    WHERE is_current = true
    ORDER BY user_id
""").show(truncate=False)

def apply_late_arriving_scd2(late_event_user_id: str, late_event_city: str,
                              late_event_date: str):
    """
    Обрабатывает запоздавшее событие для SCD Type 2.
    Вставляет новую версию строго в нужный временно́й слот, не нарушая хронологию.
    """
    # Находим строку, чей интервал содержит дату опоздавшего события
    target_row = spark.sql(f"""
        SELECT user_id, full_name, city, segment, valid_from, valid_to, is_current
        FROM lakehouse.gold.dim_users_scd2
        WHERE user_id = '{late_event_user_id}'
          AND DATE '{late_event_date}' >= valid_from
          AND DATE '{late_event_date}' < valid_to
    """).collect()

    if not target_row:
        print(f"Нет строки для user_id={late_event_user_id} на дату {late_event_date}")
        return

    row = target_row[0]
    original_valid_to = row["valid_to"]
    original_is_current = row["is_current"]

    # Шаг 1: Закрываем найденную строку на дату опоздавшего события
    spark.sql(f"""
        UPDATE lakehouse.gold.dim_users_scd2
        SET valid_to = DATE '{late_event_date}', is_current = false
        WHERE user_id = '{late_event_user_id}'
          AND valid_from = DATE '{row["valid_from"]}'
          AND valid_to = DATE '{original_valid_to}'
    """)

    # Шаг 2: Вставляем новую строку для опоздавшего события
    # valid_to берём из СТАРОЙ строки (может быть как 9999-12-31, так и промежуточной датой)
    spark.sql(f"""
        INSERT INTO lakehouse.gold.dim_users_scd2 VALUES (
            '{late_event_user_id}',
            '{row["full_name"]}',
            '{late_event_city}',
            '{row["segment"]}',
            DATE '{late_event_date}',
            DATE '{original_valid_to}',
            {str(original_is_current).lower()}
        )
    """)
    print(f"Late-arriving event обработан: {late_event_user_id} {late_event_city} {late_event_date}")


# Применяем: московский переезд (15 мая) пришёл после питерского (20 мая)
apply_late_arriving_scd2("usr-001", "Москва", "2026-05-15")

print("\n=== Временна́я шкала после обработки запоздавшего события ===")
spark.sql("""
    SELECT user_id, city, segment, valid_from, valid_to, is_current
    FROM lakehouse.gold.dim_users_scd2
    WHERE user_id = 'usr-001'
    ORDER BY valid_from
""").show(truncate=False)
# Ожидаем: Казань → Москва → Питер (хронологически правильно)

# Запускается после нескольких SCD-обновлений, например раз в сутки
spark.sql("""
    CALL lakehouse.system.rewrite_data_files(
        table => 'lakehouse.gold.dim_users_scd2',
        strategy => 'binpack',
        options => map('target-file-size-bytes', '134217728')
    )
""")

# Очищаем старые snapshot'ы (метаданные прошлых версий)
spark.sql("""
    CALL lakehouse.system.expire_snapshots(
        table => 'lakehouse.gold.dim_users_scd2',
        older_than => TIMESTAMP '2026-05-01 00:00:00',
        retain_last => 5
    )
""")