fix: Fix categorical (#106)

pola-rs · Sep 12, 2024 · 7a353c3 · 7a353c3
1 parent 632776a
commit 7a353c3
Showing 1 changed file with 7 additions and 8 deletions.
diff --git a/pyo3-polars/src/types.rs b/pyo3-polars/src/types.rs
@@ -2,8 +2,6 @@ use super::*;
 use crate::error::PyPolarsErr;
 use crate::ffi::to_py::to_py_array;
 use polars::export::arrow;
-#[cfg(feature = "dtype-categorical")]
-use polars_core::datatypes::create_enum_data_type;
 use polars_core::datatypes::{CompatLevel, DataType};
 use polars_core::prelude::*;
 use polars_core::utils::materialize_dyn_int;
@@ -343,7 +341,7 @@ impl IntoPy<PyObject> for PyExpr {
     }
 }
 
-#[cfg(feature = "dtype-full")]
+#[cfg(feature = "dtype-categorical")]
 pub(crate) fn to_series(py: Python, s: PySeries) -> PyObject {
     let series = SERIES.bind(py);
     let constructor = series
@@ -575,11 +573,12 @@ impl<'py> FromPyObject<'py> for PyDataType {
             #[cfg(feature = "dtype-categorical")]
             "Categorical" => {
                 let ordering = ob.getattr(intern!(py, "ordering")).unwrap();
-
-                let ordering = match ordering.extract::<&str>()? {
-                    "physical" => CategoricalOrdering::Physical,
-                    "lexical" => CategoricalOrdering::Lexical,
+                let ordering = ordering.extract::<PyBackedStr>()?;
+                let ordering = match  ordering.as_bytes() {
+                    b"physical" => CategoricalOrdering::Physical,
+                    b"lexical" => CategoricalOrdering::Lexical,
                     ordering => {
+                        let ordering = std::str::from_utf8(ordering).unwrap();
                         return Err(PyValueError::new_err(format!("invalid ordering argument: {ordering}")))
                     }
                 };
@@ -592,7 +591,7 @@ impl<'py> FromPyObject<'py> for PyDataType {
                 let s = get_series(&categories.as_borrowed())?;
                 let ca = s.str().map_err(PyPolarsErr::from)?;
                 let categories = ca.downcast_iter().next().unwrap().clone();
-                create_enum_data_type(categories)
+                DataType::Enum(Some(Arc::new(RevMapping::build_local(categories))), Default::default())
             },
             "Date" => DataType::Date,
             "Time" => DataType::Time,