Python Polars 滚动计数与时间窗口

发布于03月28日

我正在try 为返回Polars表达式的功能管道编写一个方法.该方法应将列名作为字符串和整数天数.我想使用等于天数的窗口对该列执行滚动计数.

似乎没有rolling_count的表达式，所以我try 使用rolling_sum，但没有用.

def temporal_rolling_count(col: str, days: int) -> pl.Expr:
    return (
        pl.lit(1)
        .rolling_sum(window_size=f"{days}d", by="date_time")
        .over(col)
        .fill_null(0)
    )

我也try 了这种方法，它更接近，但仍然不是在所有情况下都有效

def temporal_rolling_count(col: str, days: int) -> pl.Expr:
    return (
        pl.col(col)
        .cumcount()
        .over(col, (pl.col("date_time") - pl.col("date_time").min()).dt.days() % days == 0)
        .fill_null(0)
    )

有什么办法可以通过返回一个表达式来实现这一点吗？或者我必须直接对DataFrame执行操作，比如使用groupby_rolling？

推荐答案

根据@jqury的建议，通过使用.clip，我能够在不对DataFrame执行操作的情况下实现预期的结果.

def temporal_rolling_count(col: str, days: int) -> pl.Expr:
    return (
        pl.col(col).clip(1,1)
        .rolling_sum(window_size=f"{days}d", by="date_time")
        .over(col)
        .fill_null(0)
    )

EDIT个

通过执行以下操作，我成功地执行了相同的操作，但只有在条件为真时才进行计数.

def temporal_rolling_count(col: str, days: int) -> pl.Expr:
    return (
        pl.when(condition).then(1).otherwise(0)
        .rolling_sum(window_size=f"{days}d", by="date_time")
        .over(col)
        .fill_null(0)
    )