如何在Polars DataFrame中按ID分组查找跨日期的字段变更

需求回顾 #

• 按id分组，仅保留存在多个update_time（整数日期YYYYMMDD格式）的ID
• 追踪每个ID相邻日期间的字段变化
• 多字段变更时拆分为单独行展示

示例输入数据 #

import polars as pl

df = pl.DataFrame({
    "id": ["AAPL", "AAPL", "AAPL", "TSLA", "MSFT", "MSFT"],
    "update_time": [20240101, 20240102, 20240103, 20240101, 20240101, 20240102],
    "price": [150, 152, 152, 200, 300, 305],
    "volume": [1000000, 1200000, 1100000, 500000, 800000, 800000],
    "rating": ["Buy", "Buy", "Hold", "Buy", "Hold", "Hold"]
})

预期输出 #

shape: (4, 6)
┌──────┬────────────┬────────────┬────────┬───────────┬───────────┐
│ id   ┆ prev_date  ┆ curr_date  ┆ field  ┆ old_value ┆ new_value │
│ ---  ┆ ---        ┆ ---        ┆ ---    ┆ ---       ┆ ---       │
│ str  ┆ i64        ┆ i64        ┆ str    ┆ str       ┆ str       │
╞══════╪════════════╪════════════╪════════╪═══════════╪═══════════╡
│ AAPL ┆ 20240101   ┆ 20240102   ┆ volume ┆ 1000000   ┆ 1200000   │
│ AAPL ┆ 20240102   ┆ 20240103   ┆ volume ┆ 1200000   ┆ 1100000   │
│ AAPL ┆ 20240102   ┆ 20240103   ┆ rating ┆ Buy       ┆ Hold      │
│ MSFT ┆ 20240101   ┆ 20240102   ┆ price  ┆ 300       ┆ 305       │
└──────┴────────────┴────────────┴────────┴───────────┴───────────┘

解决方案代码 #

动态字段版本（推荐，无需硬编码字段）

import polars as pl

# 定义需要检查变更的字段
fields_to_check = ["price", "volume", "rating"]

# 1. 过滤仅存在单个日期的ID
valid_ids = df.group_by("id").agg(pl.count("update_time") > 1).filter(pl.col("update_time")).select("id")
filtered_df = df.join(valid_ids, on="id")

# 2. 分组排序后，获取每个行的前置日期和字段值
shift_exprs = [
    pl.col("update_time").alias("curr_date"),
    pl.col("update_time").shift().alias("prev_date")
]
for field in fields_to_check:
    shift_exprs.extend([
        pl.col(field).alias(f"curr_{field}"),
        pl.col(field).shift().alias(f"prev_{field}")
    ])

comparison_df = filtered_df.group_by("id").sort("update_time").agg(shift_exprs).explode(pl.exclude("id"))

# 3. 将字段转成行，仅保留有变更的记录
unpivot_pairs = [(f"prev_{f}", f"curr_{f}") for f in fields_to_check]
result = comparison_df.filter(pl.col("prev_date").is_not_null()).unpivot(
    index=["id", "prev_date", "curr_date"],
    variable_name="field",
    value_name=["old_value", "new_value"],
    on=unpivot_pairs
).with_columns(
    pl.col("field").str.replace("prev_", "")
).filter(pl.col("old_value") != pl.col("new_value"))

print(result)

代码说明 #

1. 过滤有效ID：通过分组统计每个ID的日期数量，只保留有多个日期记录的ID，排除TSLA这类单日期数据。
2. 获取前置数据：按ID分组并按日期排序后，用shift()获取每个记录的前一天日期及对应字段值，再通过explode展开分组后的列表。
3. 转成行并过滤变更：用unpivot把每个字段的前后值拆分为单独行，清理字段名后，只保留前后值不同的变更记录。

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