主页 / coding / 问题 / 77878728
Accepted
FredMaster
Asked: 2024-01-25 17:15:53 +0800 CST

Polars 数据框：重叠组

  • 772

我目前正在从 pandas“转换”为 Polars，因为我真的很喜欢这个 api。这个问题是我之前的问题的一个更普遍的问题（请参阅此处）

我有以下数据框


# Dummy data
df = pl.DataFrame({
    "Buy_Signal": [1, 0, 1, 0, 1, 0, 0], 
    "Returns": [0.01, 0.02, 0.03, 0.02, 0.01, 0.00, -0.01],
})

我想最终Returns对不同间隔的列条件进行聚合 - 由 column 给出Buy_Signal。在上述情况下，长度是从每个 1 到数据帧的末尾。因此，生成的数据框应如下所示

| group     | Returns   |
|------:    |--------:  |
|   u32     |     f64   |
|     1     |    0.01   |
|     1     |    0.02   |
|     1     |    0.03   |
|     1     |    0.02   |
|     1     |    0.01   |
|     1     |     0.0   |
|     1     |   -0.01   |
|     2     |    0.03   |
|     2     |    0.02   |
|     2     |    0.01   |
|     2     |     0.0   |
|     2     |   -0.01   |
|     3     |    0.01   |
|     3     |     0.0   |
|     3     |   -0.01   |

作为我上一个问题的答案发布的一种方法如下：


# Build overlapping group index
idx = df.select(index=
          pl.when(pl.col("Buy_Signal") == 1)
          .then(pl.int_ranges(pl.int_range(pl.len()), pl.len()  ))
).explode(pl.col("index")).drop_nulls().cast(pl.UInt32)

# Join index with original data
df = (df.with_row_index()
    .join(idx, on="index")
    .with_columns(group = (pl.col("index") == pl.col("index").max())
                .shift().cum_sum().backward_fill() + 1)
    .select(["group", "Returns"])
)
df

问：这个问题还有其他好的解决办法吗？

我所说的“好”是指（i）可读和/或（ii）快速。

我的实际问题包含更大的数据集。

谢谢

python

3 个回答

  1. Best Answer

    为了完整起见，这里有一个不依赖实验功能的替代解决方案。

    (
    df
    .with_columns(
        pl.col("Buy_Signal").cum_sum().alias("group")
    )
    .with_columns(
        pl.int_ranges(pl.col("group").min(), pl.col("group")+1)
    )
    .explode("group")
    .sort("group")
)

    输出。

    shape: (15, 3)
┌────────────┬─────────┬───────┐
│ Buy_Signal ┆ Returns ┆ group │
│ ---        ┆ ---     ┆ ---   │
│ i64        ┆ f64     ┆ i64   │
╞════════════╪═════════╪═══════╡
│ 1          ┆ 0.01    ┆ 1     │
│ 0          ┆ 0.02    ┆ 1     │
│ 1          ┆ 0.03    ┆ 1     │
│ 0          ┆ 0.02    ┆ 1     │
│ 1          ┆ 0.01    ┆ 1     │
│ …          ┆ …       ┆ …     │
│ 0          ┆ 0.0     ┆ 2     │
│ 0          ┆ -0.01   ┆ 2     │
│ 1          ┆ 0.01    ┆ 3     │
│ 0          ┆ 0.0     ┆ 3     │
│ 0          ┆ -0.01   ┆ 3     │
└────────────┴─────────┴───────┘
    • 2

  2. 我认为这是使用 的一个很好的案例cumulative_eval()。

    乍一看，您需要对 DataFrame 进行 2 次传递。首先，您需要计算cum_sum以获得正确的分组：

    df = df.with_columns(
        Holdings = pl.col('Buy_Signal').cum_sum()
    ).group_by("Holdings").agg(
        pl.col(["Returns"])).sort("Holdings", descending=True
    )

print(df)

shape: (3, 2)
┌──────────┬────────────────────┐
│ Holdings ┆ Returns            │
│ ---      ┆ ---                │
│ i64      ┆ list[f64]          │
╞══════════╪════════════════════╡
│ 3        ┆ [0.01, 0.0, -0.01] │
│ 2        ┆ [0.03, 0.02]       │
│ 1        ┆ [0.01, 0.02]       │
└──────────┴────────────────────┘

    请注意，目前我们的 DataFrame 是用 排序的descending=True，这很重要，因为我们希望我们的列表从最大组到最小组聚合。

    现在cumulative_eval()，连同implode()：

    df = df.with_columns(
        pl.col("Returns").cumulative_eval(pl.element().explode().implode())
    )

print(df.explode("Returns"))

shape: (15, 2)
┌──────────┬─────────┐
│ Holdings ┆ Returns │
│ ---      ┆ ---     │
│ i64      ┆ f64     │
╞══════════╪═════════╡
│ 3        ┆ 0.01    │
│ 3        ┆ 0.0     │
│ 3        ┆ -0.01   │
│ 2        ┆ 0.01    │
│ …        ┆ …       │
│ 1        ┆ 0.03    │
│ 1        ┆ 0.02    │
│ 1        ┆ 0.01    │
│ 1        ┆ 0.02    │
└──────────┴─────────┘

    您可以看到结果与您的预期相符。但是，您必须小心性能，因为cumulative_eval文档中有注释：

    警告

    此功能是实验性的，可能会发生变化，但不会被视为重大更改。

    这可能非常慢，因为它的复杂度可能是 O(n^2)。不要将其用于访问所有元素的操作。

    • 1

  3. 只是添加一些更多信息，因为我认为问题标题是基于我之前的糟糕措辞：

    我找到了问题的现有定义（这可能是问题的更好标题？）

    Polars 相当于 SQL 的ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING 窗口框架。

    ┌────────────┬─────────┐
│ Buy_Signal ┆ Returns │
│ ---        ┆ ---     │
│ i64        ┆ f64     │ # New "group/window" starts when Buy_Signal == 1
╞════════════╪═════════╡ # 
│ 1          ┆ 0.01    │ # ┌─ group 1 ─┐
│ 0          ┆ 0.02    │ # │           │ 
│ 1          ┆ 0.03    │ # │           │┌─ group 2 ─┐
│ 0          ┆ 0.02    │ # │           ││           │ 
│ 1          ┆ 0.01    │ # │           ││           │┌─ group 3 ─┐
│ 0          ┆ 0.0     │ # │           ││           ││           │
│ 0          ┆ -0.01   │ # └───────────┘└───────────┘└───────────┘
└────────────┴─────────┘

    例如使用DuckDB来演示：

    duckdb.sql("""
from df
select 
   buy_signal, 
   returns, 
   case when buy_signal = 1
   then
      list(returns) over (rows between current row and unbounded following)
   end as group
""")

    ┌────────────┬─────────┬────────────────────────────────────────────┐
│ Buy_Signal │ Returns │                   group                    │
│   int64    │ double  │                  double[]                  │
├────────────┼─────────┼────────────────────────────────────────────┤
│          1 │    0.01 │ [0.01, 0.02, 0.03, 0.02, 0.01, 0.0, -0.01] │
│          0 │    0.02 │ NULL                                       │
│          1 │    0.03 │ [0.03, 0.02, 0.01, 0.0, -0.01]             │
│          0 │    0.02 │ NULL                                       │
│          1 │    0.01 │ [0.01, 0.0, -0.01]                         │
│          0 │     0.0 │ NULL                                       │
│          0 │   -0.01 │ NULL                                       │
└────────────┴─────────┴────────────────────────────────────────────┘
    • 0

相关问题