Polars为何不复用同一上下文内重复表达式的已计算结果？

Polars用户指南提到：

Common subplan elimination - 缓存查询计划中被多个子树复用的子树/文件扫描。

我不清楚此处的“子树”具体指什么，但原本期望Polars会尝试复用同一上下文内重复子表达式的计算结果。不过以下示例显示实际并非如此。

# polars==0.18.3

import polars as pl
import time
import os


def compute_heavily(initial_value: pl.Expr) -> pl.Expr:
    """Waste some CPU power to compute the golden ratio"""
    result = initial_value
    for i in range(100):
        result = (1 + result) ** 0.5
    return result


# Block parallel execution of polars expressions
os.environ["POLARS_MAX_THREADS"] = "1"
# Create a dataframe
df = pl.DataFrame(pl.repeat(0.0, 10_000_000, eager=True).alias("value"))

# Case A
start_time = time.perf_counter()
df_a = (
    df.lazy().with_columns(compute_heavily(pl.col("value")).alias("result_1")).collect()
)
delta_time = time.perf_counter() - start_time
print(f"Case A: function called once: {delta_time:.3f} s.")

# Case B
start_time = time.perf_counter()
df_b = (
    df.lazy()
    .with_columns(compute_heavily(pl.col("value")).alias("result_1"))
    .with_columns((pl.col("result_1")).alias("result_2"))
    .collect()
)
delta_time = time.perf_counter() - start_time
print(
    f"Case B: function called once in one context "
    f"and the result is used in an updated context: {delta_time:.3f} s."
)

# Case C
start_time = time.perf_counter()
df_c = (
    df.lazy()
    .with_columns(compute_heavily(pl.col("value")).alias("result_1"))
    .with_columns(compute_heavily(pl.col("value")).alias("result_2"))
    .collect()
)
delta_time = time.perf_counter() - start_time
print(
    f"Case C: function called twice, "
    f"but each time in a different context: {delta_time:.3f} s."
)

# Case D
start_time = time.perf_counter()
df_d = (
    df.lazy()
    .with_columns(
        (compute_heavily(pl.col("value")) + compute_heavily(pl.col("value"))).alias(
            "result_1"
        )
    )
    .collect()
)
delta_time = time.perf_counter() - start_time
print(
    f"Case D: function called twice within one context "
    f"to make one column: {delta_time:.3f} s."
)

# Case E
start_time = time.perf_counter()
df_e = (
    df.lazy()
    .with_columns(
        [
            compute_heavily(pl.col("value")).alias("result_1"),
            compute_heavily(pl.col("value")).alias("result_2"),
        ]
    )
    .collect()
)
delta_time = time.perf_counter() - start_time
print(
    f"Case E: function called twice within one context, "
    f"but each time for a different column: {delta_time:.3f} s."
)

运行结果：

Case A: function called once: 2.124 s.
Case B: function called once in one context and the result is used in an updated context: 2.208 s.
Case C: function called twice, but each time in a different context: 4.454 s.
Case D: function called twice within one context to make one column: 4.390 s.
Case E: function called twice within one context, but each time for a different column: 4.299 s.

案例A、B、C的表现符合预期：A和B耗时约1个时间单位（函数仅被显式调用一次），C耗时约2个时间单位（函数在不同上下文被显式调用两次）。但案例D和E耗时也为2个时间单位，我原本以为借助优化它们应该仅耗时1个时间单位。

问题解答 #

1. Polars的这种表现是设计如此吗？
   是的，这是Polars v0.18.3版本的设计行为。当前Polars的Common Subplan Elimination（CSE）并不针对单个表达式的重复计算做自动复用，尤其是在同一个with_columns调用中多次出现相同表达式时，Polars会重复执行该计算。

这种设计主要是为了平衡查询计划的构建效率与优化收益：表达式的哈希校验和复用需要额外开销，同时Polars要避免误判（比如某些看似相同的表达式可能隐含不同的状态或副作用）。如果需要复用计算结果，建议显式将表达式结果保存为中间列，再基于中间列操作，就像案例B的实现方式。

2. Common subplan elimination具体指什么？
   这里的“子树”指的是查询计划中一个独立的操作分支，通常对应完整的数据源读取（如文件扫描）或一组连贯的操作序列（如过滤+聚合的组合）。Common Subplan Elimination的核心是识别查询计划中完全相同的子树，缓存其执行结果，避免重复执行。

举个典型场景：如果查询中多次对同一个数据源执行相同的过滤操作，Polars会复用该过滤子树的结果，只执行一次过滤：

df = pl.scan_parquet("large_file.parquet")
q1 = df.filter(pl.col("age") > 30).select("name")
q2 = df.filter(pl.col("age") > 30).select("salary")
combined = pl.concat([q1, q2], how="horizontal")
combined.collect()

上述代码中，filter(pl.col("age") > 30)对应的子树会被复用，Polars不会重复执行过滤操作。而你案例中的compute_heavily(pl.col("value"))属于单个表达式，不属于“子树”范畴，因此不会触发CSE优化。

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