nested

pola-rs · Apr 16, 2024 · 27512cb · 27512cb
1 parent 299eb05
commit 27512cb
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 15 deletions.
diff --git 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
@@ -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();