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::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 cast(&self, _to_type: &ConcreteDataType) -> Result<VectorRef> {
147        UnsupportedOperationSnafu {
148            op: "cast",
149            vector_type: self.vector_type_name(),
150        }
151        .fail()
152    }
153
154    fn filter(&self, filter: &vectors::BooleanVector) -> Result<VectorRef> {
155        let filtered =
156            datafusion_common::arrow::compute::filter(&self.array, filter.as_boolean_array())
157                .context(ArrowComputeSnafu)
158                .and_then(Helper::try_into_vector)?;
159        Ok(filtered)
160    }
161
162    fn take(&self, indices: &vectors::UInt32Vector) -> Result<VectorRef> {
163        let take_result = datafusion_common::arrow::compute::take(
164            &self.array,
165            indices.as_arrow(),
166            Some(TakeOptions { check_bounds: true }),
167        )
168        .context(ArrowComputeSnafu)
169        .and_then(Helper::try_into_vector)?;
170        Ok(take_result)
171    }
172}
173
174impl Serializable for StructVector {
175    fn serialize_to_json(&self) -> Result<Vec<serde_json::Value>> {
176        let vectors = self
177            .array
178            .columns()
179            .iter()
180            .map(|value_array| Helper::try_into_vector(value_array))
181            .collect::<Result<Vec<_>>>()?;
182
183        (0..self.array.len())
184            .map(|idx| {
185                let mut result = serde_json::Map::with_capacity(vectors.len());
186                for (field, vector) in self.fields.fields().iter().zip(vectors.iter()) {
187                    let field_value = vector.get(idx);
188                    result.insert(
189                        field.name().to_string(),
190                        field_value.try_into().context(SerializeSnafu)?,
191                    );
192                }
193                Ok(result.into())
194            })
195            .collect::<Result<Vec<serde_json::Value>>>()
196    }
197}
198
199impl TryFrom<StructArray> for StructVector {
200    type Error = Error;
201
202    fn try_from(array: StructArray) -> Result<Self> {
203        let fields = match array.data_type() {
204            ArrowDataType::Struct(fields) => StructType::from(fields),
205            other => ConversionSnafu {
206                from: other.to_string(),
207            }
208            .fail()?,
209        };
210        Ok(Self { array, fields })
211    }
212}
213
214impl ScalarVector for StructVector {
215    type OwnedItem = StructValue;
216    type RefItem<'a> = StructValueRef<'a>;
217    type Iter<'a> = StructIter<'a>;
218    type Builder = StructVectorBuilder;
219
220    fn get_data(&self, idx: usize) -> Option<Self::RefItem<'_>> {
221        if self.array.is_valid(idx) {
222            Some(StructValueRef::Indexed { vector: self, idx })
223        } else {
224            None
225        }
226    }
227
228    fn iter_data(&self) -> Self::Iter<'_> {
229        StructIter::new(self)
230    }
231}
232
233pub struct StructIter<'a> {
234    vector: &'a StructVector,
235    index: usize,
236}
237
238impl<'a> StructIter<'a> {
239    pub fn new(vector: &'a StructVector) -> Self {
240        Self { vector, index: 0 }
241    }
242}
243
244impl<'a> Iterator for StructIter<'a> {
245    type Item = Option<StructValueRef<'a>>;
246
247    fn next(&mut self) -> Option<Self::Item> {
248        if self.index < self.vector.len() {
249            let idx = self.index;
250            self.index += 1;
251
252            if self.vector.is_null(idx) {
253                Some(None)
254            } else {
255                let value = StructValueRef::Indexed {
256                    vector: self.vector,
257                    idx,
258                };
259
260                Some(Some(value))
261            }
262        } else {
263            None
264        }
265    }
266
267    fn size_hint(&self) -> (usize, Option<usize>) {
268        (self.vector.len(), Some(self.vector.len()))
269    }
270}
271
272pub struct StructVectorBuilder {
273    value_builders: Vec<Box<dyn MutableVector>>,
274    null_buffer: NullBufferBuilder,
275    fields: StructType,
276}
277
278impl StructVectorBuilder {
279    pub fn with_type_and_capacity(fields: StructType, capacity: usize) -> Self {
280        let value_builders = fields
281            .fields()
282            .iter()
283            .map(|f| f.data_type().create_mutable_vector(capacity))
284            .collect();
285        Self {
286            value_builders,
287            null_buffer: NullBufferBuilder::new(capacity),
288            fields,
289        }
290    }
291
292    fn push_struct_value(&mut self, struct_value: &StructValue) -> Result<()> {
293        for (index, value) in struct_value.items().iter().enumerate() {
294            self.value_builders[index].try_push_value_ref(&value.as_value_ref())?;
295        }
296        self.null_buffer.append_non_null();
297
298        Ok(())
299    }
300
301    pub(crate) fn push_struct_value_ref(&mut self, struct_value: StructValueRef<'_>) -> Result<()> {
302        match struct_value {
303            StructValueRef::Indexed { vector, idx } => match vector.get(idx).as_struct()? {
304                Some(struct_value) => self.push_struct_value(struct_value)?,
305                None => self.push_null_struct_value(),
306            },
307            StructValueRef::Ref(value) => self.push_struct_value(value)?,
308            StructValueRef::RefList { val, fields } => {
309                ensure!(
310                    val.len() == self.value_builders.len(),
311                    InconsistentStructFieldsAndItemsSnafu {
312                        field_len: self.value_builders.len(),
313                        item_len: val.len(),
314                    }
315                );
316                ensure!(
317                    fields.fields().len() == self.value_builders.len(),
318                    InconsistentStructFieldsAndItemsSnafu {
319                        field_len: self.value_builders.len(),
320                        item_len: fields.fields().len(),
321                    }
322                );
323                for (builder, value) in self.value_builders.iter_mut().zip(val) {
324                    builder.try_push_value_ref(&value)?;
325                }
326                self.null_buffer.append_non_null();
327            }
328        }
329        Ok(())
330    }
331
332    fn push_null_struct_value(&mut self) {
333        for builder in &mut self.value_builders {
334            builder.push_null();
335        }
336        self.null_buffer.append_null();
337    }
338}
339
340impl MutableVector for StructVectorBuilder {
341    fn data_type(&self) -> ConcreteDataType {
342        ConcreteDataType::struct_datatype(self.fields.clone())
343    }
344
345    fn len(&self) -> usize {
346        self.null_buffer.len()
347    }
348
349    fn as_any(&self) -> &dyn Any {
350        self
351    }
352
353    fn as_mut_any(&mut self) -> &mut dyn Any {
354        self
355    }
356
357    fn to_vector(&mut self) -> VectorRef {
358        Arc::new(self.finish())
359    }
360
361    fn to_vector_cloned(&self) -> VectorRef {
362        Arc::new(self.finish_cloned())
363    }
364
365    fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> {
366        if let Some(struct_ref) = value.try_into_struct()? {
367            self.push_struct_value_ref(struct_ref)?;
368        } else {
369            self.push_null();
370        }
371
372        Ok(())
373    }
374
375    fn extend_slice_of(&mut self, vector: &dyn Vector, offset: usize, length: usize) -> Result<()> {
376        for idx in offset..offset + length {
377            let value = vector.get_ref(idx);
378            self.try_push_value_ref(&value)?;
379        }
380
381        Ok(())
382    }
383
384    fn push_null(&mut self) {
385        self.push_null_struct_value();
386    }
387}
388
389impl ScalarVectorBuilder for StructVectorBuilder {
390    type VectorType = StructVector;
391
392    fn with_capacity(_capacity: usize) -> Self {
393        panic!("Must use StructVectorBuilder::with_type_capacity()");
394    }
395
396    fn push(&mut self, value: Option<<Self::VectorType as ScalarVector>::RefItem<'_>>) {
397        self.try_push_value_ref(&value.map(ValueRef::Struct).unwrap_or(ValueRef::Null))
398            .unwrap_or_else(|e| {
399                panic!(
400                    "Failed to push value, expect value type {:?}, err:{}",
401                    self.fields, e
402                );
403            });
404    }
405
406    fn finish(&mut self) -> Self::VectorType {
407        let arrays = self
408            .value_builders
409            .iter_mut()
410            .map(|b| b.to_vector().to_arrow_array())
411            .collect::<Vec<_>>();
412
413        let struct_array = if arrays.is_empty() {
414            StructArray::new_empty_fields(self.len(), self.null_buffer.finish())
415        } else {
416            StructArray::new(
417                self.fields.as_arrow_fields(),
418                arrays,
419                self.null_buffer.finish(),
420            )
421        };
422
423        StructVector::try_new(self.fields.clone(), struct_array).unwrap()
424    }
425
426    fn finish_cloned(&self) -> Self::VectorType {
427        let arrays = self
428            .value_builders
429            .iter()
430            .map(|b| b.to_vector_cloned().to_arrow_array())
431            .collect();
432
433        let struct_array = StructArray::new(
434            self.fields.as_arrow_fields(),
435            arrays,
436            self.null_buffer.finish_cloned(),
437        );
438        StructVector::try_new(self.fields.clone(), struct_array).unwrap()
439    }
440}
441
442#[cfg(test)]
443mod tests {
444    use super::*;
445    use crate::json::JsonSettings;
446    use crate::schema::{ColumnDefaultConstraint, ColumnSchema};
447    use crate::types::StructField;
448    use crate::value::ListValue;
449    use crate::value::tests::*;
450
451    #[test]
452    fn test_struct_vector_builder() {
453        let struct_type = build_struct_type();
454
455        let struct_values = (0..10).map(|_| build_struct_value());
456        let mut builder = StructVectorBuilder::with_type_and_capacity(struct_type.clone(), 20);
457        for value in struct_values {
458            builder.push(Some(StructValueRef::Ref(&value)));
459        }
460
461        builder.push_nulls(5);
462
463        let vector = builder.finish();
464        assert_eq!(
465            vector.data_type(),
466            ConcreteDataType::struct_datatype(struct_type.clone())
467        );
468        assert_eq!(vector.len(), 15);
469        assert_eq!(vector.null_count(), 5);
470
471        let mut null_count = 0;
472        for item in vector.iter_data() {
473            if let Some(value) = item.as_ref() {
474                assert_eq!(value.struct_type(), &struct_type);
475            } else {
476                null_count += 1;
477            }
478        }
479        assert_eq!(5, null_count);
480
481        let value = vector.get(2);
482        if let Value::Struct(struct_value) = value {
483            assert_eq!(struct_value.struct_type(), &struct_type);
484            let mut items = struct_value.items().iter();
485            assert_eq!(items.next(), Some(&Value::Int32(1)));
486            assert_eq!(items.next(), Some(&Value::String("tom".into())));
487            assert_eq!(items.next(), Some(&Value::UInt8(25)));
488            assert_eq!(items.next(), Some(&Value::String("94038".into())));
489            assert_eq!(items.next(), Some(&Value::List(build_list_value())));
490            assert_eq!(items.next(), None);
491        } else {
492            panic!("Expected a struct value");
493        }
494    }
495
496    #[test]
497    fn test_struct_vector_builder_push_ref_list() {
498        let struct_type = StructType::new(Arc::new(vec![
499            StructField::new("id".to_string(), ConcreteDataType::int64_datatype(), true),
500            StructField::new(
501                "name".to_string(),
502                ConcreteDataType::string_datatype(),
503                true,
504            ),
505        ]));
506        let mut builder = StructVectorBuilder::with_type_and_capacity(struct_type.clone(), 2);
507        builder
508            .push_struct_value_ref(StructValueRef::RefList {
509                val: vec![ValueRef::Int64(1), ValueRef::String("foo")],
510                fields: struct_type.clone(),
511            })
512            .unwrap();
513        builder.push_null();
514
515        let vector = builder.finish();
516        assert_eq!(vector.len(), 2);
517        assert_eq!(vector.null_count(), 1);
518        assert_eq!(
519            vector.get(0),
520            Value::Struct(StructValue::new(
521                vec![Value::Int64(1), Value::String("foo".into())],
522                struct_type,
523            ))
524        );
525        assert_eq!(vector.get(1), Value::Null);
526    }
527
528    #[test]
529    fn test_default_vector_preserves_json2_identity() {
530        let json = JsonSettings::default()
531            .encode(serde_json::json!({"answer": 42}))
532            .unwrap();
533        let fields = StructType::new(Arc::new(vec![StructField::new(
534            "payload",
535            json.data_type(),
536            true,
537        )]));
538        let data_type = ConcreteDataType::struct_datatype(fields.clone());
539        let value = Value::Struct(StructValue::new(vec![json], fields));
540        let schema = ColumnSchema::new("nested", data_type.clone(), true)
541            .with_default_constraint(Some(ColumnDefaultConstraint::Value(value)))
542            .unwrap();
543
544        let replicated = schema.create_default_vector(2).unwrap().unwrap();
545
546        assert_eq!(replicated.data_type(), data_type);
547        assert_eq!(replicated.len(), 2);
548    }
549
550    #[test]
551    fn test_deep_nested_struct_list() {
552        // level 1: struct
553        let struct_type = ConcreteDataType::struct_datatype(build_struct_type());
554        let struct_value = build_struct_value();
555        // level 2: list
556        let struct_type_ref = Arc::new(struct_type);
557        let list_type = ConcreteDataType::list_datatype(struct_type_ref.clone());
558        let list_value = ListValue::new(
559            vec![
560                Value::Struct(struct_value.clone()),
561                Value::Struct(struct_value.clone()),
562            ],
563            struct_type_ref.clone(),
564        );
565        // level 3: struct
566        let root_type = StructType::new(Arc::new(vec![StructField::new(
567            "items".to_string(),
568            list_type,
569            false,
570        )]));
571        let root_value = StructValue::new(vec![Value::List(list_value)], root_type.clone());
572
573        let mut builder = StructVectorBuilder::with_type_and_capacity(root_type.clone(), 20);
574        builder.push(Some(StructValueRef::Ref(&root_value)));
575
576        let vector = builder.finish();
577        assert_eq!(vector.len(), 1);
578        assert_eq!(vector.null_count(), 0);
579        assert_eq!(
580            vector.data_type(),
581            ConcreteDataType::struct_datatype(root_type)
582        );
583        assert_eq!(vector.get(0), Value::Struct(root_value));
584    }
585}
datatypes/vectors/struct_vector.rs

datatypes/vectors/
struct_vector.rs
