如何在pyspark中按列对数据帧进行分区以进行进一步处理？

问题描述

我需要按列对数据框进行分区。我知道可以保存在单独的文件中。但是我需要分区做进一步的处理（我需要按照一定的顺序对分区进行排序，并对有序的分区应用udf）。

我的代码是：

df = spark.createDataFrame([(2,),(1,(2,)],("name",)) \
    .repartitionByRange(2,"name") \
    .rdd.glom().collect()
print(df)

# [[Row(name=2),Row(name=1),Row(name=2),Row(name=2)],[]]

我需要得到这样的东西：

[[(2,[(1,)]]

解决方法

您可以使用 repartition 代替 repartitionByRange：

df = spark.createDataFrame([(2,),(1,(2,)],("name",)) \
    .repartition(2,"name") \
    .rdd.glom().collect()

print(df)
# [[Row(name=2),Row(name=2),Row(name=2)],[Row(name=1),Row(name=1)]]

repartitionByRange 使用抽样来估计范围，并且可能会导致您观察到的错误。