根据条件将 pyspark 数据帧拆分为多个数据帧答案

【问题标题】：split pyspark dataframe into multiple dataframes based on a condition根据条件将 pyspark 数据帧拆分为多个数据帧
【发布时间】：2021-04-18 12:03:33
【问题描述】：

我有一个 pyspark 数据框，其中包含类似于以下的数据：

id  class price  place
1   A      10      US
2   B      5       US
3   B      5       MEXICO
4   A     -20      CANADA
5   C     -15      US
6   C     -5       US
7   D      20      MEXICO
8   A      10      CANADA
9   A     -30      CANADA

我想求 price 列相对于 'class' 列的总和，这在一定程度上可以通过应用 groupby 来实现

      df.groupby('class ').agg({'price': 'sum'}).show()

output: class   sum(price)
        A       -30
        B        10
        C       -20
        D        20

现在我想根据获得的总和（价格）拆分数据。如果相对于“类”的总和（价格）大于“零”，则该数据应放入一个数据帧。（在这种情况下为 B，D 类）

id  class price place
2   B     5      US
3   B     5      MEXICO
7   D     20     MEXICO

如果“类”的总和（价格）小于“零”，则该数据应放入一个数据帧中。（在这种情况下为 A、C 类）

id  class price place
1   A   10     US
8   A   10     CANADA
4   A   -20     CANADA
9   A   -30    CANADA
5   C   -15    US
6   C   -5     US

使用 pyspark 将数据进一步写入两个不同的 csv 文件。

df.write.format('csv').option('header', 'true').save(destination_location)

如何将 groupby 结果存储到数据框中？以及如何根据上述条件实现将单个数据帧拆分为两个不同的数据帧？

【问题讨论】：

标签： python dataframe apache-spark pyspark conditional-statements

【解决方案1】：

您可以在窗口上使用总和，并使用两个filters 将数据框分成两部分。您可能需要处理sum = 0 的情况。

from pyspark.sql import functions as F, Window

summed = df.withColumn('sum', F.sum('price').over(Window.partitionBy('class')))
df1 = summed.filter('sum > 0')
df2 = summed.filter('sum < 0')

【讨论】：