如何在Polars中筛选出日期连续的首段行

我明白你想要的是从排序后的日期列里，提取最开头的一段连续月度数据，直到遇到第一个不连续的月份就停止对吧？之前用group_by_dynamic没成功很正常，因为那个函数是用来按时间窗口分组统计的，和找连续序列的需求不匹配。

给你两种简单的实现方法，都能轻松达到目的：

方法一：通过计算月份差值找断点 #

这个思路是先把日期转换成连续的月份数值（比如2021年1月就是2021*12+1=24253），然后看相邻行的这个数值差是不是1——如果是1就说明月份连续，否则就是断点。找到第一个断点后，取断点之前的所有行即可：

import polars as pl
from datetime import datetime

# 你的原始数据生成和处理代码
df = pl.DataFrame(
    {
        "time": pl.datetime_range(
            start=datetime(2020, 12, 1),
            end=datetime(2023, 12, 1),
            interval="1mo",
            eager=True,
        ),
        "n": range(37),
    }
)
df = df.sample(n=10, seed=0)  
df = df.sort(["time"], descending=False)

# 步骤1：将日期转为连续的月份数值
df = df.with_columns(month_num = pl.col("time").dt.year() * 12 + pl.col("time").dt.month())

# 步骤2：计算相邻月份的差值，第一行填充为1（视为连续起始）
df = df.with_columns(month_diff = pl.col("month_num").diff().fill_null(1))

# 步骤3：找到第一个不连续的位置
break_idx = df.select(pl.arg_where(pl.col("month_diff") != 1).first()).item()

# 步骤4：提取连续的首段数据
if break_idx is not None:
    result = df.slice(0, break_idx).drop(["month_num", "month_diff"])
else:
    # 如果所有行都是连续的，直接取全部
    result = df.drop(["month_num", "month_diff"])

print(result)

方法二：用分组ID标记连续段 #

这个方法是给连续的行分配同一个组ID，一旦遇到不连续的月份就切换组ID，最后取第一个组的所有行：

# 基于你已经排序好的df
df = df.with_columns(
    # 标记每一行是否和前一行连续，第一行默认连续
    is_consecutive = pl.col("time").dt.month_diff(pl.col("time").shift(1)).fill_null(1) == 1
)

# 生成连续段的组ID：不连续时组号+1，连续时保持原组号
df = df.with_columns(
    group_id = pl.when(~pl.col("is_consecutive")).then(1).otherwise(0).fill_null(1).cum_sum()
)

# 提取第一个组的所有行
result = df.filter(pl.col("group_id") == df["group_id"][0]).drop(["is_consecutive", "group_id"])

print(result)

这两种方法都能精准提取开头的连续月度行，在你的示例数据里，应该能得到你想要的前5行（假设前5行的月份都是连续的）。

备注：内容来源于stack exchange，提问作者pinpss