from pyspark.sql import SparkSession

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

# Create a simple DataFrame
data = [(1, "Alice"), (2, "Bob"), (3, "Charlie"),
        (4, "Dave"), (5, "Eve")]
df = spark.createDataFrame(data, ["id", "name"])

# Check logical partitions
print(f"Number of logical partitions: {df.rdd.getNumPartitions()}")
# Output: Number of logical partitions: 8  (depends on default parallelism)

# Reading CSV: each file chunk = one logical partition
df_csv = spark.read.csv("s3://bucket/data/*.csv", header=True)
print(df_csv.rdd.getNumPartitions())  # e.g., 12 (one per file or per 128MB block)

# Reading Parquet: one logical partition per row group
df_parquet = spark.read.parquet("s3://bucket/data/sales.parquet")
print(df_parquet.rdd.getNumPartitions())  # e.g., 4

# Writing physical partitions to disk
df.write.partitionBy("country").parquet("s3://bucket/output/")
# Creates: output/country=US/, output/country=IN/, output/country=UK/ ...

# Default shuffle partitions = 200. Change based on data size:
spark.conf.set("spark.sql.shuffle.partitions", "50")
# For small data: reduce to 50 (avoids 200 tiny tasks)
# For large data (TB scale): increase to 1000+

# See current value
print(spark.conf.get("spark.sql.shuffle.partitions"))

# Check partition count after groupBy (triggers shuffle)
result = df.groupBy("country").count()
print(result.rdd.getNumPartitions())  # Will be 200 by default

# Rule of thumb: partition size = total data size / num partitions
# Target: 128MB - 256MB per partition
# Example: 10GB data → ~50 to 80 partitions

from pyspark.sql.functions import spark_partition_id, col

data = [(1, 100), (2, 200), (3, 50), (4, 400), (5, 350),
        (6, 75), (7, 275), (8, 450)]
df = spark.createDataFrame(data, ["id", "amount"])

# Range repartition: divides "amount" into 3 sorted ranges
df_range = df.repartitionByRange(3, "amount")

df_range.withColumn("partition_id", spark_partition_id()) \
        .orderBy("amount") \
        .show()

# Approximate output:
# partition 0: amount 50, 75, 100        ← low values
# partition 1: amount 200, 275           ← mid values
# partition 2: amount 350, 400, 450      ← high values

# Also supports descending order
df_range_desc = df.repartitionByRange(3, col("amount").desc())

# Great before writing sorted data to files
df_range.write.parquet("s3://bucket/sorted-output/")
# Each file will have sorted amounts → better predicate pushdown on reads

# Round robin — NO column specified
# Rows distributed evenly, no grouping by key
df_rr = df.repartition(4)  # 4 partitions, balanced row count

# Check balance across partitions
df_rr.withColumn("partition_id", spark_partition_id()) \
     .groupBy("partition_id") \
     .count() \
     .orderBy("partition_id") \
     .show()
# Output: Each partition has ~equal number of rows
# +------------+-----+
# |partition_id|count|
# +------------+-----+
# |           0|    2|
# |           1|    1|
# |           2|    2|
# |           3|    1|
# +------------+-----+

# Use case: fixing skewed partitions before a write
df_skewed.repartition(100).write.parquet("output/")
# Creates 100 balanced files on disk

# Enable AQE (required for DPP)
spark.conf.set("spark.sql.adaptive.enabled", "true")

# Large fact table partitioned by date
fact_sales = spark.read.parquet("s3://bucket/sales/")
# Partition structure: sales/date=2024-01-01/, sales/date=2024-01-02/, ...

# Small dimension table — just dates we care about
dim_dates = spark.read.parquet("s3://bucket/dim_dates/") \
                 .filter(col("quarter") == "Q4")
# dim_dates now has ONLY Q4 dates

# JOIN: Spark first scans dim_dates, builds a list of Q4 dates,
# THEN only reads the matching partitions from fact_sales
result = fact_sales.join(dim_dates, on="date", how="inner")

result.explain("formatted")
# Look for: "DynamicPruning" in the plan
# This means DPP is working — fact_sales only reads Q4 partitions!

# Config to tune DPP
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true")
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.useStats", "true")

# ✅ GOOD — Region has 4-5 values
df.write.partitionBy("region").parquet("output/")
# Creates: region=US/, region=EU/, region=APAC/, region=LATAM/
# 4 partitions — lean and efficient

# ✅ GOOD — Product category has 10-20 values
df.write.partitionBy("category").parquet("output/")

# ❌ BAD — user_id has millions of distinct values
df.write.partitionBy("user_id").parquet("output/")
# Creates MILLIONS of tiny folders — kills metadata operations!

# Rule: target partition size 128MB - 1GB each
# If a partition is < 10MB, you have too many distinct values

# Hybrid: Date + Region — enables 4 query patterns with pruning
df_ymd.write \
    .partitionBy("year", "month", "region") \
    .parquet("s3://bucket/sales/")

# Directory structure:
# year=2024/month=1/region=US/  ← 4 files
# year=2024/month=1/region=EU/  ← 3 files
# year=2024/month=2/region=US/  ← 4 files
# ...

# Query pattern 1: All of 2024 → prunes to year=2024/ only
df.filter(col("year") == 2024)

# Query pattern 2: Jan 2024 → prunes year + month
df.filter((col("year") == 2024) & (col("month") == 1))

# Query pattern 3: US in Jan 2024 → prunes all three levels
df.filter((col("year") == 2024) &
          (col("month") == 1) &
          (col("region") == "US"))

# Query pattern 4: All US data → Spark still reads all year/month
# but filters region=US at each year/month level
df.filter(col("region") == "US")
# Note: Region without year/month filter = less pruning efficiency

# Data: revenue column is skewed
# Most orders are $1-$100, but a few are $1M+ (outliers)

# Hash partitioning on revenue will cause skew
df.repartition(10, "revenue")  # ❌ Most rows in a few partitions

# Range partitioning samples the data and creates even-sized buckets
df.repartitionByRange(10, "revenue")  # ✅ Each partition has similar row count

# Spark internally does quantile estimation:
# partition 0: revenue $0    - $45
# partition 1: revenue $45   - $89
# ...each with ~10% of total rows

# Verify balance
df.repartitionByRange(10, "revenue") \
  .withColumn("pid", spark_partition_id()) \
  .groupBy("pid").count().orderBy("pid").show()
# Should show roughly equal counts per partition

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

# Tuning: a partition is "skewed" if it's larger than:
# max(256MB, median_size * skewedPartitionFactor)
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256MB")

# With AQE: Spark automatically splits skewed partitions during joins
# No code changes needed! Just enable the config above.
result = df_large.join(df_small, on="country")
# AQE detects "US" partition is 5x the median → splits it → handles in parallel

# Also: use SKEW hint for explicit control
result = df_large.hint("SKEW", "country") \
                 .join(df_small, on="country")

from pyspark import SparkContext
from pyspark.rdd import RDD

# Custom partitioner: partition customers by first letter of name
# A-H → partition 0, I-P → partition 1, Q-Z → partition 2

class AlphaPartitioner:
    def __init__(self, num_partitions):
        self.num_partitions = num_partitions

    def __call__(self, key):
        # key is the name (string)
        first_letter = key[0].upper()
        if first_letter <= 'H':
            return 0
        elif first_letter <= 'P':
            return 1
        else:
            return 2

sc = spark.sparkContext
data = [("Alice", 100), ("Bob", 200), ("Zara", 300),
        ("Mike", 150), ("Kate", 250), ("Yuki", 175)]

# Create pair RDD (key, value)
rdd = sc.parallelize(data)  # RDD of (name, amount) tuples

# Apply custom partitioner via partitionBy
partitioner = AlphaPartitioner(3)
rdd_partitioned = rdd.partitionBy(3, partitioner)

# Inspect partitions
for i, partition in enumerate(rdd_partitioned.glom().collect()):
    print(f"Partition {i}: {partition}")
# Partition 0: [('Alice', 100), ('Bob', 200)]       ← A-H
# Partition 1: [('Mike', 150), ('Kate', 250)]        ← I-P
# Partition 2: [('Zara', 300), ('Yuki', 175)]        ← Q-Z

from pyspark.sql import functions as F

# Real use case: Route transactions to partitions by transaction_type
# We have: PURCHASE, REFUND, TRANSFER, DEPOSIT, WITHDRAWAL
# Business rule: each type gets its own partition for downstream processing

def business_partitioner(transaction_type):
    mapping = {
        "PURCHASE": 0,
        "REFUND": 1,
        "TRANSFER": 2,
        "DEPOSIT": 3,
        "WITHDRAWAL": 4
    }
    return mapping.get(transaction_type, 0)  # Default to 0 for unknown

sc = spark.sparkContext
NUM_PARTITIONS = 5

# Convert DataFrame to RDD pair (key, whole_row)
rdd = df_transactions.rdd.map(lambda row: (row.transaction_type, row))

# Partition by business rule
rdd_partitioned = rdd.partitionBy(NUM_PARTITIONS, business_partitioner)

# Convert back to DataFrame
df_partitioned = rdd_partitioned.map(lambda kv: kv[1]).toDF(df_transactions.schema)

# Verify: check which rows are in which partition
df_partitioned.withColumn("pid", spark_partition_id()) \
              .groupBy("transaction_type", "pid") \
              .count() \
              .orderBy("pid") \
              .show()
# All PURCHASE → partition 0
# All REFUND   → partition 1
# etc.