如何计算Polars中不同ACQDISP类型(A/D)间的日期差？

在某股票交易数据集中，需计算ACQDISP字段两类记录（A代表买入、D代表卖出）间的日期差，涵盖A-A、D-D、A-D、D-A四种间隔天数。目前已通过Polars代码实现同类型记录的日期差计算，现需实现不同类型（A与D、D与A）之间的日期差计算。

现有实现代码 #

import polars as pl

df = pl.from_repr("""
┌─────────┬────────────┬────────────┬────────────┐
│ ACQDISP ┆ TRANDATE   ┆ SHARES_ADJ ┆ TPRICE_ADJ │
│ ---     ┆ ---        ┆ ---        ┆ ---        │
│ str     ┆ str        ┆ f64        ┆ f64        │
╞═════════╪════════════╪════════════╪════════════╡
│ D       ┆ 2001-12-31 ┆ 2000.0     ┆ 6.7        │
│ D       ┆ 2001-12-31 ┆ 6000.0     ┆ 6.45       │
│ D       ┆ 2001-12-31 ┆ 3000.0     ┆ 6.52       │
│ A       ┆ 2004-08-09 ┆ 5.4e6      ┆ 60.75      │
│ A       ┆ 2004-08-24 ┆ 2.1954e6   ┆ null       │
│ D       ┆ 2004-08-24 ┆ 4.63418e6  ┆ 52.435     │
│ A       ┆ 2014-12-17 ┆ 3.84581e6  ┆ null       │
└─────────┴────────────┴────────────┴────────────┘
""")

df = df.with_columns(TRANDATE=pl.col("TRANDATE").str.to_date())  
df = df.sort("TRANDATE", descending=False) 

days_difference_subsquent_trans = df.with_columns(pl.col("TRANDATE").diff().over("ACQDISP").alias("date_diff_subsequent_trans"))

代码运行结果 #

shape: (7, 5)
┌─────────┬────────────┬────────────┬────────────┬────────────────────────────┐
│ ACQDISP ┆ TRANDATE   ┆ SHARES_ADJ ┆ TPRICE_ADJ ┆ date_diff_subsequent_trans │
│ ---     ┆ ---        ┆ ---        ┆ ---        ┆ ---                        │
│ str     ┆ date       ┆ f64        ┆ f64        ┆ duration[ms]               │
╞═════════╪════════════╪════════════╪════════════╪════════════════════════════╡
│ D       ┆ 2001-12-31 ┆ 2000.0     ┆ 6.7        ┆ null                       │
│ D       ┆ 2001-12-31 ┆ 6000.0     ┆ 6.45       ┆ 0ms                        │
│ D       ┆ 2001-12-31 ┆ 3000.0     ┆ 6.52       ┆ 0ms                        │
│ A       ┆ 2004-08-09 ┆ 5.4e6      ┆ 60.75      ┆ null                       │
│ A       ┆ 2004-08-24 ┆ 2.1954e6   ┆ null       ┆ 15d                        │
│ D       ┆ 2004-08-24 ┆ 4.63418e6  ┆ 52.435     ┆ 967d                       │
│ A       ┆ 2014-12-17 ┆ 3.84581e6  ┆ null       ┆ 3767d                      │
└─────────┴────────────┴────────────┴────────────┴────────────────────────────┘

不同类型日期差的实现方案 #

要计算不同类型（A-D、D-A）的日期差，不能再按ACQDISP分组计算差值，而是需要基于整条时间序列，获取每条记录与上一条记录的日期差，并同时标记前后记录的类型组合。具体步骤如下：

实现代码 #

# 基于已排序、转换日期后的df继续处理
df = df.with_columns(
    # 获取上一条记录的交易类型和日期
    prev_acqdisp=pl.col("ACQDISP").shift(1),
    prev_trandate=pl.col("TRANDATE").shift(1)
).with_columns(
    # 计算当前记录与上一条记录的日期差
    date_diff=pl.col("TRANDATE") - pl.col("prev_trandate"),
    # 标记类型组合（如A→D、D→A等）
    type_pair=pl.col("prev_acqdisp") + "→" + pl.col("ACQDISP")
)

# 可选：按类型组合拆分日期差为单独字段
df = df.with_columns(
    date_diff_A_A=pl.when(pl.col("type_pair") == "A→A").then(pl.col("date_diff")),
    date_diff_D_D=pl.when(pl.col("type_pair") == "D→D").then(pl.col("date_diff")),
    date_diff_A_D=pl.when(pl.col("type_pair") == "A→D").then(pl.col("date_diff")),
    date_diff_D_A=pl.when(pl.col("type_pair") == "D→A").then(pl.col("date_diff"))
)

print(df)

代码说明 #

1. shift(1)：获取每条记录的上一条记录的ACQDISP和TRANDATE，第一条记录的shift结果为null，对应日期差也为null，符合逻辑。
2. 日期差计算：直接用当前日期减去上一条日期，得到间隔时长。
3. 类型组合标记：通过字符串拼接明确前后交易的类型关系，便于后续筛选或分组。
4. 拆分字段（可选）：用when...then语法将不同类型组合的日期差拆分为单独字段，方便查看和使用。

运行结果示例 #

shape: (7, 10)
┌─────────┬────────────┬────────────┬────────────┬────────────────────────────┬──────────────┬──────────────┬──────────┬───────────┬────────────┬────────────┬────────────┬────────────┐
│ ACQDISP ┆ TRANDATE   ┆ SHARES_ADJ ┆ TPRICE_ADJ ┆ date_diff_subsequent_trans ┆ prev_acqdisp ┆ prev_trandate┆ date_diff┆ type_pair ┆ date_diff_A_A ┆ date_diff_D_D ┆ date_diff_A_D ┆ date_diff_D_A │
│ ---     ┆ ---        ┆ ---        ┆ ---        ┆ ---                        ┆ ---          ┆ ---          ┆ ---      ┆ ---       ┆ ---          ┆ ---          ┆ ---          ┆ ---          │
│ str     ┆ date       ┆ f64        ┆ f64        ┆ duration[ms]               ┆ str          ┆ date         ┆ duration ┆ str       ┆ duration     ┆ duration     ┆ duration     ┆ duration     │
╞═════════╪════════════╪════════════╪════════════╪════════════════════════════╪══════════════╪══════════════╪══════════╪═══════════╪══════════════╪══════════════╪══════════════╪══════════════╡
│ D       ┆ 2001-12-31 ┆ 2000.0     ┆ 6.7        ┆ null                       ┆ null         ┆ null         ┆ null     ┆ null      ┆ null         ┆ null         ┆ null         ┆ null         │
│ D       ┆ 2001-12-31 ┆ 6000.0     ┆ 6.45       ┆ 0ms                        ┆ D            ┆ 2001-12-31   ┆ 0ms      ┆ D→D       ┆ null         ┆ 0ms          ┆ null         ┆ null         │
│ D       ┆ 2001-12-31 ┆ 3000.0     ┆ 6.52       ┆ 0ms                        ┆ D            ┆ 2001-12-31   ┆ 0ms      ┆ D→D       ┆ null         ┆ 0ms          ┆ null         ┆ null         │
│ A       ┆ 2004-08-09 ┆ 5.4e6      ┆ 60.75      ┆ null                       ┆ D            ┆ 2001-12-31   ┆ 967d     ┆ D→A       ┆ null         ┆ null         ┆ null         ┆ 967d         │
│ A       ┆ 2004-08-24 ┆ 2.1954e6   ┆ null       ┆ 15d                        ┆ A            ┆ 2004-08-09   ┆ 15d      ┆ A→A       ┆ 15d          ┆ null         ┆ null         ┆ null         │
│ D       ┆ 2004-08-24 ┆ 4.63418e6  ┆ 52.435     ┆ 967d                       ┆ A            ┆ 2004-08-24   ┆ 0ms      ┆ A→D       ┆ null         ┆ null         ┆ 0ms          ┆ null         │
│ A       ┆ 2014-12-17 ┆ 3.84581e6  ┆ null       ┆ 3767d                      ┆ D            ┆ 2004-08-24   ┆ 3767d    ┆ D→A       ┆ null         ┆ null         ┆ null         ┆ 3767d        │
└─────────┴────────────┴────────────┴────────────┴────────────────────────────┴──────────────┴──────────────┴──────────┴───────────┴──────────────┴──────────────┴──────────────┴──────────────┘

内容的提问来源于stack exchange，提问作者lpt