PySpark DataFrame将double列转二进制并计算位和新增列

原有代码问题说明 #

• 第一段代码报错核心原因：F.col()返回的是PySpark列的引用对象，不是逐行的实际数值。直接用Python原生的字符串格式化、map、sum处理列对象时，原生函数无法识别PySpark列类型，仅会在Driver端执行，无法分布式处理DataFrame的每行数据，必然抛出类型错误。
• 第二段代码已经通过内置F.bin()正确生成了二进制字符串，也通过F.lpad()完成了补零操作，但缺少对字符串中'1'的计数逻辑，且逐列硬编码withColumn在列数较多时维护成本很高。

实现方案 #

方案1：内置bit_count函数实现（推荐，性能最优） #

Spark 3.0及以上版本自带bit_count内置函数，可直接计算整数对应二进制表达式中1的个数，无需做字符串转换，无额外序列化开销，分布式计算性能最好。
由于列是double类型，且存储值均为整数，先强转为bigint类型即可避免计算异常。

from pyspark.sql import functions as F

# 配置需要处理的列名列表
bit_columns = ['bit_1', 'bit_2', 'bit_3', 'bit_4', 'bit_5', 'bit_6']
result_df = df

# 逐列计算单值的二进制1的个数
for col_name in bit_columns:
    result_df = result_df.withColumn(
        f'{col_name}_bit_cnt',
        F.bit_count(F.col(col_name).cast('bigint'))
    )

# 如需计算所有列的位计数总和，新增总求和列
result_df = result_df.withColumn(
    'total_bit_sum',
    sum(F.col(f'{col_name}_bit_cnt') for col_name in bit_columns)
)

方案2：二进制字符串计数实现（兼容低版本Spark） #

如果Spark版本低于3.0，没有内置bit_count函数，可以基于已有的二进制字符串生成逻辑，通过替换所有0字符后统计剩余字符串长度的方式，得到二进制中1的个数。

from pyspark.sql import functions as F

bit_columns = ['bit_1', 'bit_2', 'bit_3', 'bit_4', 'bit_5', 'bit_6']
pad_len = 12 # 二进制字符串补零长度
result_df = df

for col_name in bit_columns:
    bi_col_name = f'{col_name}_bi'
    cnt_col_name = f'{col_name}_bit_cnt'
    result_df = result_df.withColumn(
        bi_col_name,
        F.lpad(F.bin(F.col(col_name).cast('bigint')), pad_len, '0')
    ).withColumn(
        cnt_col_name,
        # 移除二进制串中所有0字符，剩余字符串长度即为1的个数
        F.length(F.regexp_replace(F.col(bi_col_name), '0', ''))
    )

# 计算所有列的位计数总和
result_df = result_df.withColumn(
    'total_bit_sum',
    sum(F.col(f'{col_name}_bit_cnt') for col_name in bit_columns)
)

计算结果验证 #

以第一行样例数据为例：

• bit_1=0 → 二进制000000000000 → 1的计数为0
• bit_2=2 → 二进制000000000010 → 1的计数为1
• bit_3=8 → 二进制000000001000 → 1的计数为1
• 其余列值为0，计数均为0
• 最终total_bit_sum结果为2，和预期一致。

注意事项 #

• 处理double类型列前先强转为bigint，避免浮点数参与二进制转换出现精度异常。
• 优先使用PySpark内置函数处理列数据，不要直接用Python原生函数操作F.col()对象，内置函数比自定义UDF性能高1~2个数量级。

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