risingwave_common/array/
jsonb_array.rs

1// Copyright 2025 RisingWave Labs
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 risingwave_common_estimate_size::EstimateSize;
16use risingwave_pb::data::{PbArray, PbArrayType};
17
18use super::{Array, ArrayBuilder, ArrayImpl, ArrayResult};
19use crate::bitmap::{Bitmap, BitmapBuilder};
20use crate::types::{DataType, JsonbRef, JsonbVal, Scalar};
21
22#[derive(Debug, Clone, EstimateSize)]
23pub struct JsonbArrayBuilder {
24    bitmap: BitmapBuilder,
25    builder: jsonbb::Builder,
26}
27
28#[derive(Debug, Clone, PartialEq, Eq, EstimateSize)]
29pub struct JsonbArray {
30    bitmap: Bitmap,
31    /// Elements are stored as a single JSONB array value.
32    data: jsonbb::Value,
33}
34
35impl ArrayBuilder for JsonbArrayBuilder {
36    type ArrayType = JsonbArray;
37
38    fn new(capacity: usize) -> Self {
39        let mut builder = jsonbb::Builder::with_capacity(capacity);
40        builder.begin_array();
41        Self {
42            bitmap: BitmapBuilder::with_capacity(capacity),
43            builder,
44        }
45    }
46
47    fn with_type(capacity: usize, ty: DataType) -> Self {
48        assert_eq!(ty, DataType::Jsonb);
49        Self::new(capacity)
50    }
51
52    fn append_n(&mut self, n: usize, value: Option<<Self::ArrayType as Array>::RefItem<'_>>) {
53        match value {
54            Some(x) => {
55                self.bitmap.append_n(n, true);
56                for _ in 0..n {
57                    self.builder.add_value(x.0);
58                }
59            }
60            None => {
61                self.bitmap.append_n(n, false);
62                for _ in 0..n {
63                    self.builder.add_null();
64                }
65            }
66        }
67    }
68
69    fn append_array(&mut self, other: &Self::ArrayType) {
70        for bit in other.bitmap.iter() {
71            self.bitmap.append(bit);
72        }
73        for value in other.data.as_array().unwrap().iter() {
74            self.builder.add_value(value);
75        }
76    }
77
78    fn pop(&mut self) -> Option<()> {
79        self.bitmap.pop()?;
80        self.builder.pop();
81        Some(())
82    }
83
84    fn len(&self) -> usize {
85        self.bitmap.len()
86    }
87
88    fn finish(mut self) -> Self::ArrayType {
89        self.builder.end_array();
90        Self::ArrayType {
91            bitmap: self.bitmap.finish(),
92            data: self.builder.finish(),
93        }
94    }
95}
96
97impl JsonbArray {
98    /// Loads a `JsonbArray` from a protobuf array.
99    ///
100    /// See also `JsonbArray::to_protobuf`.
101    pub fn from_protobuf(array: &PbArray) -> ArrayResult<ArrayImpl> {
102        ensure!(
103            array.values.len() == 1,
104            "Must have exactly 1 buffer in a jsonb array"
105        );
106        let arr = JsonbArray {
107            bitmap: array.get_null_bitmap()?.into(),
108            data: jsonbb::Value::from_bytes(&array.values[0].body),
109        };
110        Ok(arr.into())
111    }
112}
113
114impl Array for JsonbArray {
115    type Builder = JsonbArrayBuilder;
116    type OwnedItem = JsonbVal;
117    type RefItem<'a> = JsonbRef<'a>;
118
119    unsafe fn raw_value_at_unchecked(&self, idx: usize) -> Self::RefItem<'_> {
120        JsonbRef(self.data.as_array().unwrap().get(idx).unwrap())
121    }
122
123    fn len(&self) -> usize {
124        self.bitmap.len()
125    }
126
127    fn to_protobuf(&self) -> PbArray {
128        use risingwave_pb::common::Buffer;
129        use risingwave_pb::common::buffer::CompressionType;
130
131        PbArray {
132            null_bitmap: Some(self.null_bitmap().to_protobuf()),
133            values: vec![Buffer {
134                compression: CompressionType::None as i32,
135                body: self.data.as_bytes().to_vec(),
136            }],
137            array_type: PbArrayType::Jsonb as i32,
138            struct_array_data: None,
139            list_array_data: None,
140        }
141    }
142
143    fn null_bitmap(&self) -> &Bitmap {
144        &self.bitmap
145    }
146
147    fn into_null_bitmap(self) -> Bitmap {
148        self.bitmap
149    }
150
151    fn set_bitmap(&mut self, bitmap: Bitmap) {
152        self.bitmap = bitmap;
153    }
154
155    fn data_type(&self) -> DataType {
156        DataType::Jsonb
157    }
158}
159
160impl FromIterator<Option<JsonbVal>> for JsonbArray {
161    fn from_iter<I: IntoIterator<Item = Option<JsonbVal>>>(iter: I) -> Self {
162        let iter = iter.into_iter();
163        let mut builder = <Self as Array>::Builder::new(iter.size_hint().0);
164        for i in iter {
165            match i {
166                Some(x) => builder.append(Some(x.as_scalar_ref())),
167                None => builder.append(None),
168            }
169        }
170        builder.finish()
171    }
172}
173
174impl FromIterator<JsonbVal> for JsonbArray {
175    fn from_iter<I: IntoIterator<Item = JsonbVal>>(iter: I) -> Self {
176        iter.into_iter().map(Some).collect()
177    }
178}
risingwave_common/array/jsonb_array.rs

risingwave_common/array/
jsonb_array.rs
