from pyspark.sql.functions import col

# Method 1: String column names
df.select("name", "dept", "salary").show()

# Method 2: col() expressions — more powerful
df.select(col("name"), col("salary")).show()

# Method 3: dot notation on DataFrame
df.select(df.name, df.salary).show()

# Method 4: Select ALL columns (rarely needed but useful)
df.select("*").show()

# Computed column inside select
df.select(
    col("name"),
    col("salary") * 1.1   # 10% raise
).show()

# selectExpr: write SQL-like expressions as strings
df.selectExpr(
    "name",
    "salary",
    "salary * 1.1 AS salary_with_raise",   # computed + renamed
    "UPPER(dept) AS department",              # SQL function
    "age > 30 AS is_senior"                   # boolean expression
).show()

# Also useful: check a condition
df.selectExpr("*", "salary > 70000 AS high_earner").show()

from pyspark.sql.functions import col, upper, round

# Without alias — ugly column name
df.select(col("salary") * 1.1).show(2)
# Column is named: (salary * 1.1)

# With alias — clean name
df.select(
    col("name"),
    (col("salary") * 1.1).alias("new_salary"),
    upper(col("dept")).alias("department"),
    round(col("salary") / 12, 2).alias("monthly_pay")
).show()

# You can also alias an existing column
df.select(col("name").alias("employee_name")).show()

# alias() on grouped column reference
df.select("name", col("salary").alias("annual_salary")).printSchema()

# Method 1: SQL string expression
df.filter("salary > 70000").show()
df.where("salary > 70000").show()   # identical

# Method 2: col() expression
df.filter(col("salary") > 70000).show()

# Method 3: dot notation
df.filter(df.salary > 70000).show()

# Multiple conditions with & (AND) and | (OR)
df.filter(
    (col("salary") > 70000) & (col("dept") == "Engineering")
).show()

df.filter(
    (col("dept") == "HR") | (col("age") < 28)
).show()

# NOT condition with ~
df.filter(~(col("dept") == "HR")).show()

# Filter for null values
df.filter(col("salary").isNull()).show()
df.filter(col("salary").isNotNull()).show()

# Filter with isin()
df.filter(col("dept").isin(["Engineering", "HR"])).show()

# Chain multiple filters (more readable)
df.filter(col("dept") == "Engineering") \
  .filter(col("age") > 28) \
  .show()

# between() — range filter
df.filter(col("salary").between(60000, 90000)).show()

# like() — pattern matching (SQL LIKE)
df.filter(col("name").like("A%")).show()   # starts with A
df.filter(col("name").like("%e")).show()   # ends with e

# rlike() — regex
df.filter(col("name").rlike("^[AB]")).show()  # starts with A or B

from pyspark.sql.functions import col, lit, when, upper, round

# Add a new column: constant value
df.withColumn("country", lit("India")).show()

# Add a computed column
df.withColumn("monthly_salary", col("salary") / 12).show()

# Overwrite existing column (modify in-place)
df.withColumn("salary", col("salary") * 1.1).show()

# Add derived column using when/otherwise (conditional)
df.withColumn(
    "grade",
    when(col("salary") >= 90000, "A")
    .when(col("salary") >= 70000, "B")
    .otherwise("C")
).show()

# Uppercase a column
df.withColumn("dept_upper", upper(col("dept"))).show()

# Chain multiple withColumn calls
df2 = df \
    .withColumn("monthly_salary", round(col("salary") / 12, 2)) \
    .withColumn("is_senior", col("age") > 30) \
    .withColumn("dept_upper", upper(col("dept")))
df2.show()

# withColumns — Spark 3.3+
df.withColumns({
    "monthly_salary": round(col("salary") / 12, 2),
    "is_senior":      col("age") > 30,
    "dept_upper":     upper(col("dept")),
    "country":        lit("India")
}).show()

# Drop a single column
df.drop("age").show()

# Drop multiple columns
df.drop("age", "id").show()

# Using a list
cols_to_drop = ["age", "id"]
df.drop(*cols_to_drop).show()

# drop() is safe — no error if column doesn't exist
df.drop("nonexistent_column").show()  # works fine

# Rename a single column
df.withColumnRenamed("name", "employee_name").show()

# Chain to rename multiple columns
df.withColumnRenamed("name", "employee_name") \
  .withColumnRenamed("dept", "department") \
  .withColumnRenamed("salary", "annual_salary") \
  .show()

# Rename many columns at once using toDF()
new_cols = ["employee_id", "employee_name", "department", "annual_salary", "age"]
df.toDF(*new_cols).show()

# Dynamic rename: make all columns lowercase
from functools import reduce
renamed = reduce(
    lambda acc, c: acc.withColumnRenamed(c, c.lower()),
    df.columns, df
)
renamed.show()

from pyspark.sql.types import IntegerType, DoubleType, StringType, LongType

# cast using type string
df.withColumn("salary", col("salary").cast("double")).printSchema()
df.withColumn("age",    col("age").cast("string")).printSchema()

# cast using type objects (more explicit)
df.withColumn("salary", col("salary").cast(DoubleType())).printSchema()
df.withColumn("id",     col("id").cast(LongType())).printSchema()

# Cast multiple columns at once using select
df.select(
    col("name"),
    col("salary").cast("double").alias("salary"),
    col("age").cast("long").alias("age")
).printSchema()

from pyspark.sql.functions import col, desc, asc

# Ascending sort (default)
df.orderBy("salary").show()
df.sort("salary").show()         # same thing

# Descending sort
df.orderBy(col("salary").desc()).show()
df.orderBy(desc("salary")).show()   # equivalent

# Sort by multiple columns
df.orderBy(
    col("dept").asc(),
    col("salary").desc()    # within each dept, highest salary first
).show()

# Nulls placement: null first or null last
df.orderBy(col("salary").asc_nulls_last()).show()   # nulls go to end
df.orderBy(col("salary").asc_nulls_first()).show()  # nulls come first

# Create a DataFrame with duplicates
dup_data = [
    (1, "Alice", "Engineering"),
    (1, "Alice", "Engineering"),  # exact duplicate
    (2, "Bob",   "Marketing"),
    (3, "Bob",   "HR"),           # same name, different dept
]
dup_df = spark.createDataFrame(dup_data, ["id", "name", "dept"])

# distinct() — removes exact duplicates
dup_df.distinct().show()
# Result: 3 rows (row 1 duplicate removed)

dup_df.distinct().count()  # 3

# dropDuplicates() with no args = same as distinct()
dup_df.dropDuplicates().show()

# Deduplicate on specific columns only
# Keep only the first row per unique "name"
dup_df.dropDuplicates(["name"]).show()

# Real-world: deduplicate on composite key
orders = spark.createDataFrame([
    ("ORD001", "2024-01-01", 100),
    ("ORD001", "2024-01-01", 100),  # duplicate event
    ("ORD002", "2024-01-02", 200),
], ["order_id", "order_date", "amount"])

orders.dropDuplicates(["order_id", "order_date"]).show()

# Drop rows where ANY column is null
df.dropna().show()

# Drop rows where ALL columns are null
df.dropna(how="all").show()

# Drop rows where specific columns are null
df.dropna(subset=["salary", "dept"]).show()

# Drop rows with fewer than N non-null values
df.dropna(thresh=4).show()  # keep rows with at least 4 non-null values

# Fill ALL null columns with a value (type must match)
df.fillna(0).show()            # fills numeric nulls with 0
df.fillna("Unknown").show()   # fills string nulls with "Unknown"

# Fill specific columns with specific values (dict)
df.fillna({
    "salary": 50000,
    "dept":   "Unknown",
    "age":    0
}).show()

# Fill only specific columns
df.fillna(0, subset=["salary", "age"]).show()

# Using na.fill() — same as fillna()
df.na.fill(0).show()

# Replace specific values
df.replace("Engineering", "Eng", subset=["dept"]).show()

# Replace multiple values at once using a dict
df.replace({
    "Engineering": "Eng",
    "Marketing":   "Mkt",
    "HR":          "Human Resources"
}, subset=["dept"]).show()

# Replace null-like sentinel values
df.replace(-1, None, subset=["salary"]).show()   # -1 → null

# Using na.replace() — same as replace()
df.na.replace("Engineering", "Eng").show()

from pyspark.sql.functions import count, sum, avg, min, max, col

# Basic groupBy + single aggregation
df.groupBy("dept").count().show()

# groupBy + avg
df.groupBy("dept").avg("salary").show()

# Multiple aggregations with agg()
df.groupBy("dept").agg(
    count("*").alias("headcount"),
    sum("salary").alias("total_salary"),
    avg("salary").alias("avg_salary"),
    min("salary").alias("min_salary"),
    max("salary").alias("max_salary"),
    min("age").alias("youngest"),
    max("age").alias("oldest")
).show()

# Group by multiple columns
df.groupBy("dept", col("age") > 30).agg(
    count("*").alias("count"),
    avg("salary").alias("avg_salary")
).show()

# groupBy + filter result (HAVING clause equivalent)
df.groupBy("dept").agg(
    count("*").alias("headcount"),
    avg("salary").alias("avg_salary")
).filter(col("headcount") > 1).show()   # HAVING headcount > 1

from pyspark.sql.functions import (
    count, countDistinct, approx_count_distinct
)

df.groupBy("dept").agg(
    count("*").alias("total_rows"),             # counts ALL rows
    count("salary").alias("non_null_salaries"),   # skips nulls
    countDistinct("age").alias("unique_ages"),     # distinct count (exact)
    approx_count_distinct("age").alias("approx_ages")  # fast approximate
).show()

from pyspark.sql.functions import (
    sum, avg, mean, min, max, stddev, variance
)

df.groupBy("dept").agg(
    sum("salary").alias("total"),
    avg("salary").alias("average"),
    mean("salary").alias("mean"),       # same as avg
    min("salary").alias("minimum"),
    max("salary").alias("maximum"),
    stddev("salary").alias("std_dev"),   # standard deviation
    variance("salary").alias("variance")  # variance
).show()

from pyspark.sql.functions import collect_list, collect_set

# Collect all names per department into an array
df.groupBy("dept").agg(
    collect_list("name").alias("all_employees"),  # with duplicates
    collect_set("age").alias("unique_ages")        # no duplicates
).show(truncate=False)

# Employees table
employees = spark.createDataFrame([
    (1, "Alice", 10),
    (2, "Bob",   20),
    (3, "Carol", 30),
    (4, "Dave",  99),  # dept_id 99 doesn't exist
], ["emp_id", "name", "dept_id"])

# Departments table
departments = spark.createDataFrame([
    (10, "Engineering"),
    (20, "Marketing"),
    (30, "HR"),
    (40, "Finance"),    # no employee in Finance
], ["dept_id", "dept_name"])

employees.join(departments, "dept_id", "inner").show()
# or simply: employees.join(departments, "dept_id") — inner is default

# Join on expression (for different column names)
employees.join(
    departments,
    employees.dept_id == departments.dept_id,
    "inner"
).show()

employees.join(departments, "dept_id", "left").show()
# Also: "left_outer" works too

# Right Join — all from right, matching from left
employees.join(departments, "dept_id", "right").show()
# Finance appears with null employees

# Full / Full Outer Join — all rows from both
employees.join(departments, "dept_id", "full").show()
# Dave (no dept) AND Finance (no emp) both appear with nulls

# Semi Join — like inner, but returns ONLY left columns
# "Give me employees who have a matching department"
employees.join(departments, "dept_id", "left_semi").show()

# Anti Join — opposite of semi — rows with NO match
# "Give me employees who do NOT have a department"
employees.join(departments, "dept_id", "left_anti").show()

# Cross Join — every row × every row (Cartesian product)
# 4 employees × 4 departments = 16 rows
employees.crossJoin(departments).show()

# Self Join — join a DataFrame with itself
# Example: find employees in the same department
emp1 = employees.alias("e1")
emp2 = employees.alias("e2")
emp1.join(emp2,
    (col("e1.dept_id") == col("e2.dept_id")) &
    (col("e1.emp_id") != col("e2.emp_id"))
).select(
    col("e1.name").alias("employee"),
    col("e2.name").alias("colleague")
).show()

from pyspark.sql.functions import broadcast

# Manual broadcast — wrap the small DataFrame
employees.join(
    broadcast(departments),   # <-- departments is small, broadcast it
    "dept_id",
    "inner"
).show()

# Spark auto-broadcasts when table is below threshold
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "50m")   # 50 MB

# Disable auto-broadcast (force sort-merge join)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")

# Enable AQE for automatic skew join handling
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

# Manual salting for skew join (when AQE is unavailable)
from pyspark.sql.functions import expr, floor, rand

SALT_BUCKETS = 5

# Add salt to the skewed (large) table
large_df = employees.withColumn(
    "salt_key",
    (floor(rand() * SALT_BUCKETS)).cast("int")
).withColumn("dept_id_salted", expr(f"concat(dept_id, '_', salt_key)"))

# Explode the small table to match all salt values
from pyspark.sql.functions import array, explode
small_df = departments.withColumn(
    "salt_array", array([lit(i) for i in range(SALT_BUCKETS)])
).withColumn("salt_key", explode(col("salt_array"))) \
 .withColumn("dept_id_salted", expr(f"concat(dept_id, '_', salt_key)"))

large_df.join(small_df, "dept_id_salted").show()

# Join hints via SQL
spark.sql("""
    SELECT /*+ BROADCAST(d) */ e.name, d.dept_name
    FROM employees e
    JOIN departments d ON e.dept_id = d.dept_id
""")

# Available hints:
# BROADCAST(table)    — force broadcast
# MERGE(table)        — force sort-merge join
# SHUFFLE_HASH(table) — force shuffle hash join
# SHUFFLE_REPLICATE_NL(table) — force nested loop

# Via DataFrame API (hint on a table reference)
employees.join(
    departments.hint("broadcast"),
    "dept_id"
).show()

q1_sales = spark.createDataFrame([
    ("Alice", 1000, "Q1"),
    ("Bob",   2000, "Q1"),
], ["name", "amount", "quarter"])

q2_sales = spark.createDataFrame([
    ("Carol", 1500, "Q2"),
    ("Dave",  3000, "Q2"),
], ["name", "amount", "quarter"])

# Stack rows — does NOT remove duplicates (like UNION ALL)
all_sales = q1_sales.union(q2_sales)
all_sales.show()

# To deduplicate, add distinct()
all_sales.distinct().show()

# Different column order — union() would mix up data!
df_a = spark.createDataFrame([("Alice", 1000)], ["name", "amount"])
df_b = spark.createDataFrame([(2000, "Bob")],   ["amount", "name"])

# union() — WRONG! amount and name get mixed
df_a.union(df_b).show()   # Bob's amount (2000) goes in name column!

# unionByName() — CORRECT! matches by column name
df_a.unionByName(df_b).show()

# allowMissingColumns=True — Spark 3.1+
# Fills null for missing columns
df_c = spark.createDataFrame([("Carol", 1500, "Engineering")], ["name", "amount", "dept"])
df_a.unionByName(df_c, allowMissingColumns=True).show()
# Alice has null for dept column

batch1 = spark.createDataFrame([(1,"Alice"),(2,"Bob"),(3,"Carol")], ["id","name"])
batch2 = spark.createDataFrame([(2,"Bob"),(3,"Carol"),(4,"Dave")],  ["id","name"])

# Rows present in BOTH DataFrames
batch1.intersect(batch2).show()
# Returns: Bob, Carol

# except() — rows in batch1 but NOT in batch2 (removes duplicates)
batch1.exceptAll(batch2).show()
# Returns: Alice

# exceptAll() vs except()
# except() deduplicates, exceptAll() preserves all occurrences
dup = spark.createDataFrame([(1,"Alice"),(1,"Alice"),(2,"Bob")], ["id","name"])
small = spark.createDataFrame([(1,"Alice")], ["id","name"])

dup.subtract(small).show()     # returns: [Bob] (except removes both Alices)
dup.exceptAll(small).show()    # returns: [Alice, Bob] (only removes one Alice)

# sample(withReplacement, fraction, seed)
# withReplacement=False: each row can appear at most once
# fraction: between 0.0 and 1.0 (10% = 0.1)
# seed: for reproducibility

# Sample ~30% of rows without replacement
df.sample(False, 0.3, seed=42).show()

# Sample with replacement (bootstrap sampling)
df.sample(True, 0.5, seed=99).show()

# Count to see approximate result
df.count()                    # 7
df.sample(False, 0.3).count() # ~2 (not exact)

# Sample 50% from Engineering, 100% from HR, 30% from Marketing
fractions = {
    "Engineering": 0.5,
    "HR":          1.0,
    "Marketing":   0.3
}
df.sampleBy("dept", fractions, seed=42).show()

# Real-world: train/test split by category
train = df.sampleBy("dept", {"Engineering":0.8, "HR":0.8, "Marketing":0.8})
test  = df.subtract(train)

# Check current partition count
print(df.rdd.getNumPartitions())   # likely 8 or 200 default

# Repartition to a specific number
df.repartition(10).rdd.getNumPartitions()  # 10

# Repartition by column (hash partitioning)
# All rows with same dept value go to same partition
df.repartition(3, "dept").show()

# Repartition by multiple columns
df.repartition(10, "dept", "age").show()

# Good practice: set shuffle partitions before groupBy/join
spark.conf.set("spark.sql.shuffle.partitions", "200")  # default
spark.conf.set("spark.sql.shuffle.partitions", "10")   # for small data

# After a large filter, you might have 200 nearly-empty partitions
filtered = df.filter(col("salary") > 90000)
filtered.rdd.getNumPartitions()  # still 200 but most are empty

# coalesce to reduce — no shuffle, just merge
filtered.coalesce(5).rdd.getNumPartitions()  # 5

# Typical pattern before writing to disk
# Avoids thousands of tiny output files
df.coalesce(1).write.parquet("/output/path")   # single file output