Polars中列重命名与延迟执行导致的连接问题分析

在使用Polars进行数据处理时，开发人员经常会遇到需要将宽表转换为长表的需求。本文通过一个典型案例，深入分析在Polars延迟执行(Lazy)模式下进行列重命名和连接操作时可能遇到的问题。

问题现象

假设我们有一个包含不同交通工具颜色信息的DataFrame，结构如下：

test_df = pl.DataFrame({
    'id':[1,2,3], 
    'car_colour':['red', 'green', 'red'], 
    'boat_colour':['blue', 'red', 'green'], 
    'plane_colour':['red', 'blue', 'green']
}).lazy()

目标是将其转换为长格式，包含三列：id、vehicle(交通工具类型)和colour(颜色)。实现思路是为每种交通工具创建一个子DataFrame，重命名颜色列后垂直连接。

两种实现方式的对比

方式一：直接延迟执行

vehicles = ['car', 'boat', 'plane']
test_dfs = []
for vehicle in vehicles:
    tmp_df = test_df.select(
        pl.col('id'),
        pl.lit(vehicle).alias('vehicle'),
        pl.selectors.contains(vehicle).name.map(lambda x: x.replace(f"{vehicle}_", ""))
    )
    test_dfs.append(tmp_df)

使用对角连接(diagonal)时，发现只有最后一个交通工具(plane)的颜色列被正确重命名，其他列保留了原始名称。

方式二：在循环中调用collect_schema()

for vehicle in vehicles:
    tmp_df = test_df.select(...)  # 同上
    tmp_df.collect_schema()  # 新增此行
    test_dfs.append(tmp_df)

这种方式下，car和plane的颜色列被正确重命名，但boat列仍保留原名。尝试垂直连接时，虽然schema显示正确，但实际收集数据时会报错。

问题根源分析

这个问题实际上是由Python的lambda闭包特性与Polars的延迟执行机制共同作用导致的。

在Python中，lambda表达式中的变量是在执行时而非定义时绑定的。当在循环中使用lambda时，所有lambda都会绑定到循环变量的最终值。在Polars的延迟执行模式下，lambda的实际执行被推迟到收集数据时，此时循环变量vehicle的值已经是最后一个值"plane"。

当在循环中调用collect_schema()时，会强制立即执行lambda，此时vehicle的值是正确的，因此部分列能正确重命名。但由于不是所有列都触发了立即执行，导致行为不一致。

解决方案

正确的做法是在lambda定义时绑定循环变量的当前值：

pl.selectors.contains(vehicle).name.map(
    lambda x, vehicle=vehicle: x.replace(f"{vehicle}_", "")
)

通过将vehicle作为默认参数传入，我们确保了每个lambda都会绑定到定义时的vehicle值，而不是执行时的最终值。

最佳实践建议

1. 在Polars延迟模式下使用循环变量时，特别注意lambda的绑定时机
2. 对于需要立即确定schema的操作，考虑使用eager模式或显式调用collect_schema()
3. 复杂的列转换操作可以先在小数据集上测试，确保行为符合预期
4. 垂直连接前确保所有子DataFrame的列名完全一致

理解这些底层机制有助于开发者更好地利用Polars的强大功能，同时避免常见的陷阱。