Allow Index Column of Type Duration in DataFrame.rolling

crates/polars-time/src/group_by/dynamic.rs

@@ -152,6 +152,14 @@ impl Wrap<&DataFrame> {
                 TimeUnit::Milliseconds,
                 None,
             ),
+            Duration(tu) => {
+                let time_type_dt = Datetime(*tu, None);
+                let dt = time.cast(&time_type_dt).unwrap();
+                let (out, by, gt) =
+                    self.impl_rolling(dt, group_by, options, *tu, None, &time_type_dt)?;
+                let out = out.cast(&Duration(*tu)).unwrap();
+                return Ok((out, by, gt));
+            },
             UInt32 | UInt64 | Int32 => {
                 let time_type_dt = Datetime(TimeUnit::Nanoseconds, None);
                 let dt = time.cast(&Int64).unwrap().cast(&time_type_dt).unwrap();
@@ -182,7 +190,7 @@ impl Wrap<&DataFrame> {
             },
             dt => polars_bail!(
                 ComputeError:
-                "expected any of the following dtypes: {{ Date, Datetime, Int32, Int64, UInt32, UInt64 }}, got {}",
+                "expected any of the following dtypes: {{ Date, Datetime, Duration, Int32, Int64, UInt32, UInt64 }}, got {}",
                 dt
             ),
         };

crates/polars/tests/it/lazy/queries.rs

@@ -66,6 +66,68 @@ fn test_special_group_by_schemas() -> PolarsResult<()> {
         &[3, 5, 7, 9, 5]
     );
 
+    // Duration index column - period have different units
+    let out = df
+        .clone()
+        .lazy()
+        .with_column(
+            col("a")
+                .cast(DataType::Duration(TimeUnit::Milliseconds))
+                .set_sorted_flag(IsSorted::Ascending),
+        )
+        .rolling(
+            col("a"),
+            [],
+            RollingGroupOptions {
+                period: Duration::parse("2ms"),
+                offset: Duration::parse("0ms"),
+                closed_window: ClosedWindow::Left,
+                ..Default::default()
+            },
+        )
+        .agg([col("b").sum().alias("sum")])
+        .select([col("a"), col("sum")])
+        .collect()?;
+
+    assert_eq!(
+        out.column("sum")?
+            .i32()?
+            .into_no_null_iter()
+            .collect::<Vec<_>>(),
+        &[3, 5, 7, 9, 5]
+    );
+
+    // Datetime index column - period have same units
+    let out = df
+        .clone()
+        .lazy()
+        .with_column(
+            col("a")
+                .cast(DataType::Datetime(TimeUnit::Milliseconds, None))
+                .set_sorted_flag(IsSorted::Ascending),
+        )
+        .rolling(
+            col("a"),
+            [],
+            RollingGroupOptions {
+                period: Duration::parse("2ms"),
+                offset: Duration::parse("0ms"),
+                closed_window: ClosedWindow::Left,
+                ..Default::default()
+            },
+        )
+        .agg([col("b").sum().alias("sum")])
+        .select([col("a"), col("sum")])
+        .collect()?;
+
+    assert_eq!(
+        out.column("sum")?
+            .i32()?
+            .into_no_null_iter()
+            .collect::<Vec<_>>(),
+        &[3, 5, 7, 9, 5]
+    );
+
     let out = df
         .lazy()
         .with_column(col("a").set_sorted_flag(IsSorted::Ascending))

py-polars/tests/unit/operations/rolling/test_rolling.py

@@ -955,3 +955,44 @@ def test_rolling_invalid() -> None:
             .rolling("index", period="3000d")
             .agg(pl.col("values").sum().alias("sum"))
         )
+
+
+@pytest.mark.parametrize("time_unit", ["ms", "us", "ns"])
+def test_rolling_duration(time_unit):
+    # Note: Both datetime with Unit != ns and Duration have weird behavior.
+    # Here we only test for consistency.
+    df = pl.DataFrame(
+        {
+            "index_column": [1, 2, 3, 4, 5],
+            "value": [
+                1,
+                10,
+                100,
+                1000,
+                10000,
+            ],
+        }
+    )
+    df_duration = df.select(
+        pl.col("index_column").cast(pl.Duration(time_unit=time_unit)).set_sorted(),
+        "value",
+    )
+
+    df_datetime = df.select(
+        pl.col("index_column").cast(pl.Datetime(time_unit=time_unit)).set_sorted(),
+        "value",
+    )
+
+    res_duration = df_duration.rolling(
+        index_column="index_column", period=f"2{time_unit}"
+    ).agg(pl.col("value").sum())
+
+    res_datetime = df_datetime.rolling(
+        index_column="index_column", period=f"2{time_unit}"
+    ).agg(pl.col("value").sum())
+
+    assert (
+        res_duration["value"].to_list() == res_datetime["value"].to_list()
+    ), f"{res_duration["value"].to_list()=}, {res_datetime["value"].to_list()=}"
+
+    assert res_duration["index_column"].dtype == pl.Duration(time_unit=time_unit)