Repository: pola-rs/polars
File: crates/polars-ops/src/chunked_array/gather/chunked.rs (Rust)

Code before:
        avoid_sharing: bool,
    ) -> Self {
        use DataType::*;
        match self.dtype() {
            dt if dt.is_primitive_numeric() => {
                with_match_physical_numeric_polars_type!(self.dtype(), |$T| {
                    let ca: &ChunkedArray<$T> = self.as_ref().as_ref().as_ref();
                    ca.take_chunked_unchecked(by, sorted, avoid_sharing).into_series()
                })
            },
            Boolean => {
                let ca = self.bool().unwrap();
                ca.take_chunked_unchecked(by, sorted, avoid_sharing)
                    .into_series()
            },
            Binary => {
                let ca = self.binary().unwrap();
                take_chunked_unchecked_binview(ca, by, sorted, avoid_sharing).into_series()
            },
            String => {
                let ca = self.str().unwrap();
                take_chunked_unchecked_binview(ca, by, sorted, avoid_sharing).into_series()
            },
            List(_) => {
                let ca = self.list().unwrap();
                let taken = ca.take_chunked_unchecked(by, sorted, avoid_sharing);
                // Preserve the original dtype (e.g., List(Categorical) instead of List(UInt32))
                unsafe {
                    Series::from_chunks_and_dtype_unchecked(
                        taken.name().clone(),
                        taken.chunks().clone(),
                        self.dtype(),
                    )
                }
            },
            #[cfg(feature = "dtype-array")]
            Array(_, _) => {
                let ca = self.array().unwrap();
                ca.take_chunked_unchecked(by, sorted, avoid_sharing)
                    .into_series()
            },
            #[cfg(feature = "dtype-struct")]
            Struct(_) => {
                let ca = self.struct_().unwrap();
                take_chunked_unchecked_struct(ca, by, sorted, avoid_sharing).into_series()
            },
            #[cfg(feature = "object")]
            Object(_) => take_unchecked_object(self, by, sorted),
            #[cfg(feature = "dtype-decimal")]
            Decimal(_, _) => {
                let ca = self.decimal().unwrap();

Review comment (suggestion):
I don't think this is the right place to fix it - could we instead update `take_chunked_unchecked()` impl for `ChunkedArray` to ensure the original type is restored there? -

https://github.com/pola-rs/polars/blob/b7c28d866f3d99c9839707706aa14977a6d4611f/crates/polars-ops/src/chunked_array/gather/chunked.rs#L343


Code after:
        avoid_sharing: bool,
    ) -> Self {
        use DataType::*;
        match self.dtype() {
            dt if dt.is_primitive_numeric() => {
                with_match_physical_numeric_polars_type!(self.dtype(), |$T| {
                    let ca: &ChunkedArray<$T> = self.as_ref().as_ref().as_ref();
                    ca.take_chunked_unchecked(by, sorted, avoid_sharing).into_series()
                })
            },
            Boolean => {
                let ca = self.bool().unwrap();
                ca.take_chunked_unchecked(by, sorted, avoid_sharing)
                    .into_series()
            },
            Binary => {
                let ca = self.binary().unwrap();
                take_chunked_unchecked_binview(ca, by, sorted, avoid_sharing).into_series()
            },
            String => {
                let ca = self.str().unwrap();
                take_chunked_unchecked_binview(ca, by, sorted, avoid_sharing).into_series()
            },
            List(_) => {
                let ca = self.list().unwrap();
                ca.take_chunked_unchecked(by, sorted, avoid_sharing)
                    .into_series()
            },
            #[cfg(feature = "dtype-array")]
            Array(_, _) => {
                let ca = self.array().unwrap();
                ca.take_chunked_unchecked(by, sorted, avoid_sharing)
                    .into_series()
            },
            #[cfg(feature = "dtype-struct")]
            Struct(_) => {
                let ca = self.struct_().unwrap();
                take_chunked_unchecked_struct(ca, by, sorted, avoid_sharing).into_series()
            },
            #[cfg(feature = "object")]
            Object(_) => take_unchecked_object(self, by, sorted),
            #[cfg(feature = "dtype-decimal")]
            Decimal(_, _) => {
                let ca = self.decimal().unwrap();
                let out = ca.phys.take_chunked_unchecked(by, sorted, avoid_sharing);
                out.into_decimal_unchecked(ca.precision(), ca.scale())
                    .into_series()
            },
            #[cfg(feature = "dtype-date")]
            Date => {
                let ca = self.date().unwrap();