Repository: pola-rs/polars
File: crates/polars-core/src/chunked_array/logical/categorical.rs (Rust)

Code before:
                    validity.push(opt_s.is_some());
                }
            },
            DataType::Enum(fcats, mapping) => {
                assert!(fcats.physical() == T::physical());
                for opt_s in strings {
                    cat_ids.push(if let Some(cat) = opt_s.and_then(|s| mapping.get_cat(s)) {
                        validity.push(true);
                        T::Native::from_cat(cat)
                    } else {
                        validity.push(false);
                        T::Native::zero()
                    });
                }
            },
            _ => panic!("from_strings_and_dtype_strict called on non-categorical type"),
        }

        let arr = <T::PolarsPhysical as PolarsDataType>::Array::from_vec(cat_ids)
            .with_validity(validity.into_opt_validity());
        let phys = ChunkedArray::<T::PolarsPhysical>::with_chunk(name, arr);
        Ok(unsafe { Self::from_cats_and_dtype_unchecked(phys, dtype) })
    }

    pub fn to_arrow(&self, compat_level: CompatLevel) -> DictionaryArray<T::Native> {
        let keys = self.physical().rechunk();
        let keys = keys.downcast_as_array();
        let values = self
            .get_mapping()
            .to_arrow(compat_level != CompatLevel::oldest());
        let values_dtype = Box::new(values.dtype().clone());
        let dtype =
            ArrowDataType::Dictionary(<T::Native as DictionaryKey>::KEY_TYPE, values_dtype, false);
        unsafe { DictionaryArray::try_new_unchecked(dtype, keys.clone(), values).unwrap() }
    }
}

impl<T: PolarsCategoricalType> LogicalType for CategoricalChunked<T> {
    fn dtype(&self) -> &DataType {
        &self.dtype
    }

    fn get_any_value(&self, i: usize) -> PolarsResult<AnyValue<'_>> {
        polars_ensure!(i < self.len(), oob = i, self.len());
        Ok(unsafe { self.get_any_value_unchecked(i) })
    }

    unsafe fn get_any_value_unchecked(&self, i: usize) -> AnyValue<'_> {
        match self.phys.get_unchecked(i) {
            Some(i) => match &self.dtype {
                DataType::Enum(_, mapping) => AnyValue::Enum(i.as_cat(), mapping),

Review comment (suggestion):
I am really not a fan of having to remap for everything, but I can imagine it is difficult to avoid here.

Code after:
                    });
                    validity.push(opt_s.is_some());
                }
            },
            DataType::Enum(fcats, mapping) => {
                assert!(fcats.physical() == T::physical());
                for opt_s in strings {
                    cat_ids.push(if let Some(cat) = opt_s.and_then(|s| mapping.get_cat(s)) {
                        validity.push(true);
                        T::Native::from_cat(cat)
                    } else {
                        validity.push(false);
                        T::Native::zero()
                    });
                }
            },
            _ => panic!("from_strings_and_dtype_strict called on non-categorical type"),
        }

        let arr = <T::PolarsPhysical as PolarsDataType>::Array::from_vec(cat_ids)
            .with_validity(validity.into_opt_validity());
        let phys = ChunkedArray::<T::PolarsPhysical>::with_chunk(name, arr);
        Ok(unsafe { Self::from_cats_and_dtype_unchecked(phys, dtype) })
    }

    pub fn to_arrow(&self, compat_level: CompatLevel) -> DictionaryArray<T::Native> {
        let in_keys = self.physical().rechunk();
        let in_keys = in_keys.downcast_as_array();
        let mapping = self.get_mapping();
        let as_views = compat_level != CompatLevel::oldest();

        // To minimize dictionary size only include used categories.
        let mut len = 0;
        let bound = mapping.num_cats_upper_bound();
        let bound_native = <T as PolarsCategoricalType>::Native::from_cat(bound as CatSize);
        let mut new_key_idx_map = vec![bound_native; bound];

        let new_key_idxs: Vec<_>;
        let values = if as_views {
            let mut builder = Utf8ViewArrayBuilder::new(ArrowDataType::Utf8View);
            new_key_idxs = in_keys
                .iter()
                .map(|opt_k| {
                    if let Some(k) = opt_k {
                        unsafe {
                            let mut new_k = *new_key_idx_map.get_unchecked(k.as_cat() as usize);
                            if new_k == bound_native {
                                new_k = <_>::from_cat(len);
                                len += 1;

                                let s = mapping.cat_to_str(k.as_cat()).unwrap_or_default();