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

schema = StructType([
    StructField("id", IntegerType(), False),  # Not nullable
    StructField("name", StringType(), True),  # Nullable
    StructField("age", IntegerType(), True)   # Nullable
])

# Catalyst pushes filter to data source
df.filter(df.age > 30).select("name", "age")
# Generates: Scan[name, age] → Filter(age > 30)
# Instead of: Scan[all columns] → Filter → Select

from pyspark.sql.functions import broadcast
result = large_df.join(broadcast(small_df), "id")

df.write.bucketBy(100, "id").saveAsTable("bucketed_table")

from pyspark.sql import SparkSession
from pyspark.sql.types import *
from pyspark.sql.functions import *

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

# Method 1: From list of tuples with explicit schema
# StructType: Defines the schema structure
# StructField: Defines each column (name, type, nullable)
schema = StructType([
    StructField("id", IntegerType(), False),      # Not nullable
    StructField("name", StringType(), True),      # Nullable
    StructField("age", IntegerType(), True),      # Nullable
    StructField("salary", DoubleType(), True),    # Nullable
    StructField("department", StringType(), True)  # Nullable
])

data = [
    (1, "Alice", 30, 75000.0, "Engineering"),
    (2, "Bob", 25, 65000.0, "Marketing"),
    (3, "Charlie", 35, 90000.0, "Engineering"),
    (4, "Diana", 28, 70000.0, "Marketing"),
    (5, "Eve", 32, 85000.0, "Engineering")
]

# createDataFrame: Create DataFrame from data and schema
df = spark.createDataFrame(data, schema)

# Method 2: From pandas DataFrame
# Converts pandas DataFrame to Spark DataFrame
import pandas as pd
pandas_df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
df_from_pandas = spark.createDataFrame(pandas_df)

# Method 3: From CSV with schema inference
# header: First line is column names
# inferSchema: Automatically detect data types (requires extra pass)
df_from_csv = spark.read.csv("data.csv", header=True, inferSchema=True)

# Method 4: From JSON
# Each line is a JSON object
df_from_json = spark.read.json("data.json")

# Method 5: From Parquet (columnar format)
# Parquet stores schema with data, so no schema needed
df_from_parquet = spark.read.parquet("data.parquet")

# Display schema
df.printSchema()
# root
#  |-- id: integer (not null)
#  |-- name: string (nullable = true)
#  |-- age: integer (nullable = true)
#  |-- salary: double (nullable = true)
#  |-- department: string (nullable = true)

# Get schema as string
schema_str = df.schema.json()
print(f"Schema JSON: {schema_str}")

# Get column names and data types
print(f"Columns: {df.columns}")
print(f"Data types: {df.dtypes}")

# Select columns
# Parameters: *cols - column names or Column objects
selected = df.select("name", "age", "salary")
selected.show()

# Filter rows
# Parameters: condition - boolean Column expression
young = df.filter(df.age < 30)
young.show()

# Alternative: where() is an alias for filter()
young_alt = df.where(df.age < 30)

# Add computed columns
# Parameters: colName - new column name, col - Column expression
with_bonus = df.withColumn("bonus", df.salary * 0.1)
with_bonus.show()

# Complex transformations
result = df \
    .filter(col("age") > 25) \           # Keep employees over 25
    .withColumn("bonus", col("salary") * 0.1) \  # Calculate 10% bonus
    .withColumn("total_compensation", col("salary") + col("bonus")) \  # Total comp
    .select("name", "department", "total_compensation") \  # Select relevant columns
    .orderBy(desc("total_compensation"))  # Sort by total comp descending

result.show()

# Aggregations
# Parameters: *exprs - aggregation expressions
dept_stats = df.groupBy("department").agg(
    count("id").alias("employee_count"),      # Count employees per department
    avg("salary").alias("avg_salary"),        # Average salary
    max("salary").alias("max_salary"),        # Maximum salary
    min("salary").alias("min_salary"),        # Minimum salary
    sum("salary").alias("total_salary")       # Total salary
)
dept_stats.show()

# Actions
print(f"Total rows: {df.count()}")
print(f"First row: {df.first()}")
print(f"Schema: {df.dtypes}")

# limit: Return first n rows (lazy transformation)
first_5 = df.limit(5)
first_5.show()

# distinct: Remove duplicate rows
distinct_depts = df.select("department").distinct()
distinct_depts.show()

# drop: Remove columns
df_without_age = df.drop("age")
df_without_age.show()

# Create second DataFrame
department_info = spark.createDataFrame([
    ("Engineering", "San Francisco", 50),
    ("Marketing", "New York", 30),
    ("Sales", "Chicago", 40)
], ["department", "location", "headcount"])

# Inner join (default)
# Parameters: other - DataFrame to join with
#             on - join condition (column name or expression)
#             how - join type (default "inner")
inner_joined = df.join(department_info, "department", "inner")
inner_joined.show()

# Left join (keep all rows from left DataFrame)
# Parameters: how="left" or "left_outer"
left_joined = df.join(department_info, "department", "left")
left_joined.show()

# Broadcast join for performance
# broadcast(): Hint to broadcast the small DataFrame
# Avoids shuffle by sending small table to all executors
from pyspark.sql.functions import broadcast
broadcast_joined = df.join(broadcast(department_info), "department")

# Complex join with multiple conditions
# Parameters: on - list of conditions combined with &
complex_join = df.join(
    department_info,
    (df.department == department_info.department) & 
    (df.salary > 50000),
    "inner"
)

# Verify join strategy using explain()
broadcast_joined.explain()

# Right join (keep all rows from right DataFrame)
right_joined = df.join(department_info, "department", "right")
right_joined.show()

# Full outer join (keep all rows from both DataFrames)
full_joined = df.join(department_info, "department", "outer")
full_joined.show()

# Left semi join (return rows from left that have match in right)
left_semi = df.join(department_info, "department", "leftsemi")
left_semi.show()

# Left anti join (return rows from left that have NO match in right)
left_anti = df.join(department_info, "department", "leftanti")
left_anti.show()

from pyspark.sql.window import Window

# Define window specification
# partitionBy: Divide data into partitions (like GROUP BY)
# orderBy: Define ordering within partition
window_spec = Window.partitionBy("department").orderBy(desc("salary"))

# Add row number within department
# row_number(): Assigns unique sequential number to each row
df_with_rank = df.withColumn(
    "rank", 
    row_number().over(window_spec)  # .over() applies window function
)

# Add salary statistics per department
# avg().over(): Calculate average within partition
df_with_stats = df.withColumn(
    "dept_avg_salary",
    avg("salary").over(Window.partitionBy("department"))
).withColumn(
    "salary_vs_avg",
    col("salary") - col("dept_avg_salary")  # Difference from average
)

# Running total
# rowsBetween: Define frame boundaries
# unboundedPreceding to currentRow: Sum from start to current row
running_window = Window.partitionBy("department").orderBy("id").rowsBetween(
    Window.unboundedPreceding,  # Start from first row
    Window.currentRow           # End at current row
)
df_running = df.withColumn(
    "running_total",
    sum("salary").over(running_window)
)

# Dense rank: Like rank but no gaps in ranking
df_dense_rank = df.withColumn(
    "dense_rank",
    dense_rank().over(window_spec)
)

# Percent rank: Relative rank as percentage
df_percent_rank = df.withColumn(
    "percent_rank",
    percent_rank().over(window_spec)
)

# Cumulative distribution
df_cume_dist = df.withColumn(
    "cume_dist",
    cume_dist().over(window_spec)
)

# Lag: Get value from previous row
lag_window = Window.partitionBy("department").orderBy("id")
df_with_lag = df.withColumn(
    "previous_salary",
    lag("salary", 1).over(lag_window)  # 1 row before
)

# Lead: Get value from next row
df_with_lead = df.withColumn(
    "next_salary",
    lead("salary", 1).over(lag_window)  # 1 row after
)

df_with_stats.show()

# Pivot: Rotate rows into columns
# Parameters: pivot_col - column to pivot
#             values - optional list of values to pivot (default: all unique values)
pivoted = df.groupBy("name").pivot("department").sum("salary")
pivoted.show()

# Pivot with specific values (faster)
pivoted_limited = df.groupBy("name").pivot("department", ["Engineering", "Marketing"]).sum("salary")
pivoted_limited.show()

# Unpivot (melt): Rotate columns into rows
# Using stack function
df_melted = df.select(
    "name",
    F.expr("stack(2, 'salary', salary, 'bonus', salary * 0.1) as (metric, value)")
)
df_melted.show()

# Cube: Multi-dimensional aggregation
# Parameters: *cols - columns to aggregate over
# Produces subtotals for all combinations
cube_result = df.cube("department", "age").agg(
    count("*").alias("count"),
    avg("salary").alias("avg_salary")
)
cube_result.show()

# Rollup: Hierarchical aggregation
# Parameters: *cols - columns in hierarchical order
# Produces subtotals for prefixes
rollup_result = df.rollup("department", "age").agg(
    count("*").alias("count"),
    avg("salary").alias("avg_salary")
)
rollup_result.show()

# BAD: Creates Python UDF (loses Catalyst optimization)
df.withColumn("new", df.age + 1)

# GOOD: Uses Catalyst-optimized expression
df.withColumn("new", col("age") + 1)

# GOOD: Use SQL expressions for complex logic
df.withColumn("new", expr("CASE WHEN age > 30 THEN 'Senior' ELSE 'Junior' END"))

# Cache DataFrame used multiple times
df = spark.read.parquet("large_dataset.parquet")
df.cache()  # MEMORY_ONLY

# Use appropriate storage level for large datasets
from pyspark import StorageLevel
df.persist(StorageLevel.MEMORY_AND_DISK_SER)  # Serialized + disk spill

# Remember to unpersist when done
df.unpersist()

# BAD: Brings all data to driver (causes OOM)
all_data = df.collect()

# GOOD: Use take for small samples
first_100 = df.take(100)

# GOOD: Use show for display
df.show(20)

# GOOD: Use toPandas for small datasets (max 1000 rows recommended)
pandas_df = df.limit(1000).toPandas()

from pyspark.sql.functions import broadcast

# Broadcast small DataFrame (< 10MB default threshold)
result = large_df.join(broadcast(small_df), "key")

# Configure broadcast threshold
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "52428800")  # 50MB

# Check join strategy
result.explain()  # Look for BroadcastHashJoin

# Use appropriate types
from pyspark.sql.types import *

# BAD: Using StringType for numeric data (no predicate pushdown)
df = spark.createDataFrame([(1, "100")], ["id", "amount"])

# GOOD: Use proper types (enables predicate pushdown)
schema = StructType([
    StructField("id", IntegerType()),
    StructField("amount", IntegerType())
])

# Use ByteType/ShortType for small numbers (saves memory)
from pyspark.sql.types import ByteType, ShortType
small_schema = StructType([
    StructField("id", IntegerType()),
    StructField("flag", ByteType()),  # 1 byte vs 4 bytes for IntegerType
    StructField("category", ShortType())  # 2 bytes vs 4 bytes
])

# Repartition for write operations (parallel writes)
df.repartition(100, "department").write.parquet("output/")

# Coalesce to reduce partitions (no shuffle)
df.coalesce(10).write.parquet("output_small/")

# Partition by column for partition pruning on reads
df.write.partitionBy("year", "month").parquet("output/")

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

# SQL can be more readable for complex queries
result = spark.sql("""
    SELECT department,
           COUNT(*) as cnt,
           AVG(salary) as avg_sal,
           MAX(salary) as max_sal
    FROM employees
    WHERE age > 25
    GROUP BY department
    HAVING COUNT(*) > 1
    ORDER BY avg_sal DESC
""")
result.show()