perf: Use single-key optimization with Categorical (#23436)

Author: orlp

crates/polars-expr/src/groups/mod.rs

@@ -87,8 +87,8 @@ pub fn new_hash_grouper(key_schema: Arc<Schema>) -> Box<dyn Grouper> {
                 Box::new(single_key::SingleKeyHashGrouper::<Int128Type>::new())
             },
             #[cfg(feature = "dtype-categorical")]
-            DataType::Enum(fcats, _) => {
-                with_match_categorical_physical_type!(fcats.physical(), |$C| {
+            dt @ (DataType::Enum(_, _) | DataType::Categorical(_, _)) => {
+                with_match_categorical_physical_type!(dt.cat_physical().unwrap(), |$C| {
                     Box::new(single_key::SingleKeyHashGrouper::<<$C as PolarsCategoricalType>::PolarsPhysical>::new())
                 })
             },

crates/polars-expr/src/hash_keys.rs

@@ -29,7 +29,7 @@ pub fn hash_keys_variant_for_dtype(dt: &DataType) -> HashKeysVariant {
         #[cfg(feature = "dtype-decimal")]
         DataType::Decimal(_, _) => HashKeysVariant::Single,
         #[cfg(feature = "dtype-categorical")]
-        DataType::Enum(_, _) => HashKeysVariant::Single,
+        DataType::Enum(_, _) | DataType::Categorical(_, _) => HashKeysVariant::Single,
 
         DataType::String | DataType::Binary => HashKeysVariant::Binview,
 
@@ -76,8 +76,8 @@ macro_rules! downcast_single_key_ca {
             #[cfg(feature = "dtype-decimal")]
             DataType::Decimal(..) => { let $ca = $self.decimal().unwrap().physical(); $($body)* },
             #[cfg(feature = "dtype-categorical")]
-            DataType::Enum(fcats, _) => {
-                match fcats.physical() {
+            dt @ (DataType::Enum(_, _) | DataType::Categorical(_, _)) => {
+                match dt.cat_physical().unwrap() {
                     CategoricalPhysical::U8 => { let $ca = $self.cat8().unwrap().physical(); $($body)* },
                     CategoricalPhysical::U16 => { let $ca = $self.cat16().unwrap().physical(); $($body)* },
                     CategoricalPhysical::U32 => { let $ca = $self.cat32().unwrap().physical(); $($body)* },

crates/polars-expr/src/hot_groups/mod.rs

@@ -82,8 +82,8 @@ pub fn new_hash_hot_grouper(key_schema: Arc<Schema>, num_groups: usize) -> Box<d
             #[cfg(feature = "dtype-decimal")]
             DataType::Decimal(_, _) => Box::new(SK::<Int128Type>::new(dt, ng)),
             #[cfg(feature = "dtype-categorical")]
-            DataType::Enum(ref fcats, _) => {
-                with_match_categorical_physical_type!(fcats.physical(), |$C| {
+            dt @ (DataType::Enum(_, _) | DataType::Categorical(_, _)) => {
+                with_match_categorical_physical_type!(dt.cat_physical().unwrap(), |$C| {
                     Box::new(SK::<<$C as PolarsCategoricalType>::PolarsPhysical>::new(dt.clone(), ng))
                 })
             },

crates/polars-expr/src/idx_table/mod.rs

@@ -104,8 +104,8 @@ pub fn new_idx_table(key_schema: Arc<Schema>) -> Box<dyn IdxTable> {
             #[cfg(feature = "dtype-decimal")]
             DataType::Decimal(_, _) => Box::new(SKIT::<Int128Type>::new()),
             #[cfg(feature = "dtype-categorical")]
-            DataType::Enum(fcats, _) => {
-                with_match_categorical_physical_type!(fcats.physical(), |$C| {
+            dt @ (DataType::Enum(_, _) | DataType::Categorical(_, _)) => {
+                with_match_categorical_physical_type!(dt.cat_physical().unwrap(), |$C| {
                     Box::new(SKIT::<<$C as PolarsCategoricalType>::PolarsPhysical>::new())
                 })
             },