Pyspark：将函数应用于多个数据帧的匹配分区

Question

我有两个Spark DataFrame，它们具有匹配的分区对。我想将每对分区发送到不同的执行器，并对它们执行Python函数。

如果我只有一个DataFrame，则可以使用mapPartitions，但是无法访问其他DataFrame。我尝试广播第二个，但是它太大了，无法分割成我的执行器。

我认为我可能需要使用pyspark.rdd.RDD中的函数，但是我不确定是哪个函数。像groupByKey和cogroup这样的几个选项似乎可能是相关的，但我不知道如何将它们应用于DataFrame对象。

预先感谢您的帮助！

这是我的输入数据（每种颜色都是一个分区）：

这就是我想要做的：

Answer 1

您可以在加入之前嵌套每个数据框。

这可能看起来像（下面在pyspark中）：

from pyspark.sql import functions as F

grouped1 = sdf1. \
    select(partition_by, F.struct(*[col for col in sdf1.columns if col != partition_by]).alias('_data')). \
    groupBy(partition_by). \
    agg(F.collect_list('_data').alias('data'))

grouped2 = sdf2. \
    select(partition_by, F.struct(*[col for col in sdf2.columns if col != partition_by]).alias('_data')). \
    groupBy(partition_by). \
    agg(F.collect_list('_data').alias('data'))

merged = grouped1.join(grouped2, on=partition_by)

然后，您可以应用merged.rdd.map(f)之类的函数。

如果有人知道这种方法与其他方法对性能的影响，我会很好奇。