Pandas可变列数据按行提取唯一值左对齐的提速优化方法

原来的代码性能差主要来自逐行groupby+apply的高额开销，以及频繁创建Series对象的冗余操作，可通过以下方案优化，性能可提升数十倍甚至上百倍：

核心优化思路 #

• 放弃逐行处理的pandas操作，改用numpy向量化操作减少开销
• 一次性初始化结果容器，避免逐行创建pandas对象
• 减少不必要的索引、属性访问操作

优化实现代码 #

import pandas as pd
import numpy as np
from numpy import nan

# 构造测试数据的代码和你提供的一致
data = [
    (111, nan, nan, 111),
    (112, 112, nan, 115),
    (113, nan, nan, nan),
    (nan, nan, nan, nan),
    (118, 110, 117, nan),
]
df = pd.DataFrame(data, columns=[f'num{i}' for i in range(len(data[0]))])

# 优化后的处理逻辑
arr = df.values
row_count = arr.shape[0]
# 计算每行非空唯一值的最大数量，作为结果列数
max_valid_count = max(len(np.unique(row[~np.isnan(row)])) for row in arr)
# 初始化全为nan的结果数组
result_arr = np.full((row_count, max_valid_count), nan)

for idx, row in enumerate(arr):
    # 提取当前行的非空唯一值
    valid_unique = np.unique(row[~np.isnan(row)])
    # 左对齐赋值，剩余位置保持nan
    result_arr[idx, :len(valid_unique)] = valid_unique

# 转换为DataFrame，设置对应列名
result_df = pd.DataFrame(result_arr, columns=[f'num{i+1}' for i in range(max_valid_count)])

超大流量场景的进一步优化 #

如果你的数据量超过10万行，还可以引入numba加速循环，只需要给循环逻辑加上装饰器即可：

from numba import jit

@jit(nopython=True)
def fast_arrange(arr, max_col):
    n_row = arr.shape[0]
    res = np.full((n_row, max_col), np.nan)
    for i in range(n_row):
        row = arr[i]
        valid = row[~np.isnan(row)]
        valid_unique = np.unique(valid)
        res[i, :len(valid_unique)] = valid_unique
    return res

# 调用方式
max_valid_count = max(len(np.unique(row[~np.isnan(row)])) for row in df.values)
result_arr = fast_arrange(df.values, max_valid_count)
result_df = pd.DataFrame(result_arr, columns=[f'num{i+1}' for i in range(max_valid_count)])

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