refactor: Rework Categorical/Enum to use (Frozen)Categories

Cargo.toml

@@ -102,6 +102,7 @@ zstd = "0.13"
 polars = { version = "0.49.1", path = "crates/polars", default-features = false }
 polars-compute = { version = "0.49.1", path = "crates/polars-compute", default-features = false }
 polars-core = { version = "0.49.1", path = "crates/polars-core", default-features = false }
+polars-dtype = { version = "0.49.1", path = "crates/polars-dtype", default-features = false }
 polars-dylib = { version = "0.49.1", path = "crates/polars-dylib", default-features = false }
 polars-error = { version = "0.49.1", path = "crates/polars-error", default-features = false }
 polars-expr = { version = "0.49.1", path = "crates/polars-expr", default-features = false }

crates/polars-arrow/src/bitmap/builder.rs

@@ -75,6 +75,13 @@ impl BitmapBuilder {
         self.bit_cap = words_available * 64;
     }
 
+    pub fn clear(&mut self) {
+        self.buf = 0;
+        self.bit_len = 0;
+        self.set_bits_in_bytes = 0;
+        self.bytes.clear();
+    }
+
     #[inline(always)]
     pub fn push(&mut self, x: bool) {
         self.reserve(1);

crates/polars-arrow/src/datatypes/field.rs

@@ -6,8 +6,15 @@ use serde::{Deserialize, Serialize};
 
 use super::{ArrowDataType, Metadata};
 
-pub static DTYPE_ENUM_VALUES: &str = "_PL_ENUM_VALUES";
-pub static DTYPE_CATEGORICAL: &str = "_PL_CATEGORICAL";
+// These two have the same encoding, but because older versions of Polars
+// were unable to read non-u32-key arrow dictionaries while _PL_ENUM_VALUES
+// is set we switched to a new version.
+pub static DTYPE_ENUM_VALUES_LEGACY: &str = "_PL_ENUM_VALUES";
+pub static DTYPE_ENUM_VALUES_NEW: &str = "_PL_ENUM_VALUES2";
+
+// These have different encodings.
+pub static DTYPE_CATEGORICAL_LEGACY: &str = "_PL_CATEGORICAL";
+pub static DTYPE_CATEGORICAL_NEW: &str = "_PL_CATEGORICAL2";
 
 /// Represents Arrow's metadata of a "column".
 ///
@@ -71,15 +78,15 @@ impl Field {
 
     pub fn is_enum(&self) -> bool {
         if let Some(md) = &self.metadata {
-            md.get(DTYPE_ENUM_VALUES).is_some()
+            md.get(DTYPE_ENUM_VALUES_LEGACY).is_some() || md.get(DTYPE_ENUM_VALUES_NEW).is_some()
         } else {
             false
         }
     }
 
     pub fn is_categorical(&self) -> bool {
         if let Some(md) = &self.metadata {
-            md.get(DTYPE_CATEGORICAL).is_some()
+            md.get(DTYPE_CATEGORICAL_LEGACY).is_some() || md.get(DTYPE_CATEGORICAL_NEW).is_some()
         } else {
             false
         }

crates/polars-arrow/src/datatypes/mod.rs

@@ -8,7 +8,10 @@ mod schema;
 use std::collections::BTreeMap;
 use std::sync::Arc;
 
-pub use field::{DTYPE_CATEGORICAL, DTYPE_ENUM_VALUES, Field};
+pub use field::{
+    DTYPE_CATEGORICAL_LEGACY, DTYPE_CATEGORICAL_NEW, DTYPE_ENUM_VALUES_LEGACY,
+    DTYPE_ENUM_VALUES_NEW, Field,
+};
 pub use physical_type::*;
 use polars_utils::pl_str::PlSmallStr;
 pub use schema::{ArrowSchema, ArrowSchemaRef};

crates/polars-core/Cargo.toml

@@ -10,6 +10,7 @@ description = "Core of the Polars DataFrame library"
 
 [dependencies]
 polars-compute = { workspace = true, features = ["gather"] }
+polars-dtype = { workspace = true }
 polars-error = { workspace = true }
 polars-row = { workspace = true }
 polars-schema = { workspace = true }
@@ -33,7 +34,7 @@ rand = { workspace = true, optional = true, features = ["small_rng", "std"] }
 rand_distr = { workspace = true, optional = true }
 rayon = { workspace = true }
 regex = { workspace = true, optional = true }
-schemars = { workspace = true, optional = true }
+schemars = { workspace = true, optional = true, features = ["uuid1"] }
 # activate if you want serde support for Series and DataFrames
 serde = { workspace = true, optional = true }
 serde_json = { workspace = true, optional = true }
@@ -132,6 +133,7 @@ serde = [
   "polars-schema/serde",
   "polars-utils/serde",
   "polars-compute/serde",
+  "polars-dtype/serde",
   "arrow/io_ipc",
   "arrow/io_ipc_compression",
   "serde_json",
@@ -140,9 +142,10 @@ serde-lazy = ["serde", "arrow/serde", "indexmap/serde", "chrono/serde"]
 dsl-schema = [
   "serde",
   "dep:schemars",
+  "polars-compute/dsl-schema",
+  "polars-dtype/dsl-schema",
   "polars-schema/dsl-schema",
   "polars-utils/dsl-schema",
-  "polars-compute/dsl-schema",
 ]
 
 docs-selection = [

crates/polars-core/src/chunked_array/builder/categorical.rs

@@ -0,0 +1,82 @@
+use arrow::bitmap::BitmapBuilder;
+
+use crate::prelude::*;
+
+pub struct CategoricalChunkedBuilder<T: PolarsCategoricalType> {
+    name: PlSmallStr,
+    dtype: DataType,
+    mapping: Arc<CategoricalMapping>,
+    is_enum: bool,
+    cats: Vec<T::Native>,
+    validity: BitmapBuilder,
+}
+
+impl<T: PolarsCategoricalType> CategoricalChunkedBuilder<T> {
+    pub fn new(name: PlSmallStr, dtype: DataType) -> Self {
+        let (DataType::Categorical(_, mapping) | DataType::Enum(_, mapping)) = &dtype else {
+            panic!("non-Categorical/Enum dtype in CategoricalChunkedbuilder")
+        };
+        Self {
+            name,
+            mapping: mapping.clone(),
+            is_enum: matches!(dtype, DataType::Enum(_, _)),
+            dtype,
+            cats: Vec::new(),
+            validity: BitmapBuilder::new(),
+        }
+    }
+
+    pub fn dtype(&self) -> &DataType {
+        &self.dtype
+    }
+
+    pub fn reserve(&mut self, len: usize) {
+        self.cats.reserve(len);
+        self.validity.reserve(len);
+    }
+
+    pub fn append_cat(
+        &mut self,
+        cat: CatSize,
+        mapping: &Arc<CategoricalMapping>,
+    ) -> PolarsResult<()> {
+        if Arc::ptr_eq(&self.mapping, mapping) {
+            self.cats.push(T::Native::from_cat(cat));
+            self.validity.push(true);
+        } else if let Some(s) = mapping.cat_to_str(cat) {
+            self.append_str(s)?;
+        } else {
+            self.append_null();
+        }
+        Ok(())
+    }
+
+    pub fn append_str(&mut self, val: &str) -> PolarsResult<()> {
+        let cat = if self.is_enum {
+            self.mapping.get_cat(val).ok_or_else(|| {
+                polars_err!(ComputeError: "attempted to insert '{val}' into Enum which does not contain this string")
+            })?
+        } else {
+            self.mapping.insert_cat(val)?
+        };
+        self.cats.push(T::Native::from_cat(cat));
+        self.validity.push(true);
+        Ok(())
+    }
+
+    pub fn append_null(&mut self) {
+        self.cats.push(T::Native::default());
+        self.validity.push(false);
+    }
+
+    pub fn finish(self) -> CategoricalChunked<T> {
+        unsafe {
+            let phys = ChunkedArray::from_vec_validity(
+                self.name,
+                self.cats,
+                self.validity.into_opt_validity(),
+            );
+            CategoricalChunked::from_cats_and_dtype_unchecked(phys, self.dtype)
+        }
+    }
+}

crates/polars-core/src/chunked_array/builder/list/anonymous.rs

@@ -4,7 +4,7 @@ pub struct AnonymousListBuilder<'a> {
     name: PlSmallStr,
     builder: AnonymousBuilder<'a>,
     fast_explode: bool,
-    inner_dtype: DtypeMerger,
+    inner_dtype: Option<DataType>,
 }
 
 impl Default for AnonymousListBuilder<'_> {
@@ -19,7 +19,7 @@ impl<'a> AnonymousListBuilder<'a> {
             name,
             builder: AnonymousBuilder::new(capacity),
             fast_explode: true,
-            inner_dtype: DtypeMerger::new(inner_dtype),
+            inner_dtype,
         }
     }
 
@@ -59,13 +59,18 @@ impl<'a> AnonymousListBuilder<'a> {
     }
 
     pub fn append_series(&mut self, s: &'a Series) -> PolarsResult<()> {
-        match s.dtype() {
-            // Empty arrays tend to be null type and thus differ
-            // if we would push it the concat would fail.
-            DataType::Null if s.is_empty() => self.append_empty(),
-            dt => self.inner_dtype.update(dt)?,
+        match (s.dtype(), &self.inner_dtype) {
+            (DataType::Null, _) => {},
+            (dt, None) => self.inner_dtype = Some(dt.clone()),
+            (dt, Some(set_dt)) => {
+                polars_bail!(ComputeError: "dtypes don't match, got {dt}, expected: {set_dt}")
+            },
+        }
+        if s.is_empty() {
+            self.append_empty();
+        } else {
+            self.builder.push_multiple(s.chunks());
         }
-        self.builder.push_multiple(s.chunks());
         Ok(())
     }
 
@@ -76,19 +81,18 @@ impl<'a> AnonymousListBuilder<'a> {
             ListChunked::full_null_with_dtype(
                 slf.name.clone(),
                 0,
-                &slf.inner_dtype.materialize().unwrap_or(DataType::Null),
+                &slf.inner_dtype.unwrap_or(DataType::Null),
             )
         } else {
-            let inner_dtype = slf.inner_dtype.materialize();
-
-            let inner_dtype_physical = inner_dtype
+            let inner_dtype_physical = self
+                .inner_dtype
                 .as_ref()
                 .map(|dt| dt.to_physical().to_arrow(CompatLevel::newest()));
             let arr = slf.builder.finish(inner_dtype_physical.as_ref()).unwrap();
 
-            let list_dtype_logical = match inner_dtype {
+            let list_dtype_logical = match &self.inner_dtype {
                 None => DataType::from_arrow_dtype(arr.dtype()),
-                Some(dt) => DataType::List(Box::new(dt)),
+                Some(dt) => DataType::List(Box::new(dt.clone())),
             };
 
             let mut ca = ListChunked::with_chunk(PlSmallStr::EMPTY, arr);
@@ -105,7 +109,7 @@ pub struct AnonymousOwnedListBuilder {
     name: PlSmallStr,
     builder: AnonymousBuilder<'static>,
     owned: Vec<Series>,
-    inner_dtype: DtypeMerger,
+    inner_dtype: Option<DataType>,
     fast_explode: bool,
 }
 
@@ -117,11 +121,17 @@ impl Default for AnonymousOwnedListBuilder {
 
 impl ListBuilderTrait for AnonymousOwnedListBuilder {
     fn append_series(&mut self, s: &Series) -> PolarsResult<()> {
+        match (s.dtype(), &self.inner_dtype) {
+            (DataType::Null, _) => {},
+            (dt, None) => self.inner_dtype = Some(dt.clone()),
+            (dt, Some(set_dt)) => {
+                polars_ensure!(dt == set_dt, ComputeError: "dtypes don't match, got {dt}, expected: {set_dt}")
+            },
+        }
         if s.is_empty() {
             self.append_empty();
         } else {
             unsafe {
-                self.inner_dtype.update(s.dtype())?;
                 self.builder
                     .push_multiple(&*(s.chunks().as_ref() as *const [ArrayRef]));
             }
@@ -138,7 +148,7 @@ impl ListBuilderTrait for AnonymousOwnedListBuilder {
     }
 
     fn finish(&mut self) -> ListChunked {
-        let inner_dtype = std::mem::take(&mut self.inner_dtype).materialize();
+        let inner_dtype = std::mem::take(&mut self.inner_dtype);
         // Don't use self from here on out.
         let slf = std::mem::take(self);
         let inner_dtype_physical = inner_dtype
@@ -166,7 +176,7 @@ impl AnonymousOwnedListBuilder {
             name,
             builder: AnonymousBuilder::new(capacity),
             owned: Vec::with_capacity(capacity),
-            inner_dtype: DtypeMerger::new(inner_dtype),
+            inner_dtype,
             fast_explode: true,
         }
     }