feat(python): Allow pl.Int32 Series as output in eager repeat.

py-polars/polars/internals/lazy_functions.py

@@ -11,6 +11,7 @@
     Date,
     Datetime,
     Duration,
+    Int32,
     Int64,
     PolarsDataType,
     SchemaDict,
@@ -2791,6 +2792,12 @@ def repeat(
         if name is None:
             name = ""
         dtype = py_type_to_dtype(type(value))
+        if (
+            dtype == Int64
+            and isinstance(value, int)
+            and -(2**31) <= value <= 2**31 - 1
+        ):
+            dtype = Int32
         s = pli.Series._repeat(name, value, n, dtype)  # type: ignore[arg-type]
         return s
     else:

py-polars/src/lib.rs

@@ -154,7 +154,7 @@ fn repeat(value: &PyAny, n_times: PyExpr) -> PyResult<PyExpr> {
     } else if let Ok(int) = value.downcast::<PyInt>() {
         let val = int.extract::<i64>().unwrap();
 
-        if val >= 0 && val <= i32::MAX as i64 || val < 0 && val >= i32::MIN as i64 {
+        if val >= i32::MIN as i64 && val <= i32::MAX as i64 {
             Ok(polars_rs::lazy::dsl::repeat(val as i32, n_times.inner).into())
         } else {
             Ok(polars_rs::lazy::dsl::repeat(val, n_times.inner).into())

py-polars/src/series.rs

@@ -285,6 +285,12 @@ impl PySeries {
                 ca.rename(name);
                 ca.into_inner().into_series().into()
             }
+            DataType::Int32 => {
+                let val = val.extract::<i32>().unwrap();
+                let mut ca: NoNull<Int32Chunked> = (0..n).map(|_| val).collect_trusted();
+                ca.rename(name);
+                ca.into_inner().into_series().into()
+            }
             DataType::Float64 => {
                 let val = val.extract::<f64>().unwrap();
                 let mut ca: NoNull<Float64Chunked> = (0..n).map(|_| val).collect_trusted();

py-polars/tests/unit/test_functions.py

@@ -317,3 +317,41 @@ def test_fill_null_unknown_output_type() -> None:
             148.4131591025766,
         ]
     }
+
+
+def test_repeat() -> None:
+    s = pl.select(pl.repeat(2**31 - 1, 3)).to_series()
+    assert s.dtype == pl.Int32
+    assert s.len() == 3
+    assert s.to_list() == [2**31 - 1] * 3
+    s = pl.select(pl.repeat(-(2**31), 4)).to_series()
+    assert s.dtype == pl.Int32
+    assert s.len() == 4
+    assert s.to_list() == [-(2**31)] * 4
+    s = pl.select(pl.repeat(2**31, 5)).to_series()
+    assert s.dtype == pl.Int64
+    assert s.len() == 5
+    assert s.to_list() == [2**31] * 5
+    s = pl.select(pl.repeat(-(2**31) - 1, 3)).to_series()
+    assert s.dtype == pl.Int64
+    assert s.len() == 3
+    assert s.to_list() == [-(2**31) - 1] * 3
+    s = pl.select(pl.repeat("foo", 2)).to_series()
+    assert s.dtype == pl.Utf8
+    assert s.len() == 2
+    assert s.to_list() == ["foo"] * 2
+    s = pl.select(pl.repeat(1.0, 5)).to_series()
+    assert s.dtype == pl.Float64
+    assert s.len() == 5
+    assert s.to_list() == [1.0] * 5
+    s = pl.select(pl.repeat(True, 4)).to_series()
+    assert s.dtype == pl.Boolean
+    assert s.len() == 4
+    assert s.to_list() == [True] * 4
+    s = pl.select(pl.repeat(None, 7)).to_series()
+    assert s.dtype == pl.Null
+    assert s.len() == 7
+    assert s.to_list() == [None] * 7
+    s = pl.select(pl.repeat(0, 0)).to_series()
+    assert s.dtype == pl.Int32
+    assert s.len() == 0

py-polars/tests/unit/test_series.py

@@ -960,22 +960,41 @@ def test_object() -> None:
 
 
 def test_repeat() -> None:
-    s = pl.repeat(1, 10, eager=True)
+    s = pl.repeat(2**31 - 1, 3, eager=True)
+    assert s.dtype == pl.Int32
+    assert s.len() == 3
+    assert s.to_list() == [2**31 - 1] * 3
+    s = pl.repeat(-(2**31), 4, eager=True)
+    assert s.dtype == pl.Int32
+    assert s.len() == 4
+    assert s.to_list() == [-(2**31)] * 4
+    s = pl.repeat(2**31, 5, eager=True)
+    assert s.dtype == pl.Int64
+    assert s.len() == 5
+    assert s.to_list() == [2**31] * 5
+    s = pl.repeat(-(2**31) - 1, 3, eager=True)
     assert s.dtype == pl.Int64
-    assert s.len() == 10
-    s = pl.repeat("foo", 10, eager=True)
+    assert s.len() == 3
+    assert s.to_list() == [-(2**31) - 1] * 3
+    s = pl.repeat("foo", 2, eager=True)
     assert s.dtype == pl.Utf8
-    assert s.len() == 10
+    assert s.len() == 2
+    assert s.to_list() == ["foo"] * 2
     s = pl.repeat(1.0, 5, eager=True)
     assert s.dtype == pl.Float64
     assert s.len() == 5
-    assert s.to_list() == [1.0, 1.0, 1.0, 1.0, 1.0]
-    s = pl.repeat(True, 5, eager=True)
+    assert s.to_list() == [1.0] * 5
+    s = pl.repeat(True, 4, eager=True)
     assert s.dtype == pl.Boolean
-    assert s.len() == 5
+    assert s.len() == 4
+    assert s.to_list() == [True] * 4
     s = pl.repeat(None, 7, eager=True)
     assert s.dtype == pl.Null
     assert s.len() == 7
+    assert s.to_list() == [None] * 7
+    s = pl.repeat(0, 0, eager=True)
+    assert s.dtype == pl.Int32
+    assert s.len() == 0
 
 
 def test_shape() -> None: