1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

use arrow::array::DictionaryArray;
use arrow::datatypes::IntegerType;
use polars_arrow::compute::cast::cast;

use super::*;
use crate::using_string_cache;

impl From<&CategoricalChunked> for DictionaryArray<u32> {
    fn from(ca: &CategoricalChunked) -> Self {
        let keys = ca.logical().rechunk();
        let keys = keys.downcast_iter().next().unwrap();
        let map = &**ca.get_rev_map();
        let dtype = ArrowDataType::Dictionary(
            IntegerType::UInt32,
            Box::new(ArrowDataType::LargeUtf8),
            false,
        );
        match map {
            RevMapping::Local(arr) => {
                // Safety:
                // the keys are in bounds
                unsafe {
                    DictionaryArray::try_new_unchecked(dtype, keys.clone(), Box::new(arr.clone()))
                        .unwrap()
                }
            }
            RevMapping::Global(reverse_map, values, _uuid) => {
                let iter = keys
                    .into_iter()
                    .map(|opt_k| opt_k.map(|k| *reverse_map.get(k).unwrap()));
                let keys = PrimitiveArray::from_trusted_len_iter(iter);

                // Safety:
                // the keys are in bounds
                unsafe {
                    DictionaryArray::try_new_unchecked(dtype, keys, Box::new(values.clone()))
                        .unwrap()
                }
            }
        }
    }
}
impl From<&CategoricalChunked> for DictionaryArray<i64> {
    fn from(ca: &CategoricalChunked) -> Self {
        let keys = ca.logical().rechunk();
        let keys = keys.downcast_iter().next().unwrap();
        let map = &**ca.get_rev_map();
        let dtype = ArrowDataType::Dictionary(
            IntegerType::UInt32,
            Box::new(ArrowDataType::LargeUtf8),
            false,
        );
        match map {
            // Safety:
            // the keys are in bounds
            RevMapping::Local(arr) => unsafe {
                DictionaryArray::try_new_unchecked(
                    dtype,
                    cast(keys, &ArrowDataType::Int64)
                        .unwrap()
                        .as_any()
                        .downcast_ref::<PrimitiveArray<i64>>()
                        .unwrap()
                        .clone(),
                    Box::new(arr.clone()),
                )
                .unwrap()
            },
            RevMapping::Global(reverse_map, values, _uuid) => {
                let iter = keys
                    .into_iter()
                    .map(|opt_k| opt_k.map(|k| *reverse_map.get(k).unwrap() as i64));
                let keys = PrimitiveArray::from_trusted_len_iter(iter);

                // Safety:
                // the keys are in bounds
                unsafe {
                    DictionaryArray::try_new_unchecked(dtype, keys, Box::new(values.clone()))
                        .unwrap()
                }
            }
        }
    }
}

impl CategoricalChunked {
    /// # Safety
    /// The caller must ensure that index values in the `keys` are in within bounds of the `values` length.
    pub(crate) unsafe fn from_keys_and_values(
        name: &str,
        keys: &PrimitiveArray<u32>,
        values: &Utf8Array<i64>,
    ) -> Self {
        if using_string_cache() {
            let mut builder = CategoricalChunkedBuilder::new(name, keys.len());
            builder.global_map_from_local(keys, values.clone());
            builder.finish()
        } else {
            CategoricalChunked::from_chunks_original(
                name,
                vec![Box::new(keys.clone())],
                RevMapping::Local(values.clone()),
            )
        }
    }
}