在Polars多索引LazyFrame中添加衍生数据的高效方法

你现在遇到的这个问题我太懂了——来回切换Lazy和Eager模式不仅麻烦，还容易拖慢性能，尤其是数据量上去之后。其实完全不用这么折腾，咱们可以利用Polars的原生能力，全程在Lazy模式下完成衍生数据的添加，不用碰pivot/unpivot那一套。

下面结合你的场景，给你两种实用的解决方案：

方法1：用窗口函数直接按索引层级计算 #

如果你的衍生逻辑是基于某个索引分组的（比如按name分组，对不同measure的数值做运算），窗口函数就是最优解，完全不需要重塑数据结构：

import numpy as np
import polars as pl

def get_raw_data() -> pl.LazyFrame:
    names = np.array(['A', 'B', 'C'])
    measures = np.array(['height', 'width'])
    repeats = np.array([3, 3, 2])
    
    # 生成模拟多索引结构数据
    name_col = np.repeat(names, repeats * len(measures))
    measure_col = np.tile(np.repeat(measures, repeats), len(names))
    value_col = np.random.randn(len(name_col))
    
    return pl.LazyFrame({
        "name": name_col,
        "measure": measure_col,
        "value": value_col
    })

# 示例：给每个name添加height和width的差值衍生列
lazy_df = get_raw_data()

result = lazy_df.with_columns(
    # 窗口函数获取每个name下height的数值
    pl.col("value").filter(pl.col("measure") == "height")
      .over("name").alias("height_val"),
    # 窗口函数获取每个name下width的数值
    pl.col("value").filter(pl.col("measure") == "width")
      .over("name").alias("width_val")
).with_columns(
    # 计算衍生列
    (pl.col("height_val") - pl.col("width_val")).alias("height_width_diff")
)

# 最后再触发计算，全程Lazy模式
print(result.collect())

这种方式让Polars的查询优化器全程参与，能自动合并计算步骤，比手动pivot/unpivot高效得多。

方法2：用group_by+agg模拟pivot（需重塑结构时） #

如果你的计算确实需要先转成宽表再处理，也可以在Lazy模式下用group_by+agg模拟pivot的效果，处理完后用melt转回长表，全程不用切Eager：

# 模拟pivot生成宽表
wide_lazy = lazy_df.group_by("name").agg(
    pl.col("value").filter(pl.col("measure") == "height").alias("height"),
    pl.col("value").filter(pl.col("measure") == "width").alias("width")
).with_columns(
    # 添加衍生列：height和width的和
    (pl.col("height") + pl.col("width")).alias("sum_dim")
)

# 模拟unpivot转回长表
long_lazy = wide_lazy.melt(
    id_vars=["name", "sum_dim"],
    value_vars=["height", "width"],
    variable_name="measure",
    value_name="value"
)

print(long_lazy.collect())

为什么这两种方法更优？ #

• 避免了模式切换的开销：每次从Lazy切到Eager，Polars都要把数据加载到内存并执行计算，全程Lazy的话，优化器可以整体规划执行路径，减少冗余计算。
• 窗口函数和分组操作都是Polars深度优化的功能，在Lazy模式下能发挥最大性能，比手动pivot/unpivot的效率高不少。

总之，只要你的衍生计算逻辑能通过窗口函数或分组聚合实现，就完全不用来回切换模式，全程在Lazy模式下就能搞定，既简洁又高效。

备注：内容来源于stack exchange，提问作者Olibarer