在 Polars 中进行透视操作时自定义列名称并根据值添加新列？

Question

比如说，我有一个极坐标数据框，如下所示：

df = pl.from_repr("""
┌──────┬────────┬──────────────┬────────────────┬─────────────┬─────┐
│ date ┆ ticker ┆ positionType ┆ predictionType ┆ holdingBars ┆ res │
│ ---  ┆ ---    ┆ ---          ┆ ---            ┆ ---         ┆ --- │
│ str  ┆ str    ┆ str          ┆ str            ┆ i64         ┆ i64 │
╞══════╪════════╪══════════════╪════════════════╪═════════════╪═════╡
│ d1   ┆ t1     ┆ Long         ┆ p1             ┆ 1           ┆ 1   │
│ d1   ┆ t1     ┆ Long         ┆ s1             ┆ 1           ┆ 2   │
│ d1   ┆ t1     ┆ Short        ┆ p1             ┆ 1           ┆ 3   │
│ d1   ┆ t1     ┆ Short        ┆ s1             ┆ 1           ┆ 4   │
│ d1   ┆ t1     ┆ Long         ┆ p1             ┆ 2           ┆ 5   │
│ d1   ┆ t1     ┆ Short        ┆ p1             ┆ 2           ┆ 6   │
└──────┴────────┴──────────────┴────────────────┴─────────────┴─────┘
""")

我尝试如下旋转它：

df.pivot(
    on=["positionType", "holdingBars"],
    index=["date", "ticker", "predictionType"],      
    values="res"
)

输出：

shape: (2, 7)
┌──────┬────────┬────────────────┬────────────┬─────────────┬────────────┬─────────────┐
│ date ┆ ticker ┆ predictionType ┆ {"Long",1} ┆ {"Short",1} ┆ {"Long",2} ┆ {"Short",2} │
│ ---  ┆ ---    ┆ ---            ┆ ---        ┆ ---         ┆ ---        ┆ ---         │
│ str  ┆ str    ┆ str            ┆ i64        ┆ i64         ┆ i64        ┆ i64         │
╞══════╪════════╪════════════════╪════════════╪═════════════╪════════════╪═════════════╡
│ d1   ┆ t1     ┆ p1             ┆ 1          ┆ 3           ┆ 5          ┆ 6           │
│ d1   ┆ t1     ┆ s1             ┆ 2          ┆ 4           ┆ null       ┆ null        │
└──────┴────────┴────────────────┴────────────┴─────────────┴────────────┴─────────────┘

1. 我希望列名称为

   Long_1

   、

   Short_1

   等。
2. 我还想添加其他列：

   Long_1/Short_1

   和

   Long_2/Short_2

shape: (2, 9)
┌──────┬────────┬────────────────┬────────┬─────────┬────────┬─────────┬──────────────────┬──────────────────┐
│ date ┆ ticker ┆ predictionType ┆ Long_1 ┆ Short_1 ┆ Long_2 ┆ Short_2 ┆ Long_1 / Short_1 ┆ Long_2 / Short_2 │
│ ---  ┆ ---    ┆ ---            ┆ ---    ┆ ---     ┆ ---    ┆ ---     ┆ ---              ┆ ---              │
│ str  ┆ str    ┆ str            ┆ i64    ┆ i64     ┆ str    ┆ str     ┆ f64              ┆ f64              │
╞══════╪════════╪════════════════╪════════╪═════════╪════════╪═════════╪══════════════════╪══════════════════╡
│ d1   ┆ t1     ┆ p1             ┆ 1      ┆ 3       ┆ 5      ┆ 6       ┆ 0.3              ┆ 0.83             │
│ d1   ┆ t1     ┆ s1             ┆ 2      ┆ 4       ┆ null   ┆ null    ┆ 0.5              ┆ null             │
└──────┴────────┴────────────────┴────────┴─────────┴────────┴─────────┴──────────────────┴──────────────────┘

请注意，这是数据的子集，它可以具有跨多个日期的多个代码，并且 res 可以跨多个持有期。

知道如何实现这些吗？

Answer 1

您可以在

.pivot()

之前创建列名称，例如使用

pl.format()

(df.with_columns(pl.format("{}_{}", "positionType", "holdingBars").alias("on"))
   .pivot(
      on = "on",
      index = ["date", "ticker", "predictionType"],
      values = "res"
   )
)

shape: (2, 7)
┌──────┬────────┬────────────────┬────────┬─────────┬────────┬─────────┐
│ date ┆ ticker ┆ predictionType ┆ Long_1 ┆ Short_1 ┆ Long_2 ┆ Short_2 │
│ ---  ┆ ---    ┆ ---            ┆ ---    ┆ ---     ┆ ---    ┆ ---     │
│ str  ┆ str    ┆ str            ┆ i64    ┆ i64     ┆ i64    ┆ i64     │
╞══════╪════════╪════════════════╪════════╪═════════╪════════╪═════════╡
│ d1   ┆ t1     ┆ p1             ┆ 1      ┆ 3       ┆ 5      ┆ 6       │
│ d1   ┆ t1     ┆ s1             ┆ 2      ┆ 4       ┆ null   ┆ null    │
└──────┴────────┴────────────────┴────────┴─────────┴────────┴─────────┘

长_1/短_1

与名称一样，您可以计算之前的值

.pivot()

由于长/空项目是保证成对的，并且按顺序 - 您可以为每对分配一个“id”，例如使用

.cum_sum()

df.with_columns(
   (pl.col("positionType") == "Long")
      .cum_sum() 
      .over("ticker", "predictionType")
      .alias("predictionId")
)

shape: (6, 7)
┌──────┬────────┬──────────────┬────────────────┬─────────────┬─────┬──────────────┐
│ date ┆ ticker ┆ positionType ┆ predictionType ┆ holdingBars ┆ res ┆ predictionId │
│ ---  ┆ ---    ┆ ---          ┆ ---            ┆ ---         ┆ --- ┆ ---          │
│ str  ┆ str    ┆ str          ┆ str            ┆ i64         ┆ i64 ┆ u32          │
╞══════╪════════╪══════════════╪════════════════╪═════════════╪═════╪══════════════╡
│ d1   ┆ t1     ┆ Long         ┆ p1             ┆ 1           ┆ 1   ┆ 1            │
│ d1   ┆ t1     ┆ Long         ┆ s1             ┆ 1           ┆ 2   ┆ 1            │
│ d1   ┆ t1     ┆ Short        ┆ p1             ┆ 1           ┆ 3   ┆ 1            │
│ d1   ┆ t1     ┆ Short        ┆ s1             ┆ 1           ┆ 4   ┆ 1            │
│ d1   ┆ t1     ┆ Long         ┆ p1             ┆ 2           ┆ 5   ┆ 2            │
│ d1   ┆ t1     ┆ Short        ┆ p1             ┆ 2           ┆ 6   ┆ 2            │
└──────┴────────┴──────────────┴────────────────┴─────────────┴─────┴──────────────┘

并将其与

.group_by()

一起使用来

.agg()

进行计算。

(df
  .with_columns(
     (pl.col("positionType") == "Long").cum_sum().over("ticker", "predictionType")
        .alias("predictionId")
  )
  .group_by("date", "ticker", "predictionType", "predictionId", maintain_order=True)
  .agg(
      pl.col("res").first().alias("Long"),
      pl.col("res").last().alias("Short"),
      (pl.col("res").first() / pl.col("res").last()).alias("Long/Short")
  )
)

shape: (3, 7)
┌──────┬────────┬────────────────┬────────┬──────┬───────┬────────────┐
│ date ┆ ticker ┆ predictionType ┆ predId ┆ Long ┆ Short ┆ Long/Short │
│ ---  ┆ ---    ┆ ---            ┆ ---    ┆ ---  ┆ ---   ┆ ---        │
│ str  ┆ str    ┆ str            ┆ u32    ┆ i64  ┆ i64   ┆ f64        │
╞══════╪════════╪════════════════╪════════╪══════╪═══════╪════════════╡
│ d1   ┆ t1     ┆ p1             ┆ 1      ┆ 1    ┆ 3     ┆ 0.333333   │
│ d1   ┆ t1     ┆ s1             ┆ 1      ┆ 2    ┆ 4     ┆ 0.5        │
│ d1   ┆ t1     ┆ p1             ┆ 2      ┆ 5    ┆ 6     ┆ 0.833333   │
└──────┴────────┴────────────────┴────────┴──────┴───────┴────────────┘

如果我们

.unpivot()

这个结果，我们最终会得到与第一个示例相同的格式/结构。

这意味着您可以添加相同的

.format()

+

.pivot()

步骤。

(df
  .with_columns(
     (pl.col("positionType") == "Long").cum_sum().over("ticker", "predictionType")
        .alias("predictionId")
  )
  .group_by("date", "ticker", "predictionType", "predictionId", maintain_order=True)
  .agg(
      pl.col("res").first().alias("Long"),
      pl.col("res").last().alias("Short"),
      (pl.col("res").first() / pl.col("res").last()).alias("Long/Short")
  )
  .unpivot(index=["date", "ticker", "predictionType", "predictionId"])
  .with_columns(pl.format("{}_{}", "variable", "predictionId"))
  .pivot(
     on = "variable",
     index = ["date", "ticker", "predictionType"],
     values = "value"
  )
)

shape: (2, 9)
┌──────┬────────┬────────────────┬────────┬────────┬─────────┬─────────┬──────────────┬──────────────┐
│ date ┆ ticker ┆ predictionType ┆ Long_1 ┆ Long_2 ┆ Short_1 ┆ Short_2 ┆ Long/Short_1 ┆ Long/Short_2 │
│ ---  ┆ ---    ┆ ---            ┆ ---    ┆ ---    ┆ ---     ┆ ---     ┆ ---          ┆ ---          │
│ str  ┆ str    ┆ str            ┆ f64    ┆ f64    ┆ f64     ┆ f64     ┆ f64          ┆ f64          │
╞══════╪════════╪════════════════╪════════╪════════╪═════════╪═════════╪══════════════╪══════════════╡
│ d1   ┆ t1     ┆ p1             ┆ 1.0    ┆ 5.0    ┆ 3.0     ┆ 6.0     ┆ 0.333333     ┆ 0.833333     │
│ d1   ┆ t1     ┆ s1             ┆ 2.0    ┆ null   ┆ 4.0     ┆ null    ┆ 0.5          ┆ null         │
└──────┴────────┴────────────────┴────────┴────────┴─────────┴─────────┴──────────────┴──────────────┘