Как понять какую строку на какую ноду закидывать при shuffle?

Shuffle в Spark: как и почему данные распределяются по нодам

Что такое Shuffle

Shuffle — это процесс перераспределения данных между нодами кластера перед агрегацией или join-операциями. Во время shuffle данные с одной ноды отправляются на другие ноды в зависимости от ключа. Это дорогостоящая операция, которая определяет производительность Spark-приложения.

Как Spark понимает, куда отправить строку

Spark использует хеширование ключа для определения целевого partition-а:

Формула:

target_partition_id = hashCode(key) % number_of_partitions

Этот процесс в Spark реализован через ShufflePartitioner:

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("shuffle_example").getOrCreate()
df = spark.createDataFrame(
    [
        ("Alice", 1000),
        ("Bob", 2000),
        ("Alice", 1500),
        ("Charlie", 1200),
    ],
    ["name", "salary"]
)

# Группировка по имени — вызывает shuffle
grouped = df.groupBy("name").agg({"salary": "sum"})

# Все строки с name="Alice" попадут в один и тот же partition
# Все строки с name="Bob" попадут в другой partition
# И так далее

Пример: SHA-1 хеш для распределения

Python пример хеширования:

import hashlib

def get_partition_id(key, num_partitions):
    """Определяет partition для ключа"""
    # Вычисляем хеш ключа
    hash_value = hash(key)  # Встроенная функция hash() в Python
    
    # Вычисляем номер partition-а
    partition_id = hash_value % num_partitions
    return partition_id

names = ["Alice", "Bob", "Charlie", "Alice", "Bob"]
num_partitions = 4

for name in names:
    pid = get_partition_id(name, num_partitions)
    print(f"{name:10} → partition {pid}")

# Результат:
# Alice      → partition 0
# Bob        → partition 2
# Charlie    → partition 3
# Alice      → partition 0  (опять в том же partition)
# Bob        → partition 2  (опять в том же partition)

Partitioner в Spark

Spark использует разные Partitioner в зависимости от операции:

1. HashPartitioner (по умолчанию):

from pyspark import SparkContext

sc = SparkContext("local", "shuffle_example")
data = [("Alice", 1), ("Bob", 2), ("Alice", 3), ("Charlie", 4)]
rdd = sc.parallelize(data)

# Используем groupByKey с HashPartitioner
grouped = rdd.groupByKey(numPartitions=4)  # 4 partition-а

for partition_id, items in grouped.glom().zipWithIndex():
    print(f"Partition {partition_id}:")
    for key, values in items:
        print(f"  {key}: {list(values)}")

2. RangePartitioner (для сортировки):

from pyspark import RangePartitioner

data = [(1, "a"), (5, "b"), (3, "c"), (2, "d"), (4, "e")]
rdd = sc.parallelize(data)

# RangePartitioner распределяет по диапазонам
partitioner = RangePartitioner(2, rdd)  # 2 partition-а с диапазонами
sorted_rdd = rdd.partitionBy(partitioner)

# Partition 0: ключи 1-2
# Partition 1: ключи 3-5

3. CustomPartitioner (собственная логика):

from pyspark import Partitioner

class GeographicPartitioner(Partitioner):
    """Partitioner, который распределяет по географии"""
    
    def __init__(self, num_partitions):
        self._num_partitions = num_partitions
    
    def numPartitions(self):
        return self._num_partitions
    
    def getPartition(self, key):
        # Распределяем по первому символу города
        if key.startswith((A, B, C)):
            return 0
        elif key.startswith((D, E, F)):
            return 1
        else:
            return 2 % self._num_partitions

partitioner = GeographicPartitioner(3)
rdd_cities = sc.parallelize([
    ("Amsterdam", 1),
    ("Berlin", 2),
    ("Paris", 3),
    ("Denver", 4),
])

partitioned = rdd_cities.partitionBy(partitioner)

Shuffle Stage в DAG

Как Spark распределяет данные:

Stage 1: Map
├─ Partition 0: (Alice, 1000) → hash(Alice) = 42 → target partition (42 % 4) = 2
├─ Partition 1: (Bob, 2000) → hash(Bob) = 15 → target partition (15 % 4) = 3
├─ Partition 2: (Alice, 1500) → hash(Alice) = 42 → target partition (42 % 4) = 2
└─ Partition 3: (Charlie, 1200) → hash(Charlie) = 99 → target partition (99 % 4) = 3

Shuffle (по сети): данные транспортируются на целевые partition-ы

Stage 2: Reduce (после shuffle)
├─ Partition 0: (empty)
├─ Partition 1: (empty)
├─ Partition 2: [(Alice, 1000), (Alice, 1500)] → sum = 2500
└─ Partition 3: [(Bob, 2000), (Charlie, 1200)]

Практический пример: контроль распределения

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

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

df = spark.createDataFrame([
    ("Alice", 100, "NYC"),
    ("Bob", 200, "LA"),
    ("Alice", 150, "NYC"),
    ("Charlie", 120, "Chicago"),
    ("Bob", 250, "LA"),
], ["name", "amount", "city"])

# Shuffle по одному ключу (name)
result1 = df.groupBy("name").agg({"amount": "sum"})
# Все Alice → одна ноде, все Bob → одна ноде, и т.д.

# Shuffle по двум ключам (composite key)
result2 = df.groupBy("name", "city").agg({"amount": "sum"})
# (Alice, NYC), (Bob, LA), и т.д. распределяются отдельно

# Контроль количества partition-ов после shuffle
result3 = df.groupBy("name").agg({"amount": "sum"}).repartition(8)
# Явно устанавливаем 8 partition-ов

# Проверяем распределение
print(f"Число partition-ов: {result1.rdd.getNumPartitions()}")

Оптимизация Shuffle

1. Минимизируй количество shuffle-операций:

# ❌ Плохо: несколько shuffle-ов
df1 = df.groupBy("department").count()
df2 = df1.filter(col("count") > 10)
df3 = df2.groupBy("department").agg({"count": "avg"})

# ✅ Хорошо: один shuffle
result = df.groupBy("department").count().filter(col("count") > 10)

2. Установи оптимальное количество partition-ов:

# По умолчанию: spark.sql.shuffle.partitions = 200
spark.conf.set("spark.sql.shuffle.partitions", 200)

# Для малых данных (< 1GB)
spark.conf.set("spark.sql.shuffle.partitions", 50)

# Для больших данных (> 100GB)
spark.conf.set("spark.sql.shuffle.partitions", 1000)

3. Используй Skewed Join Optimization:

from pyspark.sql.functions import broadcast

# Если одна из таблиц маленькая
df_small = spark.read.csv("small_dim.csv", header=True)
df_large = spark.read.csv("large_fact.csv", header=True)

# Broadcast небольшую таблицу вместо shuffle
result = df_large.join(
    broadcast(df_small),  # Отправляем в память каждой ноды
    on="key"
)

Мониторинг Shuffle

В Spark UI:

1. Открыть http://localhost:4040 во время выполнения
2. Вкладка "Stages"
3. Посмотреть "Shuffle Read Size" и "Shuffle Write Size"
4. Большие значения = медленный shuffle = нужна оптимизация

Выводы

Shuffle распределение основано на хешировании:

Процесс:
1. Вычисляется hashCode(key)
2. Берётся остаток от деления на количество partition-ов
3. Результат = номер целевого partition-а
4. Данные отправляются на соответствующую ноду

Все строки с одним и тем же ключом всегда попадут в один partition (если не менялось число partition-ов), что позволяет корректно агрегировать и объединять данные. Оптимизация shuffle критична для производительности Spark-приложений.