# Avoid shuffle by broadcasting small table
from pyspark.sql.functions import broadcast
result = large_df.join(broadcast(small_df), "key")

# Pre-partition data to avoid shuffle
df = df.repartition(100, "key")
df.write.parquet("output/")

# Bucket data by join key
df.write.bucketBy(100, "user_id").saveAsTable("users")

# Coalesce: Reduce partitions without full shuffle (narrow)
df.coalesce(10)

# Repartition: Increase or decrease with full shuffle (wide)
df.repartition(100)

# BAD: groupByKey + mapValues (two shuffles)
rdd.groupByKey().mapValues(sum)

# GOOD: reduceByKey (single shuffle)
rdd.reduceByKey(lambda a, b: a + b)

# Join always requires shuffle (unless broadcast)
joined = rdd1.join(rdd2)

# Requires shuffle to deduplicate across partitions
distinct = rdd.distinct()

from pyspark import SparkContext

sc = SparkContext("local", "TransformationTypes")

# Create RDD with 4 partitions
# parallelize: Distributes data across partitions
rdd = sc.parallelize(range(1, 101), 4)

# Narrow transformations (no shuffle)
# map: Apply function to each element
mapped = rdd.map(lambda x: x * 2)  # Double each element

# filter: Keep elements where function returns True
filtered = rdd.filter(lambda x: x % 2 == 0)  # Keep even numbers

# flatMap: Apply function that returns iterator, flatten results
flatMapped = rdd.flatMap(lambda x: [x, x * 10])  # Expand each element

# sample: Random sample
# Parameters: withReplacement (False=no replacement), fraction (0.1=10%)
sampled = rdd.sample(False, 0.1)  # 10% sample

# mapPartitions: Apply function to each partition (more efficient for I/O)
# Parameters: func - function that takes iterator, returns iterator
mappedPartitions = rdd.mapPartitions(lambda it: [sum(it)])  # Sum per partition

# mapPartitionsWithIndex: Like mapPartitions with partition index
# Parameters: func - function taking (index, iterator)
indexed = rdd.mapPartitionsWithIndex(lambda idx, it: [(idx, sum(it))])

# Execute actions
print(f"Mapped sum: {mapped.sum()}")
print(f"Filtered count: {filtered.count()}")
print(f"FlatMapped count: {flatMapped.count()}")
print(f"Sampled count: {sampled.count()}")
print(f"Per-partition sums: {mappedPartitions.collect()}")
print(f"Indexed sums: {indexed.collect()}")

# No shuffle means these can be pipelined
# Single stage execution

# Wide transformations (require shuffle)
words = sc.parallelize(["hello world", "hello spark", "world spark"], 2)

# Split words (narrow) then create key-value pairs (narrow)
word_pairs = words.flatMap(lambda line: line.split()) \
                  .map(lambda word: (word, 1))

# reduceByKey (wide) - single shuffle
# Parameters: func - binary function to combine values
# Performs map-side combine, reducing shuffle data by 10-100x
word_counts = word_pairs.reduceByKey(lambda a, b: a + b)
print(f"Word counts: {word_counts.collect()}")

# groupByKey (wide) - less efficient (no map-side combine)
# Parameters: numPartitions (optional)
# Collects all values to memory, can cause OOM for large value sets
word_groups = word_pairs.groupByKey()
for word, counts in word_groups.collect():
    print(f"{word}: {list(counts)}")

# join (wide)
# Parameters: other - RDD to join with
rdd1 = sc.parallelize([(1, "a"), (2, "b"), (3, "c")])
rdd2 = sc.parallelize([(1, "x"), (2, "y"), (3, "z")])
joined = rdd1.join(rdd2)  # Shuffle required
print(f"Joined: {joined.collect()}")

# distinct (wide)
# Removes duplicates across all partitions (requires shuffle)
duplicated = sc.parallelize([1, 2, 2, 3, 3, 3])
distinct = duplicated.distinct()  # Shuffle required
print(f"Distinct: {distinct.collect()}")

# sortByKey (wide)
# Parameters: ascending (default True), numPartitions, keyfunc
sorted_rdd = word_counts.sortByKey(ascending=True)
print(f"Sorted: {sorted_rdd.collect()}")

# Create large dataset
large_rdd = sc.parallelize(range(1000000), 8)

# Create key-value pairs
word_pairs = large_rdd.map(lambda x: (x % 100, x))

# BAD: groupByKey + mapValues (two shuffles)
# groupByKey collects all values to memory, then mapValues processes them
grouped = word_pairs.groupByKey()
result_bad = grouped.mapValues(sum)

# GOOD: reduceByKey (single shuffle)
# Performs map-side combine, reducing shuffle data significantly
result_good = word_pairs.reduceByKey(lambda a, b: a + b)

# Compare execution plans
print("=== BAD (groupByKey) ===")
print(result_bad.toDebugString().decode()[:500])

print("\n=== GOOD (reduceByKey) ===")
print(result_good.toDebugString().decode()[:500])

# Performance comparison
import time
start = time.time()
_ = result_bad.collect()
print(f"groupByKey: {time.time() - start:.2f}s")

start = time.time()
_ = result_good.collect()
print(f"reduceByKey: {time.time() - start:.2f}s")

# BETTER: combineByKey (most flexible, best for complex aggregations)
# Parameters:
#   createCombiner: Initialize accumulator from first value
#   mergeValue: Merge value into accumulator within partition
#   mergeCombiners: Merge accumulators across partitions
combined = word_pairs.combineByKey(
    lambda v: (v, 1),  # Create: (value, count)
    lambda acc, v: (acc[0] + v, acc[1] + 1),  # Merge value
    lambda acc1, acc2: (acc1[0] + acc2[0], acc1[1] + acc2[1])  # Merge combiners
)
averages = combined.mapValues(lambda x: x[0] / x[1])
print(f"Averages: {averages.take(5)}")

# Repartition vs Coalesce
rdd = sc.parallelize(range(1000), 10)
print(f"Initial partitions: {rdd.getNumPartitions()}")

# Coalesce: Reduce partitions without full shuffle (narrow transformation)
# Parameters: numPartitions, shuffle=False (default)
# Only moves data from partitions being removed
coalesced = rdd.coalesce(5)
print(f"After coalesce(5): {coalesced.getNumPartitions()}")

# Repartition: Increase or decrease with full shuffle (wide transformation)
# Parameters: numPartitions
# Redistributes all data evenly
repartitioned = rdd.repartition(20)
print(f"After repartition(20): {repartitioned.getNumPartitions()}")

# Partition by key (for key-value RDDs)
# Parameters: numPartitions, partitionFunc (default: hash)
key_value_rdd = sc.parallelize([(i % 10, i) for i in range(100)])
partitioned = key_value_rdd.partitionBy(5, lambda k: k % 5)

# Check partition distribution
def count_per_partition(rdd):
    return rdd.mapPartitionsWithIndex(
        lambda idx, it: [(idx, sum(1 for _ in it))]
    ).collect()

print(f"Partition counts: {count_per_partition(partitioned)}")

# Custom partitioner for skewed data
from pyspark import Partitioner

class RangePartitioner(Partitioner):
    def __init__(self, boundaries, numPartitions):
        self.boundaries = boundaries
        self.numPartitions = numPartitions
    
    def getPartition(self, key):
        for i, boundary in enumerate(self.boundaries):
            if key < boundary:
                return i
        return self.numPartitions - 1
    
    def numPartitions(self):
        return self.numPartitions

# Use custom partitioner
boundaries = [10, 20, 30, 40]
custom_partitioned = key_value_rdd.partitionBy(5, RangePartitioner(boundaries, 5))
print(f"Custom partition counts: {count_per_partition(custom_partitioned)}")

# aggregateByKey: Most flexible key-based aggregation
# Parameters: zeroValue, seqFunc, combFunc
from operator import add

# Zero value: Initial accumulator (sum=0, count=0)
# seqFunc: Add value to accumulator within partition
# combFunc: Merge accumulators across partitions
result = word_pairs.aggregateByKey(
    (0, 0),  # Zero value
    lambda acc, v: (acc[0] + v, acc[1] + 1),  # Within partition
    lambda acc1, acc2: (acc1[0] + acc2[0], acc1[1] + acc2[1])  # Across partitions
).mapValues(lambda x: x[0] / x[1])  # Calculate average

print(f"Key averages: {result.take(5)}")

# cartesian: All pairs between two RDDs (very expensive!)
rdd_a = sc.parallelize([1, 2, 3])
rdd_b = sc.parallelize(["a", "b"])
cartesian = rdd_a.cartesian(rdd_b)
print(f"Cartesian product: {cartesian.collect()}")
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b'), (3, 'a'), (3, 'b')]

# intersection: Common elements between two RDDs (requires shuffle)
rdd_x = sc.parallelize([1, 2, 3, 4, 5])
rdd_y = sc.parallelize([3, 4, 5, 6, 7])
intersection = rdd_x.intersection(rdd_y)
print(f"Intersection: {intersection.collect()}")  # [3, 4, 5]

# subtract: Elements in first RDD but not in second (requires shuffle)
subtracted = rdd_x.subtract(rdd_y)
print(f"Subtract: {subtracted.collect()}")  # [1, 2]

# groupWith: Group two RDDs by key (requires shuffle)
grouped = rdd1.groupWith(rdd2)
for key, (values1, values2) in grouped.collect():
    print(f"Key {key}: {list(values1)}, {list(values2)}")

# BAD: Multiple wide transformations (multiple shuffles)
result = rdd.groupByKey().mapValues(sum).sortByKey()

# GOOD: Combine operations (single shuffle)
result = rdd.reduceByKey(lambda a, b: a + b).sortByKey()

# BETTER: Use aggregateByKey for complex aggregations
result = rdd.aggregateByKey(
    (0, 0),
    lambda acc, v: (acc[0] + v, acc[1] + 1),
    lambda acc1, acc2: (acc1[0] + acc2[0], acc1[1] + acc2[1])
).mapValues(lambda x: x[0] / x[1])

# BAD: groupByKey collects all values to memory (can OOM)
grouped = rdd.groupByKey().mapValues(sum)

# GOOD: reduceByKey combines per partition first (10-100x less shuffle)
reduced = rdd.reduceByKey(lambda a, b: a + b)

# GOOD: combineByKey for complex aggregations (most flexible)
combined = rdd.combineByKey(
    lambda v: (v, 1),
    lambda acc, v: (acc[0] + v, acc[1] + 1),
    lambda acc1, acc2: (acc1[0] + acc2[0], acc1[1] + acc2[1])
)

from pyspark.sql.functions import broadcast
# Avoids shuffle for join with small table
result = large_df.join(broadcast(small_df), "key")

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

# Increase partitions for more parallelism
rdd = rdd.repartition(100)

# Decrease partitions with coalesce (no shuffle)
rdd = rdd.coalesce(10)

# Partition by key for future joins
key_value_rdd = key_value_rdd.partitionBy(100, hash)

# Cache if RDD is reused across multiple actions
rdd.cache()  # MEMORY_ONLY
# or
rdd.persist(StorageLevel.MEMORY_AND_DISK)  # With disk spill

# Uncache when no longer needed
rdd.unpersist()

# Check shuffle metrics in Spark UI
# Look for:
# - Shuffle Read Size / Records
# - Shuffle Write Size / Records
# - Shuffle Spill (Memory/Disk)
# - Task duration distribution (identify stragglers)

# Get RDD statistics
print(f"Partition count: {rdd.getNumPartitions()}")
print(f"Storage level: {rdd.getStorageLevel()}")
print(f"Is cached: {rdd.is_cached}")

# Check lineage depth
lineage = rdd.toDebugString().decode()
print(f"Lineage depth: {len(lineage.split(chr(10)))}")

# SALTING: Add random prefix to keys to distribute data evenly
import random

# Salt the keys
salted = rdd.map(lambda x: (f"{x[0]}_{random.randint(0, 9)}", x[1]))

# Reduce with salted keys
salted_reduced = salted.reduceByKey(lambda a, b: a + b)

# Remove salt and reduce again
result = salted_reduced.map(lambda x: (x[0].split("_")[0], x[1])) \
                       .reduceByKey(lambda a, b: a + b)

# Alternative: Use AQE skew join handling (Spark 3.0+)
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")