使用极坐标查找下一个值 >= 实际值加上 50%

Question

我有以下数据框：

 df = pl.DataFrame({
        "Column A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        "Column B": [2, 3, 1, 4, 1, 3, 3, 2, 1, 0]
    })

我想创建一个新列 C，用于保存当前行的 B 值与 B 列中大于或等于 B + 50% 的下一个值之间的距离（以行为单位）。

最终结果应该是这样的：

 df = pl.DataFrame({
        "Column A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        "Column B": [2, 3, 1, 4, 1, 7, 3, 2, 12, 0],
        "Column C": [1, 4, 1, 2, 1, 3, 2, 1, None, None]
    })

如何使用 Polars 有效地实现这一目标，特别是因为我正在使用大型 DataFrame？

Answer 1

好吧，首先我应该说 - 这个看起来需要在多个列上加入不等式，并且从我发现纯极坐标不太适合它。也许可以用

join_asof

做到这一点，但我无法把它做得漂亮。

我可能会使用

duckdb

与极坐标集成 来实现结果：

import duckdb

duckdb.sql("""
    select
        d."Column A",
        d."Column B",
        (
            select tt."Column A"
            from df as tt
            where tt."Column A" > d."Column A" and tt."Column B" >= d."Column B" * 1.5
            order by tt."Column A" asc
            limit 1
        ) - d."Column A" as "Column C"
    from df as d
""").pl()

┌──────────┬──────────┬──────────┐
│ Column A ┆ Column B ┆ Column C │
│ ---      ┆ ---      ┆ ---      │
│ i64      ┆ i64      ┆ i64      │
╞══════════╪══════════╪══════════╡
│ 1        ┆ 2        ┆ 1        │
│ 2        ┆ 3        ┆ 4        │
│ 3        ┆ 1        ┆ 1        │
│ 4        ┆ 4        ┆ 2        │
│ 5        ┆ 1        ┆ 1        │
│ 6        ┆ 7        ┆ 3        │
│ 7        ┆ 3        ┆ 2        │
│ 8        ┆ 2        ┆ 1        │
│ 9        ┆ 12       ┆ null     │
│ 10       ┆ 0        ┆ null     │
└──────────┴──────────┴──────────┘