PySpark中reduce+add逻辑的Pandas等价实现方法咨询

原始数据 #

首先定义示例Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({'Day_1': [2, 4, 8, 0],
                   'Day_2': [2, 0, 0, 0],
                   'Day_3': [1, 1, 0, 0]},
                  index=['user1', 'user2', 'user3', 'user4'])

输出结果：

Day_1  Day_2  Day_3
user1      2      2      1
user2      4      0      1
user3      8      0      0
user4      0      0      0

需求说明 #

新增impression列，规则为：

若用户过去n天的曝光量之和大于0则为1，否则为0

当num_days=2时，期望输出：

Day_1  Day_2  Day_3  impression
user1      2      2      1           1
user2      4      0      1           1
user3      8      0      0           1
user4      0      0      0           0

参考的PySpark实现 #

该逻辑在PySpark中的实现代码如下：

from functools import reduce
from pyspark.sql import functions as F

num_days = 2
imp_cols = ['Day_'+str(i) for i in range(1, num_days+1)]
df = df.withColumn("impression", reduce(add, [F.col(x) for x in imp_cols]))

Pandas等价实现方法 #

方法1：直接求和后转布尔整数 #

这是最直观的实现，按行求和后判断是否大于0，再转换为整数：

num_days = 2
imp_cols = [f'Day_{i}' for i in range(1, num_days+1)]
df['impression'] = df[imp_cols].sum(axis=1).gt(0).astype(int)

• sum(axis=1)：计算每行的曝光量总和
• gt(0)：判断总和是否大于0，返回布尔值
• astype(int)：将布尔值转换为1/0

方法2：判断是否存在非零曝光（更高效） #

如果只需确认是否有曝光（无需精确求和），可以直接判断每行是否有非零值：

df['impression'] = df[imp_cols].gt(0).any(axis=1).astype(int)

• gt(0)：将每个元素转换为是否大于0的布尔值
• any(axis=1)：判断每行是否存在至少一个True（即有曝光）
• astype(int)：转换为1/0

方法3：完全对应PySpark的reduce累加逻辑 #

若要严格匹配PySpark中用reduce累加列的逻辑，可使用functools.reduce：

from functools import reduce
import operator

df['impression'] = (reduce(operator.add, [df[col] for col in imp_cols]) > 0).astype(int)

通过reduce逐列累加得到每行总和，再判断并转换为整数，逻辑与PySpark完全一致。

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