StructType: вложенные схемы, schema inference и schema evolution

from pyspark.sql.types import StructField, StringType

# Полная сигнатура:
StructField(
    name="field_name",      # 1. Имя колонки (строка)
    dataType=StringType(),  # 2. Тип данных
    nullable=True,          # 3. Может ли значение быть null (по умолчанию True)
    metadata={}             # 4. Произвольные метаданные (словарь) - по умолчанию {}
)

from pyspark.sql.types import StructField, StringType, DoubleType, Metadata

# Поле с метаданными: описание, единицы измерения, источник
amount_field = StructField(
    name="transaction_amount",
    dataType=DoubleType(),
    nullable=False,
    metadata={
        "description": "Сумма транзакции в валюте счёта",
        "unit": "currency_minor_units",   # например, копейки для RUB
        "source": "payment_gateway.v2",
        "pii": False,
    }
)

# Прочитать метаданные поля
df = spark.createDataFrame([(1.0,)], schema=StructType([amount_field]))
field = df.schema["transaction_amount"]
print(field.metadata)
# {'description': 'Сумма транзакции в валюте счёта', 'unit': 'currency_minor_units', ...}

# Метаданные сохраняются при записи в Parquet и ORC
# При чтении обратно метаданные восстанавливаются из footer файла

from pyspark.sql.types import *
from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("schema-demo").getOrCreate()

# Вложенная схема: транзакция с вложенными структурами
payment_schema = StructType([

    # Идентификатор: LongType (не Int - может быть > 2 млрд)
    StructField("transaction_id", LongType(), nullable=False),

    # Временна́я метка: TimestampType хранит UTC
    StructField("created_at", TimestampType(), nullable=False),

    # Деньги: всегда Decimal, никогда Double
    StructField("amount", DecimalType(18, 2), nullable=False),
    StructField("currency_code", StringType(), nullable=False),

    # Вложенная структура: данные пользователя
    StructField("payer", StructType([
        StructField("user_id",      LongType(),   nullable=False),
        StructField("email",        StringType(), nullable=True,
                    metadata={"pii": True}),   # Метаданные: пометить PII
        StructField("country_code", StringType(), nullable=True),
        StructField("ip_address",   StringType(), nullable=True,
                    metadata={"pii": True}),
    ]), nullable=False),

    # Вложенная структура: данные получателя
    StructField("payee", StructType([
        StructField("merchant_id",   LongType(),   nullable=False),
        StructField("merchant_name", StringType(), nullable=True),
        StructField("mcc",           IntegerType(), nullable=True,
                    metadata={"description": "Merchant Category Code"}),
    ]), nullable=False),

    # Массив тегов (строк)
    StructField("tags", ArrayType(StringType(), containsNull=False), nullable=True),

    # Массив связанных транзакций (структуры)
    StructField("related_transactions", ArrayType(StructType([
        StructField("ref_id",     LongType(),   nullable=False),
        StructField("ref_type",   StringType(), nullable=False),
    ])), nullable=True),

    # Словарь с динамическими метаданными из gateway
    StructField("gateway_metadata", MapType(StringType(), StringType()), nullable=True),

    # Статус транзакции
    StructField("status", StringType(), nullable=False),
])

payment_schema.simpleString()
# 'struct<transaction_id:bigint,created_at:timestamp,amount:decimal(18,2),...>'

# Создаём DataFrame с нашей схемой (пока пустой для просмотра структуры)
df = spark.createDataFrame([], payment_schema)

# printSchema() - человекочитаемое дерево
df.printSchema()
# root
#  |-- transaction_id: long (nullable = false)
#  |-- created_at: timestamp (nullable = false)
#  |-- amount: decimal(18,2) (nullable = false)
#  |-- currency_code: string (nullable = false)
#  |-- payer: struct (nullable = false)
#  |    |-- user_id: long (nullable = false)
#  |    |-- email: string (nullable = true)
#  |    |-- country_code: string (nullable = true)
#  |    |-- ip_address: string (nullable = true)
#  |-- payee: struct (nullable = false)
#  ...

# schema - Python-объект StructType
print(type(df.schema))   # <class 'pyspark.sql.types.StructType'>
print(df.schema.names)   # ['transaction_id', 'created_at', 'amount', ...]
print(df.schema.fields)  # список StructField-объектов

# Доступ к конкретному полю по имени
payer_field = df.schema["payer"]
print(payer_field.dataType)   # StructType([StructField('user_id', LongType()...)])
print(payer_field.nullable)   # False

# Доступ к вложенному полю через dataType
email_field = df.schema["payer"].dataType["email"]
print(email_field.metadata)   # {'pii': True}

# json() - схема в виде JSON-строки (для сериализации)
schema_json = payment_schema.json()
# '{"fields":[{"metadata":{},"name":"transaction_id","nullable":false,"type":"long"},...'

# fromJson() - восстановить схему из JSON
restored_schema = StructType.fromJson(eval(schema_json))  # через ast.literal_eval или json.loads

# Сравнение схем
print(payment_schema == restored_schema)  # True

# Опасный паттерн: inference на CSV
df_inferred = spark.read.csv("transactions.csv", header=True, inferSchema=True)
# Сегодня amount → IntegerType (все значения целые)
# Завтра в CSV появляется 99.99 → ломается join с другой таблицей

# Безопасный паттерн: явная схема
explicit_schema = StructType([
    StructField("transaction_id", LongType(),         nullable=False),
    StructField("amount",         DecimalType(18, 2), nullable=False),
    StructField("currency",       StringType(),        nullable=True),
])
df_safe = spark.read.csv("transactions.csv", header=True, schema=explicit_schema)
# Работает одинаково даже если в CSV появятся дробные числа

# Параметр samplingRatio для JSON: сколько данных читать для inference
df = spark.read.option("samplingRatio", "0.1").json("large_dataset/")
# Читает 10% файлов для inference - быстрее, но может ошибиться с редкими типами
# По умолчанию 1.0 - читает всё, надёжно но медленно

# Для production: ВСЕГДА задавайте схему явно
# Для exploration/notebook: inferSchema допустим, но осознанно

from pyspark.sql.functions import schema_of_json, lit, schema_of_csv

# schema_of_json: получить DDL-строку схемы из JSON-примера
sample_json = """
{
  "order_id": "ORD-1001",
  "user_id": 42,
  "amount": 1999.99,
  "items": [{"id": "P-100", "qty": 2}],
  "created_at": "2024-01-15T10:30:00Z"
}
"""

schema_ddl = spark.range(1) \
    .select(schema_of_json(lit(sample_json))) \
    .collect()[0][0]
print(schema_ddl)
# STRUCT<amount: DOUBLE, created_at: STRING, items: ARRAY<STRUCT<id: STRING, qty: BIGINT>>,
#        order_id: STRING, user_id: BIGINT>

# Внимание: нужно проверить и исправить:
# 1. amount: DOUBLE → DecimalType(18,2) для денег
# 2. created_at: STRING → TimestampType с форматом
# 3. Поля отсортированы по алфавиту, а не в порядке оригинала

# После проверки - зафиксировать как явный StructType в коде
# Это "one-time" операция: inference для изучения, explicit schema в production

from pyspark.sql.functions import col

# Предположим что у нас есть DataFrame с вложенным payer
# payer.user_id, payer.email, payer.country_code, payer.ip_address

# Способ 1: перечислить поля явно (рекомендуется - явность важнее краткости)
df.select(
    col("transaction_id"),
    col("payer.user_id"),
    col("payer.email"),
    col("payer.country_code"),
)

# Способ 2: star expansion - разворачивает ВСЕ поля структуры
df.select(
    col("transaction_id"),
    col("payer.*"),   # Добавляет user_id, email, country_code, ip_address как отдельные колонки
)

# Способ 3: полное разворачивание в SQL
df.createOrReplaceTempView("payments")
spark.sql("""
    SELECT transaction_id, payer.*
    FROM payments
""")

# ОСТОРОЖНО: star expansion в нестабильных схемах
# Если структура payer расширится новым полем - SELECT * без явного списка
# незаметно добавит это поле в результат, сломав downstream обработку
# Всегда перечисляйте поля явно в production-коде

from pyspark.sql.functions import struct, col, upper, lit, current_timestamp

# Пример: нужно обновить поле payer.email - привести к нижнему регистру
# и добавить новое поле payer.verified_at

# ❌ СТАРЫЙ СПОСОБ (до Spark 3.1): пересборка всей структуры
# Нужно перечислить ВСЕ поля, включая неизменившиеся
df_old_way = df.withColumn("payer",
    struct(
        col("payer.user_id"),                      # не изменилось, но нужно перечислить
        F.lower(col("payer.email")).alias("email"), # изменилось
        col("payer.country_code"),                 # не изменилось
        col("payer.ip_address"),                   # не изменилось
        current_timestamp().alias("verified_at"),  # новое поле
    )
)
# Проблема: если структура расширится - нужно обновить этот код

# ✅ НОВЫЙ СПОСОБ (Spark 3.1+): withField - хирургически точно
df_new_way = df \
    .withColumn("payer",
        col("payer")
        .withField("email", F.lower(col("payer.email")))  # обновить существующее
        .withField("verified_at", current_timestamp())    # добавить новое
    )

# Можно цепочкой: несколько withField подряд
df_updated = df.withColumn("payee",
    col("payee")
    .withField("merchant_name", F.upper(col("payee.merchant_name")))
    .withField("is_verified", lit(True))
    .withField("last_seen", current_timestamp())
)

# dropFields(): удалить поле из структуры (тоже Spark 3.1+)
# Полезно для удаления PII-полей перед передачей в downstream
df_no_pii = df.withColumn("payer",
    col("payer")
    .dropFields("email")      # убрать email
    .dropFields("ip_address") # убрать IP
)
df_no_pii.select("payer.*").printSchema()
# Схема payer без email и ip_address

# Схема: order.shipping.address.city
# Нужно обновить только city

deeply_nested = spark.createDataFrame([
    ("ORD-1", {"shipping": {"address": {"city": "москва", "zip": "101000"}}},),
], StructType([
    StructField("order_id", StringType()),
    StructField("order", StructType([
        StructField("shipping", StructType([
            StructField("address", StructType([
                StructField("city", StringType()),
                StructField("zip",  StringType()),
            ]))
        ]))
    ]))
]))

# withField поддерживает путь через точку
result = deeply_nested.withColumn("order",
    col("order").withField(
        "shipping.address.city",               # путь к полю
        F.initcap(col("order.shipping.address.city"))  # значение
    )
)

result.select("order.shipping.address.city").show()
# +-----+
# |city |
# +-----+
# |Москва|
# +-----+

from pyspark.sql.functions import col, isnull

# DataFrame с nullable=False колонкой
strict_schema = StructType([
    StructField("id",     LongType(),   nullable=False),
    StructField("amount", DoubleType(), nullable=False),
])

# Spark не выбросит ошибку, но логически это некорректно
df_with_null = spark.createDataFrame(
    [(1, None)],    # amount = None, хотя nullable=False
    strict_schema
)
df_with_null.show()
# +---+------+
# | id|amount|
# +---+------+
# |  1|  null|  ← записано! nullable=False не является constraint
# +---+------+

# Реалистичная схема для ingestion: по умолчанию nullable=True
# Только идентификаторы и технические поля - nullable=False
ingestion_schema = StructType([
    StructField("event_id",   StringType(),    nullable=False),  # PK, обязателен
    StructField("event_type", StringType(),    nullable=False),  # должен быть всегда
    StructField("user_id",    LongType(),      nullable=True),   # может быть анонимный
    StructField("session_id", StringType(),    nullable=True),   # может не быть сессии
    StructField("amount",     DecimalType(18,2), nullable=True), # не у всех событий есть сумма
    StructField("created_at", TimestampType(), nullable=False),  # время обязательно
])

from pyspark.sql.functions import count, when, col, isnan

# Проверка null-количества для всех колонок (полезно при ingestion)
null_counts = df.select([
    count(when(col(c).isNull(), c)).alias(c)
    for c in df.columns
])
null_counts.show()
# Если event_id имеет null_count > 0 - это критическая ошибка в upstream

# Схема с _corrupt_record для отлова плохих JSON-документов
schema_with_corrupt = StructType([
    StructField("transaction_id", LongType(),         nullable=True),
    StructField("amount",         DecimalType(18, 2), nullable=True),
    StructField("currency",       StringType(),        nullable=True),
    StructField("_corrupt_record", StringType(),       nullable=True),  # техническая колонка
])

# Читаем JSON с PERMISSIVE mode и _corrupt_record
raw_df = spark.read \
    .option("mode", "PERMISSIVE") \
    .option("columnNameOfCorruptRecord", "_corrupt_record") \
    .schema(schema_with_corrupt) \
    .json("payments_raw/")

# Разделяем на хорошие и плохие записи
good_records = raw_df.filter(col("_corrupt_record").isNull()).drop("_corrupt_record")
bad_records  = raw_df.filter(col("_corrupt_record").isNotNull()) \
                     .select("_corrupt_record")

print(f"Хорошие записи: {good_records.count()}")
print(f"Плохие записи:  {bad_records.count()}")

# Анализируем плохие данные
bad_records.show(5, truncate=False)
# Сохраняем для последующего разбора
bad_records.write.mode("append").text("/data/quarantine/payments/")

# Алертинг: если плохих записей > 1% - это аномалия
total = raw_df.count()
bad_rate = bad_records.count() / total
if bad_rate > 0.01:
    raise RuntimeError(f"Слишком много битых записей: {bad_rate:.1%}")

# Без mergeSchema: Spark читает схему первого файла и применяет ко всем
# Если новые файлы имеют лишние колонки - они игнорируются
# Если старые файлы не имеют колонок из новых файлов - ошибка

# С mergeSchema: Spark читает footer каждого файла и объединяет схемы
df = spark.read \
    .option("mergeSchema", "true") \
    .parquet("/data/events/year=2024/")

# Результат: объединённая схема
# Новые колонки (есть только в новых файлах): nullable=True, старые записи → null
# Старые колонки (есть только в старых файлах): nullable=True, новые записи → null

# ОСТОРОЖНО: mergeSchema читает footer каждого файла
# При тысячах файлов - тысячи HDFS/S3 запросов только для планирования
# Рекомендация: spark.sql.parquet.mergeSchema=false в production,
# mergeSchema=true только при добавлении новой партиции с обновлённой схемой

# Delta Lake: два режима эволюции схемы при записи

# 1. mergeSchema=true: автоматически добавляет новые колонки
#    Существующие данные остаются нетронутыми, новые колонки → null для старых партиций
df_new_schema.write \
    .format("delta") \
    .mode("append") \
    .option("mergeSchema", "true") \   # добавить новые колонки в таблицу
    .save("/delta/events/")

# 2. overwriteSchema=true: полностью заменяет схему
#    Используйте ТОЛЬКО при полной перезаписи данных (mode="overwrite")
#    Опасен для инкрементальных обновлений - ломает совместимость
df_completely_different.write \
    .format("delta") \
    .mode("overwrite") \
    .option("overwriteSchema", "true") \  # полная замена схемы
    .save("/delta/events/")

# Просмотр истории изменений схемы через Delta
from delta.tables import DeltaTable

delta_table = DeltaTable.forPath(spark, "/delta/events/")
history = delta_table.history()
history.select("version", "timestamp", "operation", "operationParameters") \
       .show(10, truncate=False)
# Каждая операция включает изменения схемы в operationParameters

# ALTER TABLE в Delta (SQL)
spark.sql("""
    ALTER TABLE delta.`/delta/events/`
    ADD COLUMNS (new_field STRING COMMENT 'Добавлено в v2.1')
""")

# Паттерн защиты от schema drift: сравнение схем при запуске

EXPECTED_SCHEMA = StructType([
    StructField("event_id",   StringType(),    nullable=False),
    StructField("event_type", StringType(),    nullable=False),
    StructField("user_id",    LongType(),      nullable=True),
    StructField("created_at", TimestampType(), nullable=False),
])

def check_schema_compatibility(actual: StructType, expected: StructType) -> list:
    """Возвращает список проблем совместимости схем."""
    issues = []
    expected_fields = {f.name: f for f in expected.fields}
    actual_fields   = {f.name: f for f in actual.fields}

    # Проверяем что все обязательные поля присутствуют
    for name, field in expected_fields.items():
        if not field.nullable and name not in actual_fields:
            issues.append(f"CRITICAL: Отсутствует non-nullable поле '{name}'")
        elif name in actual_fields:
            actual_type   = actual_fields[name].dataType
            expected_type = field.dataType
            if type(actual_type) != type(expected_type):
                issues.append(
                    f"WARNING: Тип поля '{name}' изменился: "
                    f"{expected_type} → {actual_type}"
                )

    # Проверяем что нет неожиданных новых полей
    new_fields = set(actual_fields) - set(expected_fields)
    if new_fields:
        issues.append(f"INFO: Новые поля в источнике: {new_fields}")

    return issues

# Использование при ingestion
incoming_df = spark.read.json("/data/events/today/")
issues = check_schema_compatibility(incoming_df.schema, EXPECTED_SCHEMA)

for issue in issues:
    if issue.startswith("CRITICAL"):
        raise RuntimeError(f"Schema drift: {issue}")
    else:
        print(f"Schema drift: {issue}")

# Паттерн 1: генерация схемы из конфигурации
COLUMN_CONFIG = [
    {"name": "id",           "type": "long",       "nullable": False},
    {"name": "name",         "type": "string",      "nullable": True},
    {"name": "score",        "type": "double",      "nullable": True},
    {"name": "category",     "type": "string",      "nullable": True},
    {"name": "created_date", "type": "date",        "nullable": False},
]

TYPE_MAP = {
    "string":    StringType(),
    "long":      LongType(),
    "int":       IntegerType(),
    "double":    DoubleType(),
    "boolean":   BooleanType(),
    "date":      DateType(),
    "timestamp": TimestampType(),
}

def build_schema_from_config(config: list) -> StructType:
    fields = []
    for col_cfg in config:
        spark_type = TYPE_MAP.get(col_cfg["type"], StringType())
        fields.append(StructField(
            name=col_cfg["name"],
            dataType=spark_type,
            nullable=col_cfg.get("nullable", True),
            metadata=col_cfg.get("metadata", {}),
        ))
    return StructType(fields)

dynamic_schema = build_schema_from_config(COLUMN_CONFIG)
dynamic_schema.simpleString()

# Паттерн 2: расширение схемы программно
def add_audit_fields(schema: StructType) -> StructType:
    """Добавить технические поля аудита к любой схеме."""
    audit_fields = [
        StructField("_ingested_at",    TimestampType(), nullable=False,
                    metadata={"technical": True}),
        StructField("_source_file",    StringType(),    nullable=True,
                    metadata={"technical": True}),
        StructField("_pipeline_version", StringType(), nullable=False,
                    metadata={"technical": True}),
    ]
    return StructType(schema.fields + audit_fields)

enriched_schema = add_audit_fields(payment_schema)

# Паттерн 3: сериализация схемы для хранения в метаданных
import json

schema_dict = json.loads(payment_schema.json())
# Сохранить в конфигурационный файл, базу данных, data catalog

# Восстановить
restored = StructType.fromJson(schema_dict)
print(payment_schema == restored)  # True

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

spark = SparkSession.builder.appName("schema-evolution-demo").getOrCreate()

# ── Шаг 1: Определяем схему явно ──────────────────────────────────────────────
# Версия 1: базовая схема пользовательского профиля

user_schema_v1 = StructType([
    StructField("user_id",     LongType(),   nullable=False),
    StructField("username",    StringType(), nullable=False),
    StructField("email",       StringType(), nullable=True,
                metadata={"pii": True}),
    StructField("created_at",  TimestampType(), nullable=False),
    StructField("profile",     StructType([
        StructField("bio",      StringType(), nullable=True),
        StructField("location", StringType(), nullable=True),
        StructField("website",  StringType(), nullable=True),
    ]), nullable=True),
    StructField("_corrupt_record", StringType(), nullable=True),
])

# ── Шаг 2: Читаем данные с PERMISSIVE mode ──────────────────────────────────
# Симулируем JSON-данные (в реальности: spark.read.json("path/"))
from pyspark.sql.functions import col

json_data = spark.createDataFrame([
    ('{"user_id": 1, "username": "alice", "created_at": "2024-01-01T10:00:00Z", "profile": {"bio": "Data Engineer", "location": "Moscow"}}',),
    ('{"user_id": 2, "username": "bob", "email": "bob@example.com", "created_at": "2024-01-02T15:00:00Z"}',),
    ('{"user_id": "INVALID", "username": "charlie"}',),  # битая запись
], ["raw"])

# В реальном коде:
raw_df = spark.read \
    .option("mode", "PERMISSIVE") \
    .option("columnNameOfCorruptRecord", "_corrupt_record") \
    .schema(user_schema_v1) \
    .json("/data/users/batch=2024-01-15/")

# Разделяем хорошие и плохие записи
good_df = raw_df.filter(col("_corrupt_record").isNull()).drop("_corrupt_record")
bad_df  = raw_df.filter(col("_corrupt_record").isNotNull())

print(f"Хорошие: {good_df.count()}, Плохие: {bad_df.count()}")
bad_df.show(truncate=False)

# ── Шаг 3: Добавляем поле в вложенную структуру через withField ──────────────
# Версия 2: добавляем followers_count в profile

enriched_df = good_df.withColumn("profile",
    F.when(col("profile").isNull(),
        F.struct(
            F.lit(None).cast(StringType()).alias("bio"),
            F.lit(None).cast(StringType()).alias("location"),
            F.lit(None).cast(StringType()).alias("website"),
            F.lit(0).cast(LongType()).alias("followers_count"),
        )
    ).otherwise(
        col("profile").withField("followers_count", F.lit(0).cast(LongType()))
    )
)

enriched_df.select("user_id", "profile.*").show()
# +-------+--------+-------+-------+---------------+
# |user_id|bio     |location|website|followers_count|
# +-------+--------+-------+-------+---------------+
# |1      |Data ...|Moscow  |null   |0              |
# |2      |null    |null    |null   |0              |
# +-------+--------+-------+-------+---------------+

# ── Шаг 4: Добавляем технические поля ────────────────────────────────────────
from pyspark.sql.functions import current_timestamp, input_file_name

final_df = enriched_df \
    .withColumn("_ingested_at",     current_timestamp()) \
    .withColumn("_pipeline_version", F.lit("v2.1.0")) \
    .withColumn("_batch_date",       F.lit("2024-01-15").cast(DateType()))

# ── Шаг 5: Запись в Delta с mergeSchema ──────────────────────────────────────
# При первой записи: создаёт таблицу
# При последующих с новыми полями: расширяет схему
final_df.write \
    .format("delta") \
    .mode("append") \
    .option("mergeSchema", "true") \
    .partitionBy("_batch_date") \
    .save("/delta/users/")

# ── Проверка: схема Delta-таблицы обновилась ─────────────────────────────────
updated_schema = spark.read.format("delta").load("/delta/users/").schema
print("Обновлённая схема:")
spark.read.format("delta").load("/delta/users/").printSchema()