lance_file/versions/v1/encoding/
dictionary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Dictionary encoding.
5//!
6
7use std::fmt;
8use std::sync::Arc;
9
10use arrow_array::cast::{as_dictionary_array, as_primitive_array};
11use arrow_array::types::{
12    ArrowDictionaryKeyType, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type,
13    UInt32Type, UInt64Type,
14};
15use arrow_array::{Array, ArrayRef, DictionaryArray, PrimitiveArray, UInt32Array};
16use arrow_schema::DataType;
17
18use lance_core::{Error, Result};
19use lance_io::{
20    ReadBatchParams,
21    traits::{Reader, Writer},
22};
23
24use super::plain::{PlainDecoder, PlainEncoder};
25
26/// Encoder for Dictionary encoding.
27pub struct DictionaryEncoder<'a> {
28    writer: &'a mut dyn Writer,
29    key_type: &'a DataType,
30}
31
32impl<'a> DictionaryEncoder<'a> {
33    pub fn new(writer: &'a mut dyn Writer, key_type: &'a DataType) -> Self {
34        Self { writer, key_type }
35    }
36
37    async fn write_typed_array<T: ArrowDictionaryKeyType>(
38        &mut self,
39        arrs: &[&dyn Array],
40    ) -> Result<usize> {
41        assert!(!arrs.is_empty());
42        let data_type = arrs[0].data_type();
43        let pos = self.writer.tell().await?;
44        let mut plain_encoder = PlainEncoder::new(self.writer, data_type);
45
46        let keys = arrs
47            .iter()
48            .map(|a| {
49                let dict_arr = as_dictionary_array::<T>(*a);
50                dict_arr.keys() as &dyn Array
51            })
52            .collect::<Vec<_>>();
53
54        plain_encoder.encode(keys.as_slice()).await?;
55        Ok(pos)
56    }
57}
58
59impl DictionaryEncoder<'_> {
60    pub async fn encode(&mut self, array: &[&dyn Array]) -> Result<usize> {
61        use DataType::*;
62
63        match self.key_type {
64            UInt8 => self.write_typed_array::<UInt8Type>(array).await,
65            UInt16 => self.write_typed_array::<UInt16Type>(array).await,
66            UInt32 => self.write_typed_array::<UInt32Type>(array).await,
67            UInt64 => self.write_typed_array::<UInt64Type>(array).await,
68            Int8 => self.write_typed_array::<Int8Type>(array).await,
69            Int16 => self.write_typed_array::<Int16Type>(array).await,
70            Int32 => self.write_typed_array::<Int32Type>(array).await,
71            Int64 => self.write_typed_array::<Int64Type>(array).await,
72            _ => Err(Error::schema(format!(
73                "DictionaryEncoder: unsupported key type: {:?}",
74                self.key_type
75            ))),
76        }
77    }
78}
79
80/// Decoder for Dictionary encoding.
81pub struct DictionaryDecoder<'a> {
82    reader: &'a dyn Reader,
83    /// The start position of the key array in the file.
84    position: usize,
85    /// Number of the rows in this batch.
86    length: usize,
87    /// The dictionary data type
88    data_type: &'a DataType,
89    /// Value array,
90    value_arr: ArrayRef,
91}
92
93impl fmt::Debug for DictionaryDecoder<'_> {
94    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95        f.debug_struct("DictionaryDecoder")
96            .field("position", &self.position)
97            .field("length", &self.length)
98            .field("data_type", &self.data_type)
99            .field("value_arr", &self.value_arr)
100            .finish()
101    }
102}
103
104impl<'a> DictionaryDecoder<'a> {
105    pub fn new(
106        reader: &'a dyn Reader,
107        position: usize,
108        length: usize,
109        data_type: &'a DataType,
110        value_arr: ArrayRef,
111    ) -> Self {
112        assert!(matches!(data_type, DataType::Dictionary(_, _)));
113        Self {
114            reader,
115            position,
116            length,
117            data_type,
118            value_arr,
119        }
120    }
121
122    async fn decode_impl(&self, params: impl Into<ReadBatchParams>) -> Result<ArrayRef> {
123        let index_type = if let DataType::Dictionary(key_type, _) = &self.data_type {
124            assert!(key_type.as_ref().is_dictionary_key_type());
125            key_type.as_ref()
126        } else {
127            return Err(Error::arrow(format!(
128                "Not a dictionary type: {}",
129                self.data_type
130            )));
131        };
132
133        let decoder = PlainDecoder::new(self.reader, index_type, self.position, self.length)?;
134        let keys = decoder.get(params.into()).await?;
135
136        match index_type {
137            DataType::Int8 => self.make_dict_array::<Int8Type>(keys).await,
138            DataType::Int16 => self.make_dict_array::<Int16Type>(keys).await,
139            DataType::Int32 => self.make_dict_array::<Int32Type>(keys).await,
140            DataType::Int64 => self.make_dict_array::<Int64Type>(keys).await,
141            DataType::UInt8 => self.make_dict_array::<UInt8Type>(keys).await,
142            DataType::UInt16 => self.make_dict_array::<UInt16Type>(keys).await,
143            DataType::UInt32 => self.make_dict_array::<UInt32Type>(keys).await,
144            DataType::UInt64 => self.make_dict_array::<UInt64Type>(keys).await,
145            _ => Err(Error::arrow(format!(
146                "Dictionary encoding does not support index type: {index_type}",
147            ))),
148        }
149    }
150
151    async fn make_dict_array<T: ArrowDictionaryKeyType + Sync + Send>(
152        &self,
153        index_array: ArrayRef,
154    ) -> Result<ArrayRef> {
155        let keys: PrimitiveArray<T> = as_primitive_array(index_array.as_ref()).clone();
156        Ok(Arc::new(DictionaryArray::try_new(
157            keys,
158            self.value_arr.clone(),
159        )?))
160    }
161}
162
163impl DictionaryDecoder<'_> {
164    pub async fn decode(&self) -> Result<ArrayRef> {
165        self.decode_impl(..).await
166    }
167
168    pub async fn take(&self, indices: &UInt32Array) -> Result<ArrayRef> {
169        self.decode_impl(indices.clone()).await
170    }
171
172    pub async fn get(&self, params: impl Into<ReadBatchParams>) -> Result<ArrayRef> {
173        self.decode_impl(params).await
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    use arrow_array::StringArray;
182    use arrow_buffer::ArrowNativeType;
183    use lance_core::utils::tempfile::TempStdFile;
184    use lance_io::local::LocalObjectReader;
185    use tokio::io::AsyncWriteExt;
186
187    async fn test_dict_decoder_for_type<T: ArrowDictionaryKeyType>() {
188        let value_array: StringArray = vec![Some("a"), Some("b"), Some("c"), Some("d")]
189            .into_iter()
190            .collect();
191        let value_array_ref = Arc::new(value_array) as ArrayRef;
192
193        let keys1: PrimitiveArray<T> = vec![T::Native::from_usize(0), T::Native::from_usize(1)]
194            .into_iter()
195            .collect();
196        let arr1: DictionaryArray<T> =
197            DictionaryArray::try_new(keys1, value_array_ref.clone()).unwrap();
198
199        let keys2: PrimitiveArray<T> = vec![T::Native::from_usize(1), T::Native::from_usize(3)]
200            .into_iter()
201            .collect();
202        let arr2: DictionaryArray<T> =
203            DictionaryArray::try_new(keys2, value_array_ref.clone()).unwrap();
204
205        let keys1_ref = arr1.keys() as &dyn Array;
206        let keys2_ref = arr2.keys() as &dyn Array;
207        let arrs: Vec<&dyn Array> = vec![keys1_ref, keys2_ref];
208
209        let path = TempStdFile::default();
210
211        let pos;
212        {
213            let mut object_writer = tokio::fs::File::create(&path).await.unwrap();
214            let mut encoder = PlainEncoder::new(&mut object_writer, arr1.keys().data_type());
215            pos = encoder.encode(arrs.as_slice()).await.unwrap();
216            AsyncWriteExt::shutdown(&mut object_writer).await.unwrap();
217        }
218
219        let reader = LocalObjectReader::open_local_path(&path, 2048, None)
220            .await
221            .unwrap();
222        let decoder = DictionaryDecoder::new(
223            reader.as_ref(),
224            pos,
225            arr1.len() + arr2.len(),
226            arr1.data_type(),
227            value_array_ref.clone(),
228        );
229
230        let decoded_data = decoder.decode().await.unwrap();
231        let expected_data: DictionaryArray<T> = vec!["a", "b", "b", "d"].into_iter().collect();
232        assert_eq!(
233            &expected_data,
234            decoded_data
235                .as_any()
236                .downcast_ref::<DictionaryArray<T>>()
237                .unwrap()
238        );
239    }
240
241    #[tokio::test]
242    async fn test_dict_decoder() {
243        test_dict_decoder_for_type::<Int8Type>().await;
244        test_dict_decoder_for_type::<Int16Type>().await;
245        test_dict_decoder_for_type::<Int32Type>().await;
246        test_dict_decoder_for_type::<Int64Type>().await;
247
248        test_dict_decoder_for_type::<UInt8Type>().await;
249        test_dict_decoder_for_type::<UInt16Type>().await;
250        test_dict_decoder_for_type::<UInt32Type>().await;
251        test_dict_decoder_for_type::<UInt64Type>().await;
252    }
253}
lance_file/versions/v1/encoding/dictionary.rs

lance_file/versions/v1/encoding/
dictionary.rs
