Skip to content

Commit

Permalink
fix: recursively check allowed streaming dtypes (#11879)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Oct 20, 2023
1 parent 7b9f10e commit 27a4fe2
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
4 changes: 4 additions & 0 deletions crates/polars-lazy/src/physical_plan/streaming/convert_alp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -357,6 +357,10 @@ pub(crate) fn insert_streaming_nodes(
DataType::Object(_) => false,
#[cfg(feature = "dtype-categorical")]
DataType::Categorical(_) => string_cache,
DataType::List(inner) => allowed_dtype(inner, string_cache),
DataType::Struct(fields) => fields
.iter()
.all(|fld| allowed_dtype(fld.data_type(), string_cache)),
_ => true,
}
}
Expand Down
14 changes: 14 additions & 0 deletions py-polars/tests/unit/streaming/test_streaming_categoricals.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
import polars as pl


def test_streaming_nested_categorical() -> None:
assert (
pl.LazyFrame({"numbers": [1, 1, 2], "cat": [["str"], ["foo"], ["bar"]]})
.with_columns(pl.col("cat").cast(pl.List(pl.Categorical)))
.group_by("numbers")
.agg(pl.col("cat").first())
.sort("numbers")
).collect(streaming=True).to_dict(False) == {
"numbers": [1, 2],
"cat": [["str"], ["bar"]],
}

0 comments on commit 27a4fe2

Please sign in to comment.