基于另一列的分组聚合向Polars LazyFrame添加列

Question

我有一个包含 时间、符号 和 中间价格 的懒惰框架（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

我想在时间维度上进行一些聚合操作（也就是说：按 符号 分组）：

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

我得到了每个唯一的 符号 值对应的 列表：

>>> aggs.head().collect()

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]

现在我想把这个结果重新合并到我最初的数据框中：

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

这样每一行都会得到完整的 变化 列表，而不是各自的值。

>>> df.head().collect()

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. 我该如何在将这些从 分组 返回的列表合并回原始数据框时，进行 解压/展开？
2. 这是从 分组 中添加新列到原始数据框的推荐方法吗（也就是说：先 分组 然后 合并）？

Answer 1

听起来你并不想把数据合并成一个值（每个符号只得到一个值），而是想要独立地计算每个符号的"变化"。

在polars中，这种行为类似于PostgreSQL中的窗口函数，可以通过pl.Expr.over来实现。

df.with_columns(
    pl.col("mid_price").diff(1).over("symbol").alias("change")
)

在一些示例数据上，结果如下所示。

import polars as pl
import numpy as np
import datetime

df = pl.DataFrame({
    "symbol": ["A"] * 3 + ["B"] * 3 + ["C"] * 3,
    "time": [datetime.datetime(2024, 3, 1, hour) for hour in range(3)] * 3,
    "mid_price": np.random.randn(9),
})

df.with_columns(
    pl.col("mid_price").diff(1).over("symbol").alias("change")
)

shape: (9, 4)
┌────────┬─────────────────────┬───────────┬───────────┐
│ symbol ┆ time                ┆ mid_price ┆ change    │
│ ---    ┆ ---                 ┆ ---       ┆ ---       │
│ str    ┆ datetime[μs]        ┆ f64       ┆ f64       │
╞════════╪═════════════════════╪═══════════╪═══════════╡
│ A      ┆ 2024-03-01 00:00:00 ┆ -0.349863 ┆ null      │
│ A      ┆ 2024-03-01 01:00:00 ┆ 0.093732  ┆ 0.443595  │
│ A      ┆ 2024-03-01 02:00:00 ┆ -1.262064 ┆ -1.355796 │
│ B      ┆ 2024-03-01 00:00:00 ┆ 1.953929  ┆ null      │
│ B      ┆ 2024-03-01 01:00:00 ┆ 0.637582  ┆ -1.316348 │
│ B      ┆ 2024-03-01 02:00:00 ┆ 1.009401  ┆ 0.37182   │
│ C      ┆ 2024-03-01 00:00:00 ┆ 0.75864   ┆ null      │
│ C      ┆ 2024-03-01 01:00:00 ┆ -0.866227 ┆ -1.624867 │
│ C      ┆ 2024-03-01 02:00:00 ┆ -0.674938 ┆ 0.191289  │
└────────┴─────────────────────┴───────────┴───────────┘