fix: Fix list gather invalid fast path (#19299)

pola-rs · Oct 18, 2024 · 8c90286 · 8c90286
1 parent 960007d
commit 8c90286
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/crates/polars-plan/src/dsl/function_expr/list.rs b/crates/polars-plan/src/dsl/function_expr/list.rs
@@ -503,7 +503,7 @@ pub(super) fn gather(args: &[Column], null_on_oob: bool) -> PolarsResult<Column>
     let idx = &args[1];
     let ca = ca.list()?;
 
-    if idx.len() == 1 && null_on_oob {
+    if idx.len() == 1 && idx.dtype().is_numeric() && null_on_oob {
         // fast path
         let idx = idx.get(0)?.try_extract::<i64>()?;
         let out = ca.lst_get(idx, null_on_oob).map(Column::from)?;

diff --git a/py-polars/tests/unit/operations/test_gather.py b/py-polars/tests/unit/operations/test_gather.py
@@ -156,3 +156,14 @@ def test_gather_str_col_18099() -> None:
         "foo": [1, 1, 2],
         "idx": [0, 0, 1],
     }
+
+
+def test_gather_list_19243() -> None:
+    df = pl.DataFrame({"a": [[0.1, 0.2, 0.3]]})
+    assert df.with_columns(pl.lit([0]).alias("c")).with_columns(
+        gather=pl.col("a").list.gather(pl.col("c"), null_on_oob=True)
+    ).to_dict(as_series=False) == {
+        "a": [[0.1, 0.2, 0.3]],
+        "c": [[0]],
+        "gather": [[0.1]],
+    }