根据其他数据帧进行过滤和聚合

Question

说我有

df1 = pl.DataFrame({'start': [1., 2., 4.], 'end': [2., 4., 6.]})
df2 = pl.DataFrame({'idx': [1., 1.7, 2.3, 2.5, 3., 4.], 'values': [3, 1, 4, 2, 3, 5]})

它们看起来像这样：

In [8]: df1
Out[8]:
shape: (3, 2)
┌───────┬─────┐
│ start ┆ end │
│ ---   ┆ --- │
│ f64   ┆ f64 │
╞═══════╪═════╡
│ 1.0   ┆ 2.0 │
│ 2.0   ┆ 4.0 │
│ 4.0   ┆ 6.0 │
└───────┴─────┘

In [9]: df2
Out[9]:
shape: (6, 2)
┌─────┬────────┐
│ idx ┆ values │
│ --- ┆ ---    │
│ f64 ┆ i64    │
╞═════╪════════╡
│ 1.0 ┆ 3      │
│ 1.7 ┆ 1      │
│ 2.3 ┆ 4      │
│ 2.5 ┆ 2      │
│ 3.0 ┆ 3      │
│ 4.0 ┆ 5      │
└─────┴────────┘

我希望得到这样的结果：

In [6]: expected = pl.DataFrame({
   ...:     'start': [1., 2., 4.],
   ...:     'end': [2., 4.5, 6.],
   ...:     'sum_values': [4, 9, 5]
   ...: })

In [7]: expected
Out[7]:
shape: (3, 3)
┌───────┬─────┬────────────┐
│ start ┆ end ┆ sum_values │
│ ---   ┆ --- ┆ ---        │
│ f64   ┆ f64 ┆ i64        │
╞═══════╪═════╪════════════╡
│ 1.0   ┆ 2.0 ┆ 4          │
│ 2.0   ┆ 4.5 ┆ 9          │
│ 4.0   ┆ 6.0 ┆ 5          │
└───────┴─────┴────────────┘

这是我想出的一种低效方法，使用

map_rows

:

(
    df1.with_columns(
        df1.map_rows(
            lambda row: df2.filter(
                pl.col("idx").is_between(row[0], row[1], closed="left")
            )["values"].sum()
        )["map"].alias("sum_values")
    )
)

它提供了正确的输出，但由于它使用

map_rows

和 Python lambda 函数，因此它的性能并不理想。

有没有办法使用 Polars 原生表达式 API 来编写这个？

Answer 1

更新： “非等值” 连接已添加到 Polars 中。

• .join_where()

(df1
  .join_where(df2,
     pl.col.idx >= pl.col.start,
     pl.col.idx < pl.col.end
  )
  .group_by("start", "end")
  .agg(pl.col.values.sum())
)

---

原答案

我不确定除了交叉连接之外是否还有其他方法：

(df1.join(df2, how="cross")
    .filter(pl.col.idx.is_between("start", "end", closed="left"))
    .group_by("start", "end")
    .agg(pl.col.values.sum())
)

shape: (3, 3)
┌───────┬─────┬────────┐
│ start ┆ end ┆ values │
│ ---   ┆ --- ┆ ---    │
│ f64   ┆ f64 ┆ i64    │
╞═══════╪═════╪════════╡
│ 1.0   ┆ 2.0 ┆ 4      │
│ 4.0   ┆ 6.0 ┆ 5      │
│ 2.0   ┆ 4.0 ┆ 9      │
└───────┴─────┴────────┘