如何在Polars中按组检测日期是否匹配指定年月

问题背景 #

在Polars中执行复杂查询时，希望将操作拆分为多个方法简化流程，需要掌握如何向自定义函数传入多列数据与外部变量。

示例数据 #

import polars as pl
from datetime import datetime

test_data = pl.DataFrame({
    "class": ['A', 'A', 'A', 'B', 'B', 'C'],
    "date": [datetime(2020, 1, 31), datetime(2020, 2, 28), datetime(2021, 1, 31),
              datetime(2022, 1, 31), datetime(2023, 2, 28),
              datetime(2020, 1, 31)],
    "status": [1,0,1,0,1,0]
})

需求说明 #

按class分组，判断每组中是否存在与参考日期年月匹配的date，并根据匹配行的status生成结果：若匹配行中存在status=1则返回1，否则返回0。期望输出如下：

┌───────┬─────────────────────┬──────────────────────┐
│ class ┆ reference_date      ┆ point_in_time_status │
│ ---   ┆ ---                 ┆ ---                  │
│ str   ┆ datetime[μs]        ┆ i64                  │
╞═══════╪═════════════════════╪══════════════════════╡
│ A     ┆ 2020-01-02 00:00:00 ┆ 1                    │
│ B     ┆ 2020-01-02 00:00:00 ┆ 0                    │
│ C     ┆ 2020-01-02 00:00:00 ┆ 0                    │
└───────┴─────────────────────┴──────────────────────┘

R语言实现参考 #

# Loading library
library(tidyverse)

# Creating dataframe
df <- data.frame(
 date = c(as.Date("2020-01-31"), 
          as.Date("2020-02-28"), as.Date("2021-01-31"), 
          as.Date("2022-01-31"), as.Date("2023-02-28"), 
          as.Date("2020-01-31")), 
 status = c(1,0,1,0,1,0), 
 class = c("A","A","A","B","B","C"))

# Finding status in overlapping months
ref_date = as.Date("2020-01-02")

df %>%
  group_by("class") %>%
  filter(format(date, "%Y-%m") == format(ref_date, "%Y-%m")) %>%
  filter(status == 1)

Polars解决方案 #

方案1：向量化操作（推荐，性能更优） #

Polars优先推荐向量化操作，无需自定义函数即可高效实现需求：

reference_date = datetime(2020, 1, 2)
ref_year_month = reference_date.strftime("%Y-%m")

result = (
    test_data
    # 提取每条记录的年月
    .with_columns(pl.col("date").dt.strftime("%Y-%m").alias("year_month"))
    .group_by("class")
    .agg(
        row_count=pl.count(),
        reference_date=pl.lit(reference_date),
        # 匹配年月后取最大status，无匹配则填充0
        point_in_time_status=pl.when(pl.col("year_month") == ref_year_month)
                               .then(pl.col("status"))
                               .max()
                               .fill_null(0)
    )
)

print(result)

方案2：自定义函数+多列传入 #

如果需要拆分逻辑到自定义函数，可通过pl.struct将多列打包，再结合map_elements传入外部变量：

def check_status(group_struct, ref_date):
    ref_ym = ref_date.strftime("%Y-%m")
    # 将struct转换为DataFrame以便操作
    group_df = group_struct.to_frame().unnest("")
    # 筛选年月匹配的行
    matched_rows = group_df.filter(pl.col("date").dt.strftime("%Y-%m") == ref_ym)
    # 判断是否存在status=1的记录
    return 1 if not matched_rows.is_empty() and (matched_rows["status"] == 1).any() else 0

reference_date = datetime(2020, 1, 2)

result_with_func = (
    test_data
    .group_by("class")
    .agg(
        row_count=pl.count(),
        reference_date=pl.lit(reference_date),
        # 用struct打包date和status，传入自定义函数
        point_in_time_status=pl.struct("date", "status")
                               .map_elements(lambda x: check_status(x, reference_date), return_dtype=pl.Int64)
    )
)

print(result_with_func)

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