Hash-ключ для составных PK: sha2/xxhash64 - dedup по 10+ колонкам без shuffle penalty

from pyspark.sql import functions as F

def dv2_hash_key(*business_keys: str) -> F.Column:
    """
    Data Vault 2.0 стандарт: SHA-256 от конкатенации бизнес-ключей.
    Колонки приводятся к UPPER CASE, NULL → 'NULL'.
    Разделитель: || (двойной пайп - стандарт DV2).
    """
    normalized = [
        F.upper(F.coalesce(F.col(k).cast("string"), F.lit("NULL")))
        for k in business_keys
    ]
    return F.sha2(F.concat_ws("||", *normalized), 256)

# Применение
df = df.withColumn(
    "hash_key",
    dv2_hash_key("company_code", "plant", "material_number")
)

from pyspark.sql import functions as F

# xxhash64 принимает несколько колонок напрямую
df = df.withColumn(
    "hash_id",
    F.xxhash64(
        F.col("company_code"),
        F.col("plant"),
        F.col("material_number"),
        F.col("sales_org"),
        F.col("fiscal_year"),
    )
)

# НЕПРАВИЛЬНО: collision гарантирован
df = df.withColumn("hash_id", F.sha2(F.concat("col_a", "col_b"), 256))

# ПРАВИЛЬНО: разделитель исключает ложные совпадения
df = df.withColumn(
    "hash_id",
    F.sha2(F.concat_ws("||", F.col("col_a"), F.col("col_b")), 256)
)

# concat_ws("||", "A", NULL, "C") → "A||C"
# concat_ws("||", "A", "B",  "C") → "A||B||C"
# Но если col_b = NULL в одной строке, 
# а col_b = "B" в другой, итоговые строки
# "A||C" vs "A||B||C" - разные. Это правильно!
# НО: если col_b = NULL и col_c = NULL:
# concat_ws("||", "A", NULL, NULL) → "A"
# concat_ws("||", NULL, "A", NULL) → "A"
# Коллизия между (A, NULL, NULL) и (NULL, A, NULL)!

from pyspark.sql import functions as F
from pyspark.sql import DataFrame
from typing import List

def make_sha256_hash_key(df: DataFrame, key_cols: List[str]) -> DataFrame:
    """
    Правильное SHA-256 hash-ключ для составного PK.

    Нормализация:
    1. Все колонки → STRING
    2. NULL → литерал 'NULL'
    3. UPPER CASE (опционально, для case-insensitive сравнений)
    4. TRIM (убираем пробелы)
    5. Конкатенация через '||' разделитель
    6. SHA-256
    """
    normalized_cols = [
        F.trim(
            F.upper(
                F.coalesce(
                    F.col(c).cast("string"),
                    F.lit("NULL")
                )
            )
        )
        for c in key_cols
    ]

    hash_input = F.concat_ws("||", *normalized_cols)

    return df.withColumn("hash_key", F.sha2(hash_input, 256))


# Использование
KEY_COLS = [
    "company_code", "plant", "material_number",
    "sales_org", "distribution_channel", "customer_group",
    "document_type", "fiscal_year", "period", "line_item"
]

df_with_hash = make_sha256_hash_key(df_raw, KEY_COLS)

def make_xxhash64_key(df: DataFrame, key_cols: List[str]) -> DataFrame:
    """
    xxhash64 для ETL дедупликации.

    ВАЖНО: xxhash64 нативно обрабатывает NULL как отдельное значение.
    (NULL, 1, "A") даёт другой хэш, чем (0, 1, "A").

    Нормализация строк для консистентности:
    - TRIM убирает случайные пробелы
    - LOWER для case-insensitive ключей (опционально)
    """
    normalized_cols = [
        F.trim(F.col(c).cast("string"))
        if str(df.schema[c].dataType) in ("StringType()", "string")
        else F.col(c)
        for c in key_cols
    ]

    return df.withColumn("hash_id", F.xxhash64(*normalized_cols))


# Для числовых колонок нормализация строк не нужна:
df = df.withColumn(
    "hash_id",
    F.xxhash64(
        F.trim(F.col("company_code")),      # STRING → trim
        F.col("plant").cast("int"),          # INT → без изменений
        F.trim(F.col("material_number")),    # STRING → trim
        F.col("fiscal_year").cast("int"),    # INT → без изменений
        F.col("amount").cast("decimal(18,2)")  # DECIMAL → без изменений
    )
)

# Подход 1: dropDuplicates по 10 бизнес-колонкам
df_dedup_v1 = df.dropDuplicates([
    "company_code", "plant", "material_number",
    "sales_org", "distribution_channel", "customer_group",
    "document_type", "fiscal_year", "period", "line_item"
])
df_dedup_v1.explain(mode="formatted")

== Physical Plan ==
HashAggregate(keys=[company_code, plant, material_number, sales_org,
                    distribution_channel, customer_group, document_type,
                    fiscal_year, period, line_item],  ← 10 колонок в ключе
             functions=[], output=[...])
+- Exchange hashpartitioning(company_code, plant, material_number, sales_org,
                              distribution_channel, customer_group,
                              document_type, fiscal_year, period, line_item,
                              200),   ← Shuffle по 10 колонкам
   ENSURE_REQUIREMENTS
   +- HashAggregate(keys=[...10 cols...], functions=[], ...)
      +- FileScan parquet [...]

# Подход 2: dropDuplicates по одному hash_id
df_with_hash = df.withColumn("hash_id", F.xxhash64(
    F.col("company_code"), F.col("plant"), F.col("material_number"),
    F.col("sales_org"), F.col("distribution_channel"), F.col("customer_group"),
    F.col("document_type"), F.col("fiscal_year"), F.col("period"), F.col("line_item")
))

df_dedup_v2 = df_with_hash.dropDuplicates(["hash_id"])
df_dedup_v2.explain(mode="formatted")

== Physical Plan ==
HashAggregate(keys=[hash_id],   ← 1 колонка в ключе (BIGINT 8 байт)
             functions=[], output=[...])
+- Exchange hashpartitioning(hash_id, 200),   ← Shuffle по 1 BIGINT!
   ENSURE_REQUIREMENTS
   +- HashAggregate(keys=[hash_id], functions=[], ...)
      +- Project [company_code, ..., xxhash64(...) AS hash_id]  ← вычисляется до shuffle
         +- FileScan parquet [...]

from pyspark.sql import SparkSession, functions as F, DataFrame
from pyspark.sql.types import StringType, LongType
from typing import List
import logging

logger = logging.getLogger(__name__)


def build_hash_key(
    df: DataFrame,
    key_cols: List[str],
    hash_type: str = "xxhash64",
    hash_col: str = "hash_id",
) -> DataFrame:
    """
    Создаёт суррогатный hash-ключ.

    Args:
        df: входной DataFrame
        key_cols: список колонок, образующих составной PK
        hash_type: 'xxhash64' (быстро, BIGINT) или 'sha256' (надёжно, STRING)
        hash_col: имя результирующей колонки

    Returns:
        DataFrame с добавленной колонкой hash_col
    """
    if hash_type == "xxhash64":
        # Для строковых колонок делаем TRIM, для числовых - без изменений
        string_types = {"StringType()", "string"}
        normalized = [
            F.trim(F.col(c))
            if str(df.schema[c].dataType) in string_types
            else F.col(c)
            for c in key_cols
        ]
        hash_expr = F.xxhash64(*normalized)

    elif hash_type == "sha256":
        # Полная нормализация: TRIM → COALESCE → UPPER → concat_ws → sha2
        normalized = [
            F.trim(
                F.upper(
                    F.coalesce(F.col(c).cast("string"), F.lit("NULL"))
                )
            )
            for c in key_cols
        ]
        hash_expr = F.sha2(F.concat_ws("||", *normalized), 256)

    else:
        raise ValueError(f"Unsupported hash_type: {hash_type}. Use 'xxhash64' or 'sha256'")

    return df.withColumn(hash_col, hash_expr)


def validate_hash_uniqueness(
    df: DataFrame,
    hash_col: str,
    key_cols: List[str],
    sample_fraction: float = 0.01,
) -> None:
    """
    Проверяет отсутствие коллизий в hash-ключе.
    Выполняется на выборке (sample_fraction) для скорости.
    При нахождении коллизий - логирует WARNING и примеры.
    """
    sample = df.sample(fraction=sample_fraction, seed=42)

    # Ищем строки с одинаковым hash_id, но разными бизнес-ключами
    collision_check = (
        sample
        .groupBy(hash_col)
        .agg(
            F.count("*").alias("cnt"),
            F.countDistinct(*[F.col(c) for c in key_cols]).alias("distinct_keys"),
        )
        .filter(F.col("distinct_keys") > 1)
    )

    collision_count = collision_check.count()
    if collision_count > 0:
        logger.warning(
            f"COLLISION DETECTED: {collision_count} hash values map to multiple business keys! "
            f"Consider switching to sha256 or adding secondary key validation."
        )
        collision_check.show(5)
    else:
        logger.info(
            f"No collisions found in {sample_fraction*100:.0f}% sample. Hash key looks safe."
        )

-- Без hash-ключа: 10 условий JOIN
MERGE INTO target t
USING source s
ON t.company_code = s.company_code
   AND t.plant = s.plant
   AND t.material_number = s.material_number
   AND t.sales_org = s.sales_org
   AND t.distribution_channel = s.distribution_channel
   AND t.customer_group = s.customer_group
   AND t.document_type = s.document_type
   AND t.fiscal_year = s.fiscal_year
   AND t.period = s.period
   AND t.line_item = s.line_item
WHEN MATCHED THEN UPDATE ...
WHEN NOT MATCHED THEN INSERT ...

-- С hash-ключом: 1 условие JOIN
MERGE INTO target t
USING source s
ON t.hash_id = s.hash_id          -- 1 BIGINT сравнение!
WHEN MATCHED AND t.hash_diff != s.hash_diff THEN UPDATE ...
WHEN NOT MATCHED THEN INSERT ...

KEY_COLS = ["company_code", "plant", "material_number", "fiscal_year", "line_item"]
VALUE_COLS = ["amount", "quantity", "unit", "currency", "updated_at"]

df_source = (
    df_raw
    .transform(lambda df: build_hash_key(df, KEY_COLS, hash_type="xxhash64", hash_col="hash_id"))
    .transform(lambda df: build_hash_key(df, VALUE_COLS, hash_type="sha256", hash_col="hash_diff"))
)

# MERGE: обновляем строку только если изменились значения
df_source.createOrReplaceTempView("source_data")

spark.sql("""
    MERGE INTO catalog.silver.erp_documents target
    USING source_data source
    ON target.hash_id = source.hash_id
    WHEN MATCHED AND target.hash_diff != source.hash_diff THEN
        UPDATE SET
            target.amount      = source.amount,
            target.quantity    = source.quantity,
            target.unit        = source.unit,
            target.currency    = source.currency,
            target.updated_at  = source.updated_at,
            target.hash_diff   = source.hash_diff,
            target._updated_at = current_timestamp()
    WHEN NOT MATCHED THEN
        INSERT (hash_id, hash_diff, company_code, plant, material_number,
                fiscal_year, line_item, amount, quantity, unit, currency,
                updated_at, _ingested_at, _updated_at)
        VALUES (source.hash_id, source.hash_diff, source.company_code, source.plant,
                source.material_number, source.fiscal_year, source.line_item,
                source.amount, source.quantity, source.unit, source.currency,
                source.updated_at, current_timestamp(), current_timestamp())
""")

# При создании таблицы: включаем Bloom Filter для hash_id
spark.sql("""
    CREATE TABLE catalog.silver.erp_documents (
        hash_id      BIGINT        NOT NULL COMMENT 'xxhash64 суррогатный ключ',
        hash_diff    STRING        NOT NULL COMMENT 'SHA-256 отпечаток значений',
        company_code STRING,
        plant        STRING,
        material_number STRING,
        fiscal_year  INT,
        line_item    INT,
        amount       DECIMAL(18, 2),
        quantity     DECIMAL(18, 4),
        unit         STRING,
        currency     STRING,
        updated_at   TIMESTAMP,
        _ingested_at TIMESTAMP,
        _updated_at  TIMESTAMP
    )
    USING iceberg
    PARTITIONED BY (fiscal_year, company_code)
    TBLPROPERTIES (
        'write.parquet.bloom-filter-enabled.column.hash_id' = 'true',
        'write.parquet.bloom-filter-fpp.column.hash_id'     = '0.01',   -- 1% ложных срабатываний
        'write.parquet.bloom-filter-enabled.column.hash_diff' = 'true',
        'write.parquet.bloom-filter-fpp.column.hash_diff'   = '0.01'
    )
""")

# Для существующей таблицы: добавляем Bloom Filter через ALTER TABLE
spark.sql("""
    ALTER TABLE catalog.silver.erp_documents
    SET TBLPROPERTIES (
        'write.parquet.bloom-filter-enabled.column.hash_id'   = 'true',
        'write.parquet.bloom-filter-fpp.column.hash_id'       = '0.01'
    )
""")

size_bytes = -n * ln(fpp) / (ln(2))^2 / 8

P(коллизия) ≈ 1 - e^(-n²/(2N))

n = 1,000,000,000
N = 2^64 = 18,446,744,073,709,551,616
P ≈ n² / (2N) = (10^9)^2 / (2 × 1.8 × 10^19) ≈ 10^18 / 3.6 × 10^19 ≈ 2.7%

P ≈ (5 × 10^8)^2 / (2 × 1.8 × 10^19) ≈ 0.7%

spark.sql("""
    MERGE INTO target t
    USING source s
    ON t.hash_id = s.hash_id
       AND t.company_code = s.company_code  -- Вторичный ключ для защиты от коллизий
       AND t.material_number = s.material_number
    WHEN MATCHED AND t.hash_diff != s.hash_diff THEN UPDATE ...
    WHEN NOT MATCHED THEN INSERT ...
""")

def check_hash_collisions_in_table(
    spark: SparkSession,
    table_name: str,
    hash_col: str,
    key_cols: List[str],
    partition_filter: str = None,
) -> int:
    """
    Проверяет наличие коллизий в target-таблице.
    Возвращает количество конфликтных hash_id.
    """
    df = spark.table(table_name)
    if partition_filter:
        df = df.filter(partition_filter)

    collisions = (
        df
        .groupBy(hash_col)
        .agg(
            F.count("*").alias("row_count"),
            F.countDistinct(*[F.col(c) for c in key_cols]).alias("distinct_business_keys")
        )
        .filter(F.col("distinct_business_keys") > 1)
    )

    count = collisions.count()
    if count > 0:
        print(f"⚠️  COLLISION: {count} hash values map to multiple business keys!")
        collisions.show(10, truncate=False)
    else:
        print(f"✅ No collisions found in {table_name}.{hash_col}")

    return count


# Запускать ежедневно как часть Data Quality Check
check_hash_collisions_in_table(
    spark,
    table_name="catalog.silver.erp_documents",
    hash_col="hash_id",
    key_cols=["company_code", "plant", "material_number", "fiscal_year", "line_item"],
    partition_filter="fiscal_year >= 2024"
)

# Одна и та же компания, разный регистр → разные хэши!
"SAP AG"     → sha256("SAP AG")     = "a3f..."
"sap ag"     → sha256("sap ag")     = "b7d..."
"Sap Ag"     → sha256("Sap Ag")     = "c1e..."

# Решение: UPPER() или LOWER() перед хэшированием

# Trailing/leading пробелы меняют хэш
"MAT-001"    → sha256("MAT-001")   = "d4a..."
"MAT-001 "   → sha256("MAT-001 ") = "e5b..."  # trailing space!
"  MAT-001"  → sha256("  MAT-001") = "f6c..."  # leading space!

# Решение: TRIM() перед хэшированием

# Одно и то же время в разных зонах → разный хэш
"2024-01-15 10:00:00+00:00"  # UTC
"2024-01-15 13:00:00+03:00"  # Moscow Time
# Физически это одно и то же событие!

# Решение: приводить к UTC перед хэшированием
from pyspark.sql import functions as F

df = df.withColumn(
    "event_time_utc",
    F.to_utc_timestamp(F.col("event_time"), "Europe/Moscow")
)

# Разные форматы одного числа → разный хэш как STRING
"1000"       → хэш от "1000"
"1,000"      → хэш от "1,000"    # European format
"1000.00"    → хэш от "1000.00"  # Decimal

# Решение: приводить к числовому типу перед хэшированием
df = df.withColumn("amount", F.col("amount").cast("decimal(18,2)"))
# Затем xxhash64(F.col("amount")) - нативно обрабатывает как число

# NULL и пустая строка - разные значения семантически
NULL → coalesce → "NULL"  → хэш от "NULL"
""   → trim     → ""      → хэш от ""
# Правильно: они дают разные хэши, что семантически корректно

from pyspark.sql import functions as F, DataFrame
from pyspark.sql.types import StringType, TimestampType, DecimalType
from typing import Dict, Optional


def normalize_for_hashing(
    df: DataFrame,
    string_cols: list,
    timestamp_cols: list,
    decimal_cols: Dict[str, tuple],   # col_name → (precision, scale)
    source_timezone: str = "UTC",
) -> DataFrame:
    """
    Стандартизирует данные перед вычислением hash-ключа.

    Args:
        df: входной DataFrame
        string_cols: колонки STRING для нормализации (TRIM + UPPER)
        timestamp_cols: колонки TIMESTAMP для нормализации в UTC
        decimal_cols: словарь {col: (precision, scale)} для нормализации числовых типов
        source_timezone: часовой пояс источника для timestamp

    Returns:
        DataFrame с нормализованными колонками
    """
    for col in string_cols:
        df = df.withColumn(col, F.trim(F.upper(F.col(col))))

    for col in timestamp_cols:
        df = df.withColumn(
            col,
            F.to_utc_timestamp(F.col(col).cast("timestamp"), source_timezone)
        )

    for col, (precision, scale) in decimal_cols.items():
        df = df.withColumn(col, F.col(col).cast(f"decimal({precision},{scale})"))

    return df


# Применение
df_normalized = normalize_for_hashing(
    df=df_raw,
    string_cols=["company_code", "plant", "material_number", "sales_org"],
    timestamp_cols=["document_date", "posting_date"],
    decimal_cols={"amount": (18, 2), "quantity": (18, 4)},
    source_timezone="Europe/Berlin",  # Источник - немецкая SAP-система
)

df_with_hash = build_hash_key(df_normalized, KEY_COLS, hash_type="xxhash64")

# Пайплайн v1
hash_key = sha2(concat_ws("||", col("plant"), col("company"), col("material")), 256)

# Пайплайн v2 (после рефакторинга порядка аргументов)
hash_key = sha2(concat_ws("||", col("company"), col("plant"), col("material")), 256)

# РЕЗУЛЬТАТ: одни и те же данные дают РАЗНЫЕ hash_key!
# Все MERGE / dedup сломаны - target думает, что все строки новые

# config.py - единый источник правды
HASH_KEY_COLS_ERP = [
    "company_code",      # 1. Сначала org-структура
    "plant",             # 2.
    "sales_org",         # 3.
    "distribution_channel",  # 4.
    "material_number",   # 5. Потом объект
    "document_type",     # 6. Потом тип документа
    "fiscal_year",       # 7. Потом время
    "period",            # 8.
    "line_item",         # 9.
]

# НИКОГДА не менять порядок без миграции данных!
# При добавлении новой колонки - всегда в конец

# shared_hash_utils.py - общая библиотека для всех пайплайнов
from pyspark.sql import functions as F, DataFrame

# Единый стандарт для всей организации
ERP_HASH_COLS_ORDERED = [
    "company_code", "plant", "material_number",
    "sales_org", "fiscal_year", "line_item"
]

def compute_erp_hash_id(df: DataFrame) -> DataFrame:
    """Единственная функция вычисления hash_id для ERP данных."""
    return df.withColumn(
        "hash_id",
        F.xxhash64(*[
            F.trim(F.upper(F.col(c)))
            for c in ERP_HASH_COLS_ORDERED
        ])
    )

# ОПАСНО: concat_ws с NULL-ами неправильно
bad_hash = F.sha2(F.concat_ws("||", F.col("a"), F.col("b"), F.col("c")), 256)
# (1, NULL, 3) → "1||3"
# (1, 3, NULL) → "1||3"
# Коллизия! Разные строки дают одинаковый хэш.

# ПРАВИЛЬНО: всегда coalesce
good_hash = F.sha2(
    F.concat_ws("||",
        F.coalesce(F.col("a").cast("string"), F.lit("NULL")),
        F.coalesce(F.col("b").cast("string"), F.lit("NULL")),
        F.coalesce(F.col("c").cast("string"), F.lit("NULL")),
    ),
    256
)
# (1, NULL, 3) → "1||NULL||3"
# (1, 3, NULL) → "1||3||NULL"
# Правильно: разные строки, разные хэши

# ПЛОХО: сортировка по случайному hash_id
spark.sql("ALTER TABLE silver.events WRITE ORDERED BY hash_id")
# Iceberg не может эффективно пропускать файлы по hash_id диапазонам:
# hash_id=1001 может быть в любом файле

# ХОРОШО: физическая сортировка по бизнес-измерениям
spark.sql("ALTER TABLE silver.events WRITE ORDERED BY event_date, company_code")
# + Bloom Filter по hash_id для точечных поисков

# ОПАСНО: хэш зависит от времени вставки
bad_hash = F.xxhash64(
    F.col("company_code"),
    F.col("material"),
    F.current_timestamp()   # ← каждый раз другое значение!
)
# Один и тот же бизнес-ключ дает разный hash_id при каждом запуске

# ПРАВИЛЬНО: хэш только от бизнес-ключей
good_hash = F.xxhash64(
    F.col("company_code"),
    F.col("material"),
    # НЕ включать: текущее время, UUID, автоинкременты
)

"""
dedup_5b_rows.py
End-to-end пайплайн дедупликации 5B строк через xxhash64.
"""

from pyspark.sql import SparkSession, functions as F, DataFrame
from pyspark.sql.window import Window
from typing import List
import logging

logger = logging.getLogger(__name__)

# ─── Конфигурация ────────────────────────────────────────────────────────────

SOURCES = [
    "catalog.bronze.erp_transactions",
    "catalog.bronze.crm_transactions",
    "catalog.bronze.legacy_transactions",
]
TARGET_TABLE = "catalog.silver.transactions"

# ВАЖНО: порядок колонок зафиксирован навсегда!
KEY_COLS = [
    "transaction_id",   # Уникальный ID транзакции в источнике
    "store_id",         # ID магазина
    "pos_id",           # ID кассового терминала
    "receipt_number",   # Номер чека
    "line_number",      # Номер строки в чеке
]

# Колонки значений для hash_diff (определение изменений)
VALUE_COLS = [
    "product_sku",
    "quantity",
    "unit_price",
    "discount_amount",
    "total_amount",
    "currency",
    "payment_method",
]

# Бизнес-дата для партиционирования и lookback window
DATE_COL = "transaction_date"
LOOKBACK_DAYS = 7  # Anti-join окно (последняя неделя)


# ─── Функции ──────────────────────────────────────────────────────────────────

def read_bronze_union(spark: SparkSession, sources: List[str], date_col: str) -> DataFrame:
    """Объединяем все 3 источника в один DataFrame."""
    dfs = []
    for source in sources:
        df = spark.table(source)
        # Добавляем метку источника для отладки
        df = df.withColumn("_source", F.lit(source.split(".")[-1]))
        dfs.append(df)

    # UNION ALL: объединяем все источники
    from functools import reduce
    combined = reduce(DataFrame.unionByName, dfs)
    logger.info(f"Union of {len(sources)} sources complete")
    return combined


def normalize_data(df: DataFrame) -> DataFrame:
    """
    Нормализует данные перед хэшированием.
    Применяем ко всем строковым ключевым колонкам.
    """
    for col in KEY_COLS:
        col_type = str(df.schema[col].dataType)
        if "StringType" in col_type:
            df = df.withColumn(col, F.trim(F.upper(F.col(col))))

    # Нормализуем дату транзакции в UTC
    df = df.withColumn(
        DATE_COL,
        F.to_date(F.col(DATE_COL))  # Убеждаемся что это DATE без времени
    )

    # Нормализуем суммы
    df = df.withColumn("total_amount", F.col("total_amount").cast("decimal(18,2)"))
    df = df.withColumn("unit_price", F.col("unit_price").cast("decimal(18,4)"))
    df = df.withColumn("quantity", F.col("quantity").cast("decimal(18,3)"))
    df = df.withColumn("discount_amount", F.col("discount_amount").cast("decimal(18,2)"))

    return df


def add_hash_columns(df: DataFrame) -> DataFrame:
    """
    Добавляет hash_id (ключ строки) и hash_diff (отпечаток значений).
    """
    # hash_id: xxhash64 - быстро, BIGINT
    df = df.withColumn("hash_id", F.xxhash64(*[F.col(c) for c in KEY_COLS]))

    # hash_diff: sha256 - детерминированный отпечаток значений
    normalized_values = [
        F.coalesce(F.col(c).cast("string"), F.lit("NULL"))
        for c in VALUE_COLS
    ]
    df = df.withColumn(
        "hash_diff",
        F.sha2(F.concat_ws("||", *normalized_values), 256)
    )

    return df


def dedup_within_batch(df: DataFrame) -> DataFrame:
    """
    Дедупликация внутри батча:
    1. По hash_id: берём одну строку на уникальный бизнес-ключ
    2. При дублях - берём строку с наибольшим updated_at (самую свежую)

    Используем Window row_number вместо dropDuplicates,
    чтобы контролировать, какую из дублирующихся строк оставить.
    """
    # Шаг 1: Window dedup - берём самую свежую строку по hash_id
    window_spec = Window.partitionBy("hash_id").orderBy(
        F.col("updated_at").desc_nulls_last(),  # Сначала самая свежая
        F.col("_source").asc()                   # При равном updated_at - предпочитаем ERP
    )

    df_ranked = df.withColumn("_rn", F.row_number().over(window_spec))
    df_deduped = df_ranked.filter(F.col("_rn") == 1).drop("_rn")

    logger.info("Intra-batch dedup complete")
    return df_deduped


def anti_join_with_target(
    spark: SparkSession,
    df_batch: DataFrame,
    target_table: str,
    lookback_days: int,
) -> DataFrame:
    """
    Фильтруем батч: оставляем только строки, которых НЕТ в target
    (за период lookback_days), чтобы избежать повторной вставки дублей.

    Это anti-join dedup для инкрементального батча.
    """
    # Читаем только последние N дней из target (lookback window)
    target_keys = (
        spark.table(target_table)
        .filter(
            F.col(DATE_COL) >= F.date_sub(F.current_date(), lookback_days)
        )
        .select("hash_id")
    )

    # LEFT ANTI JOIN: оставляем строки, которых нет в target
    new_rows = df_batch.join(
        F.broadcast(target_keys),  # Broadcast если target_keys < 10M строк
        on="hash_id",
        how="left_anti"
    )

    count = new_rows.count()
    logger.info(f"Anti-join filtered: {count} truly new rows to insert")
    return new_rows


def merge_into_target(spark: SparkSession, df_new: DataFrame, target_table: str) -> None:
    """
    MERGE: вставляем новые, обновляем изменённые.
    Условие: ON hash_id (+ вторичный ключ для защиты от коллизий)
    """
    df_new.createOrReplaceTempView("source_batch")

    spark.sql(f"""
        MERGE INTO {target_table} AS target
        USING source_batch AS source
        ON target.hash_id = source.hash_id
           AND target.transaction_id = source.transaction_id  -- защита от xxhash64 коллизий
        WHEN MATCHED AND target.hash_diff != source.hash_diff THEN
            UPDATE SET
                target.product_sku      = source.product_sku,
                target.quantity         = source.quantity,
                target.unit_price       = source.unit_price,
                target.discount_amount  = source.discount_amount,
                target.total_amount     = source.total_amount,
                target.currency         = source.currency,
                target.payment_method   = source.payment_method,
                target.hash_diff        = source.hash_diff,
                target._updated_at      = current_timestamp()
        WHEN NOT MATCHED THEN
            INSERT (hash_id, hash_diff, transaction_id, store_id, pos_id,
                    receipt_number, line_number, transaction_date,
                    product_sku, quantity, unit_price, discount_amount,
                    total_amount, currency, payment_method, _source,
                    _ingested_at, _updated_at)
            VALUES (source.hash_id, source.hash_diff, source.transaction_id,
                    source.store_id, source.pos_id, source.receipt_number,
                    source.line_number, source.transaction_date,
                    source.product_sku, source.quantity, source.unit_price,
                    source.discount_amount, source.total_amount,
                    source.currency, source.payment_method, source._source,
                    current_timestamp(), current_timestamp())
    """)

    logger.info(f"MERGE INTO {target_table} complete")


# ─── Основной пайплайн ────────────────────────────────────────────────────────

def main():
    spark = (
        SparkSession.builder
        .appName("dedup_5b_rows")
        # AQE для автоматической оптимизации Shuffle
        .config("spark.sql.adaptive.enabled", "true")
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
        # Shuffle partitions: ~200 ГБ / 128 МБ = ~1600 партиций
        .config("spark.sql.shuffle.partitions", "1600")
        # Broadcast threshold для anti-join
        .config("spark.sql.autoBroadcastJoinThreshold", "100m")
        # Iceberg-специфичные настройки
        .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
        .getOrCreate()
    )

    logger.info("=== Step 1: Read and Union sources ===")
    df_raw = read_bronze_union(spark, SOURCES, DATE_COL)

    logger.info("=== Step 2: Normalize data ===")
    df_normalized = normalize_data(df_raw)

    logger.info("=== Step 3: Compute hash keys ===")
    df_hashed = add_hash_columns(df_normalized)

    logger.info("=== Step 4: Intra-batch dedup ===")
    df_deduped = dedup_within_batch(df_hashed)

    logger.info("=== Step 5: Anti-join with target ===")
    df_new = anti_join_with_target(spark, df_deduped, TARGET_TABLE, LOOKBACK_DAYS)

    logger.info("=== Step 6: MERGE INTO target ===")
    merge_into_target(spark, df_new, TARGET_TABLE)

    logger.info("=== Pipeline complete ===")


if __name__ == "__main__":
    main()

# spark-submit конфигурация для 5B строк

spark = (
    SparkSession.builder
    .appName("dedup_5b_rows")

    # AQE - автоматически оптимизирует планы при выполнении
    .config("spark.sql.adaptive.enabled", "true")
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
    .config("spark.sql.adaptive.skewJoin.enabled", "true")      # Автоматическое лечение Data Skew

    # Shuffle: ~200 ГБ данных / 128 МБ на партицию = ~1600
    # AQE может уменьшить это число автоматически
    .config("spark.sql.shuffle.partitions", "1600")

    # Размер партиции при чтении Parquet
    .config("spark.sql.files.maxPartitionBytes", "134217728")    # 128 МБ

    # Broadcast: target_keys за 7 дней (если <1B строк, ~8 ГБ BIGINT)
    # Для 5B таблицы за 7 дней = ~700M строк = ~5.6 ГБ - слишком много для broadcast
    # Оставляем 100MB для мелких таблиц
    .config("spark.sql.autoBroadcastJoinThreshold", "104857600")  # 100 МБ

    # Executor конфигурация (50 экзекьюторов × 32 ГБ × 8 vCPU)
    .config("spark.executor.memory", "28g")     # 28g: 32g минус overhead
    .config("spark.executor.cores", "8")
    .config("spark.executor.memoryOverhead", "4g")

    # OFF HEAP для сортировки
    .config("spark.memory.offHeap.enabled", "true")
    .config("spark.memory.offHeap.size", "8g")

    .getOrCreate()
)

def log_pipeline_metrics(spark: SparkSession, df_name: str, df: DataFrame) -> None:
    """Логирует метрики на каждом шаге пайплайна."""
    # Spark UI metrics через REST API (только для уже выполненных операций)
    print(f"\n=== Metrics: {df_name} ===")
    # Считаем только если явно нужно - это триггер Action
    # В production логируйте count() через Grafana/Prometheus


# Проверка результатов после MERGE
def validate_merge_result(spark: SparkSession, target: str, date_filter: str) -> None:
    """Post-MERGE валидация: нет дублей, hash_id уникален."""
    df = spark.table(target).filter(date_filter)

    total_count = df.count()
    unique_hash_count = df.select("hash_id").distinct().count()

    if total_count != unique_hash_count:
        dup_count = total_count - unique_hash_count
        raise ValueError(
            f"DUPLICATES FOUND in {target}: "
            f"{total_count} rows but only {unique_hash_count} unique hash_ids "
            f"({dup_count} duplicates)"
        )

    print(f"✅ Validation passed: {total_count:,} unique rows in {target}")


# Вызов после MERGE
validate_merge_result(
    spark,
    target=TARGET_TABLE,
    date_filter="transaction_date >= current_date() - 1"
)

# SCD Type 2: customer_id + valid_from - уникальная строка в истории
ENTITY_KEY_COLS = ["customer_id"]                      # Бизнес-ключ сущности
ROW_KEY_COLS = ["customer_id", "valid_from"]           # Ключ конкретной строки истории
ATTRIBUTE_COLS = ["name", "email", "address", "city"]  # Атрибуты

df_scd2 = (
    df_customers
    .withColumn("entity_hash_key", F.xxhash64(F.col("customer_id")))
    .withColumn(
        "row_hash_key",
        F.xxhash64(F.col("customer_id"), F.col("valid_from").cast("long"))
    )
    .withColumn(
        "attribute_hash_diff",
        F.sha2(
            F.concat_ws("||", *[
                F.coalesce(F.col(c).cast("string"), F.lit("NULL"))
                for c in ATTRIBUTE_COLS
            ]),
            256
        )
    )
)

# MERGE для SCD Type 2: закрываем старую строку, открываем новую
spark.sql("""
    MERGE INTO dim_customers target
    USING (
        SELECT
            entity_hash_key,
            row_hash_key,
            attribute_hash_diff,
            customer_id, name, email, address, city,
            valid_from,
            '9999-12-31' AS valid_to,
            true AS is_current
        FROM source_scd2
    ) source
    ON target.row_hash_key = source.row_hash_key
    WHEN MATCHED AND target.attribute_hash_diff != source.attribute_hash_diff THEN
        -- Закрываем текущую строку
        UPDATE SET target.valid_to = source.valid_from, target.is_current = false
    WHEN NOT MATCHED THEN
        INSERT *  -- Вставляем новую строку с is_current = true
""")

def load_hub_customers(spark: SparkSession, source_df: DataFrame) -> None:
    """Data Vault 2.0: Hub загрузка с SHA-256 hash_key."""
    hub_df = (
        source_df
        .select("customer_id", "record_source")
        .distinct()
        .withColumn(
            "hash_key",
            F.sha2(
                F.concat_ws("||",
                    F.trim(F.upper(F.coalesce(F.col("customer_id").cast("string"), F.lit("NULL"))))
                ),
                256
            )
        )
        .withColumn("load_date", F.current_timestamp())
    )

    # Вставляем только новые hub записи (идемпотентно)
    existing_keys = spark.table("vault.hub_customers").select("hash_key")
    new_hubs = hub_df.join(F.broadcast(existing_keys), on="hash_key", how="left_anti")
    new_hubs.writeTo("vault.hub_customers").append()