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

datatypes/vectors/
dictionary.rs
