from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("SQL37").getOrCreate()

# Create sample data
data = [(1, "Alice", "Engineering", 90000),
        (2, "Bob",   "Engineering", 85000),
        (3, "Carol", "Marketing",   70000),
        (4, "Dave",  "Marketing",   72000),
        (5, "Eve",   "HR",          60000)]

df = spark.createDataFrame(data, ["id", "name", "dept", "salary"])
df.createOrReplaceTempView("employees")

# This goes through the SQL parser first
result = spark.sql("""
    SELECT dept, AVG(salary) as avg_salary
    FROM employees
    WHERE salary > 65000
    GROUP BY dept
    ORDER BY avg_salary DESC
""")
result.show()

# Syntax error — caught by parser immediately
try:
    spark.sql("SELECT FROM employees")  # missing column list
except Exception as e:
    print(f"Parser error: {str(e)[:100]}")

# See all 4 phases with explain()
df_result = spark.sql("""
    SELECT dept, SUM(salary) as total_salary
    FROM employees
    WHERE salary > 70000
    GROUP BY dept
""")

# Simple explain — shows physical plan
df_result.explain()

# Extended explain — shows ALL 4 phases
df_result.explain(extended=True)

# Formatted explain — most readable, shows stats
df_result.explain(mode="formatted")

# Cost-Based Optimizer explain
df_result.explain(mode="cost")

# Codegen explain — shows generated Java bytecode
df_result.explain(mode="codegen")

# Control shuffle partitions for SQL queries
spark.conf.set("spark.sql.shuffle.partitions", "50")

# Enable Adaptive Query Execution (default in Spark 3+)
spark.conf.set("spark.sql.adaptive.enabled", "true")

# Enable Cost-Based Optimizer
spark.conf.set("spark.sql.cbo.enabled", "true")

# Collect table statistics for CBO (run on managed tables)
spark.sql("ANALYZE TABLE employees COMPUTE STATISTICS FOR ALL COLUMNS")

# Now run a query — CBO will use stats for join ordering
spark.sql("""
    SELECT e.name, e.dept, d.location
    FROM employees e
    JOIN departments d ON e.dept = d.dept_name
    WHERE e.salary > 80000
""").explain(mode="cost")

# Check number of output partitions
result = spark.sql("SELECT dept, COUNT(*) FROM employees GROUP BY dept")
print(f"Partitions: {result.rdd.getNumPartitions()}")

from pyspark.sql.functions import col

# Create a DataFrame
data = [(1, "Alice", "Eng", 90000), (2, "Bob", "Eng", 85000),
        (3, "Carol", "Mktg", 70000), (4, "Dave", "Mktg", 72000),
        (5, "Eve", "HR", 60000)]
df = spark.createDataFrame(data, ["id", "name", "dept", "salary"])

# Register as temp view
df.createOrReplaceTempView("employees")

# Now query with SQL — exactly like a table
spark.sql("SELECT * FROM employees WHERE salary > 70000").show()

# createTempView (fails if view already exists)
df.createTempView("employees_v2")

# createOrReplaceTempView — safe to call repeatedly (idempotent)
df.createOrReplaceTempView("employees")  # replaces existing

# List all temp views in current session
spark.catalog.listTables()  # shows all views and tables

# Drop a temp view when done
spark.catalog.dropTempView("employees_v2")

# Check if view exists
print(spark.catalog.tableExists("employees"))  # True

# Register a global temp view
df.createOrReplaceGlobalTempView("global_employees")

# Query it — MUST use global_temp. prefix
spark.sql("SELECT * FROM global_temp.global_employees").show()

# Global temp views are accessible from other sessions
spark2 = SparkSession.builder.getOrCreate()  # same app, new session
spark2.sql("SELECT COUNT(*) FROM global_temp.global_employees").show()
# → same result! Shared between sessions

# Global temp view does NOT exist in regular temp view namespace
try:
    spark.sql("SELECT * FROM global_employees").show()  # AnalysisException!
except:
    print("Must use global_temp.global_employees")

# Drop a global temp view
spark.catalog.dropGlobalTempView("global_employees")

# Create a permanent view (requires catalog + managed table or path)
spark.sql("""
    CREATE OR REPLACE VIEW high_earners AS
    SELECT id, name, dept, salary
    FROM employees
    WHERE salary > 75000
""")

# Now queryable in any future session
spark.sql("SELECT * FROM high_earners").show()

# Create view in a specific database
spark.sql("CREATE DATABASE IF NOT EXISTS analytics")
spark.sql("""
    CREATE OR REPLACE VIEW analytics.dept_summary AS
    SELECT dept,
           COUNT(*) as headcount,
           AVG(salary) as avg_salary,
           MAX(salary) as max_salary
    FROM employees
    GROUP BY dept
""")

spark.sql("SELECT * FROM analytics.dept_summary").show()

# Show view definition
spark.sql("SHOW CREATE TABLE analytics.dept_summary").show(truncate=False)

# Drop permanent view
spark.sql("DROP VIEW IF EXISTS analytics.dept_summary")

-- Simple CTE: find departments with avg salary above company average
spark.sql("""
WITH dept_stats AS (
    SELECT dept,
           AVG(salary) AS avg_salary,
           COUNT(*)    AS headcount
    FROM employees
    GROUP BY dept
),
company_avg AS (
    SELECT AVG(salary) AS co_avg
    FROM employees
)
SELECT d.dept,
       d.avg_salary,
       d.headcount,
       c.co_avg,
       ROUND(d.avg_salary - c.co_avg, 2) AS diff_from_avg
FROM dept_stats d
CROSS JOIN company_avg c
ORDER BY d.avg_salary DESC
""").show()

-- Chained CTEs: each WITH block can reference the previous one
spark.sql("""
WITH
-- Step 1: flag high earners
flagged AS (
    SELECT *,
           CASE WHEN salary > 80000 THEN 'High'
                WHEN salary > 65000 THEN 'Mid'
                ELSE 'Low' END AS pay_band
    FROM employees
),
-- Step 2: count by dept and band (uses 'flagged' CTE)
summary AS (
    SELECT dept, pay_band, COUNT(*) AS cnt
    FROM flagged
    GROUP BY dept, pay_band
),
-- Step 3: add total per dept (uses 'flagged' CTE again)
dept_total AS (
    SELECT dept, COUNT(*) AS dept_total
    FROM flagged
    GROUP BY dept
)
-- Final query joins all CTEs
SELECT s.dept, s.pay_band, s.cnt, t.dept_total,
       ROUND(s.cnt * 100.0 / t.dept_total, 1) AS pct
FROM summary s
JOIN dept_total t ON s.dept = t.dept
ORDER BY s.dept, s.pay_band
""").show()

-- Scalar subquery: compare each salary to the company average
spark.sql("""
SELECT name, dept, salary,
       (SELECT AVG(salary) FROM employees) AS company_avg,
       salary - (SELECT AVG(salary) FROM employees) AS diff
FROM employees
ORDER BY diff DESC
""").show()
-- Note: Catalyst optimizes this — doesn't run the subquery once per row

-- Create a departments reference table
depts_data = [("Eng", "San Francisco"), ("Mktg", "New York")]
depts_df = spark.createDataFrame(depts_data, ["dept_name", "location"])
depts_df.createOrReplaceTempView("departments")

-- IN subquery: employees in departments with a known office
spark.sql("""
SELECT name, dept, salary
FROM employees
WHERE dept IN (SELECT dept_name FROM departments)
""").show()

-- NOT IN subquery: employees in unknown departments
spark.sql("""
SELECT name, dept
FROM employees
WHERE dept NOT IN (SELECT dept_name FROM departments)
""").show()

-- FROM subquery: wrap a GROUP BY result and filter on it
spark.sql("""
SELECT *
FROM (
    SELECT dept,
           AVG(salary) AS avg_sal,
           COUNT(*)    AS cnt
    FROM employees
    GROUP BY dept
) dept_agg
WHERE avg_sal > 70000
ORDER BY avg_sal DESC
""").show()

-- EXISTS correlated: employees in departments that have at least 2 people
spark.sql("""
SELECT e1.name, e1.dept, e1.salary
FROM employees e1
WHERE EXISTS (
    SELECT 1
    FROM employees e2
    WHERE e2.dept = e1.dept   -- ← references outer query's e1.dept
      AND e2.id != e1.id      -- ← at least one OTHER person in same dept
)
ORDER BY e1.dept, e1.salary DESC
""").show()

-- Correlated scalar: find highest salary in same dept for each employee
spark.sql("""
SELECT e1.name, e1.dept, e1.salary,
       (SELECT MAX(e2.salary)
        FROM employees e2
        WHERE e2.dept = e1.dept) AS dept_max_salary,  -- correlated
       e1.salary = (SELECT MAX(e2.salary)
                    FROM employees e2
                    WHERE e2.dept = e1.dept) AS is_top_earner
FROM employees e1
ORDER BY e1.dept
""").show()

-- Same result — Catalyst rewrites correlated to a join anyway
-- Better to write it as a window function (covered in 37.4):
spark.sql("""
SELECT name, dept, salary,
       MAX(salary) OVER (PARTITION BY dept) AS dept_max_salary,
       salary = MAX(salary) OVER (PARTITION BY dept) AS is_top_earner
FROM employees
ORDER BY dept
""").show()

from pyspark.sql.types import StructType, StructField, StringType, ArrayType, IntegerType

# Data with arrays — each employee has a list of skills
skills_data = [
    (1, "Alice", ["Python", "Spark", "SQL"]),
    (2, "Bob",   ["Java", "Scala"]),
    (3, "Carol", []),                 # empty array
    (4, "Dave",  None),               # null array
]
schema = StructType([
    StructField("id", IntegerType()),
    StructField("name", StringType()),
    StructField("skills", ArrayType(StringType()))
])
skills_df = spark.createDataFrame(skills_data, schema)
skills_df.createOrReplaceTempView("emp_skills")

-- LATERAL VIEW EXPLODE: one row per skill (drops empty/null arrays)
spark.sql("""
SELECT id, name, skill
FROM emp_skills
LATERAL VIEW EXPLODE(skills) skill_table AS skill
ORDER BY id, skill
""").show()

-- LATERAL VIEW OUTER: preserves rows with null/empty arrays
spark.sql("""
SELECT id, name, skill
FROM emp_skills
LATERAL VIEW OUTER EXPLODE(skills) skill_table AS skill
ORDER BY id
""").show()

-- LATERAL VIEW with POSEXPLODE: get position + value
spark.sql("""
SELECT id, name, pos, skill
FROM emp_skills
LATERAL VIEW POSEXPLODE(skills) skill_table AS pos, skill
ORDER BY id, pos
""").show()

from pyspark.sql.types import MapType

# Data with maps — employee -> {score_type: score}
scores_data = [
    (1, "Alice", {"performance": 95, "attendance": 88}),
    (2, "Bob",   {"performance": 78, "attendance": 92}),
]
schema2 = StructType([
    StructField("id", IntegerType()),
    StructField("name", StringType()),
    StructField("scores", MapType(StringType(), IntegerType()))
])
scores_df = spark.createDataFrame(scores_data, schema2)
scores_df.createOrReplaceTempView("emp_scores")

-- LATERAL VIEW EXPLODE on a map: key and value columns
spark.sql("""
SELECT id, name, score_type, score_value
FROM emp_scores
LATERAL VIEW EXPLODE(scores) score_table AS score_type, score_value
ORDER BY id, score_type
""").show()

-- All four ranking functions on salary within department
spark.sql("""
SELECT
    name, dept, salary,
    ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) AS row_num,
    RANK()       OVER (PARTITION BY dept ORDER BY salary DESC) AS rnk,
    DENSE_RANK() OVER (PARTITION BY dept ORDER BY salary DESC) AS dense_rnk,
    NTILE(2)     OVER (PARTITION BY dept ORDER BY salary DESC) AS quartile
FROM employees
ORDER BY dept, salary DESC
""").show()

-- Practical: Get the TOP-1 earner per department (most common interview pattern)
spark.sql("""
WITH ranked AS (
    SELECT *,
           ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) AS rn
    FROM employees
)
SELECT name, dept, salary
FROM ranked
WHERE rn = 1
""").show()

-- Top 2 per department using DENSE_RANK (includes ties at rank 2)
spark.sql("""
WITH ranked AS (
    SELECT *,
           DENSE_RANK() OVER (PARTITION BY dept ORDER BY salary DESC) AS dr
    FROM employees
)
SELECT name, dept, salary, dr
FROM ranked
WHERE dr <= 2
ORDER BY dept, dr
""").show()

# Time-series sales data to demonstrate LAG/LEAD
sales_data = [
    ("2024-01", "Eng", 100000), ("2024-02", "Eng", 110000),
    ("2024-03", "Eng", 105000), ("2024-04", "Eng", 115000),
    ("2024-01", "Mktg", 50000), ("2024-02", "Mktg", 48000),
    ("2024-03", "Mktg", 52000), ("2024-04", "Mktg", 55000),
]
sales_df = spark.createDataFrame(sales_data, ["month", "dept", "revenue"])
sales_df.createOrReplaceTempView("monthly_sales")

spark.sql("""
SELECT
    month, dept, revenue,
    -- LAG: previous row's value (offset=1, default=0 if no prev row)
    LAG(revenue, 1, 0) OVER (PARTITION BY dept ORDER BY month) AS prev_month_rev,
    -- LEAD: next row's value
    LEAD(revenue, 1, 0) OVER (PARTITION BY dept ORDER BY month) AS next_month_rev,
    -- Month-over-month change
    revenue - LAG(revenue, 1, 0) OVER (PARTITION BY dept ORDER BY month) AS mom_change,
    -- MoM % change
    ROUND(
        (revenue - LAG(revenue, 1) OVER (PARTITION BY dept ORDER BY month))
        * 100.0
        / LAG(revenue, 1) OVER (PARTITION BY dept ORDER BY month), 1
    ) AS mom_pct
FROM monthly_sales
ORDER BY dept, month
""").show()

spark.sql("""
SELECT
    month, dept, revenue,
    -- FIRST_VALUE: first revenue in this dept's window
    FIRST_VALUE(revenue) OVER (PARTITION BY dept ORDER BY month
        ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS first_month_rev,
    -- LAST_VALUE: last revenue (IMPORTANT: need explicit frame for last_value!)
    LAST_VALUE(revenue)  OVER (PARTITION BY dept ORDER BY month
        ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS last_month_rev
FROM monthly_sales
ORDER BY dept, month
""").show()

spark.sql("""
SELECT
    name, dept, salary,
    ROUND(PERCENT_RANK() OVER (PARTITION BY dept ORDER BY salary), 3) AS pct_rank,
    ROUND(CUME_DIST()   OVER (PARTITION BY dept ORDER BY salary), 3) AS cume_dist
FROM employees
ORDER BY dept, salary
""").show()

-- Practical: Find employees in the top 25% by salary in their dept
spark.sql("""
WITH dist AS (
    SELECT *,
           CUME_DIST() OVER (PARTITION BY dept ORDER BY salary) AS cd
    FROM employees
)
SELECT name, dept, salary, ROUND(cd, 3) AS cume_dist
FROM dist
WHERE cd >= 0.75   -- top 25% = cume_dist in [0.75, 1.0]
ORDER BY dept, salary DESC
""").show()

# Richer dataset with region
data2 = [
    ("Alice", "Eng",  "West", 90000),
    ("Bob",   "Eng",  "East", 85000),
    ("Carol", "Mktg", "West", 70000),
    ("Dave",  "Mktg", "East", 72000),
    ("Eve",   "HR",   "West", 60000),
]
df2 = spark.createDataFrame(data2, ["name", "dept", "region", "salary"])
df2.createOrReplaceTempView("emp2")

-- GROUPING SETS: 3 groupings in one query
spark.sql("""
SELECT dept, region, SUM(salary) AS total_salary, COUNT(*) AS headcount
FROM emp2
GROUP BY GROUPING SETS (
    (dept, region),   -- group 1: by dept AND region
    (dept),           -- group 2: by dept only
    ()                -- group 3: grand total (no grouping)
)
ORDER BY dept NULLS LAST, region NULLS LAST
""").show()

-- ROLLUP(dept, region) generates: (dept,region), (dept), ()
-- Perfect for subtotals
spark.sql("""
SELECT
    COALESCE(dept, 'ALL DEPTS')     AS dept,
    COALESCE(region, 'ALL REGIONS') AS region,
    SUM(salary)   AS total_salary,
    COUNT(*)      AS headcount,
    ROUND(AVG(salary), 0) AS avg_salary
FROM emp2
GROUP BY ROLLUP(dept, region)
ORDER BY dept, region
""").show()

-- GROUPING() function: tells you if a column was rolled up (1=rolled up)
spark.sql("""
SELECT
    dept,
    region,
    SUM(salary) AS total_salary,
    GROUPING(dept)   AS dept_is_subtotal,    -- 1 if dept was rolled up
    GROUPING(region) AS region_is_subtotal   -- 1 if region was rolled up
FROM emp2
GROUP BY ROLLUP(dept, region)
ORDER BY dept NULLS LAST, region NULLS LAST
""").show()

-- CUBE(dept, region) = ALL 4 groupings: (dept,region), (dept), (region), ()
spark.sql("""
SELECT
    COALESCE(dept, '★ ALL')   AS dept,
    COALESCE(region, '★ ALL') AS region,
    SUM(salary)    AS total_salary,
    COUNT(*)       AS headcount
FROM emp2
GROUP BY CUBE(dept, region)
ORDER BY dept, region
""").show()

-- GROUPING_ID: unique integer identifying which grouping level each row belongs to
spark.sql("""
SELECT
    dept, region,
    SUM(salary) AS total_salary,
    GROUPING_ID(dept, region) AS gid
    -- gid=0: (dept,region)   both grouped
    -- gid=1: (dept)          only dept grouped, region rolled up
    -- gid=2: (region)        only region grouped, dept rolled up
    -- gid=3: ()              grand total, both rolled up
FROM emp2
GROUP BY CUBE(dept, region)
ORDER BY gid, dept NULLS LAST, region NULLS LAST
""").show()

-- Static PIVOT: known column values listed explicitly
spark.sql("""
SELECT *
FROM (
    SELECT dept, region, salary
    FROM emp2
)
PIVOT (
    SUM(salary) AS total_sal, COUNT(*) AS cnt
    FOR region IN ('East' AS East, 'West' AS West)
)
ORDER BY dept
""").show()

# Dynamic pivot — discover pivot values at runtime
pivot_values = [row["region"] for row in df2.select("region").distinct().collect()]
print("Pivot values:", pivot_values)

result = (df2
    .groupBy("dept")
    .pivot("region", pivot_values)
    .agg({"salary": "sum"})
)
result.show()

# Wide format: each month is a column
wide_data = [
    ("Eng",  100000, 110000, 105000),
    ("Mktg", 50000,  48000,  52000),
]
wide_df = spark.createDataFrame(wide_data, ["dept", "jan", "feb", "mar"])
wide_df.createOrReplaceTempView("wide_sales")

-- UNPIVOT: Spark 3.4+ SQL syntax
spark.sql("""
SELECT dept, month, revenue
FROM wide_sales
UNPIVOT (revenue FOR month IN (jan AS 'January', feb AS 'February', mar AS 'March'))
ORDER BY dept, month
""").show()

# Alternative using stack() — works in older Spark versions
spark.sql("""
SELECT dept, month, revenue
FROM wide_sales
LATERAL VIEW EXPLODE(
    MAP('January', jan, 'February', feb, 'March', mar)
) t AS month, revenue
ORDER BY dept, month
""").show()

# Write a Parquet file first
df2.write.mode("overwrite").parquet("/tmp/emp_parquet")

# Read and filter
df_parquet = spark.read.parquet("/tmp/emp_parquet")
df_parquet.createOrReplaceTempView("emp_pq")

result = spark.sql("SELECT dept, salary FROM emp_pq WHERE salary > 80000")
result.explain()
# Look for: PushedFilters: [IsNotNull(salary), GreaterThan(salary,80000.0)]
# This means Parquet reader applies the filter at scan time

# Pushdown FAILS when you filter on a computed column
df_parquet.selectExpr("salary * 1.1 AS bumped_salary") \
    .where("bumped_salary > 88000") \
    .explain()
# No PushedFilters — Spark can't push down computed columns

# Even though the table has 4 columns, only 'dept' and 'salary' are read
result = spark.sql("SELECT dept, SUM(salary) FROM emp_pq GROUP BY dept")
result.explain()
# Look for: ReadSchema: struct<dept:string,salary:bigint>
# 'id' and 'name' columns are NOT in ReadSchema — they were pruned

# SELECT * disables column pruning — avoid it on wide tables in production
result_star = spark.sql("SELECT * FROM emp_pq WHERE dept = 'Eng'")
result_star.explain()
# ReadSchema: struct<id:int,name:string,dept:string,salary:bigint>
# All 4 columns read even if you only use 1

# Best practice: always name your columns explicitly
result_explicit = spark.sql("SELECT dept, name FROM emp_pq WHERE dept = 'Eng'")
result_explicit.explain()
# ReadSchema: struct<dept:string,name:string> — only 2 of 4 columns

# Enable CBO
spark.conf.set("spark.sql.cbo.enabled", True)
spark.conf.set("spark.sql.cbo.joinReorder.enabled", True)

# Collect statistics so CBO has information to work with
spark.sql("ANALYZE TABLE employees COMPUTE STATISTICS FOR ALL COLUMNS")

# CBO will reorder this join to most selective first
result = spark.sql("""
SELECT e.name, e.salary, d.location
FROM employees e
JOIN departments d ON e.dept = d.dept_name
WHERE e.salary > 85000    -- highly selective — CBO applies this filter early
""")
result.explain(mode="cost")  # Shows estimated row counts from statistics

# Manual hint to force broadcast join (overrides CBO)
spark.sql("""
SELECT /*+ BROADCAST(d) */ e.name, e.salary, d.location
FROM employees e
JOIN departments d ON e.dept = d.dept_name
""").explain()
# Look for BroadcastHashJoin in plan

# Create a partitioned parquet dataset (partition by dept)
df2.write.mode("overwrite").partitionBy("dept").parquet("/tmp/emp_partitioned")

df_part = spark.read.parquet("/tmp/emp_partitioned")
df_part.createOrReplaceTempView("emp_partitioned")

-- Static partition pruning: literal value in filter
result = spark.sql("""
SELECT name, salary
FROM emp_partitioned
WHERE dept = 'Eng'   -- ← static pruning: only reads /dept=Eng/ folder
""")
result.explain()
# Look for: PartitionFilters: [isnotnull(dept#xx), (dept#xx = Eng)]
# This confirms only the Eng partition was scanned

-- Dynamic partition pruning (Spark 3.0+):
-- Filter value comes from another table's result at runtime
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true")

-- Example: only scan partitions that have a match in departments table
result2 = spark.sql("""
SELECT e.name, e.salary, d.location
FROM emp_partitioned e
JOIN departments d ON e.dept = d.dept_name
WHERE d.location = 'San Francisco'
-- Spark evaluates the departments filter first,
-- then uses the matching dept values to prune emp_partitioned partitions
""")
result2.explain(mode="formatted")

# Enable CBO and all sub-features
spark.conf.set("spark.sql.cbo.enabled", True)
spark.conf.set("spark.sql.cbo.joinReorder.enabled", True)
spark.conf.set("spark.sql.statistics.histogram.enabled", True)

# Collect statistics (do this after writing a managed table)
spark.sql("ANALYZE TABLE employees COMPUTE STATISTICS")          # row count only
spark.sql("ANALYZE TABLE employees COMPUTE STATISTICS FOR ALL COLUMNS")  # full stats

# View collected statistics
spark.sql("DESCRIBE EXTENDED employees").show(truncate=False)

# Column-level stats
spark.sql("DESCRIBE EXTENDED employees salary").show(truncate=False)

# Use explain with 'cost' mode to see CBO estimates
spark.sql("""
SELECT e.name, e.dept
FROM employees e
JOIN departments d ON e.dept = d.dept_name
""").explain(mode="cost")
# Look for: Statistics(sizeInBytes=..., rowCount=...)
# These are CBO's estimates used to choose the join strategy

-- CTE + Window + Correlated filter — cleanest in SQL
spark.sql("""
WITH dept_stats AS (
    SELECT dept,
           AVG(salary)  AS avg_sal,
           STDDEV(salary) AS std_sal
    FROM employees
    GROUP BY dept
)
SELECT e.name, e.dept, e.salary,
       d.avg_sal,
       ROUND((e.salary - d.avg_sal) / NULLIF(d.std_sal, 0), 2) AS z_score,
       ROW_NUMBER() OVER (PARTITION BY e.dept ORDER BY e.salary DESC) AS dept_rank
FROM employees e
JOIN dept_stats d USING (dept)
WHERE e.salary > d.avg_sal
ORDER BY e.dept, dept_rank
""").show()

from pyspark.sql import functions as F

# Dynamic: apply different aggregations based on config
agg_config = {
    "salary": ["avg", "max", "min"],
    "id": ["count"]
}

# Build agg expressions dynamically — IMPOSSIBLE in pure SQL string
agg_exprs = []
for col_name, funcs in agg_config.items():
    for func in funcs:
        agg_exprs.append(getattr(F, func)(col_name).alias(f"{func}_{col_name}"))

result = df.groupBy("dept").agg(*agg_exprs)
result.show()

# Reusable transformation function
def add_salary_band(df, salary_col="salary"):
    """Add pay band column based on salary. Reusable across pipelines."""
    return df.withColumn(
        "pay_band",
        F.when(F.col(salary_col) > 85000, "High")
         .when(F.col(salary_col) > 70000, "Mid")
         .otherwise("Low")
    )

df_banded = add_salary_band(df)
df_banded.show()

from pyspark.sql import functions as F

# SQL version
sql_result = spark.sql("""
    SELECT dept, AVG(salary) AS avg_salary
    FROM employees
    WHERE salary > 70000
    GROUP BY dept
    ORDER BY avg_salary DESC
""")

# DataFrame API version — exactly equivalent
df_result = (df
    .filter(F.col("salary") > 70000)
    .groupBy("dept")
    .agg(F.avg("salary").alias("avg_salary"))
    .orderBy(F.desc("avg_salary"))
)

# Compare explain plans — they should be IDENTICAL
print("=== SQL Plan ===")
sql_result.explain()

print("=== DataFrame Plan ===")
df_result.explain()
# Both plans are identical: Filter → HashAggregate → Exchange → HashAggregate → Sort