diff --git a/crates/polars-core/src/datatypes/dtype.rs b/crates/polars-core/src/datatypes/dtype.rs index b7b941e85d8e..337612e122e8 100644 --- a/crates/polars-core/src/datatypes/dtype.rs +++ b/crates/polars-core/src/datatypes/dtype.rs @@ -283,6 +283,36 @@ impl DataType { } } + pub fn contains_categoricals(&self) -> bool { + use DataType::*; + match self { + #[cfg(feature = "dtype-categorical")] + Categorical(_, _) | Enum(_, _) => true, + List(inner) => inner.contains_categoricals(), + #[cfg(feature = "dtype-array")] + Array(inner, _) => inner.contains_categoricals(), + #[cfg(feature = "dtype-struct")] + Struct(fields) => fields + .iter() + .any(|field| field.dtype.contains_categoricals()), + _ => false, + } + } + + pub fn contains_objects(&self) -> bool { + use DataType::*; + match self { + #[cfg(feature = "object")] + Object(_, _) => true, + List(inner) => inner.contains_objects(), + #[cfg(feature = "dtype-array")] + Array(inner, _) => inner.contains_objects(), + #[cfg(feature = "dtype-struct")] + Struct(fields) => fields.iter().any(|field| field.dtype.contains_objects()), + _ => false, + } + } + /// Check if type is sortable pub fn is_ord(&self) -> bool { #[cfg(feature = "dtype-categorical")] diff --git a/crates/polars-lazy/src/physical_plan/planner/lp.rs b/crates/polars-lazy/src/physical_plan/planner/lp.rs index f465b22da218..ddb24ec48889 100644 --- a/crates/polars-lazy/src/physical_plan/planner/lp.rs +++ b/crates/polars-lazy/src/physical_plan/planner/lp.rs @@ -176,23 +176,20 @@ pub fn create_physical_plan( let input_schema = lp_arena.get(input).schema(lp_arena).into_owned(); if streamable { // This can cause problems with string caches - streamable = input_schema.iter_dtypes().all(|dt| match dt { - #[cfg(feature = "dtype-categorical")] - DataType::Categorical(_, _) => false, - #[cfg(feature = "object")] - DataType::Object(_, _) => false, - _ => true, - }) || { - #[cfg(feature = "dtype-categorical")] - { - polars_core::using_string_cache() - } + streamable = !input_schema + .iter_dtypes() + .any(|dt| dt.contains_categoricals()) + || { + #[cfg(feature = "dtype-categorical")] + { + polars_core::using_string_cache() + } - #[cfg(not(feature = "dtype-categorical"))] - { - false + #[cfg(not(feature = "dtype-categorical"))] + { + false + } } - } } let input = create_physical_plan(input, lp_arena, expr_arena)?; let mut state = ExpressionConversionState::default();