Polars DataFrame实现每日各OS累计唯一ID数的透视统计

问题描述 #

给定包含DAY、OS、ID字段的Polars DataFrame，需要按日期分组，统计截至当日各操作系统对应的累计唯一ID数量。示例数据及期望输出如下：

示例输入：

import polars as pl
df = pl.DataFrame(
    {
        "DAY": [1,1,1,2,2,2,3,3,3],
        "OS" : ["A","B","A","B","A","B","A","B","A"],
        "ID": ["X","Y","Z","W","X","J","K","L","X"]
    }
)

期望输出：

shape: (3, 3)
┌─────┬─────┬─────┐
│ DAY ┆ A   ┆ B   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 2   ┆ 1   │
│ 2   ┆ 2   ┆ 3   │
│ 3   ┆ 3   ┆ 4   │
└─────┴─────┴─────┘

原代码的问题 #

你尝试的pivot写法逻辑错误：

(df
 .pivot(
     index="DAY",
     on="OS",
     aggregate_function=(pl.col("ID").cum_sum().unique())
 )
)

• cum_sum()对字符串类型的ID会执行拼接操作，完全不符合累计唯一ID的需求；
• pivot的aggregate_function需要针对每个DAY+OS分组返回单个统计值，上述写法无法正确计算截至当日的累计唯一ID数。

正确解法 #

核心思路是：先按OS分组独立处理，维护每个OS的累计唯一ID集合，再按日期统计数量，最后转换为宽表格式。

实现代码

import polars as pl

df = pl.DataFrame(
    {
        "DAY": [1,1,1,2,2,2,3,3,3],
        "OS" : ["A","B","A","B","A","B","A","B","A"],
        "ID": ["X","Y","Z","W","X","J","K","L","X"]
    }
)

result = (
    # 第一步：去重同一天、同一OS下的重复ID，避免重复统计
    df.unique(subset=["DAY", "OS", "ID"])
    # 第二步：按OS分组，确保每个OS的统计独立
    .group_by("OS", maintain_order=True)
    .agg(
        # 按日期排序，保证时间顺序正确
        pl.col("DAY").sort().alias("DAY"),
        # 累计合并每日的唯一ID并去重，得到截至当日的所有唯一ID集合
        pl.col("ID").cumulative_eval(
            pl.element().list.concat().unique(),
            parallel=False
        ).alias("cumulative_unique_ids")
    )
    # 展开分组后的列表数据
    .explode(["DAY", "cumulative_unique_ids"])
    # 计算累计唯一ID的数量
    .with_columns(count=pl.col("cumulative_unique_ids").list.len())
    # 转换为宽表，匹配期望的输出格式
    .pivot(index="DAY", columns="OS", values="count")
)

print(result)

代码逻辑解释

1. 去重处理：先移除同一天、同一OS下的重复ID，避免后续累计时重复计算；
2. 按OS分组：每个操作系统的ID统计独立进行，互不干扰；
3. 累计唯一ID：利用cumulative_eval对每日的ID列表进行累计合并，并自动去重，得到截至当日的所有唯一ID集合；
4. 统计数量：计算每个累计集合的长度，即为截至当日的唯一ID总数；
5. 转宽表：通过pivot将长表转换为以DAY为索引、各OS为列的宽表，匹配期望输出格式。

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