fix(rust, python): Fix struct get field by index out of bounds error. (…

…#10097)
pola-rs · Jul 27, 2023 · e50e1d9 · e50e1d9
1 parent ff56742
commit e50e1d9
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs
@@ -177,9 +177,15 @@ impl FunctionExpr {
                 match s {
                     FieldByIndex(index) => {
                         let (index, _) = slice_offsets(*index, 0, fields.len());
-                        fields.get(index).cloned().ok_or_else(
-                            || polars_err!(ComputeError: "index out of bounds in `struct.field`"),
-                        )
+                        if let DataType::Struct(flds) = &fields[0].dtype {
+                            flds.get(index).cloned().ok_or_else(
+                                || polars_err!(ComputeError: "index out of bounds in `struct.field`")
+                            )
+                        } else {
+                            polars_bail!(
+                                ComputeError: "expected struct dtype, got: `{}`", &fields[0].dtype
+                            )
+                        }
                     }
                     FieldByName(name) => {
                         if let DataType::Struct(flds) = &fields[0].dtype {

diff --git a/py-polars/tests/unit/datatypes/test_struct.py b/py-polars/tests/unit/datatypes/test_struct.py
@@ -857,3 +857,9 @@ def test_struct_null_count_strict_cast() -> None:
     s = pl.Series([{"a": None}]).cast(pl.Struct({"a": pl.Categorical}))
     assert s.dtype == pl.Struct([pl.Field("a", pl.Categorical)])
     assert s.to_list() == [{"a": None}]
+
+
+def test_struct_get_field_by_index() -> None:
+    df = pl.DataFrame({"val": [{"a": 1, "b": 2}]})
+    expected = {"b": [2]}
+    assert df.select(pl.all().struct[1]).to_dict(as_series=False) == expected