from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, IntegerType
from pyspark.sql.window import Window
import logging

logger = logging.getLogger(__name__)


def create_spark_session(app_name: str = "ETL-Pipeline") -> SparkSession:
    """Create a production-optimized SparkSession."""
    return (
        SparkSession.builder
        .appName(app_name)
        .config("spark.sql.adaptive.enabled", "true")
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
        .config("spark.sql.adaptive.skewJoin.enabled", "true")
        .config("spark.sql.autoBroadcastJoinThreshold", str(10 * 1024 * 1024))  # 10MB
        .config("spark.sql.shuffle.partitions", "200")
        .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
        .config("spark.sql.parquet.compression.codec", "snappy")
        .config("spark.sql.sources.partitionOverwriteMode", "dynamic")
        .config("spark.dynamicAllocation.enabled", "true")
        .config("spark.dynamicAllocation.minExecutors", "2")
        .config("spark.dynamicAllocation.maxExecutors", "50")
        .config("spark.sql.execution.arrow.pyspark.enabled", "true")
        .getOrCreate()
    )


def run_optimized_etl(spark: SparkSession, input_path: str, output_path: str):
    """Run an optimized ETL pipeline with best practices."""
    # Read source data with schema enforcement
    schema = StructType([
        StructField("transaction_id", StringType(), False),
        StructField("customer_id", StringType(), False),
        StructField("product_id", StringType(), False),
        StructField("amount", DoubleType(), True),
        StructField("quantity", IntegerType(), True),
        StructField("event_date", StringType(), False),
    ])

    raw_df = (
        spark.read
        .schema(schema)
        .parquet(input_path)
        .repartition(F.col("customer_id"))  # Partition by join key
    )

    # Cache for reuse across multiple transformations
    raw_df.cache()
    logger.info(f"Raw record count: {raw_df.count()}")

    # Filter early to reduce data volume before joins
    filtered_df = (
        raw_df
        .filter(F.col("amount").isNotNull())
        .filter(F.col("amount") > 0)
        .filter(F.col("event_date") >= "2024-01-01")
    )

    # Broadcast join with small dimension table
    customer_df = (
        spark.read
        .parquet("s3://data-lake/dimensions/customers")
        .select("customer_id", "name", "tier", "segment")
    )
    enriched_df = filtered_df.join(
        F.broadcast(customer_df), on="customer_id", how="left"
    )

    # Windowed aggregation without shuffle (if partitioned by customer_id)
    window_spec = Window.partitionBy("customer_id").orderBy("event_date")
    result_df = (
        enriched_df
        .withColumn(
            "running_total",
            F.sum("amount").over(window_spec),
        )
        .withColumn(
            "rank_in_segment",
            F.row_number().over(
                Window.partitionBy("segment").orderBy(F.desc("amount"))
            ),
        )
    )

    # Write with dynamic partition overwrite
    (
        result_df
        .write
        .mode("overwrite")
        .partitionBy("event_date")
        .parquet(output_path)
    )

    raw_df.unpersist()
    logger.info(f"ETL completed. Output: {output_path}")


if __name__ == "__main__":
    spark = create_spark_session("production-etl")
    run_optimized_etl(spark, "s3://raw/transactions", "s3://curated/transactions_enriched")
    spark.stop()

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, pandas_udf, col
from pyspark.sql.types import StringType, DoubleType
import pandas as pd
import numpy as np
from typing import Iterator

spark = create_spark_session("udf-benchmark")

# Row-at-a-time Python UDF (slow - serialization overhead)
@udf(returnType=StringType())
def categorize_amount_slow(amount: float) -> str:
    """Categorize transaction amount (row-at-a-time UDF)."""
    if amount is None:
        return "unknown"
    elif amount < 10:
        return "micro"
    elif amount < 100:
        return "small"
    elif amount < 1000:
        return "medium"
    else:
        return "large"


# Vectorized Pandas UDF (fast - Arrow serialization)
@pandas_udf(StringType())
def categorize_amount_fast(amounts: pd.Series) -> pd.Series:
    """Categorize transaction amount (vectorized Pandas UDF)."""
    def _categorize(a):
        if pd.isna(a):
            return "unknown"
        elif a < 10:
            return "micro"
        elif a < 100:
            return "small"
        elif a < 1000:
            return "medium"
        else:
            return "large"

    return amounts.apply(_categorize)


# Grouped Map Pandas UDF for complex transformations
@pandas_udf(
    schema="transaction_id string  customer_id string  amount double  z_score double",
    functionType=pandas_udf.GROUPED_MAP,
)
def compute_z_score(pdf: pd.DataFrame) -> pd.DataFrame:
    """Compute Z-score within each customer group."""
    pdf["z_score"] = (pdf["amount"] - pdf["amount"].mean()) / pdf["amount"].std()
    return pdf


# Benchmark: Compare UDF performance
df = spark.read.parquet("s3://raw/transactions")

# Slow: Row-at-a-time UDF
result_slow = df.withColumn("category", categorize_amount_slow(col("amount")))

# Fast: Vectorized Pandas UDF
result_fast = df.withColumn("category", categorize_amount_fast(col("amount")))

# Grouped transformation
result_zscore = df.groupBy("customer_id").apply(compute_z_score)

# Performance comparison (via Spark UI or time measurement)
import time

start = time.time()
result_slow.write.mode("overwrite").parquet("/tmp/output_slow")
slow_time = time.time() - start

start = time.time()
result_fast.write.mode("overwrite").parquet("/tmp/output_fast")
fast_time = time.time() - start

print(f"Row UDF time: {slow_time:.2f}s")
print(f"Pandas UDF time: {fast_time:.2f}s")
print(f"Speedup: {slow_time / fast_time:.1f}x")