pyspark按列将数据帧拆分为少量数据帧

Question

我正在使用Python（pyspark）

我有一个包含少数用户数据的数据框

IMSI   X        Y

1      139620   179284

1      139620   179284

3      138620   178284

3      138620   178284

3      138620   178284

我想创建按IMSI列分割的RDD-DF对象

这样我就可以为每个DF操作一个Python函数

SplitDF.map(PythonFunctionForEachUser)

Answer 1

您可以使用两种方法。如果您想使用RDD，可以使用groupByKey：

df.rdd.keyBy(lambda r: r.IMSI).groupByKey().map(lambda a: PythonFunctionForEachUser(a[1]))

如果您更喜欢Dataframe，可以按IMSI列分组，然后collect_list，但首先需要将这些x和y合并为一列：

>>> from pyspark.sql.functions import array
>>> df.withColumn('data', array('x', 'y')).groupBy('IMSI') \
         .agg(collect_list('data').alias('userdata')) \
         .rdd.map(PythonFunctionForEachUser)