将 PySpark Dataframe 列转换为 Python 列表，具体取决于另一列中的值

Question

我有一个 dataframe "dfClean" 有 2 列：

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

我需要创建一个 Python 列表，其中“som”中的值在同一行的“ano”列中为 1。 所以预期 output 是： pyLst = [1,3,4]

在 Pandas 我用过： pyLst = dfClean.som[dfClean.ano == 1].tolist()

如何在 PySpark 或 Scala 中执行此操作，我需要导入哪些其他库？

Answer 1

您可以先进行过滤，然后 select som列，将行收集到 Python 列表中，并从每一行中获取元素：

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

print(pyLst)
[1, 3, 4]

另一种方法是先转换为 pandas dataframe ，然后像您在问题中所做的那样调用.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()

Answer 2

您可以使用过滤器，然后仅将所需的列收集为 rdd，然后转换为列表。 您可以使用列表推导来执行此操作。

它看起来像这样，

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

Answer 3

你可以做一个groupby并使用spark的collect_list function

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']