pandas使用apply从关联表查询数据新增两列的报错问题求解

问题原因 #

报错的核心原因是：apply沿axis=1逐行运行时，返回的每个结果是长度为2的numpy数组，最终apply整体输出的是一维的对象数组（元素为长度2的数组），而给多列赋值时，pandas期望接收的是「行数与table一致、列数与新增列数一致的二维结构」，结构不匹配就触发了长度错误。
另外原写法存在潜在风险：如果table里的某个Color在lookup表里不存在，values[0]会直接报索引越界错误，建议补充兜底处理。

解决方案 #

方法1：最小改动修复原有代码 #

在apply返回结果后面加.tolist()，把嵌套的数组转成嵌套列表，pandas就能自动对齐多列：

lookup_columns = ["Source", "Lead Days"]
table[lookup_columns] = table.apply(
    lambda row: 
          lookup.query('`Color` == "{color}"'.format(color=row["Color"])).loc[:, lookup_columns].values[0]
      , axis = 1).tolist()

方法2：更稳妥的优化写法（兼容缺省值、效率更高） #

提前把lookup转成字典映射，不需要逐行执行query遍历整张表，查询效率从O(n)降到O(1)，同时兼容匹配不到的场景自动填充NaN：

lookup_columns = ["Source", "Lead Days"]
# 提前构建Color到关联信息的映射字典
color_map = lookup.set_index("Color")[lookup_columns].to_dict("index")

def get_lookup_info(row):
    # 匹配不到返回空Series，自动填充NaN
    return pd.Series(color_map.get(row["Color"], {}))

table[lookup_columns] = table.apply(get_lookup_info, axis=1)

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