Polars：在同一分组上“执行”许多操作

Question

显示一个玩具示例，K=2但问题主要与高g基数和相关K>>1：

\n

df = pl.DataFrame(dict(\n    g=[1, 2, 1, 2, 1, 2],\n    v=[1, 2, 3, 4, 5, 6],\n))\n\nK = 2\n\ndf.with_columns((col.v.shift(k+1).over(\'g\').alias(f\'s{k}\') for k in range(K)))\n

\n

\xe2\x95\xad\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x95\xae\n\xe2\x94\x82   g \xe2\x94\x86   v \xe2\x94\x86   s0 \xe2\x94\x86   s1 \xe2\x94\x82\n\xe2\x94\x82 i64 \xe2\x94\x86 i64 \xe2\x94\x86  i64 \xe2\x94\x86  i64 \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82   1 \xe2\x94\x86   1 \xe2\x94\x86 null \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82   2 \xe2\x94\x86   2 \xe2\x94\x86 null \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82   1 \xe2\x94\x86   3 \xe2\x94\x86    1 \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82   2 \xe2\x94\x86   4 \xe2\x94\x86    2 \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82   1 \xe2\x94\x86   5 \xe2\x94\x86    3 \xe2\x94\x86    1 \xe2\x94\x82\n\xe2\x94\x82   2 \xe2\x94\x86   6 \xe2\x94\x86    4 \xe2\x94\x86    2 \xe2\x94\x82\n\xe2\x95\xb0\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x95\xaf\n

\n

如何确保分组依据g只完成一次？

\n

Polars 似乎没有在查询计划中对此进行优化。

\n

我希望它的运行速度与以下速度一样快：

\n

\xe2\x95\xad\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x95\xae\n\xe2\x94\x82   g \xe2\x94\x86   v \xe2\x94\x86   s0 \xe2\x94\x86   s1 \xe2\x94\x82\n\xe2\x94\x82 i64 \xe2\x94\x86 i64 \xe2\x94\x86  i64 \xe2\x94\x86  i64 \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82   1 \xe2\x94\x86   1 \xe2\x94\x86 null \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82   2 \xe2\x94\x86   2 \xe2\x94\x86 null \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82   1 \xe2\x94\x86   3 \xe2\x94\x86    1 \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82   2 \xe2\x94\x86   4 \xe2\x94\x86    2 \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82   1 \xe2\x94\x86   5 \xe2\x94\x86    3 \xe2\x94\x86    1 \xe2\x94\x82\n\xe2\x94\x82   2 \xe2\x94\x86   6 \xe2\x94\x86    4 \xe2\x94\x86    2 \xe2\x94\x82\n\xe2\x95\xb0\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x95\xaf\n

\n

Answer 1

Polars 缓存窗口表达式。虽然您可能在查询计划中看不到这一点，但over分组只完成一次。

\n

尽管如此，您的查询的运行速度仍不会像您的查询over一样快。group_by这是因为over必须将结果添加回原始数据框中。

\n

更公平的比较是下面的查询，它将与over查询的结果匹配。如您所见，join需要额外的。

\n

import polars as pl\n\ndf = pl.DataFrame(\n    {\n        "g": [1, 2, 1, 2, 1, 2],\n        "v": [1, 2, 3, 4, 5, 6],\n    }\n)\nK = 2\n\ndf_shift = (\n    df.group_by("g")\n    .agg([pl.col("v")] + [pl.col("v").shift(k + 1).alias(f"s{k}") for k in range(K)])\n    .explode(["v"] + [f"s{k}" for k in range(K)])\n)\nresult = df.join(df_shift, on=["g", "v"], how="left")\nprint(result)\n

\n

shape: (6, 4)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 g   \xe2\x94\x86 v   \xe2\x94\x86 s0   \xe2\x94\x86 s1   \xe2\x94\x82\n\xe2\x94\x82 --- \xe2\x94\x86 --- \xe2\x94\x86 ---  \xe2\x94\x86 ---  \xe2\x94\x82\n\xe2\x94\x82 i64 \xe2\x94\x86 i64 \xe2\x94\x86 i64  \xe2\x94\x86 i64  \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 1   \xe2\x94\x86 1   \xe2\x94\x86 null \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82 2   \xe2\x94\x86 2   \xe2\x94\x86 null \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82 1   \xe2\x94\x86 3   \xe2\x94\x86 1    \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82 2   \xe2\x94\x86 4   \xe2\x94\x86 2    \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82 1   \xe2\x94\x86 5   \xe2\x94\x86 3    \xe2\x94\x86 1    \xe2\x94\x82\n\xe2\x94\x82 2   \xe2\x94\x86 6   \xe2\x94\x86 4    \xe2\x94\x86 2    \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n