Polars LazyFrame分组聚合后关联原表的问题咨询

我有一个包含time、symbols和mid_price字段的Polars LazyFrame，示例数据如下：

time                symbols             mid_price
datetime[ns]        str                 f64
2024-03-01 00:01:00 "PERP_SOL_USDT@…    126.1575
2024-03-01 00:01:00 "PERP_WAVES_USD…    2.71235
2024-03-01 00:01:00 "SOL_USDT@BINAN…    126.005
2024-03-01 00:01:00 "WAVES_USDT@BIN…    2.7085
2024-03-01 00:02:00 "PERP_SOL_USDT@…    126.3825

我想要按symbols分组，对mid_price做时间维度的差分，于是写了如下聚合代码：

aggs = (
    df
        .group_by('symbols')
        .agg([
            pl.col('mid_price').diff(1).alias("change"),
        ])
)

得到的结果是每个唯一symbols对应一个change列表：

symbols             change
str                 list[f64]
"SOL_USDT@BINAN…    [null, 0.25, … -0.55]
"PERP_SOL_USDT@…    [null, 0.225, … -0.605]
"WAVES_USDT@BIN…    [null, -0.002, … -0.001]
"PERP_WAVES_USD…    [null, -0.00255, … 0.0001]

之后我把这个结果关联回原DataFrame：

df = df.join(
    aggs,
    on='symbols',
    how='left',
)

但结果里每行都带上了完整的change列表，而不是对应行的单个值：

time                symbols             mid_price   change
datetime[ns]        str                 f64         list[f64]
2024-03-01 00:01:00 "PERP_SOL_USDT@…    126.1575    [null, 0.225, … -0.605]
2024-03-01 00:01:00 "PERP_WAVES_USD…    2.71235     [null, -0.00255, … 0.0001]
2024-03-01 00:01:00 "SOL_USDT@BINAN…    126.005     [null, 0.25, … -0.55]
2024-03-01 00:01:00 "WAVES_USDT@BIN…    2.7085      [null, -0.002, … -0.001]
2024-03-01 00:02:00 "PERP_SOL_USDT@…    126.3825    [null, 0.225, … -0.605]

我有两个问题：

1. 如何在关联回原DataFrame时，将分组返回的列表拆分为对应行的单个值？
2. 这种先分组聚合再关联的方式，是给原DataFrame添加新列的推荐做法吗？

问题1：拆分列表到对应行的正确姿势 #

你不需要先分组聚合再关联，Polars提供了更直接的方式——窗口函数（over子句），可以直接在原表上按分组计算差分并生成对应行的单个值，代码如下：

df = df.with_columns(
    pl.col('mid_price').diff().over('symbols').alias('change')
)

这样处理后，每个symbols分组内的mid_price会按时间顺序计算差分，结果直接对应到原表的每一行，不会生成列表。

如果一定要基于你现有的aggs结果来处理，需要确保每个symbols的列表长度和原表中该分组的行数一致，然后通过添加分组内索引来匹配，但这种方法繁琐且容易出错，不推荐：

# 不推荐的写法，仅作示例
aggs_exploded = aggs.explode('change')
# 给原表和拆分后的聚合表添加分组内的行号
df = df.with_row_index().group_by('symbols').with_row_index("group_idx")
aggs_exploded = aggs_exploded.group_by('symbols').with_row_index("group_idx")
# 按symbols和分组内索引关联
df = df.join(aggs_exploded, on=['symbols', 'group_idx'], how='left').drop('row_index', 'group_idx')

问题2：先分组聚合再关联不是推荐做法 #

这种做法不仅代码繁琐，还会带来额外的性能开销（比如生成列表、执行关联操作），完全没必要。Polars中给原表添加分组计算的新列，首选**transform方法或者窗口函数（over子句）**，两者都能保留原表的行数，直接将分组计算的结果映射到每一行，代码更简洁，执行效率也更高。

比如用transform的写法：

df = df.group_by('symbols').transform(
    pl.col('mid_price').diff().alias('change')
)

效果和窗口函数完全一致，两种写法可以根据个人习惯选择。

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