datatypes/vectors/
dictionary.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::fmt;
17use std::sync::Arc;
18
19use arrow::array::{
20    Array, ArrayBuilder, ArrayRef, DictionaryArray, PrimitiveArray, StringDictionaryBuilder,
21};
22use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, UInt32Type};
23use serde_json::Value as JsonValue;
24use snafu::ResultExt;
25
26use crate::data_type::ConcreteDataType;
27use crate::error::{self, Result};
28use crate::serialize::Serializable;
29use crate::types::DictionaryType;
30use crate::value::{Value, ValueRef};
31use crate::vectors::operations::VectorOp;
32use crate::vectors::{self, Helper, MutableVector, Validity, Vector, VectorRef};
33
34/// Builder for `Dictionary<UInt32, Utf8>` vectors.
35pub(crate) struct StringDictionaryVectorBuilder {
36    builder: StringDictionaryBuilder<UInt32Type>,
37}
38
39impl StringDictionaryVectorBuilder {
40    pub(crate) fn with_capacity(capacity: usize) -> Self {
41        Self {
42            builder: StringDictionaryBuilder::with_capacity(capacity, 0, 0),
43        }
44    }
45
46    fn vector(array: DictionaryArray<UInt32Type>) -> VectorRef {
47        Arc::new(DictionaryVector::new(array, ConcreteDataType::string_datatype()).unwrap())
48    }
49}
50
51impl MutableVector for StringDictionaryVectorBuilder {
52    fn data_type(&self) -> ConcreteDataType {
53        ConcreteDataType::dictionary_datatype(
54            ConcreteDataType::uint32_datatype(),
55            ConcreteDataType::string_datatype(),
56        )
57    }
58
59    fn len(&self) -> usize {
60        self.builder.len()
61    }
62
63    fn as_any(&self) -> &dyn Any {
64        self
65    }
66
67    fn as_mut_any(&mut self) -> &mut dyn Any {
68        self
69    }
70
71    fn to_vector(&mut self) -> VectorRef {
72        Self::vector(self.builder.finish())
73    }
74
75    fn to_vector_cloned(&self) -> VectorRef {
76        Self::vector(self.builder.finish_cloned())
77    }
78
79    fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> {
80        match value.try_into_string()? {
81            Some(value) => {
82                self.builder
83                    .append(value)
84                    .context(error::ArrowComputeSnafu)?;
85            }
86            None => self.builder.append_null(),
87        }
88        Ok(())
89    }
90
91    fn push_null(&mut self) {
92        self.builder.append_null();
93    }
94
95    fn extend_slice_of(&mut self, vector: &dyn Vector, offset: usize, length: usize) -> Result<()> {
96        for index in offset..offset + length {
97            self.try_push_value_ref(&vector.get_ref(index))?;
98        }
99        Ok(())
100    }
101}
102
103/// Vector of dictionaries, basically backed by Arrow's `DictionaryArray`.
104pub struct DictionaryVector<K: ArrowDictionaryKeyType> {
105    array: DictionaryArray<K>,
106    /// The datatype of the keys in the dictionary.
107    key_type: ConcreteDataType,
108    /// The datatype of the items in the dictionary.
109    item_type: ConcreteDataType,
110    /// The vector of items in the dictionary.
111    item_vector: VectorRef,
112}
113
114impl<K: ArrowDictionaryKeyType> fmt::Debug for DictionaryVector<K> {
115    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
116        f.debug_struct("DictionaryVector")
117            .field("array", &self.array)
118            .field("key_type", &self.key_type)
119            .field("item_type", &self.item_type)
120            .finish()
121    }
122}
123
124impl<K: ArrowDictionaryKeyType> PartialEq for DictionaryVector<K> {
125    fn eq(&self, other: &DictionaryVector<K>) -> bool {
126        self.array == other.array
127            && self.key_type == other.key_type
128            && self.item_type == other.item_type
129    }
130}
131
132impl<K: ArrowDictionaryKeyType> DictionaryVector<K> {
133    /// Create a new instance of `DictionaryVector` from a dictionary array and item type
134    pub fn new(array: DictionaryArray<K>, item_type: ConcreteDataType) -> Result<Self> {
135        let key_type = ConcreteDataType::try_from(&K::DATA_TYPE)?;
136        let item_vector = Helper::try_into_vector(array.values())?;
137
138        Ok(Self {
139            array,
140            key_type,
141            item_type,
142            item_vector,
143        })
144    }
145
146    /// Returns the underlying Arrow dictionary array
147    pub fn array(&self) -> &DictionaryArray<K> {
148        &self.array
149    }
150
151    /// Returns the keys array of this dictionary
152    pub fn keys(&self) -> &arrow_array::PrimitiveArray<K> {
153        self.array.keys()
154    }
155
156    /// Returns the values array of this dictionary
157    pub fn values(&self) -> &ArrayRef {
158        self.array.values()
159    }
160
161    pub fn as_arrow(&self) -> &dyn Array {
162        &self.array
163    }
164}
165
166impl<K: ArrowDictionaryKeyType> Vector for DictionaryVector<K> {
167    fn data_type(&self) -> ConcreteDataType {
168        ConcreteDataType::Dictionary(DictionaryType::new(
169            self.key_type.clone(),
170            self.item_type.clone(),
171        ))
172    }
173
174    fn vector_type_name(&self) -> String {
175        "DictionaryVector".to_string()
176    }
177
178    fn as_any(&self) -> &dyn Any {
179        self
180    }
181
182    fn len(&self) -> usize {
183        self.array.len()
184    }
185
186    fn to_arrow_array(&self) -> ArrayRef {
187        Arc::new(self.array.clone())
188    }
189
190    fn to_boxed_arrow_array(&self) -> Box<dyn Array> {
191        Box::new(self.array.clone())
192    }
193
194    fn validity(&self) -> Validity {
195        self.array
196            .logical_nulls()
197            .map(Validity::from_null_buffer)
198            .unwrap_or_else(|| Validity::all_valid(self.len()))
199    }
200
201    fn memory_size(&self) -> usize {
202        self.array.get_buffer_memory_size()
203    }
204
205    fn null_count(&self) -> usize {
206        self.array.logical_null_count()
207    }
208
209    fn is_null(&self, row: usize) -> bool {
210        self.array
211            .key(row)
212            .is_none_or(|key| self.item_vector.is_null(key))
213    }
214
215    fn slice(&self, offset: usize, length: usize) -> VectorRef {
216        Arc::new(Self {
217            array: self.array.slice(offset, length),
218            key_type: self.key_type.clone(),
219            item_type: self.item_type.clone(),
220            item_vector: self.item_vector.clone(),
221        })
222    }
223
224    fn get(&self, index: usize) -> Value {
225        if !self.array.is_valid(index) {
226            return Value::Null;
227        }
228
229        let key = self.array.keys().value(index);
230        self.item_vector.get(key.as_usize())
231    }
232
233    fn get_ref(&self, index: usize) -> ValueRef<'_> {
234        if !self.array.is_valid(index) {
235            return ValueRef::Null;
236        }
237
238        let key = self.array.keys().value(index);
239        self.item_vector.get_ref(key.as_usize())
240    }
241}
242
243impl<K: ArrowDictionaryKeyType> Serializable for DictionaryVector<K> {
244    fn serialize_to_json(&self) -> Result<Vec<JsonValue>> {
245        // Convert the dictionary array to JSON, where each element is either null or
246        // the value it refers to in the dictionary
247        let mut result = Vec::with_capacity(self.len());
248
249        let keys = self.array.keys();
250        let key_values = &keys.values()[..self.len()];
251        for (i, &key) in key_values.iter().enumerate() {
252            if self.is_null(i) {
253                result.push(JsonValue::Null);
254            } else {
255                let value = self.item_vector.get(key.as_usize());
256                let json_value = serde_json::to_value(value).context(error::SerializeSnafu)?;
257                result.push(json_value);
258            }
259        }
260
261        Ok(result)
262    }
263}
264
265impl<K: ArrowDictionaryKeyType> TryFrom<DictionaryArray<K>> for DictionaryVector<K> {
266    type Error = crate::error::Error;
267
268    fn try_from(array: DictionaryArray<K>) -> Result<Self> {
269        let key_type = ConcreteDataType::try_from(array.keys().data_type())?;
270        let item_type = ConcreteDataType::try_from(array.values().data_type())?;
271        let item_vector = Helper::try_into_vector(array.values())?;
272
273        Ok(Self {
274            array,
275            key_type,
276            item_type,
277            item_vector,
278        })
279    }
280}
281
282pub struct DictionaryIter<'a, K: ArrowDictionaryKeyType> {
283    vector: &'a DictionaryVector<K>,
284    idx: usize,
285}
286
287impl<'a, K: ArrowDictionaryKeyType> DictionaryIter<'a, K> {
288    pub fn new(vector: &'a DictionaryVector<K>) -> DictionaryIter<'a, K> {
289        DictionaryIter { vector, idx: 0 }
290    }
291}
292
293impl<'a, K: ArrowDictionaryKeyType> Iterator for DictionaryIter<'a, K> {
294    type Item = Option<ValueRef<'a>>;
295
296    #[inline]
297    fn next(&mut self) -> Option<Self::Item> {
298        if self.idx >= self.vector.len() {
299            return None;
300        }
301
302        let idx = self.idx;
303        self.idx += 1;
304
305        if self.vector.is_null(idx) {
306            return Some(None);
307        }
308
309        Some(Some(self.vector.get_ref(idx)))
310    }
311
312    #[inline]
313    fn size_hint(&self) -> (usize, Option<usize>) {
314        (
315            self.vector.len() - self.idx,
316            Some(self.vector.len() - self.idx),
317        )
318    }
319}
320
321impl<K: ArrowDictionaryKeyType> VectorOp for DictionaryVector<K> {
322    fn filter(&self, filter: &vectors::BooleanVector) -> Result<VectorRef> {
323        let key_array: ArrayRef = Arc::new(self.array.keys().clone());
324        let key_vector = Helper::try_into_vector(&key_array)?;
325        let filtered_key_vector = key_vector.filter(filter)?;
326        let filtered_key_array = filtered_key_vector.to_arrow_array();
327        let filtered_key_array = filtered_key_array
328            .as_any()
329            .downcast_ref::<PrimitiveArray<K>>()
330            .unwrap();
331
332        let new_array = DictionaryArray::try_new(filtered_key_array.clone(), self.values().clone())
333            .expect("Failed to create filtered dictionary array");
334
335        Ok(Arc::new(Self {
336            array: new_array,
337            key_type: self.key_type.clone(),
338            item_type: self.item_type.clone(),
339            item_vector: self.item_vector.clone(),
340        }))
341    }
342
343    fn cast(&self, to_type: &ConcreteDataType) -> Result<VectorRef> {
344        let new_items = self.item_vector.cast(to_type)?;
345        let new_array =
346            DictionaryArray::try_new(self.array.keys().clone(), new_items.to_arrow_array())
347                .expect("Failed to create casted dictionary array");
348        Ok(Arc::new(Self {
349            array: new_array,
350            key_type: self.key_type.clone(),
351            item_type: to_type.clone(),
352            item_vector: self.item_vector.clone(),
353        }))
354    }
355
356    fn take(&self, indices: &vectors::UInt32Vector) -> Result<VectorRef> {
357        let key_array: ArrayRef = Arc::new(self.array.keys().clone());
358        let key_vector = Helper::try_into_vector(&key_array)?;
359        let new_key_vector = key_vector.take(indices)?;
360        let new_key_array = new_key_vector.to_arrow_array();
361        let new_key_array = new_key_array
362            .as_any()
363            .downcast_ref::<PrimitiveArray<K>>()
364            .unwrap();
365
366        let new_array = DictionaryArray::try_new(new_key_array.clone(), self.values().clone())
367            .expect("Failed to create filtered dictionary array");
368
369        Ok(Arc::new(Self {
370            array: new_array,
371            key_type: self.key_type.clone(),
372            item_type: self.item_type.clone(),
373            item_vector: self.item_vector.clone(),
374        }))
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    use std::sync::Arc;
381
382    use arrow::array::{Int64Array, StringArray, UInt32Array};
383    use arrow::datatypes::{Int64Type, UInt32Type};
384
385    use super::*;
386
387    // Helper function to create a test dictionary vector with string values
388    fn create_test_dictionary() -> DictionaryVector<Int64Type> {
389        // Dictionary values: ["a", "b", "c", "d"]
390        // Keys: [0, 1, 2, null, 1, 3]
391        // Resulting in: ["a", "b", "c", null, "b", "d"]
392        let values = StringArray::from(vec!["a", "b", "c", "d"]);
393        let keys = Int64Array::from(vec![Some(0), Some(1), Some(2), None, Some(1), Some(3)]);
394        let dict_array = DictionaryArray::new(keys, Arc::new(values));
395        DictionaryVector::<Int64Type>::try_from(dict_array).unwrap()
396    }
397
398    #[test]
399    fn test_dictionary_vector_basics() {
400        let dict_vec = create_test_dictionary();
401
402        // Test length and null count
403        assert_eq!(dict_vec.len(), 6);
404        assert_eq!(dict_vec.null_count(), 1);
405
406        // Test data type
407        let data_type = dict_vec.data_type();
408        if let ConcreteDataType::Dictionary(dict_type) = data_type {
409            assert_eq!(*dict_type.value_type(), ConcreteDataType::string_datatype());
410        } else {
411            panic!("Expected Dictionary data type");
412        }
413
414        // Test is_null
415        assert!(!dict_vec.is_null(0));
416        assert!(dict_vec.is_null(3));
417
418        // Test get values
419        assert_eq!(dict_vec.get(0), Value::String("a".to_string().into()));
420        assert_eq!(dict_vec.get(1), Value::String("b".to_string().into()));
421        assert_eq!(dict_vec.get(3), Value::Null);
422        assert_eq!(dict_vec.get(4), Value::String("b".to_string().into()));
423    }
424
425    #[test]
426    fn test_dictionary_vector_logical_nulls() {
427        let values = StringArray::from(vec![Some("a"), None]);
428        let keys = Int64Array::from(vec![Some(0), Some(1), None, Some(1)]);
429        let dict_array = DictionaryArray::new(keys, Arc::new(values));
430        let dict_vec = DictionaryVector::<Int64Type>::try_from(dict_array).unwrap();
431
432        assert_eq!(3, dict_vec.null_count());
433        assert!(!dict_vec.is_null(0));
434        assert!(dict_vec.is_null(1));
435        assert!(dict_vec.is_null(2));
436        assert!(dict_vec.is_null(3));
437
438        let validity = dict_vec.validity();
439        assert_eq!(3, validity.null_count());
440        assert!(validity.is_set(0));
441        assert!(!validity.is_set(1));
442        assert!(!validity.is_set(2));
443        assert!(!validity.is_set(3));
444    }
445
446    #[test]
447    fn test_slice() {
448        let dict_vec = create_test_dictionary();
449        let sliced = dict_vec.slice(1, 3);
450
451        assert_eq!(sliced.len(), 3);
452        assert_eq!(sliced.get(0), Value::String("b".to_string().into()));
453        assert_eq!(sliced.get(1), Value::String("c".to_string().into()));
454        assert_eq!(sliced.get(2), Value::Null);
455    }
456
457    #[test]
458    fn test_filter() {
459        let dict_vec = create_test_dictionary();
460
461        // Keep only indices 0, 2, 4
462        let filter_values = vec![true, false, true, false, true, false];
463        let filter = vectors::BooleanVector::from(filter_values);
464
465        let filtered = dict_vec.filter(&filter).unwrap();
466        assert_eq!(filtered.len(), 3);
467
468        // Check the values
469        assert_eq!(filtered.get(0), Value::String("a".to_string().into()));
470        assert_eq!(filtered.get(1), Value::String("c".to_string().into()));
471        assert_eq!(filtered.get(2), Value::String("b".to_string().into()));
472    }
473
474    #[test]
475    fn test_cast() {
476        let dict_vec = create_test_dictionary();
477
478        // Cast to the same type should return an equivalent vector
479        let casted = dict_vec.cast(&ConcreteDataType::string_datatype()).unwrap();
480
481        // The returned vector should have string values
482        assert_eq!(
483            casted.data_type(),
484            ConcreteDataType::Dictionary(DictionaryType::new(
485                ConcreteDataType::int64_datatype(),
486                ConcreteDataType::string_datatype(),
487            ))
488        );
489        assert_eq!(casted.len(), dict_vec.len());
490
491        // Values should match the original dictionary lookups
492        assert_eq!(casted.get(0), Value::String("a".to_string().into()));
493        assert_eq!(casted.get(1), Value::String("b".to_string().into()));
494        assert_eq!(casted.get(2), Value::String("c".to_string().into()));
495        assert_eq!(casted.get(3), Value::Null);
496        assert_eq!(casted.get(4), Value::String("b".to_string().into()));
497        assert_eq!(casted.get(5), Value::String("d".to_string().into()));
498    }
499
500    #[test]
501    fn test_take() {
502        let dict_vec = create_test_dictionary();
503
504        // Take indices 2, 0, 4
505        let indices_vec = vec![Some(2u32), Some(0), Some(4)];
506        let indices = vectors::UInt32Vector::from(indices_vec);
507
508        let taken = dict_vec.take(&indices).unwrap();
509        assert_eq!(taken.len(), 3);
510
511        // Check the values
512        assert_eq!(taken.get(0), Value::String("c".to_string().into()));
513        assert_eq!(taken.get(1), Value::String("a".to_string().into()));
514        assert_eq!(taken.get(2), Value::String("b".to_string().into()));
515    }
516
517    #[test]
518    fn test_other_type() {
519        let values = StringArray::from(vec!["a", "b", "c", "d"]);
520        let keys = UInt32Array::from(vec![Some(0), Some(1), Some(2), None, Some(1), Some(3)]);
521        let dict_array = DictionaryArray::new(keys, Arc::new(values));
522        let dict_vec = DictionaryVector::<UInt32Type>::try_from(dict_array).unwrap();
523        assert_eq!(
524            ConcreteDataType::dictionary_datatype(
525                ConcreteDataType::uint32_datatype(),
526                ConcreteDataType::string_datatype()
527            ),
528            dict_vec.data_type()
529        );
530    }
531
532    #[test]
533    fn test_string_dictionary_vector_builder() {
534        let mut builder = StringDictionaryVectorBuilder::with_capacity(4);
535        builder.push_value_ref(&ValueRef::String("a"));
536        builder.push_value_ref(&ValueRef::String("b"));
537        builder.push_value_ref(&ValueRef::String("a"));
538        builder.push_null();
539
540        let vector = builder.to_vector();
541        assert_eq!(vector.data_type(), builder.data_type());
542        assert_eq!(vector.get(0), Value::String("a".to_string().into()));
543        assert_eq!(vector.get(1), Value::String("b".to_string().into()));
544        assert_eq!(vector.get(2), Value::String("a".to_string().into()));
545        assert_eq!(vector.get(3), Value::Null);
546
547        let array = vector
548            .to_arrow_array()
549            .as_any()
550            .downcast_ref::<DictionaryArray<UInt32Type>>()
551            .unwrap()
552            .clone();
553        assert_eq!(array.values().len(), 2);
554    }
555}
datatypes/vectors/dictionary.rs

datatypes/vectors/
dictionary.rs
