serde_arrow 0.14.0

Convert sequences of Rust objects to Arrow arrays and back again
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use std::collections::{BTreeMap, HashMap};

use marrow::{
    array::{Array, BooleanArray},
    datatypes::FieldMeta,
};
use serde::{Serialize, Serializer};

use crate::internal::{
    error::{set_default, try_, Context, ContextSupport, Result},
    serialization::utils::impl_serializer,
    utils::array_ext::{reserve_bits, set_bit_buffer, set_validity, set_validity_default},
};

use super::array_builder::ArrayBuilder;

#[derive(Debug, Clone)]
pub struct BoolBuilder {
    pub name: String,
    array: BooleanArray,
    metadata: HashMap<String, String>,
}

impl BoolBuilder {
    pub fn new(name: String, is_nullable: bool, metadata: HashMap<String, String>) -> Self {
        Self {
            name,
            array: BooleanArray {
                len: 0,
                validity: is_nullable.then(Vec::new),
                values: Vec::new(),
            },
            metadata,
        }
    }

    pub fn take(&mut self) -> ArrayBuilder {
        ArrayBuilder::Bool(Self {
            name: self.name.clone(),
            metadata: self.metadata.clone(),
            array: BooleanArray {
                len: std::mem::take(&mut self.array.len),
                validity: self.array.validity.as_mut().map(std::mem::take),
                values: std::mem::take(&mut self.array.values),
            },
        })
    }

    pub fn is_nullable(&self) -> bool {
        self.array.validity.is_some()
    }

    #[inline]
    pub fn into_array_and_field_meta(self) -> Result<(Array, FieldMeta)> {
        let meta = FieldMeta {
            name: self.name,
            nullable: self.array.validity.is_some(),
            metadata: self.metadata,
        };
        let array = Array::Boolean(self.array);
        Ok((array, meta))
    }

    pub fn reserve(&mut self, additional: usize) {
        if let Some(validity) = &mut self.array.validity {
            reserve_bits(validity, self.array.len, additional);
        }
        reserve_bits(&mut self.array.values, self.array.len, additional);
    }

    pub fn serialize_default_value(&mut self) -> Result<()> {
        try_(|| {
            set_validity_default(self.array.validity.as_mut(), self.array.len);
            set_bit_buffer(&mut self.array.values, self.array.len, false);
            self.array.len += 1;
            Ok(())
        })
        .ctx(self)
    }

    pub fn serialize_value<V: Serialize>(&mut self, value: V) -> Result<()> {
        value.serialize(&mut *self).ctx(self)
    }
}

impl Context for BoolBuilder {
    fn annotate(&self, annotations: &mut BTreeMap<String, String>) {
        set_default(annotations, "field", &self.name);
        set_default(annotations, "data_type", "Boolean");
    }
}

impl<'a> Serializer for &'a mut BoolBuilder {
    impl_serializer!(
        'a, BoolBuilder;
        override serialize_none,
        override serialize_bool,
    );

    fn serialize_none(self) -> Result<()> {
        set_validity(self.array.validity.as_mut(), self.array.len, false)?;
        set_bit_buffer(&mut self.array.values, self.array.len, false);
        self.array.len += 1;
        Ok(())
    }

    fn serialize_bool(self, v: bool) -> Result<()> {
        set_validity(self.array.validity.as_mut(), self.array.len, true)?;
        set_bit_buffer(&mut self.array.values, self.array.len, v);
        self.array.len += 1;
        Ok(())
    }
}