from pyspark.sql import SparkSession

# Create SparkSession with AQE enabled
spark = SparkSession.builder \
    .appName("SparkSQLEngine") \
    .config("spark.sql.adaptive.enabled", "true") \  # Enable AQE
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \  # Coalesce small partitions
    .getOrCreate()

# Create tables
employees = spark.createDataFrame([
    (1, "Alice", "Engineering", 30, 75000),
    (2, "Bob", "Marketing", 25, 65000),
    (3, "Charlie", "Engineering", 35, 90000),
    (4, "Diana", "Marketing", 28, 70000),
    (5, "Eve", "Engineering", 32, 85000),
    (6, "Frank", "Sales", 29, 60000),
    (7, "Grace", "Sales", 31, 68000)
], ["id", "name", "department", "age", "salary"])

# Create temporary view for SQL access
employees.createOrReplaceTempView("employees")

# Simple query with GROUP BY and HAVING
result = spark.sql("""
    SELECT department, 
           COUNT(*) as emp_count,
           AVG(salary) as avg_salary,
           MAX(age) as max_age
    FROM employees
    WHERE age > 25
    GROUP BY department
    HAVING COUNT(*) > 1
    ORDER BY avg_salary DESC
""")

result.show()

# Explain the query plan (True = all plans)
# Shows: Parsed → Analyzed → Optimized → Physical
result.explain(True)

# Get query execution details
print(f"Number of output rows: {result.count()}")
print(f"Query execution time: {result.head(1)}")

# Show optimized plan
df = spark.sql("""
    SELECT e.name, e.salary, d.location
    FROM employees e
    JOIN departments d ON e.department = d.name
    WHERE e.age > 30
    AND e.salary > 70000
""")

# Print physical plan with optimizations
df.explain(True)

# Output shows:
# == Parsed Logical Plan ==
# ...
# == Analyzed Logical Plan ==
# ... (resolved references)
# == Optimized Logical Plan ==
# ... (with predicate pushdown, column pruning)
# == Physical Plan ==
# ... (with join strategy, exchange)

# Verify predicate pushdown
# Look for "PushedFilters" in Parquet scan
df.write.parquet("/tmp/employees")
parquet_df = spark.read.parquet("/tmp/employees")
filtered = parquet_df.filter((parquet_df.age > 30) & (parquet_df.salary > 70000))
filtered.explain()  # Should show PushedFilters in scan

# Enable AQE
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.targetPartitionSize", "64MB")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256MB")

# Complex query that benefits from AQE
result = spark.sql("""
    SELECT 
        department,
        COUNT(*) as cnt,
        AVG(salary) as avg_sal,
        PERCENTILE(salary, 0.5) as median_sal
    FROM employees
    GROUP BY department
    HAVING COUNT(*) >= 2
""")

# AQE will:
# 1. Coalesce partitions after shuffle (reduces small files)
# 2. Handle data skew automatically (splits large partitions)
# 3. Optimize join strategies based on runtime stats
result.show()

# Check AQE optimizations in UI
print(f"Number of partitions: {result.rdd.getNumPartitions()}")

# AQE runtime statistics
spark.sql("SET -v").filter("key like '%adaptive%'").show(truncate=False)

# Read Parquet with predicate pushdown
parquet_df = spark.read.parquet("employee_data.parquet")

# Filter pushed to Parquet reader (PushedFilters in plan)
filtered = parquet_df.filter(
    (parquet_df.age > 30) & 
    (parquet_df.salary > 70000)
)

# Only reads relevant row groups (Parquet statistics)
filtered.explain()

# Write with bucketing for future joins
# bucketBy: Partition data by hash of column
# sortBy: Sort within each bucket
employees.write \
    .bucketBy(10, "department") \
    .sortBy("id") \
    .saveAsTable("bucketed_employees")

# Bucketed join avoids shuffle (same bucketing scheme)
spark.sql("""
    SELECT /*+ MAPJOIN(e2) */
        e1.name, e2.name as manager
    FROM bucketed_employees e1
    JOIN employees e2 ON e1.department = e2.department
""").show()

# Write with partitioning for partition pruning
employees.write \
    .partitionBy("department") \
    .parquet("partitioned_employees/")

# Read with partition pruning
# Only reads partitions matching filter
pruned = spark.read.parquet("partitioned_employees/") \
    .filter("department = 'Engineering'")
pruned.explain()  # Shows partition pruning

# CTEs (Common Table Expressions)
spark.sql("""
    WITH dept_stats AS (
        SELECT department, AVG(salary) as avg_sal
        FROM employees
        GROUP BY department
    )
    SELECT e.name, e.salary, d.avg_sal,
           e.salary - d.avg_sal as diff_from_avg
    FROM employees e
    JOIN dept_stats d ON e.department = d.department
""").show()

# Window functions in SQL
spark.sql("""
    SELECT name, department, salary,
           ROW_NUMBER() OVER (PARTITION BY department ORDER BY salary DESC) as rank,
           LAG(salary) OVER (PARTITION BY department ORDER BY salary) as prev_salary
    FROM employees
""").show()

# Subqueries
spark.sql("""
    SELECT name, salary
    FROM employees
    WHERE salary > (SELECT AVG(salary) FROM employees)
""").show()

# Correlated subqueries
spark.sql("""
    SELECT e.name, e.salary
    FROM employees e
    WHERE e.salary > (
        SELECT AVG(salary) 
        FROM employees 
        WHERE department = e.department
    )
""").show()

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.targetPartitionSize", "64MB")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.localShuffleReader.enabled", "true")

from pyspark.sql.functions import broadcast
result = large_df.join(broadcast(small_df), "key")
# Or configure threshold
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "52428800")  # 50MB

df.cache()  # For DataFrames reused multiple times
# Or persist with specific storage level
from pyspark import StorageLevel
df.persist(StorageLevel.MEMORY_AND_DISK_SER)  # Serialized + disk spill

# Parquet for columnar storage (recommended)
df.write.parquet("output/", compression="snappy")

# ORC for Hive integration
df.write.orc("output/")

# Avoid CSV/JSON for large datasets (no predicate pushdown)

# Partition by frequently filtered columns (enables partition pruning)
df.write.partitionBy("year", "month").parquet("output/")

# Use bucketing for join optimization (avoids shuffle)
df.write.bucketBy(100, "user_id").saveAsTable("users_bucketed")

# Always check explain plans
df.explain(True)

# Look for:
# - BroadcastHashJoin (good for small tables)
# - FileScan with PushedFilters (predicate pushdown working)
# - Project with only needed columns (column pruning working)
# - AQE optimizations in physical plan

# BAD: Python UDF (loses Catalyst optimization)
from pyspark.sql.functions import udf
@udf(returnType=IntegerType())
def double_it(x):
    return x * 2
df.withColumn("doubled", double_it(df.age))

# GOOD: Built-in function (Catalyst optimized)
from pyspark.sql.functions import col
df.withColumn("doubled", col("age") * 2)

# GOOD: Pandas UDF (Vectorized, better than Python UDF)
from pyspark.sql.functions import pandas_udf
@pandas_udf(IntegerType())
def double_it_pandas(s: pd.Series) -> pd.Series:
    return s * 2
df.withColumn("doubled", double_it_pandas(df.age))