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
202impl Write for StringWriter<'_> {
203    fn write_str(&mut self, s: &str) -> std::fmt::Result {
204        self.bytes.write_ref(s.as_bytes());
205        Ok(())
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use std::hash::Hash;
212
213    use itertools::Itertools;
214
215    use super::*;
216    use crate::array::NULL_VAL_FOR_HASH;
217    use crate::util::iter_util::ZipEqFast;
218
219    #[test]
220    fn test_utf8_builder() {
221        let mut builder = Utf8ArrayBuilder::new(0);
222        for i in 0..100 {
223            if i % 2 == 0 {
224                builder.append(Some(&format!("{}", i)));
225            } else {
226                builder.append(None);
227            }
228        }
229        builder.finish();
230    }
231
232    #[test]
233    fn test_utf8_writer() {
234        let mut builder = Utf8ArrayBuilder::new(0);
235        {
236            let mut writer = builder.writer();
237            for _ in 0..2 {
238                writer.write_str("ran").unwrap();
239            }
240            writer.finish()
241        };
242        let array = builder.finish();
243        assert_eq!(array.len(), 1);
244        assert_eq!(array.value_at(0), Some("ranran"));
245        assert_eq!(unsafe { array.value_at_unchecked(0) }, Some("ranran"));
246    }
247
248    #[test]
249    fn test_utf8_writer_failed() {
250        let mut builder = Utf8ArrayBuilder::new(0);
251        // Write a record.
252        {
253            let mut writer = builder.writer();
254            writer.write_str("Dia").unwrap();
255            writer.write_str("na").unwrap();
256            writer.finish()
257        };
258
259        // Write a record failed.
260        {
261            let mut writer = builder.writer();
262            writer.write_str("Ca").unwrap();
263            writer.write_str("rol").unwrap();
264            // We don't finish here.
265        };
266
267        // Write a record.
268        {
269            let mut writer = builder.writer();
270            writer.write_str("Ki").unwrap();
271            writer.write_str("ra").unwrap();
272            writer.finish()
273        };
274
275        // Verify only two valid records.
276        let array = builder.finish();
277        assert_eq!(array.len(), 2);
278        assert_eq!(array.value_at(0), Some("Diana"));
279        assert_eq!(array.value_at(1), Some("Kira"));
280    }
281
282    #[test]
283    fn test_utf8_array() {
284        let input = vec![
285            Some("1"),
286            Some("22"),
287            None,
288            Some("4444"),
289            None,
290            Some("666666"),
291        ];
292
293        let array = Utf8Array::from_iter(&input);
294        assert_eq!(array.len(), input.len());
295        assert_eq!(input, array.iter().collect_vec());
296    }
297
298    #[test]
299    fn test_utf8_array_to_protobuf() {
300        let input = vec![
301            Some("1"),
302            Some("22"),
303            None,
304            Some("4444"),
305            None,
306            Some("666666"),
307        ];
308
309        let array = Utf8Array::from_iter(&input);
310        let buffers = array.to_protobuf().values;
311        assert!(buffers.len() >= 2);
312    }
313
314    #[test]
315    fn test_utf8_array_hash() {
316        use std::hash::BuildHasher;
317
318        use super::super::test_util::{hash_finish, test_hash};
319
320        const ARR_NUM: usize = 3;
321        const ARR_LEN: usize = 90;
322        let vecs: [Vec<Option<&str>>; ARR_NUM] = [
323            (0..ARR_LEN)
324                .map(|x| match x % 2 {
325                    0 => Some("1"),
326                    1 => None,
327                    _ => unreachable!(),
328                })
329                .collect_vec(),
330            (0..ARR_LEN)
331                .map(|x| match x % 3 {
332                    0 => Some("1"),
333                    1 => Some("abc"),
334                    2 => None,
335                    _ => unreachable!(),
336                })
337                .collect_vec(),
338            (0..ARR_LEN)
339                .map(|x| match x % 5 {
340                    0 => Some("1"),
341                    1 => Some("abc"),
342                    2 => None,
343                    3 => Some("ABCDEF"),
344                    4 => Some("666666"),
345                    _ => unreachable!(),
346                })
347                .collect_vec(),
348        ];
349
350        let arrs = vecs.iter().map(Utf8Array::from_iter).collect_vec();
351
352        let hasher_builder = twox_hash::xxhash64::RandomState::default();
353        let mut states = vec![hasher_builder.build_hasher(); ARR_LEN];
354        vecs.iter().for_each(|v| {
355            v.iter()
356                .zip_eq_fast(&mut states)
357                .for_each(|(x, state)| match x {
358                    Some(inner) => inner.hash(state),
359                    None => NULL_VAL_FOR_HASH.hash(state),
360                })
361        });
362        let hashes = hash_finish(&states[..]);
363
364        let count = hashes.iter().counts().len();
365        assert_eq!(count, 30);
366
367        test_hash(arrs, hashes, hasher_builder);
368    }
369}
risingwave_common/array/utf8_array.rs

risingwave_common/array/
utf8_array.rs
