相当于熊猫groupby shift()的极性

您可以使用 over

import polars as pl

df = pl.DataFrame({
    'object': [1, 1, 1, 2, 2],
    'period': [1, 2, 4, 4, 23],
    'value': [24, 67, 89, 5, 23],
})

df.with_column(
    pl.col('value').shift().over('object').alias('prev_value')
)

shape: (5, 4)
┌────────┬────────┬───────┬────────────┐
│ object ┆ period ┆ value ┆ prev_value │
│ ---    ┆ ---    ┆ ---   ┆ ---        │
│ i64    ┆ i64    ┆ i64   ┆ i64        │
╞════════╪════════╪═══════╪════════════╡
│ 1      ┆ 1      ┆ 24    ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 2      ┆ 67    ┆ 24         │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 4      ┆ 89    ┆ 67         │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 4      ┆ 5     ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 23     ┆ 23    ┆ 5          │
└────────┴────────┴───────┴────────────┘

要在多个列上执行此操作,可以在 pl.col 表达式,然后使用前缀/后缀来命名新列。例如:

df.with_columns(
    pl.col(['period', 'value']).shift().over('object').prefix("prev_")
)

shape: (5, 5)
┌────────┬────────┬───────┬─────────────┬────────────┐
│ object ┆ period ┆ value ┆ prev_period ┆ prev_value │
│ ---    ┆ ---    ┆ ---   ┆ ---         ┆ ---        │
│ i64    ┆ i64    ┆ i64   ┆ i64         ┆ i64        │
╞════════╪════════╪═══════╪═════════════╪════════════╡
│ 1      ┆ 1      ┆ 24    ┆ null        ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 2      ┆ 67    ┆ 1           ┆ 24         │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1      ┆ 4      ┆ 89    ┆ 2           ┆ 67         │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 4      ┆ 5     ┆ null        ┆ null       │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2      ┆ 23     ┆ 23    ┆ 4           ┆ 5          │
└────────┴────────┴───────┴─────────────┴────────────┘

结束

让我们使用这些数据。

df = pl.DataFrame(
    {
        "id": [1] * 5 + [2] * 5,
        "date": ["2020-01-01", "2020-01-01", "2020-02-01", "2020-02-01", "2020-02-01"] * 2,
        "value1": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        "value2": [10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
    }
).with_column(pl.col('date').str.strptime(pl.Date))
df

shape: (10, 4)
┌─────┬────────────┬────────┬────────┐
│ id  ┆ date       ┆ value1 ┆ value2 │
│ --- ┆ ---        ┆ ---    ┆ ---    │
│ i64 ┆ date       ┆ i64    ┆ i64    │
╞═════╪════════════╪════════╪════════╡
│ 1   ┆ 2020-01-01 ┆ 1      ┆ 10     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2020-01-01 ┆ 2      ┆ 20     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2020-02-01 ┆ 3      ┆ 30     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2020-02-01 ┆ 4      ┆ 40     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2020-02-01 ┆ 5      ┆ 50     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-01-01 ┆ 6      ┆ 60     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-01-01 ┆ 7      ┆ 70     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-02-01 ┆ 8      ┆ 80     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-02-01 ┆ 9      ┆ 90     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-02-01 ┆ 10     ┆ 100    │
└─────┴────────────┴────────┴────────┘

结束 表达式(以及我们的 表达式)。Polars将并行运行它们。

df.with_columns([
    pl.col(["value1", "value2"]).shift().over(['id','date']).prefix("prev_"),
    pl.col(["value1", "value2"]).diff().over(['id','date']).suffix("_diff"),
])

shape: (10, 8)
┌─────┬────────────┬────────┬────────┬─────────────┬─────────────┬─────────────┬─────────────┐
│ id  ┆ date       ┆ value1 ┆ value2 ┆ prev_value1 ┆ prev_value2 ┆ value1_diff ┆ value2_diff │
│ --- ┆ ---        ┆ ---    ┆ ---    ┆ ---         ┆ ---         ┆ ---         ┆ ---         │
│ i64 ┆ date       ┆ i64    ┆ i64    ┆ i64         ┆ i64         ┆ i64         ┆ i64         │
╞═════╪════════════╪════════╪════════╪═════════════╪═════════════╪═════════════╪═════════════╡
│ 1   ┆ 2020-01-01 ┆ 1      ┆ 10     ┆ null        ┆ null        ┆ null        ┆ null        │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2020-01-01 ┆ 2      ┆ 20     ┆ 1           ┆ 10          ┆ 1           ┆ 10          │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2020-02-01 ┆ 3      ┆ 30     ┆ null        ┆ null        ┆ null        ┆ null        │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2020-02-01 ┆ 4      ┆ 40     ┆ 3           ┆ 30          ┆ 1           ┆ 10          │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2020-02-01 ┆ 5      ┆ 50     ┆ 4           ┆ 40          ┆ 1           ┆ 10          │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-01-01 ┆ 6      ┆ 60     ┆ null        ┆ null        ┆ null        ┆ null        │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-01-01 ┆ 7      ┆ 70     ┆ 6           ┆ 60          ┆ 1           ┆ 10          │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-02-01 ┆ 8      ┆ 80     ┆ null        ┆ null        ┆ null        ┆ null        │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-02-01 ┆ 9      ┆ 90     ┆ 8           ┆ 80          ┆ 1           ┆ 10          │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2020-02-01 ┆ 10     ┆ 100    ┆ 9           ┆ 90          ┆ 1           ┆ 10          │
└─────┴────────────┴────────┴────────┴─────────────┴─────────────┴─────────────┴─────────────┘