在 agg 中使用缩减 ufunc

Question

如何在聚合上下文中使用缩减为标量的 ufunc？例如，使用 numpy.trapz:

总结 table

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.groupby('id').agg(pl.map(['t', 'y'], np.trapz))
# Segmentation fault (core dumped)

Answer 1

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

在 groupby 上下文中使用 apply（而不是 map）。

在这种情况下，numpy trapz 函数只接受一个位置参数 (y)

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

因此，我们需要在调用中明确指定 x 关键字参数。（我还假设你的意思是你的 y 列被映射为 y 参数，你的 t 列被映射为调用中的 x 参数麻木。）

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

一个额外的问题，numpy return 是一个 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 return 成为一个 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  │
└─────┴──────┘

在 agg 中使用缩减 ufunc

Using a reduction ufunc in agg

python-polars