如何实现Polars表达式的动态多步链式运算？

核心思路 #

不要每次调用辅助函数就立即执行collect()，而是保留LazyFrame的链式调用状态，要么把所有运算步骤先构建成表达式列表，一次性传入with_columns；要么通过迭代逐步扩展LazyFrame，最后统一执行collect()。

改造方案 #

1. 调整辅助函数，返回表达式而非DataFrame #

先修改原有函数，让它生成Polars表达式（Expr）而非直接输出DataFrame，这样就能灵活组合多个运算步骤：

from typing import Callable, Optional, List
import polars as pl
from polars import IntoExpr

def demo_operator_addition(var1: str, var2: str) -> IntoExpr:
    return pl.col(var1).add(pl.col(var2))

def build_calc_expr(
        target_var: str,
        operator: Callable,
        reference_var: Optional[str] = None,
        by: Optional[List[str]] = None,
        col_name: Optional[str] = None,
        ) -> pl.Expr:
    by = by if by is not None else []
    expr = operator(target_var, reference_var).over([True, *by])
    return expr.alias(col_name) if col_name is not None else expr.alias(target_var)

2. 动态构建表达式列表，一次性执行 #

如果运算步骤是动态生成的（比如从配置列表读取），可以先把所有需要的表达式收集到列表里，再一次性传给with_columns：

# 示例：定义动态运算步骤，后一步依赖前一步的结果
calc_steps = [
    # 第一步：用var1+var2生成var3
    {"target_var": "var1", "operator": demo_operator_addition, "reference_var": "var2", "col_name": "var3"},
    # 第二步：用var3+var1生成var4
    {"target_var": "var3", "operator": demo_operator_addition, "reference_var": "var1", "col_name": "var4"},
    # 可继续添加更多依赖步骤
]

# 收集所有表达式
exprs = []
for step in calc_steps:
    expr = build_calc_expr(**step)
    exprs.append(expr)

# 一次性执行所有运算
data = pl.DataFrame({"var1": [1,2,3], "var2": [4,5,6]})
result = data.lazy().with_columns(exprs).collect()
print(result)

3. 迭代式扩展LazyFrame #

如果需要更灵活的步骤控制（比如根据前一步结果动态生成下一步），可以迭代更新LazyFrame，最后统一执行计算：

data_lazy = pl.DataFrame({"var1": [1,2,3], "var2": [4,5,6]}).lazy()

# 第一步：生成var3
data_lazy = data_lazy.with_columns(
    build_calc_expr("var1", demo_operator_addition, "var2", col_name="var3")
)

# 第二步：基于var3生成var4
data_lazy = data_lazy.with_columns(
    build_calc_expr("var3", demo_operator_addition, "var1", col_name="var4")
)

# 最后统一执行计算
result = data_lazy.collect()
print(result)

关键说明 #

• 避免多次collect()：每次collect()都会触发实际计算，多次调用会大幅降低效率，保留LazyFrame直到最后再执行是Polars高效运算的核心。
• 表达式延迟解析：Polars的表达式是延迟计算的，即使后一步引用的是前一步生成的列名，只要在同一个LazyFrame的with_columns链中，Polars会自动处理依赖关系。

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