Polars从立即执行转延迟执行时触发ColumnNotFoundError问题咨询

问题原因 #

这段代码在立即执行模式正常运行，但切换到延迟执行时抛出ColumnNotFoundError，核心原因是延迟执行依赖明确的查询计划，而map_groups中动态生成的列无法被Polars提前推断。立即执行是实时计算生成列，后续步骤能直接读取；但延迟执行先构建完整查询计划，Polars无法预知lambda会生成A_lag_0/1/2这些列，导致后续with_columns引用时报错。

解决方案 #

方案1：显式指定map_groups的schema

给map_groups传入明确的schema定义，让Polars提前知晓输出列结构，即可正常构建查询计划：

import polars as pl 

df = pl.DataFrame({
     'A':[1,2,3,3,2,1],
     'B':[1,1,1,2,2,2]
})

(df
 .lazy()
 .group_by('B')
 .map_groups(
     lambda x: x.with_columns(pl.col("A").shift(i).alias(f"A_lag_{i}") for i in range(3)),
     # 显式声明输出schema：包含原列和所有新生成列
     schema={'A': pl.Int64, 'B': pl.Int64, 'A_lag_0': pl.Int64, 'A_lag_1': pl.Int64, 'A_lag_2': pl.Int64}
 )
 .with_columns(pl.col(f'A_lag_{i}') / pl.col('A') for i in range(3))
 .collect()
)

方案2：将后续计算逻辑移入map_groups内部

把生成延迟列和除法计算都封装到map_groups的lambda中，避免在外部引用动态生成的列：

import polars as pl 

df = pl.DataFrame({
     'A':[1,2,3,3,2,1],
     'B':[1,1,1,2,2,2]
})

(df
 .lazy()
 .group_by('B')
 .map_groups(
     lambda x: x.with_columns(
         pl.col("A").shift(i).alias(f"A_lag_{i}") for i in range(3)
     ).with_columns(
         pl.col(f'A_lag_{i}') / pl.col('A') for i in range(3)
     ),
     schema=None
 )
 .collect()
)

补充说明 #

延迟执行模式下使用schema=None时，Polars会尝试自动推断输出schema，但动态生成列的场景下推断逻辑失效。因此要么显式指定schema，要么把相关计算都封装到map_groups内部，即可解决该问题。

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