Как заставить определенное разбиение в PySpark DataFrame?

Question

Предположим, у меня есть DataFrame со столбцом partition_id:

n_partitions = 2

df = spark.sparkContext.parallelize([
    [1, 'A'],
    [1, 'B'],
    [2, 'A'],
    [2, 'C']
]).toDF(('partition_id', 'val'))

Как я могу перераспределить DataFrame, чтобы гарантировать, что каждое значение partition_id направляется в отдельный раздел, и что точностолько фактических разделов, сколько существует различных значений partition_id?

Если я сделаю хэш-раздел, то есть df.repartition(n_partitions, 'partition_id'), это гарантирует правильное количество разделов, но некоторые разделы могут быть пустыми, а другие могут содержатьмножественные значения partition_id из-за коллизий хешей.

Answer 1

Нет такой опции с Python и DataFrame API.API секционирования в Dataset не является подключаемым и поддерживает только предопределенные диапазоны и схемы хэширования .

. Вы можете конвертировать данные в RDD, секционировать с помощью пользовательского секционера и читать преобразование обратно вDataFrame:

from pyspark.sql.functions import col, struct, spark_partition_id

mapping = {k: i for i, k in enumerate(
    df.select("partition_id").distinct().rdd.flatMap(lambda x: x).collect()
)}

result = (df
    .select("partition_id", struct([c for c in df.columns]))
    .rdd.partitionBy(len(mapping), lambda k: mapping[k])
    .values()
    .toDF(df.schema))

result.withColumn("actual_partition_id", spark_partition_id()).show()
# +------------+---+-------------------+
# |partition_id|val|actual_partition_id|
# +------------+---+-------------------+
# |           1|  A|                  0|
# |           1|  B|                  0|
# |           2|  A|                  1|
# |           2|  C|                  1|
# +------------+---+-------------------+

Пожалуйста, помните, что это только создает определенное распределение данных и не устанавливает разделитель, который может использоваться оптимизатором Catalyst.