fix and test

pola-rs · Jan 18, 2025 · 08a58a2 · 08a58a2
1 parent f34e9a8
commit 08a58a2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/crates/polars-ops/src/chunked_array/repeat_by.rs b/crates/polars-ops/src/chunked_array/repeat_by.rs
@@ -14,13 +14,12 @@ fn check_lengths(length_srs: usize, length_by: usize) -> PolarsResult<()> {
 }
 
 fn new_by(by: &IdxCa, len: usize) -> IdxCa {
-    let values = if let Some(x) = by.get(0) {
-        std::iter::repeat(x).take(len).collect::<Vec<IdxSize>>()
+    if let Some(x) = by.get(0) {
+        let values = std::iter::repeat(x).take(len).collect::<Vec<IdxSize>>();
+        return IdxCa::new(PlSmallStr::EMPTY, values);
     } else {
-        println!("HEREE");
-        vec![0u32; len]
-    };
-    IdxCa::new(PlSmallStr::EMPTY, values)
+        return IdxCa::full_null(PlSmallStr::EMPTY, len);
+    }
 }
 
 fn repeat_by_primitive<T>(ca: &ChunkedArray<T>, by: &IdxCa) -> PolarsResult<ListChunked>

diff --git a/py-polars/tests/unit/functions/test_repeat.py b/py-polars/tests/unit/functions/test_repeat.py
@@ -207,3 +207,14 @@ def test_repeat_by_none_13053(data: list[Any], expected_data: list[list[Any]]) -
     res = df.select(repeat=pl.col("x").repeat_by("by"))
     expected = pl.Series("repeat", expected_data)
     assert_series_equal(res.to_series(), expected)
+
+
+def test_repeat_by_literal_none_20268() -> None:
+    df = pl.DataFrame({"x": ["a", "b"]})
+    expected = pl.Series("repeat", [None, None], dtype=pl.List(pl.String))
+
+    res = df.select(repeat=pl.col("x").repeat_by(pl.lit(None)))
+    assert_series_equal(res.to_series(), expected)
+
+    res = df.select(repeat=pl.col("x").repeat_by(None))
+    assert_series_equal(res.to_series(), expected)