为Pandas DataFrame添加新列并扩展行的高效实现方案

问题描述 #

现有如下Pandas DataFrame：

需求：为每行复制出一行，新增col3列；当col1为偶数时，复制行的col3值为'a'和'b'；当col1为奇数时，复制行的col3值为'c'和'd'，期望输出如下：

已通过iterrows()遍历行实现需求，但处理数十万行、数千列的大数据时速度极慢；尝试用apply()函数实现却得到异常输出，求更优实现方式。

现有实现代码：

import pandas as pd
import numpy as np

d = {'col1': [1, 2], 'col2': [3, 4]}
df = pd.DataFrame(data=d)


def add_rows(row):
    """
    This function is used in apply function. It adds a new column and adds rows to the dataframe.
    :param row:
    :return:
    """
    col3 = np.array(['a', 'b'])
    out_df = pd.DataFrame([row.tolist()], columns=row.index)
    out_df['col3'] = None
    if row['col1'] % 2 == 0:
        out_df.at[0, 'col3'] = col3
    else:
        out_df.at[0, 'col3'] = np.array(['c', 'd'])
    out = out_df.explode('col3', ignore_index=True)
    return out

cols = list(df.columns)
cols.append('col3')
result = pd.DataFrame(columns=cols)
for index, row in df.iterrows():
    rows = add_rows(row)
    result = pd.concat([result, rows])
result.reset_index(drop=True, inplace=True)
print(result)

优化实现方案 #

使用Pandas的矢量化操作替代循环，避免逐行处理的性能损耗，以下两种方案均能高效解决问题：

方案一：结合apply生成列表+explode展开 #

import pandas as pd

d = {'col1': [1, 2], 'col2': [3, 4]}
df = pd.DataFrame(data=d)

# 为每行生成对应的col3值列表
df['col3'] = df['col1'].apply(lambda x: ['c', 'd'] if x % 2 != 0 else ['a', 'b'])

# 展开列表为多行
result = df.explode('col3', ignore_index=True)

print(result)

方案二：numpy批量赋值（性能最优） #

如果处理超大规模数据，可直接用numpy构造重复行并批量赋值，完全规避循环开销：

import pandas as pd
import numpy as np

d = {'col1': [1, 2], 'col2': [3, 4]}
df = pd.DataFrame(data=d)

# 重复原数据每行2次
repeated_df = df.loc[df.index.repeat(2)].reset_index(drop=True)

# 根据col1奇偶性批量赋值col3
mask = repeated_df['col1'] % 2 == 0
repeated_df['col3'] = np.where(
    mask,
    np.tile(['a', 'b'], len(df)),
    np.tile(['c', 'd'], len(df))
)

print(repeated_df)

方案优势说明 #

• 两种方案均避免了iterrows()或逐行拼接的低效操作，矢量化处理是Pandas针对大数据集优化的核心特性，处理数十万行时性能提升可达数十倍。
• 方案二基于numpy的数组操作，底层由C实现，是性能最优的选择，适合超大规模数据场景。

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