Polars中如何在select/with_columns语句中声明中间变量？

以下是几种实现你需求的方案，无需链式调用多个select/with_columns：

1. 用Python变量存储临时表达式（最简单） #

把重复使用的表达式赋值给普通Python变量，在多个列的定义中直接引用，既避免重复代码，又提升可读性：

import polars as pl

df = pl.DataFrame({"a": [1,2,3], "b": [4,5,6]})

# 定义临时表达式，作为Python变量
temp_expr = pl.col("b") * 2

result = df.with_columns(
    # 输出到结果的列
    step_1 = pl.col("a") + 1,
    # 复用temp_expr计算step_2，同时引用刚定义的step_1
    step_2 = pl.col("step_1") + temp_expr,
    # 再次复用temp_expr计算另一列
    another_column = pl.col("a") * temp_expr
)

print(result)

2. 使用Polars的pl.let()表达式（Polars 0.19.0+） #

pl.let()允许你在表达式内部定义临时变量，该变量仅在当前表达式上下文有效，不会输出为DataFrame的列：

import polars as pl

df = pl.DataFrame({"a": [1,2,3], "b": [4,5,6]})

result = df.with_columns(
    step_1 = pl.col("a") + 1,
    # 在step_2的表达式中定义临时变量temp_var并使用
    step_2 = pl.let(
        temp_var = pl.col("b") * 2,
        then = pl.col("step_1") + temp_var
    ),
    # 在another_column中复用相同的临时表达式逻辑
    another_column = pl.let(
        temp_var = pl.col("b") * 2,
        then = pl.col("a") * temp_var
    )
)

print(result)

如果需要在多个列中复用同一个临时变量的定义，可以结合Python变量和pl.let()：

temp_expr = pl.col("b") * 2

result = df.with_columns(
    step_1 = pl.col("a") + 1,
    step_2 = pl.let(temp_var = temp_expr, then = pl.col("step_1") + temp_var),
    another_column = pl.let(temp_var = temp_expr, then = pl.col("a") * temp_var)
)

3. 嵌套pl.select子查询 #

通过嵌套的select先定义临时列，再在子查询中计算最终列，临时列不会出现在结果中：

import polars as pl

df = pl.DataFrame({"a": [1,2,3], "b": [4,5,6]})

result = df.with_columns(
    pl.select(
        # 定义临时列和需要保留的列
        step_1 = pl.col("a") + 1,
        temp_var = pl.col("b") * 2
    ).select(
        # 只保留需要的列，引用临时列计算
        step_1,
        step_2 = pl.col("step_1") + pl.col("temp_var"),
        another_column = pl.col("a") * pl.col("temp_var")
    )
)

print(result)

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