fix: Implement list.min and list.max for list[i128] (#20488)

pola-rs · Dec 29, 2024 · 5671b0f · 5671b0f
1 parent 9a44c02
commit 5671b0f
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 2 deletions.
diff --git a/crates/polars-lazy/src/physical_plan/streaming/convert_alp.rs b/crates/polars-lazy/src/physical_plan/streaming/convert_alp.rs
@@ -356,6 +356,7 @@ pub(crate) fn insert_streaming_nodes(
                         DataType::Unknown(_) => false,
                         #[cfg(feature = "dtype-decimal")]
                         DataType::Decimal(_, _) => false,
+                        DataType::Int128 => false,
                         _ => true,
                     }
                 }

diff --git a/crates/polars-ops/src/chunked_array/list/min_max.rs b/crates/polars-ops/src/chunked_array/list/min_max.rs
@@ -55,6 +55,7 @@ fn min_list_numerical(ca: &ListChunked, inner_type: &DataType) -> Series {
                 Int16 => dispatch_min::<i16>(values, offsets, arr.validity()),
                 Int32 => dispatch_min::<i32>(values, offsets, arr.validity()),
                 Int64 => dispatch_min::<i64>(values, offsets, arr.validity()),
+                Int128 => dispatch_min::<i128>(values, offsets, arr.validity()),
                 UInt8 => dispatch_min::<u8>(values, offsets, arr.validity()),
                 UInt16 => dispatch_min::<u16>(values, offsets, arr.validity()),
                 UInt32 => dispatch_min::<u32>(values, offsets, arr.validity()),
@@ -164,6 +165,7 @@ fn max_list_numerical(ca: &ListChunked, inner_type: &DataType) -> Series {
                 Int16 => dispatch_max::<i16>(values, offsets, arr.validity()),
                 Int32 => dispatch_max::<i32>(values, offsets, arr.validity()),
                 Int64 => dispatch_max::<i64>(values, offsets, arr.validity()),
+                Int128 => dispatch_max::<i128>(values, offsets, arr.validity()),
                 UInt8 => dispatch_max::<u8>(values, offsets, arr.validity()),
                 UInt16 => dispatch_max::<u16>(values, offsets, arr.validity()),
                 UInt32 => dispatch_max::<u32>(values, offsets, arr.validity()),

diff --git a/py-polars/polars/_utils/constants.py b/py-polars/polars/_utils/constants.py
@@ -5,10 +5,12 @@
 I16_MIN = -(2**15)
 I32_MIN = -(2**31)
 I64_MIN = -(2**63)
+I128_MIN = -(2**127)
 I8_MAX = 2**7 - 1
 I16_MAX = 2**15 - 1
 I32_MAX = 2**31 - 1
 I64_MAX = 2**63 - 1
+I128_MAX = 2**127 - 1
 U8_MAX = 2**8 - 1
 U16_MAX = 2**16 - 1
 U32_MAX = 2**32 - 1

diff --git a/py-polars/polars/testing/parametric/strategies/data.py b/py-polars/polars/testing/parametric/strategies/data.py
@@ -21,6 +21,8 @@
     I32_MIN,
     I64_MAX,
     I64_MIN,
+    I128_MAX,
+    I128_MIN,
     U8_MAX,
     U16_MAX,
     U32_MAX,
@@ -43,6 +45,7 @@
     Int16,
     Int32,
     Int64,
+    Int128,
     List,
     Null,
     Object,
@@ -78,6 +81,7 @@
         16: st.integers(I16_MIN, I16_MAX),
         32: st.integers(I32_MIN, I32_MAX),
         64: st.integers(I64_MIN, I64_MAX),
+        128: st.integers(I128_MIN, I128_MAX),
     },
     False: {
         8: st.integers(0, U8_MAX),
@@ -89,7 +93,7 @@
 
 
 def integers(
-    bit_width: Literal[8, 16, 32, 64] = 64, *, signed: bool = True
+    bit_width: Literal[8, 16, 32, 64, 128] = 64, *, signed: bool = True
 ) -> SearchStrategy[int]:
     """Create a strategy for generating integers."""
     return _INTEGER_STRATEGIES[signed][bit_width]
@@ -356,6 +360,7 @@ def objects() -> SearchStrategy[object]:
     Int16: integers(16, signed=True),
     Int32: integers(32, signed=True),
     Int64: integers(64, signed=True),
+    Int128: integers(128, signed=True),
     UInt8: integers(8, signed=False),
     UInt16: integers(16, signed=False),
     UInt32: integers(32, signed=False),

diff --git a/py-polars/tests/unit/conftest.py b/py-polars/tests/unit/conftest.py
@@ -22,7 +22,7 @@
 )
 
 # Data type groups
-SIGNED_INTEGER_DTYPES = [pl.Int8(), pl.Int16(), pl.Int32(), pl.Int64()]
+SIGNED_INTEGER_DTYPES = [pl.Int8(), pl.Int16(), pl.Int32(), pl.Int64(), pl.Int128()]
 UNSIGNED_INTEGER_DTYPES = [pl.UInt8(), pl.UInt16(), pl.UInt32(), pl.UInt64()]
 INTEGER_DTYPES = SIGNED_INTEGER_DTYPES + UNSIGNED_INTEGER_DTYPES
 FLOAT_DTYPES = [pl.Float32(), pl.Float64()]