fix: Fix rolling aggregations for various integer types (#20512)

pola-rs · Dec 31, 2024 · 25fab78 · 25fab78
1 parent 4c14e70
commit 25fab78
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 7 deletions.
diff --git a/crates/polars-core/src/frame/group_by/aggregations/mod.rs b/crates/polars-core/src/frame/group_by/aggregations/mod.rs
@@ -286,7 +286,7 @@ impl_take_extremum!(i8);
 impl_take_extremum!(i16);
 impl_take_extremum!(i32);
 impl_take_extremum!(i64);
-#[cfg(feature = "dtype-decimal")]
+#[cfg(any(feature = "dtype-decimal", feature = "dtype-i128"))]
 impl_take_extremum!(i128);
 impl_take_extremum!(float: f32);
 impl_take_extremum!(float: f64);

diff --git a/crates/polars-core/src/hashing/vector_hasher.rs b/crates/polars-core/src/hashing/vector_hasher.rs
@@ -167,7 +167,7 @@ vec_hash_numeric!(UInt16Chunked);
 vec_hash_numeric!(UInt8Chunked);
 vec_hash_numeric!(Float64Chunked);
 vec_hash_numeric!(Float32Chunked);
-#[cfg(feature = "dtype-decimal")]
+#[cfg(any(feature = "dtype-decimal", feature = "dtype-i128"))]
 vec_hash_numeric!(Int128Chunked);
 
 impl VecHash for StringChunked {

diff --git a/crates/polars-expr/Cargo.toml b/crates/polars-expr/Cargo.toml
@@ -49,6 +49,7 @@ dtype-full = [
   "dtype-decimal",
   "dtype-duration",
   "dtype-i16",
+  "dtype-i128",
   "dtype-i8",
   "dtype-struct",
   "dtype-time",

diff --git a/crates/polars-lazy/Cargo.toml b/crates/polars-lazy/Cargo.toml
@@ -97,6 +97,7 @@ dtype-full = [
   "dtype-decimal",
   "dtype-duration",
   "dtype-i16",
+  "dtype-i128",
   "dtype-i8",
   "dtype-struct",
   "dtype-time",
@@ -144,6 +145,7 @@ dtype-duration = [
   "polars-mem-engine/dtype-duration",
 ]
 dtype-i16 = ["polars-plan/dtype-i16", "polars-pipe?/dtype-i16", "polars-expr/dtype-i16", "polars-mem-engine/dtype-i16"]
+dtype-i128 = ["polars-plan/dtype-i128", "polars-pipe?/dtype-i128", "polars-expr/dtype-i128"]
 dtype-i8 = ["polars-plan/dtype-i8", "polars-pipe?/dtype-i8", "polars-expr/dtype-i8", "polars-mem-engine/dtype-i8"]
 dtype-struct = [
   "polars-plan/dtype-struct",

diff --git a/crates/polars/Cargo.toml b/crates/polars/Cargo.toml
@@ -265,6 +265,7 @@ dtype-full = [
   "dtype-array",
   "dtype-i8",
   "dtype-i16",
+  "dtype-i128",
   "dtype-decimal",
   "dtype-u8",
   "dtype-u16",
@@ -318,12 +319,20 @@ dtype-i8 = [
   "polars-io/dtype-i8",
   "polars-lazy?/dtype-i8",
   "polars-ops/dtype-i8",
+  "polars-time?/dtype-i8",
 ]
 dtype-i16 = [
   "polars-core/dtype-i16",
   "polars-io/dtype-i16",
   "polars-lazy?/dtype-i16",
   "polars-ops/dtype-i16",
+  "polars-time?/dtype-i16",
+]
+dtype-i128 = [
+  "polars-core/dtype-i128",
+  "polars-lazy?/dtype-i128",
+  "polars-ops/dtype-i128",
+  "polars-time?/dtype-i128",
 ]
 dtype-decimal = [
   "polars-core/dtype-decimal",
@@ -337,12 +346,14 @@ dtype-u8 = [
   "polars-io/dtype-u8",
   "polars-lazy?/dtype-u8",
   "polars-ops/dtype-u8",
+  "polars-time?/dtype-u8",
 ]
 dtype-u16 = [
   "polars-core/dtype-u16",
   "polars-io/dtype-u16",
   "polars-lazy?/dtype-u16",
   "polars-ops/dtype-u16",
+  "polars-time?/dtype-u16",
 ]
 dtype-categorical = [
   "polars-core/dtype-categorical",

diff --git a/py-polars/tests/unit/operations/rolling/test_rolling.py b/py-polars/tests/unit/operations/rolling/test_rolling.py
@@ -15,6 +15,7 @@
 from polars.testing import assert_frame_equal, assert_series_equal
 from polars.testing.parametric import column, dataframes
 from polars.testing.parametric.strategies.dtype import _time_units
+from tests.unit.conftest import INTEGER_DTYPES
 
 if TYPE_CHECKING:
     from hypothesis.strategies import SearchStrategy
@@ -739,11 +740,18 @@ def test_rolling_aggregations_with_over_11225() -> None:
     assert_frame_equal(result, expected)
 
 
-def test_rolling() -> None:
-    s = pl.Series("a", [1, 2, 3, 2, 1])
-    assert_series_equal(s.rolling_min(2), pl.Series("a", [None, 1, 2, 2, 1]))
-    assert_series_equal(s.rolling_max(2), pl.Series("a", [None, 2, 3, 3, 2]))
-    assert_series_equal(s.rolling_sum(2), pl.Series("a", [None, 3, 5, 5, 3]))
+@pytest.mark.parametrize("dtype", INTEGER_DTYPES)
+def test_rolling(dtype: PolarsDataType) -> None:
+    s = pl.Series("a", [1, 2, 3, 2, 1], dtype=dtype)
+    assert_series_equal(
+        s.rolling_min(2), pl.Series("a", [None, 1, 2, 2, 1], dtype=dtype)
+    )
+    assert_series_equal(
+        s.rolling_max(2), pl.Series("a", [None, 2, 3, 3, 2], dtype=dtype)
+    )
+    assert_series_equal(
+        s.rolling_sum(2), pl.Series("a", [None, 3, 5, 5, 3], dtype=dtype)
+    )
     assert_series_equal(s.rolling_mean(2), pl.Series("a", [None, 1.5, 2.5, 2.5, 1.5]))
 
     assert s.rolling_std(2).to_list()[1] == pytest.approx(0.7071067811865476)