Polars LazyFrame预处理函数的最优返回类型咨询

我有一个大型Polars LazyFrame，现在写了一个只能接收pl.LazyFrame（不能用pl.Expr）的预处理函数，这个函数会对LazyFrame的不同列执行一系列操作，最终只需要返回一列用来追加到原LazyFrame中。想请教：

• 哪种返回类型能让查询规划器最优地整合到原数据框？
• 目前用.collect().to_series()再通过.with_columns()追加的方式会触发collect，损失优化潜力，有没有办法避免返回Series和执行collect？
• 当前方案的主要问题是如果步骤改变顺序或者遗漏maintain_order参数会出问题。我还考虑过两种替代方案：找ID列关联；或者在with_columns表达式上下文里直接返回完整LazyFrame不调用.collect()。

附上伪代码：

def preprocessing(
    df: pl.LazyFrame | pl.DataFrame,
) -> pl.Series : #? 待确定返回类型
    return (
        df.lazy()
        .with_columns(
            [
                pl.col("some_col")
                .cast(pl.Int64)
                .cast(pl.Utf8)
                .str.pad_start(5, "0")
                .alias("some_col_preprocessed"),
                pl.col("other_col")
                .cast(pl.Int64)
                .cast(pl.Utf8)
                .str.pad_start(7, "0")
                .alias("other_col_preprocessed"),
            ],
        )
# ... 更多处理步骤
        .select(
            pl.concat_str(
                [
                    pl.col("some_col_preprocessed"),
                    pl.col("other_col_preprocessed"),
                    pl.lit("42133723"),
                ],
            )
        ).collect().to_series()

starting_df = starting_df.with_columns(final_col=preprocessing(starting_df))

最优方案：返回仅含目标列的pl.LazyFrame #

直接返回包含目标列的LazyFrame，全程保持延迟计算状态，让Polars的查询规划器把预处理逻辑和原LazyFrame的后续操作完全合并优化，这是效率最高的方式。

具体修改步骤 #

1. 去掉.collect().to_series()，让函数返回pl.LazyFrame类型
2. 给最终生成的列指定明确别名，方便后续合并
3. 调用时直接用.with_columns()引用返回的LazyFrame即可

修改后的代码示例 #

def preprocessing(
    df: pl.LazyFrame | pl.DataFrame,
) -> pl.LazyFrame:
    return (
        df.lazy()
        .with_columns(
            [
                pl.col("some_col")
                .cast(pl.Int64)
                .cast(pl.Utf8)
                .str.pad_start(5, "0")
                .alias("some_col_preprocessed"),
                pl.col("other_col")
                .cast(pl.Int64)
                .cast(pl.Utf8)
                .str.pad_start(7, "0")
                .alias("other_col_preprocessed"),
            ],
        )
        # ... 更多处理步骤
        .select(
            pl.concat_str(
                [
                    pl.col("some_col_preprocessed"),
                    pl.col("other_col_preprocessed"),
                    pl.lit("42133723"),
                ],
            ).alias("final_col")  # 给目标列指定明确别名
        )

# 合并到原LazyFrame
starting_df = starting_df.with_columns(preprocessing(starting_df))

方案优势 #

• 完全保留优化潜力：全程延迟计算，Polars可以自动做谓词下推、列裁剪等优化，处理超大型数据集时性能提升明显
• 无顺序风险：只要预处理是逐行操作，返回的LazyFrame列顺序和原数据完全一致，不用依赖maintain_order或ID列
• 查询计划完全整合：原数据和预处理逻辑会合并成一个查询计划，执行时一次性计算，没有中间数据落地的额外开销

其他方案的弊端对比 #

• ID列关联：需要额外维护ID字段，join操作会增加查询复杂度，还可能因ID重复导致数据膨胀，完全没必要
• 返回Series：触发提前collect不仅丢失优化机会，还可能因内存限制无法处理超大型数据集，顺序问题也难以保障

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