Polars - 如何有效地在其他行上运行计算

Question

我正在使用 Polars DataFrame，需要使用其他行的值对每一行执行计算。目前，我正在使用map_elements方法，但效率不高。

在以下示例中，我向 DataFrame 添加两个新列：

1. sum_lower：所有小于当前元素的元素之和。
2. max_other：DataFrame 中的最大值，不包括当前元素。

这是我当前的实现：

import polars as pl

COL_VALUE = "value"

def fun_sum_lower(current_row, df):
    tmp_df = df.filter(pl.col(COL_VALUE) < current_row[COL_VALUE])
    sum_lower = tmp_df.select(pl.sum(COL_VALUE)).item()
    return sum_lower

def fun_max_other(current_row, df):
    tmp_df = df.filter(pl.col(COL_VALUE) != current_row[COL_VALUE])
    max_other = tmp_df.select(pl.col(COL_VALUE)).max().item()
    return max_other

if __name__ == '__main__':
    df = pl.DataFrame({COL_VALUE: [3, 7, 1, 9, 4]})

    df = df.with_columns(
        pl.struct([COL_VALUE])
        .map_elements(lambda row: fun_sum_lower(row, df), return_dtype=pl.Int64)
        .alias("sum_lower")
    )

    df = df.with_columns(
        pl.struct([COL_VALUE])
        .map_elements(lambda row: fun_max_other(row, df), return_dtype=pl.Int64)
        .alias("max_other")
    )

    print(df)

上述代码的输出为：

shape: (5, 3)
┌───────┬───────────┬───────────┐
│ value ┆ sum_lower ┆ max_other │
│ ---   ┆ ---       ┆ ---       │
│ i64   ┆ i64       ┆ i64       │
╞═══════╪═══════════╪═══════════╡
│ 3     ┆ 1         ┆ 9         │
│ 7     ┆ 8         ┆ 9         │
│ 1     ┆ 0         ┆ 9         │
│ 9     ┆ 15        ┆ 7         │
│ 4     ┆ 4         ┆ 9         │
└───────┴───────────┴───────────┘

虽然此代码有效，但由于使用了 lambda 和逐行操作，效率不高。

问题： 在 Polars 中是否有更有效的方法来实现此目的，而不使用 lambda、迭代行或运行 Python 代码？

我还尝试使用 Polars 方法：

cum_sum

、

group_by_dynamic

和

rolling

，但我认为这些方法不能用于此任务。

Answer 1

您可以使用

.join_where()

df_sum = (
   df.join_where(df, pl.col.value > pl.col.value_right)
     .group_by("value")
     .sum()
)

df_max = (
   df.join_where(df, pl.col.value != pl.col.value_right)
     .group_by("value")
     .max()
)

(df.select("value")
   .join(df_sum, on="value", how="left")
   .join(df_max, on="value", how="left")
)

shape: (5, 3)
┌───────┬─────────────┬───────────────────┐
│ value ┆ value_right ┆ value_right_right │
│ ---   ┆ ---         ┆ ---               │
│ i64   ┆ i64         ┆ i64               │
╞═══════╪═════════════╪═══════════════════╡
│ 3     ┆ 1           ┆ 9                 │
│ 7     ┆ 8           ┆ 9                 │
│ 1     ┆ null        ┆ 9                 │
│ 9     ┆ 15          ┆ 7                 │
│ 4     ┆ 4           ┆ 9                 │
└───────┴─────────────┴───────────────────┘