如何借助Polars Expressions优化数据分析管道性能？

核心需求 #

• 用Polars Expressions替代多步数据对象转换，减少中间表/冗余列
• 实现分组聚合结果的复用，避免重复计算
• 优化last_cum_sum_min列生成逻辑及recent_price_data的调用
• 减少流程启停损耗，提升管道性能，最终输出指定格式的价格数据

现有流程痛点 #

• 多步拆分的DataFrame转换（多次创建中间对象、调用collect()）导致性能损耗
• 中间列冗余，增加内存占用
• 聚合结果与后续计算的关联逻辑拆分，未利用Expr链式调用的优势

针对性优化方案 #

1. 合并多步转换为链式Expr调用 #

将原流程中拆分的「去趋势收益率标准差计算→RS聚合→数据关联→价格区间计算→价格调整」合并为单链式操作，避免多次创建DataFrame对象：

# 优化前（拆分多步）
df1 = df.with_columns(pl.col("return").std().over("group").alias("std_dev"))
df2 = df1.group_by("group").agg(pl.col("rs").sum().alias("rs_sum"))
df3 = df.join(df2, on="group")

# 优化后（链式Lazy调用）
result = df.lazy()
.with_columns(
    # 分组内计算去趋势收益率标准差
    pl.col("return").std().over("group").alias("std_dev")
)
.group_by("group")
.agg(
    pl.col("rs").sum().alias("rs_sum"),
    # 一步生成last_cum_sum_min
    pl.col("value").cumsum().min().last().alias("last_cum_sum_min"),
    # 直接提取分组最新价格，无需单独创建recent_price_data
    pl.col("price").last().alias("recent_price")
)
.with_columns(
    # 计算价格区间，直接复用聚合列
    (pl.col("rs_sum") * pl.col("std_dev")).alias("price_range"),
    # 价格调整逻辑转为Expr（替代单独函数）
    (pl.col("recent_price") * 1.2 + 5).alias("adjusted_price")
)
.collect()

2. 优化last_cum_sum_min列生成 #

避免分步计算cum_sum、取min、取last的冗余操作，直接用Polars函数链式组合：

# 错误分步写法（易出错且低效）
df = df.with_columns(pl.col("value").cumsum().alias("cum_sum"))
df = df.group_by("group").agg(pl.col("cum_sum").min().alias("cum_sum_min"))
df = df.with_columns(pl.col("cum_sum_min").last().alias("last_cum_sum_min"))

# 优化后Expr写法（一步生成）
pl.col("value").cumsum().min().over("group").last().alias("last_cum_sum_min")

3. 复用recent_price_data避免重复关联 #

无需单独提取最新价格表再做join，直接用窗口函数在原表中生成最新价格列：

# 原逻辑（冗余join）
recent_price_data = df.group_by("group").agg(pl.col("price").last().alias("recent_price"))
df = df.join(recent_price_data, on="group")

# 优化后（直接生成列）
df = df.with_columns(pl.col("price").last().over("group").alias("recent_price"))

4. 调整_modify_mandelbrot_prices为Expr兼容逻辑 #

尽量用Polars内置向量化函数替代Python自定义函数，若必须保留自定义逻辑，用map_elements但优先向量化：

# 原自定义函数
def _modify_mandelbrot_prices(price):
    return price * 1.2 + 5

# 优化为原生Expr（推荐）
pl.col("recent_price") * 1.2 + 5.alias("adjusted_price")

# 保留自定义逻辑的Expr写法
pl.col("recent_price").map_elements(_modify_mandelbrot_prices).alias("adjusted_price")

5. 减少启停损耗：延迟collect()调用 #

全程用LazyFrame操作，直到最终需要结果时再调用collect()，让Polars查询优化器自动合并操作：

# 优化前（多次collect）
df_lazy = df.lazy()
df1 = df_lazy.with_columns(...).collect()
df2 = df1.group_by(...).collect()

# 优化后（单次collect）
result = df.lazy().with_columns(...).group_by(...).agg(...).with_columns(...).collect()

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