如何在Polars DataFrame中创建标记最新记录的哑列？

在包含服务使用日期信息的数据表中，查找每个用户的最新记录。需实现需求：在Polars DataFrame中创建哑列，标记每条记录是否为对应用户的最新记录。

示例数据 #

sample_pl = pl.DataFrame(
    {
        'id':['id1', 'id1', 'id1', 'id1',
              'id2', 'id2', 'id2', 'id2', 
              'id3', 'id3', 'id3', 'id3'],

        'dt_event': ['2023-06-01', '2023-06-04', '2023-06-10', '2023-06-28',
                     '2023-06-01', '2023-06-04', '2023-06-10', '2023-06-28',
                     '2023-06-01', '2023-06-04', '2023-06-10', '2023-06-28']
    }, 
    schema={
        'id': pl.Utf8, 'dt_event': pl.Utf8
    }
)

可以通过Polars的窗口函数直接实现需求，无需额外的分组连接操作，步骤如下：

实现逻辑 #

1. 先将字符串类型的日期列转换为Polars日期类型，确保日期比较的准确性
2. 使用窗口函数over('id')对每个用户分组，计算该组内的最大日期
3. 对比每条记录的日期与组内最大日期，生成布尔类型的哑列标记是否为最新记录

代码实现 #

import polars as pl

# 处理示例数据
result = sample_pl.with_columns(
    # 转换日期格式
    pl.col('dt_event').str.to_date().alias('dt_event')
).with_columns(
    # 生成是否为最新记录的哑列
    (pl.col('dt_event') == pl.col('dt_event').max().over('id')).alias('is_latest')
)

# 输出结果
print(result)

结果说明 #

运行代码后会新增is_latest列：

• 该列值为True时，表示对应行是该用户的最新记录
• 示例数据中每个用户的最后一条记录（日期为2023-06-28）的is_latest值为True，其余均为False

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