在agg中使用归约ufunc

Question

如何在聚合上下文中使用 ufunc 简化为标量？例如，使用

numpy.trapz

:

总结表格

import polars as pl
import numpy as np

df = pl.DataFrame(dict(id=[0, 0, 0, 1, 1, 1], t=[2, 4, 5, 10, 11, 14], y=[0, 1, 1, 2, 3, 4]))
df.group_by('id').agg(pl.map_groups(['t', 'y'], np.trapz))
# Segmentation fault (core dumped)

Answer 1

编辑：从 Polars 0.13.18 开始，apply 方法将 Numpy 数据类型转换为 Polars 数据类型，而不需要 Numpy item 方法。

在 groupby 上下文中使用

apply

（而不是

map

）。

在这种情况下，numpy trapz 函数仅采用一个位置参数 (y)

numpy.trapz(y, x=None, dx=1.0, axis=- 1)

因此，我们需要在调用中显式指定 x 关键字参数。 （我还假设您的意思是将

y

列映射为

y

参数，并将

t

列映射为 numpy 调用中的

x

参数。）

系列“y”和“t”将作为系列列表传递给 lambda 函数，因此我们将使用索引来指示哪一列映射到哪个 numpy 参数。

还有一个问题，numpy 返回类型为

numpy.float64

的值，而不是 Python 浮点数。

type(np.trapz([0, 1, 1], x=[2, 4, 5]))

<class 'numpy.float64'>

目前，Polars 中的

apply

功能不会自动将

numpy.float64

转换为

polars.Float64

。 为了解决这个问题，我们将使用 numpy

item

方法让 numpy 返回 Python 浮点数，而不是

numpy.float64

。

type(np.trapz([0, 1, 1], x=[2, 4, 5]).item())

<class 'float'>

有了这个，我们现在可以编写我们的

apply

声明。

df.groupby("id").agg(
    pl.apply(
        ["y", "t"],
        lambda lst: np.trapz(y=lst[0], x=lst[1]).item()
    )
)

shape: (2, 2)
┌─────┬──────┐
│ id  ┆ y    │
│ --- ┆ ---  │
│ i64 ┆ f64  │
╞═════╪══════╡
│ 1   ┆ 13.0 │
├╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 0   ┆ 2.0  │
└─────┴──────┘