datatypes/vectors/
struct_vector.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::sync::Arc;
17
18use arrow::array::{MutableArrayData, NullBufferBuilder};
19use arrow::compute::TakeOptions;
20use arrow::datatypes::DataType as ArrowDataType;
21use arrow_array::{Array, ArrayRef, StructArray};
22use datafusion_common::ScalarValue;
23use snafu::{ResultExt, ensure};
24
25use crate::error::{
26    ArrowComputeSnafu, ConversionSnafu, Error, InconsistentStructFieldsAndItemsSnafu, Result,
27    SerializeSnafu, UnsupportedOperationSnafu,
28};
29use crate::prelude::{ConcreteDataType, DataType, ScalarVector, ScalarVectorBuilder};
30use crate::serialize::Serializable;
31use crate::types::StructType;
32use crate::value::{StructValue, StructValueRef, Value, ValueRef};
33use crate::vectors::operations::VectorOp;
34use crate::vectors::{self, Helper, MutableVector, Validity, Vector, VectorRef};
35
36/// A simple wrapper around `StructArray` to represent a vector of structs in GreptimeDB.
37#[derive(Debug, PartialEq)]
38pub struct StructVector {
39    array: StructArray,
40    fields: StructType,
41}
42
43impl StructVector {
44    pub fn try_new(fields: StructType, array: StructArray) -> Result<Self> {
45        ensure!(
46            fields.fields().len() == array.fields().len(),
47            InconsistentStructFieldsAndItemsSnafu {
48                field_len: fields.fields().len(),
49                item_len: array.fields().len(),
50            }
51        );
52        Ok(StructVector { array, fields })
53    }
54
55    pub fn array(&self) -> &StructArray {
56        &self.array
57    }
58
59    pub fn as_arrow(&self) -> &dyn Array {
60        &self.array
61    }
62
63    pub fn struct_type(&self) -> &StructType {
64        &self.fields
65    }
66}
67
68impl Vector for StructVector {
69    fn data_type(&self) -> ConcreteDataType {
70        ConcreteDataType::struct_datatype(self.fields.clone())
71    }
72
73    fn vector_type_name(&self) -> String {
74        "StructVector".to_string()
75    }
76
77    fn as_any(&self) -> &dyn std::any::Any {
78        self
79    }
80
81    fn len(&self) -> usize {
82        self.array.len()
83    }
84
85    fn to_arrow_array(&self) -> ArrayRef {
86        Arc::new(self.array.clone())
87    }
88
89    fn to_boxed_arrow_array(&self) -> Box<dyn Array> {
90        Box::new(self.array.clone())
91    }
92
93    fn validity(&self) -> Validity {
94        vectors::impl_validity_for_vector!(self.array)
95    }
96
97    fn memory_size(&self) -> usize {
98        self.array.get_buffer_memory_size()
99    }
100
101    fn null_count(&self) -> usize {
102        self.array.null_count()
103    }
104
105    fn is_null(&self, row: usize) -> bool {
106        self.array.is_null(row)
107    }
108
109    fn slice(&self, offset: usize, length: usize) -> VectorRef {
110        Arc::new(StructVector {
111            array: self.array.slice(offset, length),
112            fields: self.fields.clone(),
113        })
114    }
115
116    fn get(&self, index: usize) -> Value {
117        if !self.array.is_valid(index) {
118            return Value::Null;
119        }
120
121        let values = (0..self.fields.fields().len())
122            .map(|i| {
123                let field_array = &self.array.column(i);
124
125                if field_array.is_null(index) {
126                    Value::Null
127                } else {
128                    let scalar_value = ScalarValue::try_from_array(field_array, index).unwrap();
129                    Value::try_from(scalar_value).unwrap()
130                }
131            })
132            .collect();
133
134        Value::Struct(StructValue::try_new(values, self.fields.clone()).unwrap())
135    }
136
137    fn get_ref(&self, index: usize) -> ValueRef<'_> {
138        ValueRef::Struct(StructValueRef::Indexed {
139            vector: self,
140            idx: index,
141        })
142    }
143}
144
145impl VectorOp for StructVector {
146    fn replicate(&self, offsets: &[usize]) -> VectorRef {
147        assert_eq!(offsets.len(), self.len());
148        assert!(offsets.is_sorted(), "offsets must be non-decreasing");
149
150        let Some(&output_len) = offsets.last() else {
151            return self.slice(0, 0);
152        };
153
154        let source = self.array.to_data();
155        let mut output = MutableArrayData::new(vec![&source], false, output_len);
156        let mut previous_offset = 0;
157
158        for (index, &offset) in offsets.iter().enumerate() {
159            for _ in previous_offset..offset {
160                output.extend(0, index, index + 1);
161            }
162            previous_offset = offset;
163        }
164
165        Arc::new(StructVector {
166            array: StructArray::from(output.freeze()),
167            fields: self.fields.clone(),
168        })
169    }
170
171    fn cast(&self, _to_type: &ConcreteDataType) -> Result<VectorRef> {
172        UnsupportedOperationSnafu {
173            op: "cast",
174            vector_type: self.vector_type_name(),
175        }
176        .fail()
177    }
178
179    fn filter(&self, filter: &vectors::BooleanVector) -> Result<VectorRef> {
180        let filtered =
181            datafusion_common::arrow::compute::filter(&self.array, filter.as_boolean_array())
182                .context(ArrowComputeSnafu)
183                .and_then(Helper::try_into_vector)?;
184        Ok(filtered)
185    }
186
187    fn take(&self, indices: &vectors::UInt32Vector) -> Result<VectorRef> {
188        let take_result = datafusion_common::arrow::compute::take(
189            &self.array,
190            indices.as_arrow(),
191            Some(TakeOptions { check_bounds: true }),
192        )
193        .context(ArrowComputeSnafu)
194        .and_then(Helper::try_into_vector)?;
195        Ok(take_result)
196    }
197}
198
199impl Serializable for StructVector {
200    fn serialize_to_json(&self) -> Result<Vec<serde_json::Value>> {
201        let vectors = self
202            .array
203            .columns()
204            .iter()
205            .map(|value_array| Helper::try_into_vector(value_array))
206            .collect::<Result<Vec<_>>>()?;
207
208        (0..self.array.len())
209            .map(|idx| {
210                let mut result = serde_json::Map::with_capacity(vectors.len());
211                for (field, vector) in self.fields.fields().iter().zip(vectors.iter()) {
212                    let field_value = vector.get(idx);
213                    result.insert(
214                        field.name().to_string(),
215                        field_value.try_into().context(SerializeSnafu)?,
216                    );
217                }
218                Ok(result.into())
219            })
220            .collect::<Result<Vec<serde_json::Value>>>()
221    }
222}
223
224impl TryFrom<StructArray> for StructVector {
225    type Error = Error;
226
227    fn try_from(array: StructArray) -> Result<Self> {
228        let fields = match array.data_type() {
229            ArrowDataType::Struct(fields) => StructType::from(fields),
230            other => ConversionSnafu {
231                from: other.to_string(),
232            }
233            .fail()?,
234        };
235        Ok(Self { array, fields })
236    }
237}
238
239impl ScalarVector for StructVector {
240    type OwnedItem = StructValue;
241    type RefItem<'a> = StructValueRef<'a>;
242    type Iter<'a> = StructIter<'a>;
243    type Builder = StructVectorBuilder;
244
245    fn get_data(&self, idx: usize) -> Option<Self::RefItem<'_>> {
246        if self.array.is_valid(idx) {
247            Some(StructValueRef::Indexed { vector: self, idx })
248        } else {
249            None
250        }
251    }
252
253    fn iter_data(&self) -> Self::Iter<'_> {
254        StructIter::new(self)
255    }
256}
257
258pub struct StructIter<'a> {
259    vector: &'a StructVector,
260    index: usize,
261}
262
263impl<'a> StructIter<'a> {
264    pub fn new(vector: &'a StructVector) -> Self {
265        Self { vector, index: 0 }
266    }
267}
268
269impl<'a> Iterator for StructIter<'a> {
270    type Item = Option<StructValueRef<'a>>;
271
272    fn next(&mut self) -> Option<Self::Item> {
273        if self.index < self.vector.len() {
274            let idx = self.index;
275            self.index += 1;
276
277            if self.vector.is_null(idx) {
278                Some(None)
279            } else {
280                let value = StructValueRef::Indexed {
281                    vector: self.vector,
282                    idx,
283                };
284
285                Some(Some(value))
286            }
287        } else {
288            None
289        }
290    }
291
292    fn size_hint(&self) -> (usize, Option<usize>) {
293        (self.vector.len(), Some(self.vector.len()))
294    }
295}
296
297pub struct StructVectorBuilder {
298    value_builders: Vec<Box<dyn MutableVector>>,
299    null_buffer: NullBufferBuilder,
300    fields: StructType,
301}
302
303impl StructVectorBuilder {
304    pub fn with_type_and_capacity(fields: StructType, capacity: usize) -> Self {
305        let value_builders = fields
306            .fields()
307            .iter()
308            .map(|f| f.data_type().create_mutable_vector(capacity))
309            .collect();
310        Self {
311            value_builders,
312            null_buffer: NullBufferBuilder::new(capacity),
313            fields,
314        }
315    }
316
317    fn push_struct_value(&mut self, struct_value: &StructValue) -> Result<()> {
318        for (index, value) in struct_value.items().iter().enumerate() {
319            self.value_builders[index].try_push_value_ref(&value.as_value_ref())?;
320        }
321        self.null_buffer.append_non_null();
322
323        Ok(())
324    }
325
326    pub(crate) fn push_struct_value_ref(&mut self, struct_value: StructValueRef<'_>) -> Result<()> {
327        match struct_value {
328            StructValueRef::Indexed { vector, idx } => match vector.get(idx).as_struct()? {
329                Some(struct_value) => self.push_struct_value(struct_value)?,
330                None => self.push_null_struct_value(),
331            },
332            StructValueRef::Ref(value) => self.push_struct_value(value)?,
333            StructValueRef::RefList { val, fields } => {
334                ensure!(
335                    val.len() == self.value_builders.len(),
336                    InconsistentStructFieldsAndItemsSnafu {
337                        field_len: self.value_builders.len(),
338                        item_len: val.len(),
339                    }
340                );
341                ensure!(
342                    fields.fields().len() == self.value_builders.len(),
343                    InconsistentStructFieldsAndItemsSnafu {
344                        field_len: self.value_builders.len(),
345                        item_len: fields.fields().len(),
346                    }
347                );
348                for (builder, value) in self.value_builders.iter_mut().zip(val) {
349                    builder.try_push_value_ref(&value)?;
350                }
351                self.null_buffer.append_non_null();
352            }
353        }
354        Ok(())
355    }
356
357    fn push_null_struct_value(&mut self) {
358        for builder in &mut self.value_builders {
359            builder.push_null();
360        }
361        self.null_buffer.append_null();
362    }
363}
364
365impl MutableVector for StructVectorBuilder {
366    fn data_type(&self) -> ConcreteDataType {
367        ConcreteDataType::struct_datatype(self.fields.clone())
368    }
369
370    fn len(&self) -> usize {
371        self.null_buffer.len()
372    }
373
374    fn as_any(&self) -> &dyn Any {
375        self
376    }
377
378    fn as_mut_any(&mut self) -> &mut dyn Any {
379        self
380    }
381
382    fn to_vector(&mut self) -> VectorRef {
383        Arc::new(self.finish())
384    }
385
386    fn to_vector_cloned(&self) -> VectorRef {
387        Arc::new(self.finish_cloned())
388    }
389
390    fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> {
391        if let Some(struct_ref) = value.try_into_struct()? {
392            self.push_struct_value_ref(struct_ref)?;
393        } else {
394            self.push_null();
395        }
396
397        Ok(())
398    }
399
400    fn extend_slice_of(&mut self, vector: &dyn Vector, offset: usize, length: usize) -> Result<()> {
401        for idx in offset..offset + length {
402            let value = vector.get_ref(idx);
403            self.try_push_value_ref(&value)?;
404        }
405
406        Ok(())
407    }
408
409    fn push_null(&mut self) {
410        self.push_null_struct_value();
411    }
412}
413
414impl ScalarVectorBuilder for StructVectorBuilder {
415    type VectorType = StructVector;
416
417    fn with_capacity(_capacity: usize) -> Self {
418        panic!("Must use StructVectorBuilder::with_type_capacity()");
419    }
420
421    fn push(&mut self, value: Option<<Self::VectorType as ScalarVector>::RefItem<'_>>) {
422        self.try_push_value_ref(&value.map(ValueRef::Struct).unwrap_or(ValueRef::Null))
423            .unwrap_or_else(|e| {
424                panic!(
425                    "Failed to push value, expect value type {:?}, err:{}",
426                    self.fields, e
427                );
428            });
429    }
430
431    fn finish(&mut self) -> Self::VectorType {
432        let arrays = self
433            .value_builders
434            .iter_mut()
435            .map(|b| b.to_vector().to_arrow_array())
436            .collect::<Vec<_>>();
437
438        let struct_array = if arrays.is_empty() {
439            StructArray::new_empty_fields(self.len(), self.null_buffer.finish())
440        } else {
441            StructArray::new(
442                self.fields.as_arrow_fields(),
443                arrays,
444                self.null_buffer.finish(),
445            )
446        };
447
448        StructVector::try_new(self.fields.clone(), struct_array).unwrap()
449    }
450
451    fn finish_cloned(&self) -> Self::VectorType {
452        let arrays = self
453            .value_builders
454            .iter()
455            .map(|b| b.to_vector_cloned().to_arrow_array())
456            .collect();
457
458        let struct_array = StructArray::new(
459            self.fields.as_arrow_fields(),
460            arrays,
461            self.null_buffer.finish_cloned(),
462        );
463        StructVector::try_new(self.fields.clone(), struct_array).unwrap()
464    }
465}
466
467#[cfg(test)]
468mod tests {
469    use arrow::array::{DictionaryArray, Int8Array, StringArray};
470    use arrow::buffer::NullBuffer;
471    use arrow::datatypes::Int8Type;
472
473    use super::*;
474    use crate::json::JsonSettings;
475    use crate::schema::{ColumnDefaultConstraint, ColumnSchema};
476    use crate::types::StructField;
477    use crate::value::ListValue;
478    use crate::value::tests::*;
479
480    #[test]
481    fn test_struct_vector_builder() {
482        let struct_type = build_struct_type();
483
484        let struct_values = (0..10).map(|_| build_struct_value());
485        let mut builder = StructVectorBuilder::with_type_and_capacity(struct_type.clone(), 20);
486        for value in struct_values {
487            builder.push(Some(StructValueRef::Ref(&value)));
488        }
489
490        builder.push_nulls(5);
491
492        let vector = builder.finish();
493        assert_eq!(
494            vector.data_type(),
495            ConcreteDataType::struct_datatype(struct_type.clone())
496        );
497        assert_eq!(vector.len(), 15);
498        assert_eq!(vector.null_count(), 5);
499
500        let mut null_count = 0;
501        for item in vector.iter_data() {
502            if let Some(value) = item.as_ref() {
503                assert_eq!(value.struct_type(), &struct_type);
504            } else {
505                null_count += 1;
506            }
507        }
508        assert_eq!(5, null_count);
509
510        let value = vector.get(2);
511        if let Value::Struct(struct_value) = value {
512            assert_eq!(struct_value.struct_type(), &struct_type);
513            let mut items = struct_value.items().iter();
514            assert_eq!(items.next(), Some(&Value::Int32(1)));
515            assert_eq!(items.next(), Some(&Value::String("tom".into())));
516            assert_eq!(items.next(), Some(&Value::UInt8(25)));
517            assert_eq!(items.next(), Some(&Value::String("94038".into())));
518            assert_eq!(items.next(), Some(&Value::List(build_list_value())));
519            assert_eq!(items.next(), None);
520        } else {
521            panic!("Expected a struct value");
522        }
523    }
524
525    #[test]
526    fn test_struct_vector_builder_push_ref_list() {
527        let struct_type = StructType::new(Arc::new(vec![
528            StructField::new("id".to_string(), ConcreteDataType::int64_datatype(), true),
529            StructField::new(
530                "name".to_string(),
531                ConcreteDataType::string_datatype(),
532                true,
533            ),
534        ]));
535        let mut builder = StructVectorBuilder::with_type_and_capacity(struct_type.clone(), 2);
536        builder
537            .push_struct_value_ref(StructValueRef::RefList {
538                val: vec![ValueRef::Int64(1), ValueRef::String("foo")],
539                fields: struct_type.clone(),
540            })
541            .unwrap();
542        builder.push_null();
543
544        let vector = builder.finish();
545        assert_eq!(vector.len(), 2);
546        assert_eq!(vector.null_count(), 1);
547        assert_eq!(
548            vector.get(0),
549            Value::Struct(StructValue::new(
550                vec![Value::Int64(1), Value::String("foo".into())],
551                struct_type,
552            ))
553        );
554        assert_eq!(vector.get(1), Value::Null);
555    }
556
557    #[test]
558    fn test_replicate_preserves_json2_identity() {
559        let json = JsonSettings::default()
560            .encode(serde_json::json!({"answer": 42}))
561            .unwrap();
562        let fields = StructType::new(Arc::new(vec![StructField::new(
563            "payload",
564            json.data_type(),
565            true,
566        )]));
567        let data_type = ConcreteDataType::struct_datatype(fields.clone());
568        let value = Value::Struct(StructValue::new(vec![json], fields));
569        let schema = ColumnSchema::new("nested", data_type.clone(), true)
570            .with_default_constraint(Some(ColumnDefaultConstraint::Value(value)))
571            .unwrap();
572
573        let replicated = schema.create_default_vector(2).unwrap().unwrap();
574
575        assert_eq!(replicated.data_type(), data_type);
576        assert_eq!(replicated.len(), 2);
577    }
578
579    #[test]
580    fn test_replicate_preserves_dictionary_and_nulls() {
581        let fields = StructType::new(Arc::new(vec![StructField::new(
582            "label",
583            ConcreteDataType::dictionary_datatype(
584                ConcreteDataType::int8_datatype(),
585                ConcreteDataType::string_datatype(),
586            ),
587            true,
588        )]));
589        let dictionary = DictionaryArray::<Int8Type>::new(
590            Int8Array::from(vec![Some(0), Some(1)]),
591            Arc::new(StringArray::from(vec!["a", "b"])),
592        );
593        let array = StructArray::new(
594            fields.as_arrow_fields(),
595            vec![Arc::new(dictionary)],
596            Some(NullBuffer::from(vec![true, false])),
597        );
598        let vector = StructVector::try_new(fields.clone(), array).unwrap();
599
600        let replicated = vector.replicate(&[2, 3]);
601
602        assert_eq!(
603            replicated.data_type(),
604            ConcreteDataType::struct_datatype(fields)
605        );
606        assert_eq!(replicated.len(), 3);
607        assert_eq!(replicated.null_count(), 1);
608        assert_eq!(replicated.get(0), replicated.get(1));
609        assert!(replicated.is_null(2));
610    }
611
612    #[test]
613    fn test_deep_nested_struct_list() {
614        // level 1: struct
615        let struct_type = ConcreteDataType::struct_datatype(build_struct_type());
616        let struct_value = build_struct_value();
617        // level 2: list
618        let struct_type_ref = Arc::new(struct_type);
619        let list_type = ConcreteDataType::list_datatype(struct_type_ref.clone());
620        let list_value = ListValue::new(
621            vec![
622                Value::Struct(struct_value.clone()),
623                Value::Struct(struct_value.clone()),
624            ],
625            struct_type_ref.clone(),
626        );
627        // level 3: struct
628        let root_type = StructType::new(Arc::new(vec![StructField::new(
629            "items".to_string(),
630            list_type,
631            false,
632        )]));
633        let root_value = StructValue::new(vec![Value::List(list_value)], root_type.clone());
634
635        let mut builder = StructVectorBuilder::with_type_and_capacity(root_type.clone(), 20);
636        builder.push(Some(StructValueRef::Ref(&root_value)));
637
638        let vector = builder.finish();
639        assert_eq!(vector.len(), 1);
640        assert_eq!(vector.null_count(), 0);
641        assert_eq!(
642            vector.data_type(),
643            ConcreteDataType::struct_datatype(root_type)
644        );
645        assert_eq!(vector.get(0), Value::Struct(root_value));
646    }
647}
datatypes/vectors/struct_vector.rs

datatypes/vectors/
struct_vector.rs
