Skip to content

Commit

Permalink
fix: Struct filter by index (#18778)
Browse files Browse the repository at this point in the history
  • Loading branch information
barak1412 authored Sep 22, 2024
1 parent 8276084 commit 545117c
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 0 deletions.
5 changes: 5 additions & 0 deletions crates/polars-plan/src/plans/conversion/dsl_to_ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -811,6 +811,11 @@ fn expand_filter(
| Expr::DtypeColumn(_)
| Expr::IndexColumn(_)
| Expr::Nth(_) => true,
#[cfg(feature = "dtype-struct")]
Expr::Function {
function: FunctionExpr::StructExpr(StructFunction::FieldByIndex(_)),
..
} => true,
_ => false,
}) {
let mut rewritten = rewrite_projections(vec![predicate], &schema, &[], opt_flags)?;
Expand Down
22 changes: 22 additions & 0 deletions py-polars/tests/unit/operations/namespaces/test_struct.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,12 @@
import datetime
from collections import OrderedDict

import pytest

import polars as pl
from polars.exceptions import (
OutOfBoundsError,
)
from polars.testing import assert_frame_equal


Expand Down Expand Up @@ -95,3 +100,20 @@ def test_empty_list_eval_schema_5734() -> None:
assert df.filter(False).select(
pl.col("a").list.eval(pl.element().struct.field("b"))
).schema == {"a": pl.List(pl.Int64)}


def test_field_by_index_18732() -> None:
df = pl.DataFrame({"foo": [{"a": 1, "b": 2}, {"a": 2, "b": 1}]})

# illegal upper bound
with pytest.raises(OutOfBoundsError, match=r"index 2 for length: 2"):
df.filter(pl.col.foo.struct[2] == 1)

# legal
expected_df = pl.DataFrame({"foo": [{"a": 1, "b": 2}]})
result_df = df.filter(pl.col.foo.struct[0] == 1)
assert_frame_equal(expected_df, result_df)

expected_df = pl.DataFrame({"foo": [{"a": 2, "b": 1}]})
result_df = df.filter(pl.col.foo.struct[-1] == 1)
assert_frame_equal(expected_df, result_df)

0 comments on commit 545117c

Please sign in to comment.