如何在Python Polars中统计特定时间范围内的匹配值数量？

问题场景 #

已有两个Polars DataFrame：

• events：包含时间范围（startdatetime/enddatetime）和允许的数值列表event_values
• test_data：包含时间戳datetime和数值test_values

需求：统计test_data中满足以下条件的记录数，添加为events的新列numEvents，最终导出带该列的CSV：

1. datetime落在对应events行的startdatetime到enddatetime区间内（左闭右开，符合示例逻辑）
2. test_values存在于对应events行的event_values列表中

修正后的示例数据构造 #

先修正原始代码的语法错误，确保DataFrame能正常创建：

import polars as pl

# 构造events数据集
events = pl.DataFrame({
    "startdatetime": ["2021-01-01 00:00:00", "2021-01-01 00:10:00", "2021-01-01 00:20:00"],
    "enddatetime": ["2021-01-01 00:15:00", "2021-01-01 00:25:00", "2021-01-01 00:35:00"],
    "event_values": [[1, 2, 3], [2, 3, 4], [4, 5, 6]]
})

# 构造test_data数据集
test_data = pl.DataFrame({
    "datetime": ["2021-01-01 00:02:00", "2021-01-01 00:07:00", "2021-01-01 00:13:00",
                 "2021-01-01 00:15:00", "2021-01-01 00:16:00", "2021-01-01 00:19:00",
                 "2021-01-01 00:22:00", "2021-01-01 00:25:00", "2021-01-01 00:28:00",
                 "2021-01-01 00:37:00"],
    "test_values": [1, 3, 2, 7, 3, 4, 5, 6, 4, 5]
})

解决方案 #

步骤1：转换时间列类型 #

将字符串格式的时间列转为Polars datetime类型，才能进行范围比较：

events = events.with_columns(
    pl.col("startdatetime").str.to_datetime(),
    pl.col("enddatetime").str.to_datetime()
)

test_data = test_data.with_columns(
    pl.col("datetime").str.to_datetime()
)

步骤2：计算匹配记录数 #

提供两种方案，根据数据量选择：

方案一：行级映射（小数据量友好）

利用map_elements对events每行进行条件过滤统计，代码简洁直观：

events = events.with_columns(
    numEvents=pl.struct(["startdatetime", "enddatetime", "event_values"]).map_elements(
        lambda row: (
            test_data
            .filter(
                # 左闭右开区间匹配
                pl.col("datetime").is_between(row["startdatetime"], row["enddatetime"], closed="left")
                # 值匹配
                & pl.col("test_values").is_in(row["event_values"])
            )
            .height  # 获取匹配记录数
        ),
        return_dtype=pl.Int64
    )
)

方案二：关联聚合（大数据量友好）

通过explode展开event_values，再用关联+聚合的方式统计，避免行级循环，效率更高：

# 展开event_values并添加行索引
exploded_events = events.explode("event_values").with_row_index("event_id")

# 关联test_data并过滤时间范围
matched_records = test_data.join(
    exploded_events,
    left_on="test_values",
    right_on="event_values",
    how="inner"
).filter(
    pl.col("datetime").is_between(pl.col("startdatetime"), pl.col("enddatetime"), closed="left")
)

# 按event_id统计匹配数，合并回原events
match_counts = matched_records.group_by("event_id").agg(numEvents=pl.count())
events = (
    events.with_row_index("event_id")
    .join(match_counts, on="event_id", how="left")
    .drop("event_id")
    .fill_null(0)  # 无匹配的行填充0
)

步骤3：导出CSV #

将处理后的events导出为CSV文件：

events.write_csv("events_with_numEvents.csv")

验证结果 #

处理后的events将包含numEvents列，与示例预期一致：

print(events)
# 输出：
# shape: (3, 4)
# ┌─────────────────────┬─────────────────────┬──────────────┬──────────┐
# │ startdatetime       ┆ enddatetime         ┆ event_values ┆ numEvents│
# │ ---                 ┆ ---                 ┆ ---          ┆ ---      │
# │ datetime[μs]        ┆ datetime[μs]        ┆ list[i64]    ┆ i64      │
# ╞═════════════════════╪═════════════════════╪══════════════╪══════════╡
# │ 2021-01-01 00:00:00 ┆ 2021-01-01 00:15:00 ┆ [1, 2, 3]    ┆ 3        │
# │ 2021-01-01 00:10:00 ┆ 2021-01-01 00:25:00 ┆ [2, 3, 4]    ┆ 4        │
# │ 2021-01-01 00:20:00 ┆ 2021-01-01 00:35:00 ┆ [4, 5, 6]    ┆ 3        │
# └─────────────────────┴─────────────────────┴──────────────┴──────────┘

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