Python 如何使用pyspark有条件地求和几列？

我试图找出一种方法来求多个列的和，但每个和的条件不同

这是我在数据帧中拥有的数据：

order_id    article_id  article_name         nr_of_items price       is_black is_fabric
----------- ----------- -------------------- ----------- ----------- -------- ---------
1           567         batteries            6           5           0        0
1           645         pants                1           20          1        1
2           876         tent                 1           40          0        1
2           434         socks                10          5           1        1

这就是我想要的：

order_id    total_order_amount black_order_amount fabric_order_amount
----------- ------------------ ------------------ -------------------
1           50                 20                 20
2           90                 50                 90

这是如何在SQL中实现的：

选择
订单号：，
总订单金额（项目数量*价格），
作为黑色订单金额的总和（如果黑色=1，则价格*其他项目的数量为0），
总和（如果面料=1，则价格*n个其他物品0结束）作为面料订单金额
从订单行
按订单分组\u id
;

如何使用pyspark进行相同的操作？也就是说，我想知道的是如何在不同的条件下聚合多个列

我准备了一个pyspark数据帧，以防有人想尝试：

from pyspark.sql.types import *

cSchema = StructType([StructField("order_id", IntegerType())\
                      ,StructField("article_id", IntegerType())\
                      ,StructField("article_name", StringType())\
                      ,StructField("nr_of_items", IntegerType())\
                      ,StructField("price", IntegerType())\
                      ,StructField("is_black", BooleanType())\
                      ,StructField("is_fabric", BooleanType())])

test_list = [[1, 567, 'batteries', 6, 5, False, False],
             [1, 645, 'pants', 1, 20, True, True],
             [2, 876, 'tent', 1, 40, False, True],
             [2, 434, 'socks', 10, 5, True, True]]

df = spark.createDataFrame(test_list,schema=cSchema) 

我使用的是spark版本2.4.4和python版本3.7.3。

Johanrex

下面是一段代码：

from pyspark.sql.functions import *

df.groupBy("order_id").agg(
    sum(col("nr_of_items")*col("price")).alias("total_order_amount"),
    sum(when(col("is_black") == lit(1), col("price")*col("nr_of_items")).otherwise(lit(0))).alias("black_order_amount"),
    sum(when(col("is_fabric") == lit(1), col("price")*col("nr_of_items")).otherwise(lit(0))).alias("fabric_order_amount")
).limit(100).toPandas()

输出：

order_id    total_order_amount  black_order_amount  fabric_order_amount
1               50                  20                  20
2               90                  50                  90

---

您可以使用sum和when内置函数来聚合数据。在scala中，这看起来是这样的：df.groupBy（“key”）.agg（sum（当（$”条件“，$”val”）。否则（lit（0）））