Polars LazyFrame迭代使用pl.struct+map_elements触发KeyError问题

问题重现 #

在Polars 0.19.13版本中，对多列应用自定义逻辑时出现以下差异：

• DataFrame立即执行模式：循环迭代列，通过pl.struct传递列到map_elements的lambda中访问字段，代码可正常运行。
• LazyFrame延迟执行模式：相同逻辑循环迭代时，lambda访问struct字段会触发KeyError；但将每一步迭代单独写出时，代码又能正常运行。

正常运行的DataFrame代码 #

import polars as pl

my_df = pl.DataFrame(
    {
        "foo": ["a", "b", "c", "d"], 
        "bar": ["w", "x", "y", "z"], 
        "notes": ["1", "2", "3", "4"]
    }
)

cols_to_validate = ("foo", "bar")

def validate_stuff(value, notes):
    if value not in ["a", "b", "x"]:
        return f"FAILED {value} - PREVIOUS ({notes})"
    else:
        return notes

for col in cols_to_validate:
    my_df = my_df.with_columns(
        pl.struct([col, "notes"]).map_elements(
            lambda row: validate_stuff(row[col], row["notes"])
        ).alias("notes")
    )

print(my_df)

触发KeyError的LazyFrame代码 #

import polars as pl

my_lf = pl.DataFrame(
    {
        "foo": ["a", "b", "c", "d"], 
        "bar": ["w", "x", "y", "z"], 
        "notes": ["1", "2", "3", "4"]
    }
).lazy()

def validate_stuff(value, notes):
    if value not in ["a", "b", "x"]:
        return f"FAILED {value} - PREVIOUS ({notes})"
    else:
        return notes

cols_to_validate = ("foo", "bar")

for col in cols_to_validate:
    my_lf = my_lf.with_columns(
        pl.struct([col, "notes"]).map_elements(
            lambda row: validate_stuff(row[col], row["notes"])
        ).alias("notes")
    )

print(my_lf.collect())  # 执行时触发KeyError

可正常运行的单独迭代写法 #

# 接上述LazyFrame定义
my_lf = my_lf.with_columns(
    pl.struct(["foo", "notes"]).map_elements(
        lambda row: validate_stuff(row["foo"], row["notes"])
    ).alias("notes")
)

my_lf = my_lf.with_columns(
    pl.struct(["bar", "notes"]).map_elements(
        lambda row: validate_stuff(row["bar"], row["notes"])
    ).alias("notes")
)

print(my_lf.collect())

原因分析 #

这不是Polars的bug，而是Python闭包特性结合LazyFrame延迟执行机制导致的问题：

• LazyFrame的map_elements不会立即执行lambda，而是将其作为计算计划的一部分延迟到collect()时执行。
• 循环中的lambda是闭包，它捕获的col变量是对循环变量的引用，而非每次迭代的具体值。当collect()执行时，所有lambda中的row[col]都会使用循环最后一次迭代的col值（即"bar"）。
• 第一次迭代生成的struct是["foo", "notes"]，不存在"bar"字段，因此访问row["bar"]时触发KeyError。
• DataFrame是立即执行，每次迭代的lambda都会捕获当前循环的col值；单独写步骤时lambda用的是硬编码的字段名，不存在闭包变量引用问题，因此都能正常运行。

解决方案 #

方案1：绑定循环变量到函数（避免闭包引用） #

使用functools.partial将当前迭代的col值绑定到处理函数，确保每个lambda使用的是当前迭代的列名：

import polars as pl
from functools import partial

my_lf = pl.DataFrame(
    {
        "foo": ["a", "b", "c", "d"], 
        "bar": ["w", "x", "y", "z"], 
        "notes": ["1", "2", "3", "4"]
    }
).lazy()

def validate_stuff(value, notes):
    if value not in ["a", "b", "x"]:
        return f"FAILED {value} - PREVIOUS ({notes})"
    else:
        return notes

def validate_row(row, col):
    return validate_stuff(row[col], row["notes"])

cols_to_validate = ("foo", "bar")

for col in cols_to_validate:
    my_lf = my_lf.with_columns(
        pl.struct([col, "notes"]).map_elements(
            partial(validate_row, col=col)
        ).alias("notes")
    )

print(my_lf.collect())

方案2：直接传递列给map_elements（无需struct） #

跳过pl.struct，直接用pl.map_elements接收多列参数，避免struct字段访问问题：

import polars as pl

my_lf = pl.DataFrame(
    {
        "foo": ["a", "b", "c", "d"], 
        "bar": ["w", "x", "y", "z"], 
        "notes": ["1", "2", "3", "4"]
    }
).lazy()

def validate_stuff(value, notes):
    if value not in ["a", "b", "x"]:
        return f"FAILED {value} - PREVIOUS ({notes})"
    else:
        return notes

cols_to_validate = ("foo", "bar")

for col in cols_to_validate:
    my_lf = my_lf.with_columns(
        pl.map_elements(
            lambda val, note: validate_stuff(val, note),
            pl.col(col), pl.col("notes"),
            return_dtype=pl.String
        ).alias("notes")
    )

print(my_lf.collect())

方案3：硬编码每一步迭代 #

即用户提到的单独写出每列处理逻辑的方式，适合列数量较少的场景。

结论 #

该问题是Python闭包特性与Polars LazyFrame延迟执行结合导致的用法问题，而非Polars 0.19.13的bug。通过上述方案避免闭包对循环变量的引用，即可解决KeyError问题。

内容的提问来源于stack exchange，提问作者J. Alvarez