极地：过度表达效率低下

Question

我发现至少对于下面的场景，做

over

比做

group_by/agg

+

explode

慢得多（2~3 倍）。而且，结果是完全一样的。

根据这一发现，我有以下问题：

• 这样的行为符合预期吗？如果是这样，我们是否应该始终执行两步过程（

  group_by/agg

  +

  explode

  ）而不是直接使用

  over

  ？
• 或者，这是否意味着可能还有一些优化的空间

  over

  ？
• 或者，这两种方法之间的性能实际上取决于问题设置，用户应该尝试看看哪种方法更适合？

import time

import numpy as np
import polars as pl
from polars.testing import assert_frame_equal

## setup
rng = np.random.default_rng(1)

nrows = 20_000_000
df = pl.DataFrame(
    dict(
        id=rng.integers(1, 50, nrows),
        id2=rng.integers(1, 500, nrows),
        v=rng.normal(0, 1, nrows),
        v1=rng.normal(0, 1, nrows),
        v2=rng.normal(0, 1, nrows),
        v3=rng.normal(0, 1, nrows),
        v4=rng.normal(0, 1, nrows),
        v5=rng.normal(0, 1, nrows),
        v6=rng.normal(0, 1, nrows),
        v7=rng.normal(0, 1, nrows),
        v8=rng.normal(0, 1, nrows),
        v9=rng.normal(0, 1, nrows),
        v10=rng.normal(0, 1, nrows),
    )
)

## over
start = time.perf_counter()
res = (
    df.lazy()
    .select(
        "id",
        "id2",
        *[
             (pl.col(f"v{i}") - pl.col(f"v{i}").mean().over("id", "id2"))
             / pl.col(f"v{i}").std().over("id", "id2")
             for i in range(1, 11)
        ],
    )
    .collect()
)
print(
    time.perf_counter() - start
)
# 8.541702497983351

## groupby/agg + explode
start = time.perf_counter()
res2 = (
    df.lazy()
    .group_by("id", "id2")
    .agg(
        (pl.col(f"v{i}") - pl.col(f"v{i}").mean()) / pl.col(f"v{i}").std()
        for i in range(1, 11)
    )
    .explode(pl.exclude("id", "id2"))
    .collect()
)
print(
    time.perf_counter() - start
)
# 3.1841439900454134

## compare results
assert_frame_equal(res.sort(pl.all()), res2.sort(pl.all()))

Answer 1

@ritchie46 最近在 Github 上回答了这个问题。

在此发布用于文档目的。

这是预期的，因为默认情况下，Polars 中的窗口函数有一个限制，即它必须按照输入帧的顺序返回数据。这是一项成本高昂的手术。

如果你写

over(.., mapping_strategy=explode)

，窗口函数没有这个约束，但这应该只在选择上下文中使用，因为返回顺序不会被绑定到到输入帧。

可以这样写：

df.select(
        pl.all().over('id', mapping_strategy='explode'),
        pl.col(..).shift(l).over('id', mapping_strategy='explode')
)

• https://github.com/pola-rs/polars/issues/18556#issuecomment-2335192323