risingwave_common/array/
proto_reader.rs

1// Copyright 2023 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::io::{Cursor, Read};
16
17use anyhow::Context;
18use byteorder::{BigEndian, ReadBytesExt};
19use risingwave_pb::data::PbArrayType;
20
21use super::*;
22
23impl ArrayImpl {
24    pub fn from_protobuf(array: &PbArray, cardinality: usize) -> ArrayResult<Self> {
25        let array = match array.array_type() {
26            PbArrayType::Unspecified => unreachable!(),
27            PbArrayType::Int16 => read_primitive_array::<i16>(array, cardinality)?,
28            PbArrayType::Int32 => read_primitive_array::<i32>(array, cardinality)?,
29            PbArrayType::Int64 => read_primitive_array::<i64>(array, cardinality)?,
30            PbArrayType::Serial => read_primitive_array::<Serial>(array, cardinality)?,
31            PbArrayType::Float32 => read_primitive_array::<F32>(array, cardinality)?,
32            PbArrayType::Float64 => read_primitive_array::<F64>(array, cardinality)?,
33            PbArrayType::Bool => read_bool_array(array, cardinality)?,
34            PbArrayType::Utf8 => read_var_sized_array::<Utf8ValueReader>(array, cardinality)?,
35            PbArrayType::Decimal => read_primitive_array::<Decimal>(array, cardinality)?,
36            PbArrayType::Date => read_primitive_array::<Date>(array, cardinality)?,
37            PbArrayType::Time => read_primitive_array::<Time>(array, cardinality)?,
38            PbArrayType::Timestamp => read_primitive_array::<Timestamp>(array, cardinality)?,
39            PbArrayType::Timestamptz => read_primitive_array::<Timestamptz>(array, cardinality)?,
40            PbArrayType::Interval => read_primitive_array::<Interval>(array, cardinality)?,
41            PbArrayType::Jsonb => JsonbArray::from_protobuf(array)?,
42            PbArrayType::Variant => read_var_sized_array::<VariantValueReader>(array, cardinality)?,
43            PbArrayType::Struct => StructArray::from_protobuf(array)?,
44            PbArrayType::List => ListArray::from_protobuf(array)?,
45            PbArrayType::Bytea => read_var_sized_array::<BytesValueReader>(array, cardinality)?,
46            PbArrayType::Int256 => Int256Array::from_protobuf(array, cardinality)?,
47            PbArrayType::Map => MapArray::from_protobuf(array)?,
48            PbArrayType::Vector => VectorArray::from_protobuf(array)?,
49        };
50        Ok(array)
51    }
52}
53
54// TODO: Use techniques like apache arrow flight RPC to eliminate deserialization.
55// https://arrow.apache.org/docs/format/Flight.html
56
57fn read_primitive_array<T: PrimitiveArrayItemType>(
58    array: &PbArray,
59    cardinality: usize,
60) -> ArrayResult<ArrayImpl> {
61    ensure!(
62        array.get_values().len() == 1,
63        "Must have only 1 buffer in a numeric array"
64    );
65
66    let buf = array.get_values()[0].get_body().as_slice();
67
68    let mut builder = PrimitiveArrayBuilder::<T>::new(cardinality);
69    let bitmap: Bitmap = array.get_null_bitmap()?.into();
70    let mut cursor = Cursor::new(buf);
71    for not_null in bitmap.iter() {
72        if not_null {
73            let v = T::from_protobuf(&mut cursor)?;
74            builder.append(Some(v));
75        } else {
76            builder.append(None);
77        }
78    }
79    let arr = builder.finish();
80    ensure_eq!(arr.len(), cardinality);
81
82    Ok(arr.into())
83}
84
85fn read_bool_array(array: &PbArray, cardinality: usize) -> ArrayResult<ArrayImpl> {
86    ensure!(
87        array.get_values().len() == 1,
88        "Must have only 1 buffer in a bool array"
89    );
90
91    let data = (&array.get_values()[0]).into();
92    let bitmap: Bitmap = array.get_null_bitmap()?.into();
93
94    let arr = BoolArray::new(data, bitmap);
95    ensure_eq!(arr.len(), cardinality);
96
97    Ok(arr.into())
98}
99
100fn read_offset(offset_cursor: &mut Cursor<&[u8]>) -> ArrayResult<i64> {
101    let offset = offset_cursor
102        .read_i64::<BigEndian>()
103        .context("failed to read i64 from offset buffer")?;
104    // `read_var_sized_array` uses a negative `prev_offset` as its "not yet read" sentinel, so a
105    // negative offset here would silently lengthen the first value instead of being rejected.
106    ensure!(
107        offset >= 0,
108        "invalid offset buffer: offsets must not be negative"
109    );
110
111    Ok(offset)
112}
113
114trait VarSizedValueReader {
115    type AB: ArrayBuilder;
116    fn new_builder(capacity: usize) -> Self::AB;
117    fn read(buf: &[u8], builder: &mut Self::AB) -> ArrayResult<()>;
118}
119
120struct Utf8ValueReader;
121
122impl VarSizedValueReader for Utf8ValueReader {
123    type AB = Utf8ArrayBuilder;
124
125    fn new_builder(capacity: usize) -> Self::AB {
126        Utf8ArrayBuilder::new(capacity)
127    }
128
129    fn read(buf: &[u8], builder: &mut Utf8ArrayBuilder) -> ArrayResult<()> {
130        let s = std::str::from_utf8(buf).context("failed to read utf8 string from bytes")?;
131        builder.append(Some(s));
132        Ok(())
133    }
134}
135
136struct BytesValueReader;
137
138impl VarSizedValueReader for BytesValueReader {
139    type AB = BytesArrayBuilder;
140
141    fn new_builder(capacity: usize) -> Self::AB {
142        BytesArrayBuilder::new(capacity)
143    }
144
145    fn read(buf: &[u8], builder: &mut BytesArrayBuilder) -> ArrayResult<()> {
146        builder.append(Some(buf));
147        Ok(())
148    }
149}
150
151struct VariantValueReader;
152
153impl VarSizedValueReader for VariantValueReader {
154    type AB = VariantArrayBuilder;
155
156    fn new_builder(capacity: usize) -> Self::AB {
157        VariantArrayBuilder::new(capacity)
158    }
159
160    fn read(buf: &[u8], builder: &mut VariantArrayBuilder) -> ArrayResult<()> {
161        let variant =
162            VariantRef::from_serialized(buf).context("failed to read variant from bytes")?;
163        builder.append(Some(variant));
164        Ok(())
165    }
166}
167
168fn read_var_sized_array<R: VarSizedValueReader>(
169    array: &PbArray,
170    cardinality: usize,
171) -> ArrayResult<ArrayImpl> {
172    ensure!(
173        array.get_values().len() == 2,
174        "Must have exactly 2 buffers in a variable-sized array"
175    );
176    let offset_buff = array.get_values()[0].get_body().as_slice();
177    let data_buf = array.get_values()[1].get_body().as_slice();
178
179    let mut builder = R::new_builder(cardinality);
180    let bitmap: Bitmap = array.get_null_bitmap()?.into();
181    let mut offset_cursor = Cursor::new(offset_buff);
182    let mut data_cursor = Cursor::new(data_buf);
183    let mut prev_offset: i64 = -1;
184
185    let mut buf = Vec::new();
186    for not_null in bitmap.iter() {
187        if not_null {
188            if prev_offset < 0 {
189                prev_offset = read_offset(&mut offset_cursor)?;
190            }
191            let offset = read_offset(&mut offset_cursor)?;
192            // A decreasing offset would wrap to a huge length and drive the `resize` below.
193            let length = offset
194                .checked_sub(prev_offset)
195                .and_then(|length| usize::try_from(length).ok())
196                .context("invalid offset buffer: offsets must be non-decreasing")?;
197            prev_offset = offset;
198            buf.resize(length, Default::default());
199            data_cursor
200                .read_exact(buf.as_mut_slice())
201                .with_context(|| {
202                    format!(
203                        "failed to read value from data buffer [length={}, offset={}]",
204                        length, offset
205                    )
206                })?;
207            R::read(buf.as_slice(), &mut builder)?;
208        } else {
209            builder.append(None);
210        }
211    }
212    let arr = builder.finish();
213    ensure_eq!(arr.len(), cardinality);
214
215    Ok(arr.into())
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    // Convert a column to protobuf, then convert it back to column, and ensures the two are
223    // identical.
224    #[test]
225    fn test_column_protobuf_conversion() {
226        let cardinality = 2048;
227        let mut builder = I32ArrayBuilder::new(cardinality);
228        for i in 0..cardinality {
229            if i % 2 == 0 {
230                builder.append(Some(i as i32));
231            } else {
232                builder.append(None);
233            }
234        }
235        let col: ArrayImpl = builder.finish().into();
236        let new_col = ArrayImpl::from_protobuf(&col.to_protobuf(), cardinality).unwrap();
237        assert_eq!(new_col.len(), cardinality);
238        let arr: &I32Array = new_col.as_int32();
239        arr.iter().enumerate().for_each(|(i, x)| {
240            if i % 2 == 0 {
241                assert_eq!(i as i32, x.unwrap());
242            } else {
243                assert!(x.is_none());
244            }
245        });
246    }
247
248    #[test]
249    fn test_bool_column_protobuf_conversion() {
250        let cardinality = 2048;
251        let mut builder = BoolArrayBuilder::new(cardinality);
252        for i in 0..cardinality {
253            match i % 3 {
254                0 => builder.append(Some(false)),
255                1 => builder.append(Some(true)),
256                _ => builder.append(None),
257            }
258        }
259        let col: ArrayImpl = builder.finish().into();
260        let new_col = ArrayImpl::from_protobuf(&col.to_protobuf(), cardinality).unwrap();
261        assert_eq!(new_col.len(), cardinality);
262        let arr: &BoolArray = new_col.as_bool();
263        arr.iter().enumerate().for_each(|(i, x)| match i % 3 {
264            0 => assert_eq!(Some(false), x),
265            1 => assert_eq!(Some(true), x),
266            _ => assert_eq!(None, x),
267        });
268    }
269
270    #[test]
271    fn test_utf8_column_conversion() {
272        let cardinality = 2048;
273        let mut builder = Utf8ArrayBuilder::new(cardinality);
274        for i in 0..cardinality {
275            if i % 2 == 0 {
276                builder.append(Some("abc"));
277            } else {
278                builder.append(None);
279            }
280        }
281        let col: ArrayImpl = builder.finish().into();
282        let new_col = ArrayImpl::from_protobuf(&col.to_protobuf(), cardinality).unwrap();
283        let arr: &Utf8Array = new_col.as_utf8();
284        arr.iter().enumerate().for_each(|(i, x)| {
285            if i % 2 == 0 {
286                assert_eq!("abc", x.unwrap());
287            } else {
288                assert!(x.is_none());
289            }
290        });
291    }
292
293    #[test]
294    fn test_decimal_protobuf_conversion() {
295        let cardinality = 2048;
296        let mut builder = DecimalArrayBuilder::new(cardinality);
297        for i in 0..cardinality {
298            if i % 2 == 0 {
299                builder.append(Some(Decimal::from(i)));
300            } else {
301                builder.append(None);
302            }
303        }
304        let col: ArrayImpl = builder.finish().into();
305        let new_col = ArrayImpl::from_protobuf(&col.to_protobuf(), cardinality).unwrap();
306        assert_eq!(new_col.len(), cardinality);
307        let arr: &DecimalArray = new_col.as_decimal();
308        arr.iter().enumerate().for_each(|(i, x)| {
309            if i % 2 == 0 {
310                assert_eq!(Decimal::from(i), x.unwrap());
311            } else {
312                assert!(x.is_none());
313            }
314        });
315    }
316
317    #[test]
318    fn test_date_protobuf_conversion() {
319        let cardinality = 2048;
320        let mut builder = DateArrayBuilder::new(cardinality);
321        for i in 0..cardinality {
322            if i % 2 == 0 {
323                builder.append(Date::with_days_since_ce(i as i32).ok());
324            } else {
325                builder.append(None);
326            }
327        }
328        let col: ArrayImpl = builder.finish().into();
329        let new_col = ArrayImpl::from_protobuf(&col.to_protobuf(), cardinality).unwrap();
330        assert_eq!(new_col.len(), cardinality);
331        let arr: &DateArray = new_col.as_date();
332        arr.iter().enumerate().for_each(|(i, x)| {
333            if i % 2 == 0 {
334                assert_eq!(Date::with_days_since_ce(i as i32).ok().unwrap(), x.unwrap());
335            } else {
336                assert!(x.is_none());
337            }
338        });
339    }
340
341    #[test]
342    fn test_time_protobuf_conversion() {
343        let cardinality = 2048;
344        let mut builder = TimeArrayBuilder::new(cardinality);
345        for i in 0..cardinality {
346            if i % 2 == 0 {
347                builder.append(Time::with_secs_nano(i as u32, i as u32 * 1000).ok());
348            } else {
349                builder.append(None);
350            }
351        }
352        let col: ArrayImpl = builder.finish().into();
353        let new_col = ArrayImpl::from_protobuf(&col.to_protobuf(), cardinality).unwrap();
354        assert_eq!(new_col.len(), cardinality);
355        let arr: &TimeArray = new_col.as_time();
356        arr.iter().enumerate().for_each(|(i, x)| {
357            if i % 2 == 0 {
358                assert_eq!(
359                    Time::with_secs_nano(i as u32, i as u32 * 1000)
360                        .ok()
361                        .unwrap(),
362                    x.unwrap()
363                );
364            } else {
365                assert!(x.is_none());
366            }
367        });
368    }
369
370    #[test]
371    fn test_timestamp_protobuf_conversion() {
372        let cardinality = 2048;
373        let mut builder = TimestampArrayBuilder::new(cardinality);
374        for i in 0..cardinality {
375            if i % 2 == 0 {
376                builder.append(Timestamp::with_secs_nsecs(i as i64, i as u32 * 1000).ok());
377            } else {
378                builder.append(None);
379            }
380        }
381        let col: ArrayImpl = builder.finish().into();
382        let new_col = ArrayImpl::from_protobuf(&col.to_protobuf(), cardinality).unwrap();
383        assert_eq!(new_col.len(), cardinality);
384        let arr: &TimestampArray = new_col.as_timestamp();
385        arr.iter().enumerate().for_each(|(i, x)| {
386            if i % 2 == 0 {
387                assert_eq!(
388                    Timestamp::with_secs_nsecs(i as i64, i as u32 * 1000)
389                        .ok()
390                        .unwrap(),
391                    x.unwrap()
392                );
393            } else {
394                assert!(x.is_none());
395            }
396        });
397    }
398}
risingwave_common/array/proto_reader.rs

risingwave_common/array/
proto_reader.rs
