pyspark 按 条件 计算 行数

72qzrwbm  于 2022-11-21  发布在  Spark
关注(0)|答案(4)|浏览(152)

我有一个 Dataframe

test = spark.createDataFrame([('bn', 12452, 221), ('mb', 14521, 330),('bn',2,220),('mb',14520,331)],['x','y','z'])
test.show()

我需要根据一个条件计算行数:

test.groupBy("x").agg(count(col("y")>12453),count(col("z")>230)).show()

其给出了

+---+------------------+----------------+
 |  x|count((y > 12453))|count((z > 230))|
 +---+------------------+----------------+
 | bn|                 2|               2|
 | mb|                 2|               2|
 +---+------------------+----------------+

它只是行数而不是特定条件下的行数。

piok6c0g

piok6c0g1#

count不会对 Trues求和,它只计算非空值的数量。要计算 True 值,需要将条件转换为1 / 0,然后sum

import pyspark.sql.functions as F

cnt_cond = lambda cond: F.sum(F.when(cond, 1).otherwise(0))
test.groupBy('x').agg(
    cnt_cond(F.col('y') > 12453).alias('y_cnt'), 
    cnt_cond(F.col('z') > 230).alias('z_cnt')
).show()
+---+-----+-----+
|  x|y_cnt|z_cnt|
+---+-----+-----+
| bn|    0|    0|
| mb|    2|    2|
+---+-----+-----+
92dk7w1h

92dk7w1h2#

根据@Psidom的答案,我的答案如下

from pyspark.sql.functions import col,when,count

test.groupBy("x").agg(
    count(when(col("y") > 12453, True)),
    count(when(col("z") > 230, True))
).show()
0mkxixxg

0mkxixxg3#

由于Spark 3.0.0存在count_if(exp),请参见Spark函数文档

ct3nt3jp

ct3nt3jp4#

count函数跳过null值,因此您可以尝试执行以下操作:

import pyspark.sql.functions as F

def count_with_condition(cond):
    return F.count(F.when(cond, True))

且还在此存储库中起作用:戈朗

相关问题