python - 如何更好地使用 Polars 中的 apply？

Question

我有一个极坐标数据框，如下所示。

import polars as pl

df = pl.DataFrame(
    {
        "a": [1, 4, 3, 2, 8, 4, 5, 6],
        "b": [2, 3, 1, 3, 9, 7, 6, 8],
        "c": [1, 1, 1, 1, 2, 2, 2, 2],
    }
)

我的任务是

1. 按列“c”分组
2. 对于每个组，检查“a”列中的所有数字是否都小于“b”列中的相应值。
   • 如果是这样，只需在 groupby 上下文中返回与“a”相同的列。
   • 否则，应用一个名为“convert”的第三方函数，它接受两个 numpy 数组并返回一个具有相同大小的 numpy 数组，所以在我的例子中，我可以首先将列“a "和 "b"到 numpy 数组，并将它们作为输入提供给 "convert"。最后，在 groupby 上下文中返回从 "convert" 返回的数组(可能需要在返回之前将其转换为极坐标系)。

因此，对于上面的示例，我想要的输出如下(为了更好地说明，在 groupby 之后进行了分解)。

shape: (8, 2)
┌─────┬─────┐
│ c   ┆ a   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 1   │
│ 1   ┆ 3   │
│ 1   ┆ 1   │
│ 1   ┆ 2   │
│ 2   ┆ 8   │
│ 2   ┆ 4   │
│ 2   ┆ 5   │
│ 2   ┆ 6   │
└─────┴─────┘

根据假设，

>>> import numpy as np
>>> convert(np.array([1, 4, 3, 2]), np.array([2, 3, 1, 3]))
np.array([1, 3, 1, 2])

# [1, 4, 3, 2] is from column a of df when column c is 1, and [2, 3, 1, 3] comes from column b of df when column c is 1.
# I have to apply my custom python function 'convert' for the c == 1 group, because not all values in a are smaller than those in b according to the task description above.

我的问题是我应该如何以高性能或极性惯用的方式实现这个逻辑，而不牺牲运行 Rust 代码和并行化所获得的如此多的速度？

我之所以问这个问题，是因为根据我的理解，将 apply 与自定义 python 函数一起使用会减慢程序速度，但就我而言，在某些情况下，我不需要求助于第三方函数来帮助。那么，有什么办法可以让我获得最好的世界呢？ (对于不需要第三方功能的场景，充分利用 Polars 的优势，仅在必要时应用第三方功能)。

Answer 1

听起来您想找到匹配的组:

(
   df
   .with_row_count()
   .filter(
      (pl.col("a") >= pl.col("b"))
      .any()
      .over("c"))
)

shape: (4, 4)
┌────────┬─────┬─────┬─────┐
│ row_nr | a   | b   | c   │
│ ---    | --- | --- | --- │
│ u32    | i64 | i64 | i64 │
╞════════╪═════╪═════╪═════╡
│ 0      | 1   | 2   | 1   │
│ 1      | 4   | 3   | 1   │
│ 2      | 3   | 1   | 1   │
│ 3      | 2   | 3   | 1   │
└────────┴─────┴─────┴─────┘

并将您的自定义函数应用于每个组。

(      
   df
   .with_row_count()
   .filter(
      (pl.col("a") >= pl.col("b"))
      .any()
      .over("c"))
   .select(
      pl.col("row_nr"),
      pl.apply(
         ["a", "b"], # np.minimum is just for example purposes
         lambda s: np.minimum(s[0], s[1]))
      .over("c"))
)

shape: (4, 2)
┌────────┬─────┐
│ row_nr | a   │
│ ---    | --- │
│ u32    | i64 │
╞════════╪═════╡
│ 0      | 1   │
│ 1      | 3   │
│ 2      | 1   │
│ 3      | 2   │
└────────┴─────┘

(注意:How to Write Poisson CDF as Python Polars Expression 中可能有一些关于 scipy/numpy ufunc 的有用信息，并可能避免使用 .apply())

然后您可以.join()将结果返回到原始数据。

(
   df
   .with_row_count()
   .join(
      df
      .with_row_count()
      .filter(
         (pl.col("a") >= pl.col("b"))
         .any()
         .over("c"))
      .select(
         pl.col("row_nr"),
         pl.apply(
            ["a", "b"],
            lambda s: np.minimum(s[0], s[1]))
         .over("c")),
      on="row_nr",
      how="left")
)

shape: (8, 5)
┌────────┬─────┬─────┬─────┬─────────┐
│ row_nr | a   | b   | c   | a_right │
│ ---    | --- | --- | --- | ---     │
│ u32    | i64 | i64 | i64 | i64     │
╞════════╪═════╪═════╪═════╪═════════╡
│ 0      | 1   | 2   | 1   | 1       │
│ 1      | 4   | 3   | 1   | 3       │
│ 2      | 3   | 1   | 1   | 1       │
│ 3      | 2   | 3   | 1   | 2       │
│ 4      | 8   | 9   | 2   | null    │
│ 5      | 4   | 7   | 2   | null    │
│ 6      | 5   | 6   | 2   | null    │
│ 7      | 6   | 8   | 2   | null    │
└────────┴─────┴─────┴─────┴─────────┘

然后您可以填写空值。

.with_columns(
   pl.col("a_right").fill_null(pl.col("a")))