Блок 4 - Python и расширенные возможности

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

def clean_email(email):
    return email.strip().lower() if email else None

clean_email_udf = udf(clean_email, StringType())
df = df.withColumn("email", clean_email_udf("raw_email"))

# Или через декоратор
@udf("string")
def clean_email(email):
    return email.strip().lower() if email else None

from pyspark.sql.functions import pandas_udf
import pandas as pd

# SCALAR: Series → Series (одна колонка → одна колонка)
@pandas_udf("double")
def multiply_by_pi(s: pd.Series) -> pd.Series:
    return s * 3.14159

df = df.withColumn("area", multiply_by_pi("radius"))

# GROUPED_MAP: группа строк → DataFrame (для sessionization, custom agg)
from pyspark.sql.functions import PandasUDFType

@pandas_udf("id long, value double", PandasUDFType.GROUPED_MAP)
def normalize_group(df: pd.DataFrame) -> pd.DataFrame:
    df["value"] = (df["value"] - df["value"].mean()) / df["value"].std()
    return df

df.groupby("group_id").apply(normalize_group)

# Включить Arrow для toPandas/createDataFrame
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
pd_df = spark_df.toPandas()  # Теперь через Arrow

# 1. addPyFile - простой py-файл или zip
spark.sparkContext.addPyFile("utils.py")
spark.sparkContext.addPyFile("my_lib.zip")

# 2. --py-files при spark-submit
# spark-submit --py-files libs.zip job.py

# 3. conda-pack / venv-pack - полная упаковка окружения
# conda-pack -o env.tar.gz
# spark-submit --archives env.tar.gz#env --conf spark.pyspark.python=./env/bin/python job.py

# 4. Docker image - для Kubernetes
# Поместить библиотеки в Dockerfile Executor'а

import logging

# Метод 1: стандартный Python logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.info("Processing started")

# Метод 2: log4j через SparkContext (пишет в тот же поток, что Spark)
log4j = spark._jvm.org.apache.log4j
logger = log4j.LogManager.getLogger("MyApp")
logger.info("Processing started")
logger.warn("Warning message")

# Уровень логирования Spark
spark.sparkContext.setLogLevel("WARN")  # ERROR, WARN, INFO, DEBUG

# Чтение
df = spark.read.format("jdbc") \
    .option("url", "jdbc:postgresql://host:5432/db") \
    .option("dbtable", "public.users") \
    .option("user", "user") \
    .option("password", "pass") \
    .option("driver", "org.postgresql.Driver") \
    .option("numPartitions", 16) \
    .option("partitionColumn", "id") \
    .option("lowerBound", 1) \
    .option("upperBound", 1000000) \
    .load()

# Запись
df.write.format("jdbc") \
    .option("url", "jdbc:postgresql://host:5432/db") \
    .option("dbtable", "public.result") \
    .option("batchsize", 10000) \
    .mode("append") \
    .save()

# Чтение из Kafka
stream_df = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "kafka:9092") \
    .option("subscribe", "events") \
    .load()

# Обработка
result = stream_df.selectExpr("CAST(value AS STRING) as json") \
    .select(from_json(col("json"), schema).alias("data")) \
    .select("data.*") \
    .groupBy(window("timestamp", "5 minutes"), "user_id") \
    .count()

# Запись
query = result.writeStream \
    .format("parquet") \
    .option("checkpointLocation", "s3://bucket/checkpoints/") \
    .trigger(processingTime="1 minute") \
    .start()

df.withWatermark("event_timestamp", "10 minutes") \
  .groupBy(window("event_timestamp", "5 minutes"), "user_id") \
  .count()

# Tumbling
window("timestamp", "5 minutes")

# Sliding
window("timestamp", "10 minutes", "5 minutes")

# Session (через mapGroupsWithState - нет встроенного синтаксиса)

query = df.writeStream \
    .option("checkpointLocation", "s3://bucket/checkpoints/query1/") \
    .start()

from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, VectorAssembler, StandardScaler
from pyspark.ml.classification import LogisticRegression

# Этапы пайплайна
indexer   = StringIndexer(inputCol="category", outputCol="category_idx")
assembler = VectorAssembler(
    inputCols=["age", "salary", "category_idx"], outputCol="features")
scaler    = StandardScaler(inputCol="features", outputCol="scaled_features")
lr        = LogisticRegression(featuresCol="scaled_features", labelCol="label")

pipeline = Pipeline(stages=[indexer, assembler, scaler, lr])

# Обучение
model = pipeline.fit(train_df)

# Предсказание
predictions = model.transform(test_df)

# Сохранение
model.save("s3://bucket/models/lr_model_v1")
pipeline.save("s3://bucket/models/pipeline_v1")

# Загрузка
from pyspark.ml.classification import LogisticRegressionModel
from pyspark.ml import PipelineModel

loaded_model    = LogisticRegressionModel.load("s3://bucket/models/lr_model_v1")
loaded_pipeline = PipelineModel.load("s3://bucket/models/pipeline_v1")

# Предсказание
predictions = loaded_pipeline.transform(new_df)

from pyspark.ml.feature import VectorAssembler

assembler = VectorAssembler(
    inputCols=["age", "salary", "tenure", "score"],
    outputCol="features",
    handleInvalid="skip"  # "error" или "keep" (для null)
)
df = assembler.transform(df)
# Новая колонка features: DenseVector([25, 50000, 3, 0.87])

from pyspark.ml.feature import StringIndexer, OneHotEncoder

# Шаг 1: Категориальные строки → числа
indexer = StringIndexer(inputCol="country", outputCol="country_idx")
indexed = indexer.fit(df).transform(df)

# Шаг 2: Числа → бинарный вектор
encoder = OneHotEncoder(inputCol="country_idx", outputCol="country_ohe")
encoded = encoder.fit(indexed).transform(indexed)

# В Pipeline - автоматически
pipeline = Pipeline(stages=[indexer, encoder])

import pytest
from pyspark.sql import SparkSession
from pyspark.testing import assertDataFrameEqual

@pytest.fixture(scope="session")
def spark():
    return SparkSession.builder.master("local[2]").appName("test").getOrCreate()

def test_dedup(spark):
    input_df = spark.createDataFrame([
        (1, "a"), (1, "a"), (2, "b")
    ], ["id", "value"])

    result = input_df.dropDuplicates(["id"])

    expected = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "value"])
    assertDataFrameEqual(result, expected, checkRowOrder=False)

# Append: для Parquet/Iceberg - записать новые события, не трогать старые
query = result.writeStream.outputMode("append").format("parquet").start(path)

# Complete: для live-dashboard - перезаписать текущий топ целиком
query = top_n.writeStream.outputMode("complete").format("memory").queryName("top").start()

# Update: для Delta/Cassandra - обновить только изменившиеся агрегаты
query = counts.writeStream.outputMode("update").format("delta").start(path)

query = df.writeStream \
    .format("delta") \
    .option("checkpointLocation", "s3://bucket/checkpoints/stream1/") \
    .outputMode("append") \
    .start("s3://bucket/output/")

# Статический справочник (читается один раз при старте)
products = spark.read.parquet("s3://bucket/products/")

# Поток событий из Kafka
events = spark.readStream \
    .format("kafka") \
    .option("subscribe", "purchase_events") \
    .load() \
    .selectExpr("CAST(value AS STRING) AS json")

parsed = events.select(from_json(col("json"), schema).alias("d")).select("d.*")

# Join: каждый batch событий × статический справочник (без shuffle)
enriched = parsed.join(broadcast(products), "product_id", "left")

enriched.writeStream \
    .format("delta") \
    .option("checkpointLocation", "s3://checkpoints/") \
    .start("s3://output/enriched/")

# 1. Агрегация с window - нарастающие счётчики по временным окнам
df.withWatermark("event_time", "10 minutes") \
  .groupBy(window("event_time", "5 minutes"), "user_id") \
  .count()

# 2. Дедупликация - запоминает уже виденные ключи
df.withWatermark("event_time", "1 hour") \
  .dropDuplicates(["event_id"])

# 3. Stream-Stream Join - state обеих сторон до появления совпадения
events.join(clicks, ["session_id"], "inner")  # Требует watermark на обоих

# 4. Произвольный stateful processing (сессии, автоматы)
df.groupBy("user_id").flatMapGroupsWithState(
    outputMode=OutputMode.Append,
    timeoutConf=GroupStateTimeout.EventTimeTimeout,
    func=update_session_state
)

# Один поток чтения с несколькими партициями Kafka
stream = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "kafka:9092") \
    .option("subscribe", "events") \
    .option("maxOffsetsPerTrigger", 100_000)  # Ограничить батч-размер
    .load()

# Вариант 1: явная схема + PERMISSIVE режим - новые колонки → null, лишние - игнорируются
schema = StructType([...])  # Зафиксированная схема
stream = spark.readStream \
    .format("json") \
    .schema(schema) \
    .option("mode", "PERMISSIVE") \
    .load(path)

# Вариант 2: Для Kafka+JSON - schema-on-read с from_json
stream.select(from_json(col("value"), schema).alias("d"))
# Новые поля в JSON просто игнорируются; удалённые → null

# Вариант 3: Delta Lake как sink - schema merging при записи
stream.writeStream \
    .format("delta") \
    .option("mergeSchema", "true") \  # Добавляет новые колонки в таблицу
    .start(path)

# Включение WAL в Spark Streaming
ssc = StreamingContext(sc, batchDuration=5)

# 1. Обязательно - checkpoint directory
ssc.checkpoint("hdfs:///checkpoints/myapp")

# 2. Включить WAL
spark.conf.set("spark.stream.receiver.writeAheadLog.enable", "true")

# 3. Для Kafka - использовать Direct API (без Receiver), тогда WAL не нужен:
stream = KafkaUtils.createDirectStream(ssc, topics, kafkaParams)