Window Functions Advanced: сложные примеры в DWH и Data Lake

import pyspark.sql.functions as F
from pyspark.sql import Window
from pyspark.sql.functions import col, lead, row_number, coalesce, lit

# Сырой поток изменений (Bronze/Silver)
changes = spark.read.parquet("/data/silver/price_changes/")

# Окно: для каждого товара, хронологически
scd_w = Window.partitionBy("product_id").orderBy("change_date")

scd2 = changes \
    .withColumn("valid_from", col("change_date")) \
    .withColumn(
        "valid_to_raw",
        lead("change_date", 1).over(scd_w)
    ) \
    .withColumn(
        "valid_to",
        coalesce(
            F.date_sub(col("valid_to_raw"), 1),  # предыдущий день до следующей версии
            lit("9999-12-31").cast("date")        # текущая запись
        )
    ) \
    .withColumn(
        "is_current",
        col("valid_to_raw").isNull()
    ) \
    .drop("valid_to_raw", "change_date")

scd2.orderBy("product_id", "valid_from").show(truncate=False)
# +-----------+-------+----------+----------+----------+
# |product_id |price  |valid_from|valid_to  |is_current|
# +-----------+-------+----------+----------+----------+
# |42         |100.00 |2024-01-01|2024-03-14|false     |
# |42         |120.00 |2024-03-15|2024-05-31|false     |
# |42         |115.00 |2024-06-01|9999-12-31|true      |
# +-----------+-------+----------+----------+----------+

# Найти цену товара на 2024-04-10
query_date = "2024-04-10"

current_price = scd2.filter(
    (col("valid_from") <= query_date) &
    (col("valid_to") >= query_date)
)
# Результат: product_id=42, price=120 (версия с 2024-03-15 по 2024-05-31)

# Задача: из нескольких версий одной записи оставить:
# 1. Самую свежую по updated_at
# 2. При равенстве updated_at - ту, где больше непустых полей

from pyspark.sql.functions import row_number, col, when

# Считаем количество непустых полей как proxy для "полноты"
completeness = (
    when(col("email").isNotNull(), 1).otherwise(0) +
    when(col("phone").isNotNull(), 1).otherwise(0) +
    when(col("address").isNotNull(), 1).otherwise(0) +
    when(col("birth_date").isNotNull(), 1).otherwise(0)
)

dedup_w = Window.partitionBy("customer_id").orderBy(
    col("updated_at").desc_nulls_last(),
    completeness.desc()
)

best_version = customers \
    .withColumn("completeness", completeness) \
    .withColumn("rn", row_number().over(dedup_w)) \
    .filter(col("rn") == 1) \
    .drop("rn", "completeness")

# Bronze содержит все версии, включая поздние корректировки
bronze_w = Window.partitionBy("event_id").orderBy(
    col("ingested_at").desc()  # самая свежая запись о событии - победитель
)

silver = (
    bronze
    .withColumn("rn", row_number().over(bronze_w))
    .filter(col("rn") == 1)
    .drop("rn", "ingested_at")
)

# Задача: найти непрерывные периоды, когда цена товара была выше 100
# Данные: price history по дням

from pyspark.sql.functions import row_number, sum as spark_sum, col

price_w   = Window.partitionBy("product_id").orderBy("date")
all_w     = Window.partitionBy("product_id").orderBy("date")

# Шаг 1: пронумеровать все строки и строки с условием отдельно
df = prices \
    .withColumn("rn_all",        row_number().over(price_w)) \
    .withColumn("above_100",     col("price") > 100) \
    .withColumn("rn_above",      row_number().over(
        Window.partitionBy("product_id", (col("price") > 100)).orderBy("date")
    ))

# Шаг 2: разность rn_all - rn_above - постоянна внутри одного острова
# (если строки не прерываются, разность одна и та же)
df = df.withColumn("island_id", col("rn_all") - col("rn_above"))

# Шаг 3: агрегировать острова
islands = df.filter(col("above_100")) \
    .groupBy("product_id", "island_id") \
    .agg(
        F.min("date").alias("island_start"),
        F.max("date").alias("island_end"),
        F.count("*").alias("island_days"),
        F.avg("price").alias("avg_price"),
    )

from pyspark.sql.functions import lag, col, unix_timestamp, sum as spark_sum, count, first, last

SESSION_GAP = 30 * 60  # 30 минут

events = spark.read.parquet("/data/silver/clickstream/")

user_w = Window.partitionBy("user_id").orderBy("ts")

# Шаг 1: пометить начала сессий
with_flags = events \
    .withColumn("prev_ts", lag("ts", 1).over(user_w)) \
    .withColumn("gap_sec", col("ts") - col("prev_ts")) \
    .withColumn(
        "is_new_session",
        col("prev_ts").isNull() | (col("gap_sec") > SESSION_GAP)
    )

# Шаг 2: ID сессии через накопленную сумму флагов
with_session = with_flags.withColumn(
    "session_id",
    spark_sum(col("is_new_session").cast("int")).over(user_w)
)

# Шаг 3: метрики сессии через full-partition window
sess_w = Window.partitionBy("user_id", "session_id")
full_w = sess_w.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

enriched = with_session \
    .withColumn("session_start",    first("ts").over(full_w)) \
    .withColumn("session_end",      last("ts").over(full_w)) \
    .withColumn("session_events",   count("*").over(sess_w)) \
    .withColumn("session_dur_sec",  last("ts").over(full_w) - first("ts").over(full_w)) \
    .withColumn("has_purchase",     (spark_sum(
        (col("event_type") == "purchase").cast("int")
    ).over(sess_w) > 0))

from pyspark.sql.functions import last, first

# Генерируем полный calendar spine (все даты в диапазоне)
date_range = spark.sql("""
    SELECT explode(sequence(
        date('2024-01-01'), date('2024-12-31'), interval 1 day
    )) AS date
""")

# Cross join с уникальными product_id
products = prices.select("product_id").distinct()
spine = products.crossJoin(date_range)

# Left join с реальными данными
full = spine.join(prices, ["product_id", "date"], "left")

# Forward fill: last non-null value в накопленном окне
ffill_w = Window.partitionBy("product_id").orderBy("date") \
    .rowsBetween(Window.unboundedPreceding, Window.currentRow)

filled = full.withColumn(
    "price_filled",
    last("price", ignorenulls=True).over(ffill_w)
)

# Backward fill (если нужно заполнить и в обратную сторону)
bfill_w = Window.partitionBy("product_id").orderBy("date") \
    .rowsBetween(Window.currentRow, Window.unboundedFollowing)

filled = filled.withColumn(
    "price_bfill",
    first("price_filled", ignorenulls=True).over(bfill_w)
)

from pyspark.sql.functions import lag, round as spark_round

# Месячная выручка уже агрегирована: (year, month, region, revenue)
monthly = spark.read.parquet("/data/gold/monthly_revenue/")

# Окно: внутри региона, хронологически по месяцам
mom_w = Window.partitionBy("region").orderBy("year", "month")

# Window 12 месяцев назад - lag с offset 12 для Year-over-Year
enriched = monthly \
    .withColumn("prev_month_rev",  lag("revenue", 1).over(mom_w)) \
    .withColumn("prev_year_rev",   lag("revenue", 12).over(mom_w)) \
    .withColumn(
        "mom_growth_pct",
        spark_round(
            (col("revenue") - col("prev_month_rev")) / col("prev_month_rev") * 100,
            1
        )
    ) \
    .withColumn(
        "yoy_growth_pct",
        spark_round(
            (col("revenue") - col("prev_year_rev")) / col("prev_year_rev") * 100,
            1
        )
    )

from pyspark.sql.functions import max as spark_max, min as spark_min, when

# Шаг 1: определить первое появление каждого события у пользователя в сессии
session_w = Window.partitionBy("user_id", "session_id")

events_with_first = events_sessionized.withColumn(
    "first_view",     spark_min(when(col("event") == "page_view",  col("ts"))).over(session_w)
).withColumn(
    "first_cart",     spark_min(when(col("event") == "add_to_cart", col("ts"))).over(session_w)
).withColumn(
    "first_checkout", spark_min(when(col("event") == "checkout",   col("ts"))).over(session_w)
).withColumn(
    "first_purchase", spark_min(when(col("event") == "purchase",   col("ts"))).over(session_w)
)

# Шаг 2: взять одну строку на сессию и проверить порядок событий
funnel = events_with_first \
    .groupBy("user_id", "session_id", "first_view", "first_cart", "first_checkout", "first_purchase") \
    .agg(F.count("*").alias("events")) \
    .withColumn(
        "reached_cart",
        col("first_cart").isNotNull() & (col("first_cart") > col("first_view"))
    ) \
    .withColumn(
        "reached_checkout",
        col("first_checkout").isNotNull() & (col("first_checkout") > col("first_cart"))
    ) \
    .withColumn(
        "converted",
        col("first_purchase").isNotNull() & (col("first_purchase") > col("first_checkout"))
    )

# Шаг 3: конверсия по воронке
funnel_stats = funnel.agg(
    F.count("*").alias("sessions"),
    F.sum(col("reached_cart").cast("int")).alias("to_cart"),
    F.sum(col("reached_checkout").cast("int")).alias("to_checkout"),
    F.sum(col("converted").cast("int")).alias("purchased"),
)

from pyspark.sql.functions import min as spark_min, months_between, floor

orders = spark.read.parquet("/data/silver/orders/").filter(col("status") == "COMPLETED")

# Шаг 1: первая покупка каждого клиента (когорта)
first_purchase = orders \
    .groupBy("customer_id") \
    .agg(spark_min("order_date").alias("cohort_date")) \
    .withColumn("cohort_month", F.date_trunc("month", col("cohort_date")))

# Шаг 2: join с полной историей заказов
orders_with_cohort = orders.join(first_purchase, "customer_id") \
    .withColumn("order_month", F.date_trunc("month", col("order_date"))) \
    .withColumn(
        "months_since_first",
        floor(months_between(col("order_month"), col("cohort_month"))).cast("int")
    )

# Шаг 3: когортная матрица
cohort_matrix = orders_with_cohort \
    .groupBy("cohort_month", "months_since_first") \
    .agg(F.countDistinct("customer_id").alias("active_customers"))

# Шаг 4: retention rate через window
cohort_size_w = Window.partitionBy("cohort_month")

cohort_with_size = cohort_matrix \
    .withColumn(
        "cohort_size",
        spark_max(when(col("months_since_first") == 0, col("active_customers"))).over(cohort_size_w)
    ) \
    .withColumn(
        "retention_rate",
        spark_round(col("active_customers") / col("cohort_size"), 3)
    )

# Pivot для читаемой таблицы: строки = когорта, колонки = месяц
cohort_pivot = cohort_with_size \
    .groupBy("cohort_month") \
    .pivot("months_since_first", list(range(13))) \
    .agg(spark_max("retention_rate"))

from pyspark.sql.functions import lag, col, when, sum as spark_sum

# История статусов заказа: (order_id, status, changed_at)
status_w = Window.partitionBy("order_id").orderBy("changed_at")

status_history = orders_events \
    .withColumn("prev_status", lag("status", 1).over(status_w)) \
    .withColumn(
        "is_status_change",
        col("prev_status").isNull() | (col("status") != col("prev_status"))
    ) \
    .withColumn(
        "state_id",
        spark_sum(col("is_status_change").cast("int")).over(status_w)
    )

# Длительность каждого состояния
state_duration_w = Window.partitionBy("order_id", "state_id")
full_state_w = state_duration_w.rowsBetween(
    Window.unboundedPreceding, Window.unboundedFollowing
)

with_duration = status_history \
    .withColumn("state_start", F.first("changed_at").over(full_state_w)) \
    .withColumn("state_end",   F.last("changed_at").over(full_state_w)) \
    .withColumn(
        "state_duration_hrs",
        (F.last("changed_at").over(full_state_w).cast("long") -
         F.first("changed_at").over(full_state_w).cast("long")) / 3600
    )

from pyspark.sql.functions import count, sum as spark_sum, avg, stddev

# Окно: транзакции за последние 10 минут для каждого счёта
ten_min_sec = 10 * 60

fraud_w = Window \
    .partitionBy("account_id") \
    .orderBy(col("ts").cast("long")) \
    .rangeBetween(-ten_min_sec, Window.currentRow)

# Окно за 24 часа для расчёта базовой линии
day_sec = 24 * 60 * 60
baseline_w = Window \
    .partitionBy("account_id") \
    .orderBy(col("ts").cast("long")) \
    .rangeBetween(-day_sec, Window.currentRow)

flagged = transactions \
    .withColumn("tx_count_10m",  count("*").over(fraud_w)) \
    .withColumn("sum_10m",       spark_sum("amount").over(fraud_w)) \
    .withColumn("avg_24h",       avg("amount").over(baseline_w)) \
    .withColumn("stddev_24h",    stddev("amount").over(baseline_w)) \
    .withColumn(
        "is_suspicious",
        (col("tx_count_10m") >= 5) |
        (col("sum_10m") > col("avg_24h") * 3) |
        (col("amount") > col("avg_24h") + 3 * col("stddev_24h"))  # 3σ выброс
    )

alerts = flagged.filter(col("is_suspicious"))

# ПЛОХО: 3 отдельных withColumn с одним WindowSpec → 3 WindowExec
w = Window.partitionBy("customer_id").orderBy("order_date")
df = df.withColumn("seq", row_number().over(w))
df = df.withColumn("prev_date", lag("order_date").over(w))
df = df.withColumn("running", F.sum("amount").over(w))

# ХОРОШО: один select → Spark может объединить в один WindowExec
df = df.select(
    "*",
    row_number().over(w).alias("seq"),
    lag("order_date").over(w).alias("prev_date"),
    F.sum("amount").over(w.rowsBetween(Window.unboundedPreceding, 0)).alias("running"),
)

# Сохранить отсортированные данные в Parquet с партиционированием по ключу
orders.repartition(col("customer_id")) \
      .sortWithinPartitions("customer_id", "order_date") \
      .write.mode("overwrite").parquet("/data/sorted_orders/")

# При последующем чтении + window с тем же partitionBy/orderBy
# Spark может пропустить Exchange и использовать уже отсортированные данные
sorted_orders = spark.read.parquet("/data/sorted_orders/")
sorted_orders.withColumn("seq", row_number().over(w))  # нет Exchange в physical plan

# Паттерн: накопленные итоги инкрементально
# Шаг 1: загрузить предыдущие агрегаты (Gold)
prev_gold = spark.read.parquet("/data/gold/customer_stats/")

# Шаг 2: обработать только новые данные (сегодняшний Bronze)
new_silver = spark.read.parquet("/data/silver/orders/today/")

# Шаг 3: join + пересчёт только изменившихся ключей
updated_customers = new_silver.select("customer_id").distinct()

affected_gold = prev_gold.join(updated_customers, "customer_id")
unchanged_gold = prev_gold.join(updated_customers, "customer_id", "left_anti")

# Пересчёт только для изменившихся клиентов
new_stats = (
    spark.read.parquet("/data/silver/orders/full/")
         .filter(col("customer_id").isin(updated_customers.rdd.flatMap(lambda x: x).collect()))
         .withColumn("seq", row_number().over(w))
    # ... остальные метрики
)

# Merge и запись
new_gold = unchanged_gold.union(new_stats)
new_gold.write.mode("overwrite").parquet("/data/gold/customer_stats/")

import pyspark.sql.functions as F
from pyspark.sql import Window
from pyspark.sql.functions import (
    col, row_number, lag, lead, first, last,
    sum as spark_sum, count, avg, min as spark_min, max as spark_max,
    datediff, months_between, floor, coalesce, lit, when,
    percent_rank, ntile
)

orders = spark.read.parquet("/data/silver/orders/").filter(col("status") == "COMPLETED")

# === Window Specs ===
cust_time_w  = Window.partitionBy("customer_id").orderBy("order_date")
cust_full_w  = Window.partitionBy("customer_id")
cust_time_full = cust_time_w.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
cust_cum_w   = cust_time_w.rowsBetween(Window.unboundedPreceding, Window.currentRow)
cust_roll3_w = cust_time_w.rowsBetween(-2, Window.currentRow)
region_w     = Window.partitionBy("region").orderBy(col("amount").desc())

# === Обогащение ===
customer_mart = orders.select(
    "order_id", "customer_id", "region", "order_date", "amount", "category",

    # --- Порядковая аналитика ---
    row_number().over(cust_time_w).alias("order_seq"),
    lag("order_date", 1).over(cust_time_w).alias("prev_order_date"),
    lead("order_date", 1).over(cust_time_w).alias("next_order_date"),

    # --- Временные метрики ---
    datediff(col("order_date"), lag("order_date", 1).over(cust_time_w)).alias("days_since_prev"),
    datediff(lead("order_date", 1).over(cust_time_w), col("order_date")).alias("days_to_next"),

    # --- Накопленная выручка (Running LTV) ---
    spark_sum("amount").over(cust_cum_w).alias("cumulative_ltv"),
    count("*").over(cust_cum_w).alias("cumulative_orders"),

    # --- Скользящие метрики ---
    avg("amount").over(cust_roll3_w).alias("rolling_avg_3tx"),
    spark_max("amount").over(cust_roll3_w).alias("rolling_max_3tx"),

    # --- Полная статистика клиента ---
    spark_sum("amount").over(cust_full_w).alias("total_ltv"),
    count("*").over(cust_full_w).alias("total_orders"),
    spark_min("order_date").over(cust_full_w).alias("first_order_date"),
    spark_max("order_date").over(cust_full_w).alias("last_order_date"),

    # --- Флаги ---
    (row_number().over(cust_time_w) == 1).alias("is_first_order"),
    lead("order_date", 1).over(cust_time_w).isNull().alias("is_last_order"),

    # --- Ранжирование в регионе ---
    percent_rank().over(region_w).alias("region_pct_rank"),
    ntile(4).over(Window.partitionBy("region").orderBy("amount")).alias("region_quartile"),
)

# === Сегментация клиентов ===
final_mart = customer_mart \
    .withColumn(
        "customer_segment",
        when(col("total_ltv") >= 50000, "VIP")
        .when(col("total_ltv") >= 10000, "HIGH_VALUE")
        .when(col("total_orders") == 1,  "ONE_TIME")
        .otherwise("REGULAR")
    ) \
    .withColumn(
        "days_since_last_order",
        datediff(F.current_date(), col("last_order_date"))
    ) \
    .withColumn(
        "churn_risk",
        when(col("days_since_last_order") > 180, "HIGH")
        .when(col("days_since_last_order") > 90,  "MEDIUM")
        .otherwise("LOW")
    )

# Запись в Gold (последний заказ как снапшот клиента)
snapshot = final_mart \
    .filter(col("is_last_order")) \
    .select(
        "customer_id", "region", "total_ltv", "total_orders",
        "first_order_date", "last_order_date", "days_since_last_order",
        "customer_segment", "churn_risk", "region_pct_rank",
    )

snapshot.write.mode("overwrite").parquet("/data/gold/customer_360/")