Skip to content

Commit

Permalink
fix(rust, python): object filter should keep single chunk (#5913)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Dec 27, 2022
1 parent 486c7e7 commit 882fb58
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 1 deletion.
4 changes: 4 additions & 0 deletions polars/polars-core/src/chunked_array/object/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -142,4 +142,8 @@ where
length: len as IdxSize,
}
}

pub fn new_empty(name: &str) -> Self {
Self::new_from_vec(name, vec![])
}
}
2 changes: 1 addition & 1 deletion polars/polars-core/src/chunked_array/ops/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ where
if filter.len() == 1 {
return match filter.get(0) {
Some(true) => Ok(self.clone()),
_ => Ok(ObjectChunked::from_chunks(self.name(), vec![])),
_ => Ok(ObjectChunked::new_empty(self.name())),
};
}
if self.is_empty() {
Expand Down
19 changes: 19 additions & 0 deletions py-polars/tests/unit/test_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,3 +15,22 @@ def test_object_when_then_4702() -> None:
"Type": [pl.Date, pl.UInt8],
"New_Type": [pl.UInt16, pl.UInt8],
}


def test_object_empty_filter_5911() -> None:
df = pl.DataFrame(
data=[
(1, "dog", {}),
],
columns=[
("pet_id", pl.Int64),
("pet_type", pl.Categorical),
("pet_obj", pl.Object),
],
orient="row",
)

empty_df = df.filter(pl.col("pet_type") == "cat")
out = empty_df.select(["pet_obj"])
assert out.dtypes == [pl.Object]
assert out.shape == (0, 1)

0 comments on commit 882fb58

Please sign in to comment.