Polars中collect_all与LazyFrame.collect的性能差异及优化问询

先看以下示例代码：

import time
import numpy as np
import polars as pl

n_index = 1000
n_a = 10
n_b = 500
n_obs = 5000000

df = pl.DataFrame(
    {
        "id": np.random.randint(0, n_index, size=n_obs),
        "a": np.random.randint(0, n_a, size=n_obs),
        "b": np.random.randint(0, n_b, size=n_obs),
        "x": np.random.normal(0, 1, n_obs),
    }
).lazy()

dfs = [
    pl.DataFrame(
        {
            "id": np.random.randint(0, n_index, size=n_obs),
            "a": np.random.randint(0, n_a, size=n_obs),
            f"b_{i}": np.random.randint(0, n_b, size=n_obs),
            "x": np.random.normal(0, 1, n_obs),
        }
    ).lazy()
    for i in range(50)
]

res = [
    df.join(
        dfs[i], left_on=["id", "a", "b"], right_on=["id", "a", f"b_{i}"], how="inner"
    )
    .group_by("a", "b")
    .agg((pl.col("x") * pl.col("x_right")).sum().alias(f"x_{i}"))
    for i in range(50)
]

需求为处理多个DataFrame并计算，最后将所有结果拼接。上述代码构建了包含所有计算结果的列表res。

针对结果拼接，我尝试了两种方案：

方案1 #

start = time.perf_counter()
res2 = pl.collect_all(res)
res3 = res2[0]
for i in range(1, 50):
    res3 = res3.join(res2[i], on=["a", "b"])
time.perf_counter() - start

方案2 #

start = time.perf_counter()
res4 = res[0]
for i in range(1, 50):
    res4 = res4.join(res[i], on=["a", "b"])
res4 = res4.collect()
time.perf_counter() - start

方案1先执行collect_all再逐个拼接DataFrame；方案2全程采用延迟执行，最后才调用collect。

据我了解，collect会在底层做优化，原本预期两种方案性能相近，但测试结果显示方案2耗时是方案1的两倍（32核系统下21s vs 10s）。

因此想请教：

• 这种性能差异是否符合预期？我的实现方式是否存在低效问题？
• 方案2的优势是全程延迟执行，适合需要提供全延迟API、让用户后续自主操作的场景，但性能损耗过大。有没有办法在保持全延迟API的同时，达到接近方案1的性能？

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