fix(rust, python): allow None in arr.slice length (#5934)

pola-rs · Dec 29, 2022 · 8e335cf · 8e335cf
1 parent 04b96e6
commit 8e335cf
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/list.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/list.rs
@@ -47,7 +47,11 @@ pub(super) fn slice(args: &mut [Series]) -> PolarsResult<Series> {
     let mut out: ListChunked = match (offset_s.len(), length_s.len()) {
         (1, 1) => {
             let offset = offset_s.get(0).unwrap().try_extract::<i64>()?;
-            let slice_len = length_s.get(0).unwrap().try_extract::<usize>()?;
+            let slice_len = length_s
+                .get(0)
+                .unwrap()
+                .extract::<usize>()
+                .unwrap_or(usize::MAX);
             return Ok(list_ca.lst_slice(offset, slice_len).into_series());
         }
         (1, length_slice_len) => {
@@ -73,7 +77,11 @@ pub(super) fn slice(args: &mut [Series]) -> PolarsResult<Series> {
             if offset_len != list_ca.len() {
                 return Err(PolarsError::ComputeError("the length of the slice 'offset' argument does not match that of the list column".into()));
             }
-            let length_slice = length_s.get(0).unwrap().try_extract::<usize>()?;
+            let length_slice = length_s
+                .get(0)
+                .unwrap()
+                .extract::<usize>()
+                .unwrap_or(usize::MAX);
             let offset_ca = offset_s.cast(&DataType::Int64)?;
             let offset_ca = offset_ca.i64().unwrap();
             list_ca

diff --git a/py-polars/tests/unit/test_lists.py b/py-polars/tests/unit/test_lists.py
@@ -572,3 +572,9 @@ def test_list_amortized_apply_explode_5812() -> None:
     assert s.arr.max().to_list() == [None, 3, 0, 2]
     assert s.arr.arg_min().to_list() == [None, 0, 1, 0]
     assert s.arr.arg_max().to_list() == [None, 1, 0, 1]
+
+
+def test_list_slice_5866() -> None:
+    vals = [[1, 2, 3, 4], [10, 2, 1]]
+    s = pl.Series("a", vals)
+    assert s.arr.slice(1).to_list() == [[2, 3, 4], [2, 1]]