如何在Polars中为subject_id与timestamp唯一对创建唯一ID列？

核心思路 #

利用Polars的factorize方法给每个subject_id+timestamp的唯一组合分配整数ID，这个ID可作为后续合并的关联键，完全不需要依赖索引。

具体步骤 #

1. 生成唯一关联ID
   把subject_id和timestamp组合成结构体，对结构体调用factorize就能得到每个唯一组合对应的ID，这种方式比手动哈希更稳定，无哈希冲突风险。

示例代码：

import polars as pl

# 示例DataFrame
df = pl.DataFrame({
    "subject_id": [1, 1, 2, 2, 3],
    "timestamp": ["2023-01-01", "2023-01-01", "2023-01-02", "2023-01-03", "2023-01-01"],
    "event": ["A", "B", "C", "D", "E"],
    "col1": [10, 20, 30, 40, 50],
    "col2": [100, 200, 300, 400, 500]
})

# 添加唯一关联ID列
df_with_id = df.with_columns(
    pl.struct(["subject_id", "timestamp"]).factorize().alias("pair_id")
)

这里的pair_id就是每个subject_id+timestamp唯一组合的标识，相同组合会得到相同ID。

2. 拆分DataFrame
   按需求拆分为两个DataFrame，注意保留pair_id列用于后续合并：

# 第一个DF：包含subject_id、timestamp、event、pair_id
df_event = df_with_id.select(["subject_id", "timestamp", "event", "pair_id"])

# 第二个DF：包含subject_id、timestamp、col1、col2、pair_id
df_cols = df_with_id.select(["subject_id", "timestamp", "col1", "col2", "pair_id"])

3. 处理后合并
   对两个DF各自完成分组或其他处理后，直接通过pair_id列合并即可：

# 示例：对df_event按pair_id统计事件数量
df_event_processed = df_event.group_by("pair_id").agg(pl.col("event").count().alias("event_count"))

# 示例：对df_cols按pair_id计算col1和col2的均值
df_cols_processed = df_cols.group_by("pair_id").agg(
    pl.col("col1").mean().alias("col1_mean"),
    pl.col("col2").mean().alias("col2_mean")
)

# 通过pair_id合并处理结果
merged_df = df_event_processed.join(df_cols_processed, on="pair_id", how="inner")

补充说明 #

• 若担心factorize生成的整数ID不够直观，也可用pl.hash生成哈希值，但哈希存在极小冲突概率，factorize是更稳妥的选择。
• 拆分时可选择保留原始的subject_id和timestamp，方便后续核对数据。

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