fix: allow null index in list.get and array.get (#15239)

pola-rs · Mar 22, 2024 · dbbd3e8 · dbbd3e8
1 parent d214f55
commit dbbd3e8
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 6 deletions.
diff --git a/crates/polars-ops/src/chunked_array/array/get.rs b/crates/polars-ops/src/chunked_array/array/get.rs
@@ -26,7 +26,7 @@ pub fn array_get(ca: &ArrayChunked, index: &Int64Chunked) -> PolarsResult<Series
             if let Some(index) = index {
                 array_get_literal(ca, index)
             } else {
-                polars_bail!(ComputeError: "unexpected null index received in `arr.get`")
+                Ok(Series::full_null(ca.name(), ca.len(), &ca.inner_dtype()))
             }
         },
         len if len == ca.len() => {

diff --git a/crates/polars-plan/src/dsl/function_expr/list.rs b/crates/polars-plan/src/dsl/function_expr/list.rs
@@ -425,7 +425,11 @@ pub(super) fn get(s: &mut [Series]) -> PolarsResult<Option<Series>> {
             if let Some(index) = index {
                 ca.lst_get(index).map(Some)
             } else {
-                polars_bail!(ComputeError: "unexpected null index received in `list.get`")
+                Ok(Some(Series::full_null(
+                    ca.name(),
+                    ca.len(),
+                    &ca.inner_dtype(),
+                )))
             }
         },
         len if len == ca.len() => {

diff --git a/py-polars/tests/unit/namespaces/array/test_array.py b/py-polars/tests/unit/namespaces/array/test_array.py
@@ -128,22 +128,38 @@ def test_array_arg_min_max() -> None:
 
 
 def test_array_get() -> None:
-    # test index literal
     s = pl.Series(
         "a",
         [[1, 2, 3, 4], [5, 6, None, None], [7, 8, 9, 10]],
         dtype=pl.Array(pl.Int64, 4),
     )
+
+    # Test index literal.
     out = s.arr.get(1)
     expected = pl.Series("a", [2, 6, 8], dtype=pl.Int64)
     assert_series_equal(out, expected)
 
-    # test index expr
-    out = s.arr.get(pl.Series([1, -2, 4]))
+    # Null index literal.
+    out_df = s.to_frame().select(pl.col.a.arr.get(pl.lit(None)))
+    expected_df = pl.Series("a", [None, None, None], dtype=pl.Int64).to_frame()
+    assert_frame_equal(out_df, expected_df)
+
+    # Out-of-bounds index literal.
+    out = s.arr.get(100)
+    expected = pl.Series("a", [None, None, None], dtype=pl.Int64)
+    assert_series_equal(out, expected)
+
+    # Negative index literal.
+    out = s.arr.get(-2)
+    expected = pl.Series("a", [3, None, 9], dtype=pl.Int64)
+    assert_series_equal(out, expected)
+
+    # Test index expr.
+    out = s.arr.get(pl.Series([1, -2, 100]))
     expected = pl.Series("a", [2, None, None], dtype=pl.Int64)
     assert_series_equal(out, expected)
 
-    # test logical type
+    # Test logical type.
     s = pl.Series(
         "a",
         [

diff --git a/py-polars/tests/unit/namespaces/list/test_list.py b/py-polars/tests/unit/namespaces/list/test_list.py
@@ -30,6 +30,16 @@ def test_list_arr_get() -> None:
     out = pl.select(pl.lit(a).list.last()).to_series()
     assert_series_equal(out, expected)
 
+    # Out of bounds index.
+    out = a.list.get(3)
+    expected = pl.Series("a", [None, None, 9])
+    assert_series_equal(out, expected)
+
+    # Null index.
+    out_df = a.to_frame().select(pl.col.a.list.get(pl.lit(None)))
+    expected_df = pl.Series("a", [None, None, None], dtype=pl.Int64).to_frame()
+    assert_frame_equal(out_df, expected_df)
+
     a = pl.Series("a", [[1, 2, 3], [4, 5], [6, 7, 8, 9]])
     out = a.list.get(-3)
     expected = pl.Series("a", [1, None, 7])