Question

我有一个 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

我想找到关于列'class'的价格列的总和，这在一定程度上可以通过应用groupby来实现

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

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

现在我想根据获得的总和（价格）拆分数据。如果 sum(price) 相对于 'class' 大于 '0' 那么这个数据应该进入一个数据帧。（在这种情况下是 B，D 类）

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

如果 sum(price) 相对于 'class' 小于 '0' 那么这个数据应该进入一个数据帧。（在这种情况下是 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)

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

Answer 1

您可以在窗口上使用总和，并使用两个 filter 将数据帧拆分为两个。您可能需要处理 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')

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

1 个答案: