使用 groupby 在另一列上使用条件聚合行上的列

Question

假设我有以下 Pyspark 数据框：

 Country    Direction    Quantity     Price
 Belgium    In           5            10
 Belgium    Out          2            8
 Belgium    Out          3            9
 France     In           2            3
 France     Out          3            2
 France     Out          4            3
 

是否可以按“国家/地区”列对该数据框进行分组，照常聚合“价格”列的平均值，但对“数量”列使用函数“第一”，仅适用于“方向”列为“出去”？ 我想应该是这样的：

df.groupby("Country").agg(F.mean('Price'), F.first(F.col('Quantity').filter(F.col('Direction') == "Out")))

Answer 1

您可以为 Direction != 'out' 屏蔽 Quantity 并使用 ignoreNulls:

做一个 first

df.groupby("Country").agg(
    F.mean('Price'),
    F.first(
        F.when(
            F.col('Direction') == "Out",
            F.col('Quantity')
        ),
        ignoreNulls=True
    )
)