就说我有
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 │
└───────┴─────┴────────────┘
以下是我想出的一种低效方法,使用apply
:
(
df1.with_columns(
df1.apply(
lambda row: df2.filter(
pl.col("idx").is_between(row[0], row[1], closed="left")
)["values"].sum()
)["apply"].alias("sum_values")
)
)
它给出了正确的输出,但是因为它使用了apply
和一个Python lambda函数,所以它的性能不是很好.
有没有办法使用Polars原生表达式API来编写?