risingwave_common/array/
variant_array.rs

1// Copyright 2026 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 std::sync::LazyLock;
16
17use risingwave_common_estimate_size::EstimateSize;
18use risingwave_pb::data::{PbArray, PbArrayType};
19
20use super::{Array, ArrayBuilder, BytesArray, BytesArrayBuilder};
21use crate::bitmap::Bitmap;
22use crate::types::{DataType, Scalar, VariantRef, VariantVal};
23
24/// Returned by raw iteration for NULL entries, whose empty buffer is not a valid variant.
25static NULL_VARIANT_PLACEHOLDER: LazyLock<VariantVal> = LazyLock::new(VariantVal::null);
26
27/// `VariantArray` is a collection of Parquet Variant values. It's a wrapper of [`BytesArray`],
28/// and every non-null slot holds bytes accepted by [`VariantRef::from_serialized`].
29#[derive(Debug, Clone, PartialEq, Eq, EstimateSize)]
30pub struct VariantArray {
31    bytes: BytesArray,
32}
33
34impl Array for VariantArray {
35    type Builder = VariantArrayBuilder;
36    type OwnedItem = VariantVal;
37    type RefItem<'a> = VariantRef<'a>;
38
39    unsafe fn raw_value_at_unchecked(&self, idx: usize) -> Self::RefItem<'_> {
40        if unsafe { !self.bytes.null_bitmap().is_set_unchecked(idx) } {
41            return NULL_VARIANT_PLACEHOLDER.as_scalar_ref();
42        }
43        VariantRef::from_serialized_unchecked(unsafe { self.bytes.raw_value_at_unchecked(idx) })
44    }
45
46    #[inline]
47    fn len(&self) -> usize {
48        self.bytes.len()
49    }
50
51    #[inline]
52    fn to_protobuf(&self) -> PbArray {
53        PbArray {
54            array_type: PbArrayType::Variant as i32,
55            ..self.bytes.to_protobuf()
56        }
57    }
58
59    fn null_bitmap(&self) -> &Bitmap {
60        self.bytes.null_bitmap()
61    }
62
63    fn into_null_bitmap(self) -> Bitmap {
64        self.bytes.into_null_bitmap()
65    }
66
67    fn set_bitmap(&mut self, bitmap: Bitmap) {
68        self.bytes.set_bitmap(bitmap);
69    }
70
71    fn data_type(&self) -> DataType {
72        DataType::Variant
73    }
74}
75
76#[derive(Debug, Clone, EstimateSize)]
77pub struct VariantArrayBuilder {
78    bytes: BytesArrayBuilder,
79}
80
81impl ArrayBuilder for VariantArrayBuilder {
82    type ArrayType = VariantArray;
83
84    fn new(capacity: usize) -> Self {
85        Self {
86            bytes: BytesArrayBuilder::new(capacity),
87        }
88    }
89
90    fn with_type(capacity: usize, ty: DataType) -> Self {
91        assert_eq!(ty, DataType::Variant);
92        Self::new(capacity)
93    }
94
95    #[inline]
96    fn append_n(&mut self, n: usize, value: Option<VariantRef<'_>>) {
97        self.bytes.append_n(n, value.map(|v| v.as_bytes()));
98    }
99
100    #[inline]
101    fn append_array(&mut self, other: &VariantArray) {
102        self.bytes.append_array(&other.bytes);
103    }
104
105    #[inline]
106    fn pop(&mut self) -> Option<()> {
107        self.bytes.pop()
108    }
109
110    fn len(&self) -> usize {
111        self.bytes.len()
112    }
113
114    fn finish(self) -> VariantArray {
115        VariantArray {
116            bytes: self.bytes.finish(),
117        }
118    }
119}
120
121impl FromIterator<Option<VariantVal>> for VariantArray {
122    fn from_iter<I: IntoIterator<Item = Option<VariantVal>>>(iter: I) -> Self {
123        let iter = iter.into_iter();
124        let mut builder = <Self as Array>::Builder::new(iter.size_hint().0);
125        for i in iter {
126            builder.append(i.as_ref().map(|v| v.as_scalar_ref()));
127        }
128        builder.finish()
129    }
130}
131
132impl FromIterator<VariantVal> for VariantArray {
133    fn from_iter<I: IntoIterator<Item = VariantVal>>(iter: I) -> Self {
134        iter.into_iter().map(Some).collect()
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use crate::array::ArrayImpl;
142
143    fn variant(text: &str) -> VariantVal {
144        text.parse().unwrap()
145    }
146
147    #[test]
148    fn raw_iter_tolerates_null_slots() {
149        let array: VariantArray = [Some(variant("1")), None].into_iter().collect();
150
151        assert!(array.value_at(0).is_some());
152        assert!(array.value_at(1).is_none());
153
154        let texts: Vec<_> = array.raw_iter().map(|v| v.to_string()).collect();
155        assert_eq!(texts, ["1", "null"]);
156    }
157
158    #[test]
159    fn protobuf_round_trip_preserves_values_and_nulls() {
160        let array: VariantArray = [
161            Some(variant("1")),
162            None,
163            Some(variant(r#"{"a":[1,2],"b":"x"}"#)),
164            Some(variant("null")),
165        ]
166        .into_iter()
167        .collect();
168
169        let decoded = ArrayImpl::from_protobuf(&array.to_protobuf(), array.len()).unwrap();
170        assert_eq!(ArrayImpl::from(array), decoded);
171    }
172
173    #[test]
174    fn rejects_invalid_serialized_values_from_protobuf() {
175        let array: VariantArray = [Some(variant("1"))].into_iter().collect();
176        let mut proto = array.to_protobuf();
177        // Keep the original length, so the failure comes from validating the bytes rather than
178        // from the data buffer running short.
179        proto.values[1].body = vec![0xFF; proto.values[1].body.len()];
180
181        let err = ArrayImpl::from_protobuf(&proto, 1).unwrap_err();
182        assert!(
183            err.to_string()
184                .contains("failed to read variant from bytes"),
185            "{err:?}"
186        );
187    }
188
189    #[test]
190    fn append_array_concatenates() {
191        let left: VariantArray = [Some(variant("1")), None].into_iter().collect();
192        let right: VariantArray = [Some(variant(r#""x""#))].into_iter().collect();
193
194        let mut builder = VariantArrayBuilder::new(3);
195        builder.append_array(&left);
196        builder.append_array(&right);
197        let joined = builder.finish();
198
199        assert_eq!(joined.len(), 3);
200        let texts: Vec<_> = joined
201            .iter()
202            .map(|v| v.map(|v| v.to_string()))
203            .collect::<Vec<_>>();
204        assert_eq!(
205            texts,
206            [Some("1".to_owned()), None, Some("\"x\"".to_owned())]
207        );
208    }
209}
risingwave_common/array/variant_array.rs

risingwave_common/array/
variant_array.rs
