使用 struct 将带有移位版本的 Polars 数据帧合并到单个 df 中

Question

我有一个数据框，其中包含按日期列出的公司股票价格历史记录。

df = pl.from_repr("""
┌──────┬─────┬───────┬──────────┐
│ date ┆ ibm ┆ apple ┆ coinbase │
│ ---  ┆ --- ┆ ---   ┆ ---      │
│ i64  ┆ i64 ┆ i64   ┆ i64      │
╞══════╪═════╪═══════╪══════════╡
│ 1    ┆ 50  ┆ 100   ┆ 40       │
│ 2    ┆ 48  ┆ 200   ┆ 60       │
│ 3    ┆ 49  ┆ 400   ┆ 10       │
│ 4    ┆ 27  ┆ 800   ┆ 300      │
│ 5    ┆ 11  ┆ 1600  ┆ 7000     │
└──────┴─────┴───────┴──────────┘
""")

还有一个价格移动了 2 个日期单位。

df_2 = df.shift(-2)

shape: (5, 4)
┌──────┬──────┬───────┬──────────┐
│ date ┆ ibm  ┆ apple ┆ coinbase │
│ ---  ┆ ---  ┆ ---   ┆ ---      │
│ i64  ┆ i64  ┆ i64   ┆ i64      │
╞══════╪══════╪═══════╪══════════╡
│ 3    ┆ 49   ┆ 400   ┆ 10       │
│ 4    ┆ 27   ┆ 800   ┆ 300      │
│ 5    ┆ 11   ┆ 1600  ┆ 7000     │
│ null ┆ null ┆ null  ┆ null     │
│ null ┆ null ┆ null  ┆ null     │
└──────┴──────┴───────┴──────────┘

我想将 2 个

ibm

系列捆绑到单个

struct

列中的

ibm

中，其中包含价格@日期的字段和偏移（-2）价格。类似

apple

等

我该怎么做？效率会高吗？

不知道从哪里开始，因为结果

ibm

列似乎不会包含2个系列，而是由一对（原始，转移）价格组成的单个系列

Answer 1

pl.struct

接受任意数量的表达式，允许将

shift

表达式直接放入其中：

df.with_columns(
    pl.struct(col, pl.col(col).shift(-2).name.suffix('2'))
    for col in df.columns
    if col != 'date'
)

shape: (5, 4)
┌──────┬───────────┬─────────────┬─────────────┐
│ date ┆ ibm       ┆ apple       ┆ coinbase    │
│ ---  ┆ ---       ┆ ---         ┆ ---         │
│ i64  ┆ struct[2] ┆ struct[2]   ┆ struct[2]   │
╞══════╪═══════════╪═════════════╪═════════════╡
│ 1    ┆ {50,49}   ┆ {100,400}   ┆ {40,10}     │
│ 2    ┆ {48,27}   ┆ {200,800}   ┆ {60,300}    │
│ 3    ┆ {49,11}   ┆ {400,1600}  ┆ {10,7000}   │
│ 4    ┆ {27,null} ┆ {800,null}  ┆ {300,null}  │
│ 5    ┆ {11,null} ┆ {1600,null} ┆ {7000,null} │
└──────┴───────────┴─────────────┴─────────────┘

此处，

ibm

列的基础字段为

ibm2

和

ibm

（请参阅

.name.suffix

）等，您可以使用

alias

或类似名称来重命名您想要的列/字段.

df.with_columns(...).schema

{'date': Int64, 'ibm': Struct([Field('ibm', Int64), Field('ibm2', Int64)]), 'apple': Struct([Field('apple', Int64), Field('apple2', Int64)]), 'coinbase': Struct([Field('coinbase', Int64), Field('coinbase2', Int64)])}