risingwave_common/array/
utf8_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 std::fmt::{Display, Write};
16
17use risingwave_common_estimate_size::EstimateSize;
18use risingwave_pb::data::{ArrayType, PbArray};
19
20use super::bytes_array::BytesWriter;
21use super::{Array, ArrayBuilder, BytesArray, BytesArrayBuilder, DataType};
22use crate::bitmap::Bitmap;
23
24/// `Utf8Array` is a collection of Rust Utf8 `str`s. It's a wrapper of `BytesArray`.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct Utf8Array {
27    bytes: BytesArray,
28}
29
30impl EstimateSize for Utf8Array {
31    fn estimated_heap_size(&self) -> usize {
32        self.bytes.estimated_heap_size()
33    }
34}
35
36impl Array for Utf8Array {
37    type Builder = Utf8ArrayBuilder;
38    type OwnedItem = Box<str>;
39    type RefItem<'a> = &'a str;
40
41    unsafe fn raw_value_at_unchecked(&self, idx: usize) -> Self::RefItem<'_> {
42        unsafe {
43            let bytes = self.bytes.raw_value_at_unchecked(idx);
44            std::str::from_utf8_unchecked(bytes)
45        }
46    }
47
48    #[inline]
49    fn len(&self) -> usize {
50        self.bytes.len()
51    }
52
53    #[inline]
54    fn to_protobuf(&self) -> PbArray {
55        PbArray {
56            array_type: ArrayType::Utf8 as i32,
57            ..self.bytes.to_protobuf()
58        }
59    }
60
61    fn null_bitmap(&self) -> &Bitmap {
62        self.bytes.null_bitmap()
63    }
64
65    fn into_null_bitmap(self) -> Bitmap {
66        self.bytes.into_null_bitmap()
67    }
68
69    fn set_bitmap(&mut self, bitmap: Bitmap) {
70        self.bytes.set_bitmap(bitmap);
71    }
72
73    fn data_type(&self) -> DataType {
74        DataType::Varchar
75    }
76}
77
78impl<'a> FromIterator<Option<&'a str>> for Utf8Array {
79    fn from_iter<I: IntoIterator<Item = Option<&'a str>>>(iter: I) -> Self {
80        Self {
81            bytes: iter.into_iter().map(|s| s.map(|s| s.as_bytes())).collect(),
82        }
83    }
84}
85
86impl<'a> FromIterator<&'a Option<&'a str>> for Utf8Array {
87    fn from_iter<I: IntoIterator<Item = &'a Option<&'a str>>>(iter: I) -> Self {
88        iter.into_iter().cloned().collect()
89    }
90}
91
92impl<'a> FromIterator<&'a str> for Utf8Array {
93    fn from_iter<I: IntoIterator<Item = &'a str>>(iter: I) -> Self {
94        iter.into_iter().map(Some).collect()
95    }
96}
97
98impl Utf8Array {
99    pub fn into_bytes_array(self) -> BytesArray {
100        self.bytes
101    }
102
103    pub fn from_iter_display(iter: impl IntoIterator<Item = Option<impl Display>>) -> Self {
104        let iter = iter.into_iter();
105        let mut builder = Utf8ArrayBuilder::new(iter.size_hint().0);
106        for e in iter {
107            if let Some(s) = e {
108                let mut writer = builder.writer();
109                write!(writer, "{}", s).unwrap();
110                writer.finish();
111            } else {
112                builder.append_null();
113            }
114        }
115        builder.finish()
116    }
117}
118
119/// `Utf8ArrayBuilder` use `&str` to build an `Utf8Array`.
120#[derive(Debug, Clone, EstimateSize)]
121pub struct Utf8ArrayBuilder {
122    bytes: BytesArrayBuilder,
123}
124
125impl ArrayBuilder for Utf8ArrayBuilder {
126    type ArrayType = Utf8Array;
127
128    /// Creates a new `Utf8ArrayBuilder`.
129    ///
130    /// `item_capacity` is the number of items to pre-allocate. The size of the preallocated
131    /// buffer of offsets is the number of items plus one.
132    /// No additional memory is pre-allocated for the data buffer.
133    fn new(item_capacity: usize) -> Self {
134        Self {
135            bytes: BytesArrayBuilder::new(item_capacity),
136        }
137    }
138
139    fn with_type(item_capacity: usize, ty: DataType) -> Self {
140        assert_eq!(ty, DataType::Varchar);
141        Self::new(item_capacity)
142    }
143
144    #[inline]
145    fn append_n<'a>(&'a mut self, n: usize, value: Option<&'a str>) {
146        self.bytes.append_n(n, value.map(|v| v.as_bytes()));
147    }
148
149    #[inline]
150    fn append_array(&mut self, other: &Utf8Array) {
151        self.bytes.append_array(&other.bytes);
152    }
153
154    #[inline]
155    fn pop(&mut self) -> Option<()> {
156        self.bytes.pop()
157    }
158
159    fn len(&self) -> usize {
160        self.bytes.len()
161    }
162
163    fn finish(self) -> Utf8Array {
164        Utf8Array {
165            bytes: self.bytes.finish(),
166        }
167    }
168}
169
170impl Utf8ArrayBuilder {
171    pub fn writer(&mut self) -> StringWriter<'_> {
172        StringWriter {
173            bytes: self.bytes.writer(),
174        }
175    }
176
177    /// Append an element as the `Display` format to the array.
178    pub fn append_display(&mut self, value: Option<impl Display>) {
179        if let Some(s) = value {
180            let mut writer = self.writer();
181            write!(writer, "{}", s).unwrap();
182            writer.finish();
183        } else {
184            self.append_null();
185        }
186    }
187}
188
189/// Note: dropping an unfinished `StringWriter` will rollback the partial data, which is the behavior of the inner `BytesWriter`.
190pub struct StringWriter<'a> {
191    bytes: BytesWriter<'a>,
192}
193
194impl StringWriter<'_> {
195    /// `finish` will be called while the entire record is written.
196    /// Exactly one new record was appended and the `builder` can be safely used.
197    pub fn finish(self) {
198        self.bytes.finish()
199    }
200
201    /// `rollback` will be called while the entire record is abandoned.
202    /// The partial data was cleaned and the `builder` can be safely used.
203    pub fn rollback(self) {
204        self.bytes.rollback();
205    }
206}
207
208impl Write for StringWriter<'_> {
209    fn write_str(&mut self, s: &str) -> std::fmt::Result {
210        self.bytes.write_ref(s.as_bytes());
211        Ok(())
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use std::hash::Hash;
218
219    use itertools::Itertools;
220
221    use super::*;
222    use crate::array::NULL_VAL_FOR_HASH;
223    use crate::util::iter_util::ZipEqFast;
224
225    #[test]
226    fn test_utf8_builder() {
227        let mut builder = Utf8ArrayBuilder::new(0);
228        for i in 0..100 {
229            if i % 2 == 0 {
230                builder.append(Some(&format!("{}", i)));
231            } else {
232                builder.append(None);
233            }
234        }
235        builder.finish();
236    }
237
238    #[test]
239    fn test_utf8_writer() {
240        let mut builder = Utf8ArrayBuilder::new(0);
241        {
242            let mut writer = builder.writer();
243            for _ in 0..2 {
244                writer.write_str("ran").unwrap();
245            }
246            writer.finish()
247        };
248        let array = builder.finish();
249        assert_eq!(array.len(), 1);
250        assert_eq!(array.value_at(0), Some("ranran"));
251        assert_eq!(unsafe { array.value_at_unchecked(0) }, Some("ranran"));
252    }
253
254    #[test]
255    fn test_utf8_writer_failed() {
256        let mut builder = Utf8ArrayBuilder::new(0);
257        // Write a record.
258        {
259            let mut writer = builder.writer();
260            writer.write_str("Dia").unwrap();
261            writer.write_str("na").unwrap();
262            writer.finish()
263        };
264
265        // Write a record failed.
266        {
267            let mut writer = builder.writer();
268            writer.write_str("Ca").unwrap();
269            writer.write_str("rol").unwrap();
270            // We don't finish here.
271        };
272
273        // Write a record.
274        {
275            let mut writer = builder.writer();
276            writer.write_str("Ki").unwrap();
277            writer.write_str("ra").unwrap();
278            writer.finish()
279        };
280
281        // Verify only two valid records.
282        let array = builder.finish();
283        assert_eq!(array.len(), 2);
284        assert_eq!(array.value_at(0), Some("Diana"));
285        assert_eq!(array.value_at(1), Some("Kira"));
286    }
287
288    #[test]
289    fn test_utf8_array() {
290        let input = vec![
291            Some("1"),
292            Some("22"),
293            None,
294            Some("4444"),
295            None,
296            Some("666666"),
297        ];
298
299        let array = Utf8Array::from_iter(&input);
300        assert_eq!(array.len(), input.len());
301        assert_eq!(input, array.iter().collect_vec());
302    }
303
304    #[test]
305    fn test_utf8_array_to_protobuf() {
306        let input = vec![
307            Some("1"),
308            Some("22"),
309            None,
310            Some("4444"),
311            None,
312            Some("666666"),
313        ];
314
315        let array = Utf8Array::from_iter(&input);
316        let buffers = array.to_protobuf().values;
317        assert!(buffers.len() >= 2);
318    }
319
320    #[test]
321    fn test_utf8_array_hash() {
322        use std::hash::BuildHasher;
323
324        use super::super::test_util::{hash_finish, test_hash};
325
326        const ARR_NUM: usize = 3;
327        const ARR_LEN: usize = 90;
328        let vecs: [Vec<Option<&str>>; ARR_NUM] = [
329            (0..ARR_LEN)
330                .map(|x| match x % 2 {
331                    0 => Some("1"),
332                    1 => None,
333                    _ => unreachable!(),
334                })
335                .collect_vec(),
336            (0..ARR_LEN)
337                .map(|x| match x % 3 {
338                    0 => Some("1"),
339                    1 => Some("abc"),
340                    2 => None,
341                    _ => unreachable!(),
342                })
343                .collect_vec(),
344            (0..ARR_LEN)
345                .map(|x| match x % 5 {
346                    0 => Some("1"),
347                    1 => Some("abc"),
348                    2 => None,
349                    3 => Some("ABCDEF"),
350                    4 => Some("666666"),
351                    _ => unreachable!(),
352                })
353                .collect_vec(),
354        ];
355
356        let arrs = vecs.iter().map(Utf8Array::from_iter).collect_vec();
357
358        let hasher_builder = twox_hash::xxhash64::RandomState::default();
359        let mut states = vec![hasher_builder.build_hasher(); ARR_LEN];
360        vecs.iter().for_each(|v| {
361            v.iter()
362                .zip_eq_fast(&mut states)
363                .for_each(|(x, state)| match x {
364                    Some(inner) => inner.hash(state),
365                    None => NULL_VAL_FOR_HASH.hash(state),
366                })
367        });
368        let hashes = hash_finish(&states[..]);
369
370        let count = hashes.iter().counts().len();
371        assert_eq!(count, 30);
372
373        test_hash(arrs, hashes, hasher_builder);
374    }
375}
risingwave_common/array/utf8_array.rs

risingwave_common/array/
utf8_array.rs
