根据另一列中的值，将PySpark Dataframe列转换为Python列表

Question

我有一个数据帧“dfClean”，有两列：

+---+-----+
|som| ano |
+---+-----+
| 1 |  1  |
| 2 |  0  |
| 3 |  1  |
| 4 |  1  |

我需要创建一个Python列表，其中“som”中的值在同一行的“ano”列中有1。 因此，预期产出为： pyLst = [1,3,4]

在熊猫中，我使用了： pyLst = dfClean.som[dfClean.ano == 1].tolist()

如何在PySpark或Scala中实现这一点，以及需要导入哪些其他库

Answer 1

您可以执行groupby并使用spark的collect_list函数

import pyspark.sql.functions as F   
df = spark.createDataFrame([(1, 1), (2, 0), (3, 1), (4, 1), ], ['som', 'ano'])
pyLst = df.groupby('ano').agg(F.collect_list('som').alias('pyLst')).where('ano = 1').collect()[0]['pyLst']

Answer 2

您可以使用过滤器，然后将所需列仅收集为rdd，然后再转换为列表。您可以使用列表理解来完成此操作

看起来是这样的,

lst = [row[0] for row in sdf.filter("ano = 1").select('som').collect()]

Answer 3

您可以先进行筛选，然后选择som列，将行收集到Python列表中，然后从每一行获取元素：

pyLst = [r[0] for r in df.filter(df.ano == 1).select('som').collect()]

print(pyLst)
[1, 3, 4]

另一种方法是首先转换为pandas数据帧，然后调用.tolist()，就像您在问题中所做的那样：

pyLst = df.filter(df.ano == 1).select('som').toPandas()['som'].tolist()

您还可以使用RDD API：

pyLst = df.filter(df.ano == 1).select('som').rdd.map(lambda r: r[0]).collect()