如何使用numba加速pandas差异统计for循环并解决相关报错

• 你遇到的AssertionError是Numba的nopython模式无法编译代码中的不支持特性导致的：该模式完全不支持Pandas的各类API（包括applymap、loc、布尔索引、unique()等），同时你用到的lambda表达式、列表推导式内调用set.add、全局DataFrame引用、tqdm调用都是nopython模式无法识别的语法，编译阶段出现变量冲突触发了报错。
• 代码本身还存在逻辑漏洞：循环中重复调用unique()产生冗余计算、用于存储差异列的s列表没有每次重置导致不同分组的结果互相污染、函数内部直接读取全局变量df不符合Numba的传参规则。

你的需求可以直接通过Pandas内置的groupby聚合实现，所有运算都走C层优化，性能远高于你手写的Python循环加Numba的方案：

import pandas as pd

# 全列转字符串
df = df.astype(str)
# 定义需要检查差异的列：排除分组用的两个字段
check_cols = [col for col in df.columns if col not in ["client_cookie_id", "OS_name"]]

def get_diff_columns(group):
    # 分组行数小于2时无差异
    if len(group) < 2:
        return []
    # 统计每列的唯一值数量，大于1说明存在差异
    return [col for col in check_cols if group[col].nunique() > 1]

# 按client_cookie_id和OS_name分组计算差异列
group_result = df.groupby(["client_cookie_id", "OS_name"]).apply(get_diff_columns).to_dict()

# 按client_cookie_id汇总所有OS分组下的差异列
final_dict = {}
for (cid, os_name), diff_cols in group_result.items():
    if cid not in final_dict:
        final_dict[cid] = set()
    final_dict[cid].update(diff_cols)
# 转换为列表格式
final_dict = {k: list(v) for k, v in final_dict.items()}

如果数据量极大确实需要Numba加速，必须先把Pandas对象转换为Numba支持的numpy数组传入，不能在Numba函数内部操作Pandas实例：

import pandas as pd
import numpy as np
from numba import jit

# 预处理：将所有分类字段和待比较值编码为整数（Numba处理整数效率远高于字符串）
cid_codes, cid_uniques = pd.factorize(df["client_cookie_id"])
os_codes, os_uniques = pd.factorize(df["OS_name"])
val_array = np.vstack([pd.factorize(df[col])[0] for col in check_cols]).T
col_name_array = np.array(check_cols)

@jit(nopython=True)
def calc_diff_columns(cid_codes, os_codes, val_array, col_name_array, cid_uniques):
    res = dict()
    max_cid = cid_codes.max() + 1
    max_os = os_codes.max() + 1
    for cid_idx in range(max_cid):
        cid_mask = cid_codes == cid_idx
        diff_set = set()
        for os_idx in range(max_os):
            group_mask = cid_mask & (os_codes == os_idx)
            row_indexes = np.where(group_mask)[0]
            if len(row_indexes) < 2:
                continue
            # 逐列检查值是否全部相同
            for col_idx in range(val_array.shape[1]):
                col_vals = val_array[row_indexes, col_idx]
                if not (col_vals == col_vals[0]).all():
                    diff_set.add(col_name_array[col_idx])
        res[cid_uniques[cid_idx]] = list(diff_set)
    return res

final_dict = calc_diff_columns(cid_codes, os_codes, val_array, col_name_array, cid_uniques)

• Numba的nopython模式不支持tqdm，需要进度条只能关闭nopython模式，或者在外层逻辑中添加
• 所有需要在Numba函数中使用的数据必须通过参数传入，不能引用全局变量
• 尽量不要在Numba中直接处理字符串，提前编码为整数可以让性能提升数十倍

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