PySpark 根据条件计算行数

Question

我有一个数据框

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

我需要根据条件计算行数：

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|
+---+------------------+----------------+

这只是行数，而不是某些条件下的计数。

Answer 1

count

不会对 True 求和，它只计算非空值的数量。要计算 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|
+---+-----+-----+

Answer 2

根据@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()

Answer 3

从 Spark 3.0.0 开始有

count_if(exp)

，参见 Spark 功能文档

Answer 4

count 函数跳过

null

值，因此您可以尝试以下操作：

import pyspark.sql.functions as F

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

并且也在这个仓库中起作用：kolang

Answer 5

Spark 3.5+ 在 Python API 中具有

count_if

：

from pyspark.sql import functions as F

test.groupBy('x').agg(
    F.count_if(F.col('y') > 12453).alias('y_cnt'),
    F.count_if(F.col('z') > 230).alias('z_cnt')
).show()
# +---+-----+-----+
# |  x|y_cnt|z_cnt|
# +---+-----+-----+
# | bn|    0|    0|
# | mb|    2|    2|
# +---+-----+-----+

Spark 3.0+也有，但是必须使用

expr

：

test.groupBy('x').agg(
    F.expr("count_if(y > 12453) y_cnt"),
    F.expr("count_if(z > 230) z_cnt")
).show()
# +---+-----+-----+
# |  x|y_cnt|z_cnt|
# +---+-----+-----+
# | bn|    0|    0|
# | mb|    2|    2|
# +---+-----+-----+