在Databricks中统计PySpark DataFrame每行True值的数量

问题背景 #

拥有如下布尔类型的PySpark DataFrame：

columns = ["a", "b", "c"]
data = [(True, True, True), (True, True, True), (True, False, True)]
df = spark.createDataFrame(data).toDF(*columns)
df.display()

想要新增一列number_true_values统计每行中True的数量，但直接使用sum对布尔列求和会抛出类型不匹配异常：

import pyspark.sql.functions as F
df.withColumn('number_true_values', sum([F.col(column) for column in df.columns]))

异常信息：

AnalysisException: [DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES] Cannot resolve "(a + 0)" due to data type mismatch: the left and right operands of the binary operator have incompatible types ("BOOLEAN" and "INT").;

而如果是数值类型的DataFrame，上述求和语法可以正常运行：

columns = ["a", "b", "c"]
data = [(1, 0, 1), (1, 0, 0), (1, 1, 1)]
df = spark.createDataFrame(data).toDF(*columns)
df.display()

解决方案 #

方法1：将布尔列转换为整数后求和 #

利用类型转换把True转为1、False转为0，再对转换后的列求和：

import pyspark.sql.functions as F

# 方式一：用cast直接转换类型
df_with_count = df.withColumn(
    "number_true_values",
    sum(F.col(col).cast("int") for col in df.columns)
)

# 方式二：用when函数映射数值
df_with_count = df.withColumn(
    "number_true_values",
    sum(F.when(F.col(col), 1).otherwise(0) for col in df.columns)
)

df_with_count.display()

方法2：使用数组聚合函数 #

将行数据转为数组，通过聚合逻辑统计其中True的数量：

import pyspark.sql.functions as F

df_with_count = df.withColumn(
    "number_true_values",
    F.aggregate(
        F.array(*df.columns),
        F.lit(0),
        lambda acc, x: acc + F.when(x, 1).otherwise(0)
    )
)

df_with_count.display()

方法3：通过SQL表达式实现 #

直接编写SQL风格的转换与求和逻辑：

import pyspark.sql.functions as F

# 构建SQL求和表达式
sum_expr = " + ".join([f"CAST({col} AS INT)" for col in df.columns])
df_with_count = df.withColumn("number_true_values", F.expr(sum_expr))

df_with_count.display()

所有方法最终都会得到如下结果：

内容的提问来源于stack exchange，提问作者the_economist