Python Pyspark-带过滤器的groupby-优化速度

我需要使用Pyspark处理数十亿行

Dataframe如下所示：

category    value    flag
   A          10       1
   A          12       0
   B          15       0
and so on...

我需要运行两个groupby操作：一个对flag==1的行，另一个对所有行。目前我正在这样做：

frame_1 = df.filter(df.flag==1).groupBy('category').agg(F.sum('value').alias('foo1'))
frame_2 = df.groupBy('category').agg(F.sum('value').alias(foo2))
final_frame = frame1.join(frame2,on='category',how='left')

到目前为止，这段代码正在运行，但我的问题是它非常慢。

---

有没有办法提高这段代码的速度，或者这是一个限制，因为我知道PySpark的惰性计算确实需要一些时间，但这段代码是最好的方法吗？

请注意，连接操作的成本很高。 您只需执行此操作并将标志添加到您的组：

frame_1 = df.groupBy(["category", "flag"]).agg(F.sum('value').alias('foo1'))

如果您有两个以上的标志，并且希望执行

flag==1 vs其余的

，则：

import pyspark.sql.functions as F
frame_1 = df.withColumn("flag2", F.when(F.col("flag") == 1, 1).otherwise(0))
frame_1 = df.groupBy(["category", "flag2"]).agg(F.sum('value').alias('foo1'))

如果要对所有行执行groupby应用，只需创建一个新框架，在其中对类别执行另一个汇总：

frame_1=df.groupBy（“category”）.agg（F.sum（'foo1'）。alias（'foo2'））

---

不可能在一个步骤中同时完成这两项工作，因为基本上存在组重叠。

IIUC，您可以避免昂贵的连接，并使用一个

groupBy

来实现这一点

final_frame_2 = df.groupBy("category").agg(
    F.sum(F.col("value")*F.col("flag")).alias("foo1"),
    F.sum(F.col("value")).alias("foo2"),
)
final_frame_2.show()
#+--------+----+----+
#|category|foo1|foo2|
#+--------+----+----+
#|       B| 0.0|15.0|
#|       A|10.0|22.0|
#+--------+----+----+

现在比较执行计划：

首先，你的方法是：

final_frame.explain()
#== Physical Plan ==
#*(5) Project [category#0, foo1#68, foo2#75]
#+- SortMergeJoin [category#0], [category#78], LeftOuter
#   :- *(2) Sort [category#0 ASC NULLS FIRST], false, 0
#   :  +- *(2) HashAggregate(keys=[category#0], functions=[sum(cast(value#1 as double))])
#   :     +- Exchange hashpartitioning(category#0, 200)
#   :        +- *(1) HashAggregate(keys=[category#0], functions=[partial_sum(cast(value#1 as double))])
#   :           +- *(1) Project [category#0, value#1]
#   :              +- *(1) Filter (isnotnull(flag#2) && (cast(flag#2 as int) = 1))
#   :                 +- Scan ExistingRDD[category#0,value#1,flag#2]
#   +- *(4) Sort [category#78 ASC NULLS FIRST], false, 0
#      +- *(4) HashAggregate(keys=[category#78], functions=[sum(cast(value#79 as double))])
#         +- Exchange hashpartitioning(category#78, 200)
#            +- *(3) HashAggregate(keys=[category#78], functions=[partial_sum(cast(value#79 as double))])
#               +- *(3) Project [category#78, value#79]
#                  +- Scan ExistingRDD[category#78,value#79,flag#80]

现在，对于最后一帧2，也一样：

final_frame_2.explain()
#== Physical Plan ==
#*(2) HashAggregate(keys=[category#0], functions=[sum((cast(value#1 as double) * cast(flag#2 as double))), sum(cast(value#1 as double))])
#+- Exchange hashpartitioning(category#0, 200)
#   +- *(1) HashAggregate(keys=[category#0], functions=[partial_sum((cast(value#1 as double) * cast(flag#2 as double))), partial_sum(cast(value#1 as double))])
#      +- Scan ExistingRDD[category#0,value#1,flag#2]

注意：严格来说，这与您给出的示例（如下所示）的输出不完全相同，因为您的内部联接将消除没有

标志=1的行的所有类别

final_frame.show()
#+--------+----+----+
#|category|foo1|foo2|
#+--------+----+----+
#|       A|10.0|22.0|
#+--------+----+----+

您可以向sum
标志添加一个聚合，并过滤那些求和为零的值（如果这是一个要求），而对性能的影响很小

final_frame_3 = df.groupBy("category").agg(
    F.sum(F.col("value")*F.col("flag")).alias("foo1"),
    F.sum(F.col("value")).alias("foo2"),
    F.sum(F.col("flag")).alias("foo3")
).where(F.col("foo3")!=0).drop("foo3")

final_frame_3.show()
#+--------+----+----+
#|category|foo1|foo2|
#+--------+----+----+
#|       A|10.0|22.0|
#+--------+----+----+

你真是个天才！对于sum等聚合方法，您建议的方法是这样的。对于countDistinct（）等方法，有什么方法可以实现这一点吗@pault@pault，类似的，并尝试使用您的方法，在这里。