在 RDD 转换上保留 Spark DataFrame 列分区

Posted 2023-04-17

【中文标题】在 RDD 转换上保留 Spark DataFrame 列分区

【英文标题】：Preserving Spark DataFrame column partitions on RDD conversion

【发布时间】：2016-11-17 23:22:34

【问题描述】：

给定一个看起来像这样的 Spark DataFrame：

==================================
| Name | Col1 | Col2 | .. | ColN |
----------------------------------
|    A |    1 |   11 | .. |   21 |
|    A |   31 |   41 | .. |   51 |
|    B |    2 |   12 | .. |   22 |
|    B |   32 |   42 | .. |   52 |
==================================

我想运行逻辑，该逻辑对对应于特定Name 值的表分区执行聚合/计算。所述逻辑要求分区的全部内容——并且仅该分区——在执行该逻辑的节点上的内存中实现；它看起来像下面的processSegment 函数：

def processDataMatrix(dataMatrix):
    # do some number crunching on a 2-D matrix

def processSegment(dataIter):
    # "running" value of the Name column in the iterator
    dataName = None
    # as the iterator is processed, put the data in a matrix
    dataMatrix = []

    for dataTuple in dataIter:
        # separate the name column from the other columns
        (name, *values) = dataTuple
        # SANITY CHECK: ensure that all rows have same name
        if (dataName is None):
            dataName = name
        else:
            assert (dataName == name), 'row name ' + str(name) + ' does not match expected ' + str(dataName)

        # put the row in the matrix
        dataMatrix.append(values)

    # if any rows were processed, number-crunch the matrix
    if (dataName is not None):
        return processDataMatrix(dataMatrix)
    else:
        return []

我试图通过基于Name 列重新分区，然后通过底层RDD 上的mapPartitions 在每个分区上运行processSegment 来完成这项工作：

result = \
    stacksDF \
        .repartition('Name') \
        .rdd \
        .mapPartitions(processSegment) \
        .collect()

但是，该过程通常会使processSegment 中的SANITY CHECK 断言失败：

AssertionError: row name Q7 does not match expected A9

当我尝试在底层 RDD 上运行 mapPartitions 时，为什么表面上在 DataFrame 上执行的分区没有被保留？如果上述方法无效，是否有某种方法（使用 DataFrame API 或 RDD API）可以让我对 DataFrame 分区的内存再现执行聚合逻辑？

（由于我使用的是 PySpark，而我希望执行的特定数字运算逻辑是 Python，用户定义的聚合函数 (UDAF) would not appear to be an option。）

【参考方案1】：

我相信您误解了分区的工作原理。一般来说，partioner 是一个满射函数，而不是一个双射函数。虽然特定值的所有记录都将移动到单个分区，但分区可能包含具有多个不同值的记录。

DataFrame API 无法让您对分区程序进行任何控制，但可以在使用 RDD API 时定义自定义 partitionFunc。这意味着您可以使用双射的，例如：

mapping = (df
    .select("Name")
    .distinct()
    .rdd.flatMap(lambda x: x)
    .zipWithIndex()
    .collectAsMap())

def partitioner(x):
    return mapping[x]

并按如下方式使用：

df.rdd.map(lambda row: (row.Name, row)).partitionBy(len(mapping), partitioner)

尽管您可能必须记住分区不是空闲的，并且如果唯一值的数量很大，它可能会成为一个严重的性能问题。