mz_persist_types/
arrow.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! A [protobuf] representation of [Apache Arrow] arrays.
11//!
12//! # Motivation
13//!
14//! Persist can store a small amount of data inline at the consensus layer.
15//! Because we are space constrained, we take particular care to store only the
16//! data that is necessary. Other Arrow serialization formats, e.g. [Parquet]
17//! or [Arrow IPC], include data that we don't need and would be wasteful to
18//! store.
19//!
20//! [protobuf]: https://protobuf.dev/
21//! [Apache Arrow]: https://arrow.apache.org/
22//! [Parquet]: https://parquet.apache.org/docs/
23//! [Arrow IPC]: https://arrow.apache.org/docs/format/Columnar.html#serialization-and-interprocess-communication-ipc
24
25use std::cmp::Ordering;
26use std::fmt::{Debug, Display};
27use std::sync::Arc;
28
29use arrow::array::*;
30use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer};
31use arrow::datatypes::{ArrowNativeType, DataType, Field, FieldRef, Fields};
32use itertools::Itertools;
33use mz_ore::cast::CastFrom;
34use mz_ore::soft_assert_eq_no_log;
35use mz_proto::{IntoRustIfSome, ProtoType, RustType, TryFromProtoError};
36use prost::Message;
37
38#[allow(missing_docs)]
39mod proto {
40    include!(concat!(env!("OUT_DIR"), "/mz_persist_types.arrow.rs"));
41}
42use crate::arrow::proto::data_type;
43pub use proto::{DataType as ProtoDataType, ProtoArrayData};
44
45/// Extract the list of fields for our recursive datatypes.
46pub fn fields_for_type(data_type: &DataType) -> &[FieldRef] {
47    match data_type {
48        DataType::Struct(fields) => fields,
49        DataType::List(field) => std::slice::from_ref(field),
50        DataType::Map(field, _) => std::slice::from_ref(field),
51        DataType::Null
52        | DataType::Boolean
53        | DataType::Int8
54        | DataType::Int16
55        | DataType::Int32
56        | DataType::Int64
57        | DataType::UInt8
58        | DataType::UInt16
59        | DataType::UInt32
60        | DataType::UInt64
61        | DataType::Float16
62        | DataType::Float32
63        | DataType::Float64
64        | DataType::Timestamp(_, _)
65        | DataType::Date32
66        | DataType::Date64
67        | DataType::Time32(_)
68        | DataType::Time64(_)
69        | DataType::Duration(_)
70        | DataType::Interval(_)
71        | DataType::Binary
72        | DataType::FixedSizeBinary(_)
73        | DataType::LargeBinary
74        | DataType::BinaryView
75        | DataType::Utf8
76        | DataType::LargeUtf8
77        | DataType::Utf8View
78        | DataType::Decimal32(_, _)
79        | DataType::Decimal64(_, _)
80        | DataType::Decimal128(_, _)
81        | DataType::Decimal256(_, _) => &[],
82        DataType::ListView(_)
83        | DataType::FixedSizeList(_, _)
84        | DataType::LargeList(_)
85        | DataType::LargeListView(_)
86        | DataType::Union(_, _)
87        | DataType::Dictionary(_, _)
88        | DataType::RunEndEncoded(_, _) => unimplemented!("not supported"),
89    }
90}
91
92/// Encode the array into proto. If an expected data type is passed, that implies it is
93/// encoded at some higher level, and we omit it from the data.
94fn into_proto_with_type(data: &ArrayData, expected_type: Option<&DataType>) -> ProtoArrayData {
95    let data_type = match expected_type {
96        Some(expected) => {
97            // Equality is recursive, and this function is itself called recursively,
98            // skip the call in production to avoid a quadratic overhead.
99            soft_assert_eq_no_log!(
100                expected,
101                data.data_type(),
102                "actual type should match expected type"
103            );
104            None
105        }
106        None => Some(data.data_type().into_proto()),
107    };
108
109    ProtoArrayData {
110        data_type,
111        length: u64::cast_from(data.len()),
112        offset: u64::cast_from(data.offset()),
113        buffers: data.buffers().iter().map(|b| b.into_proto()).collect(),
114        children: data
115            .child_data()
116            .iter()
117            .zip_eq(fields_for_type(
118                expected_type.unwrap_or_else(|| data.data_type()),
119            ))
120            .map(|(child, expect)| into_proto_with_type(child, Some(expect.data_type())))
121            .collect(),
122        nulls: data.nulls().map(|n| n.inner().into_proto()),
123    }
124}
125
126/// Decode the array data.
127/// If the data type is omitted from the proto, we decode it as the expected type.
128fn from_proto_with_type(
129    proto: ProtoArrayData,
130    expected_type: Option<&DataType>,
131) -> Result<ArrayData, TryFromProtoError> {
132    let ProtoArrayData {
133        data_type,
134        length,
135        offset,
136        buffers,
137        children,
138        nulls,
139    } = proto;
140    let data_type: Option<DataType> = data_type.into_rust()?;
141    let data_type = match (data_type, expected_type) {
142        (Some(data_type), None) => data_type,
143        (Some(data_type), Some(expected_type)) => {
144            // Equality is recursive, and this function is itself called recursively,
145            // skip the call in production to avoid a quadratic overhead.
146            soft_assert_eq_no_log!(
147                data_type,
148                *expected_type,
149                "expected type should match actual type"
150            );
151            data_type
152        }
153        (None, Some(expected_type)) => expected_type.clone(),
154        (None, None) => {
155            return Err(TryFromProtoError::MissingField(
156                "ProtoArrayData::data_type".to_string(),
157            ));
158        }
159    };
160    let nulls = nulls
161        .map(|n| n.into_rust())
162        .transpose()?
163        .map(NullBuffer::new);
164
165    let mut builder = ArrayDataBuilder::new(data_type.clone())
166        .len(usize::cast_from(length))
167        .offset(usize::cast_from(offset))
168        .nulls(nulls);
169
170    for b in buffers.into_iter().map(|b| b.into_rust()) {
171        builder = builder.add_buffer(b?);
172    }
173    // `children` is reconstructed from the wire, so its count may disagree with
174    // what `data_type` expects. Guard the length explicitly and reject a mismatch
175    // as a decode error (corrupted/crafted parts must not crash readers); the
176    // `zip_eq` below is then infallible.
177    let fields = fields_for_type(&data_type);
178    if children.len() != fields.len() {
179        return Err(TryFromProtoError::RowConversionError(format!(
180            "ProtoArrayData for {:?} has {} children but its data type expects {}",
181            data_type,
182            children.len(),
183            fields.len(),
184        )));
185    }
186    for (c, field) in children.into_iter().zip_eq(fields) {
187        let c = from_proto_with_type(c, Some(field.data_type()))?;
188        builder = builder.add_child_data(c);
189    }
190
191    // Construct the builder which validates all inputs and aligns data.
192    builder
193        .align_buffers(true)
194        .build()
195        .map_err(|e| TryFromProtoError::RowConversionError(e.to_string()))
196}
197
198impl RustType<ProtoArrayData> for arrow::array::ArrayData {
199    fn into_proto(&self) -> ProtoArrayData {
200        into_proto_with_type(self, None)
201    }
202
203    fn from_proto(proto: ProtoArrayData) -> Result<Self, TryFromProtoError> {
204        from_proto_with_type(proto, None)
205    }
206}
207
208impl RustType<proto::DataType> for arrow::datatypes::DataType {
209    fn into_proto(&self) -> proto::DataType {
210        let kind = match self {
211            DataType::Null => proto::data_type::Kind::Null(()),
212            DataType::Boolean => proto::data_type::Kind::Boolean(()),
213            DataType::UInt8 => proto::data_type::Kind::Uint8(()),
214            DataType::UInt16 => proto::data_type::Kind::Uint16(()),
215            DataType::UInt32 => proto::data_type::Kind::Uint32(()),
216            DataType::UInt64 => proto::data_type::Kind::Uint64(()),
217            DataType::Int8 => proto::data_type::Kind::Int8(()),
218            DataType::Int16 => proto::data_type::Kind::Int16(()),
219            DataType::Int32 => proto::data_type::Kind::Int32(()),
220            DataType::Int64 => proto::data_type::Kind::Int64(()),
221            DataType::Float32 => proto::data_type::Kind::Float32(()),
222            DataType::Float64 => proto::data_type::Kind::Float64(()),
223            DataType::Utf8 => proto::data_type::Kind::String(()),
224            DataType::Binary => proto::data_type::Kind::Binary(()),
225            DataType::FixedSizeBinary(size) => proto::data_type::Kind::FixedBinary(*size),
226            DataType::List(inner) => proto::data_type::Kind::List(Box::new(inner.into_proto())),
227            DataType::Map(inner, sorted) => {
228                let map = proto::data_type::Map {
229                    value: Some(Box::new(inner.into_proto())),
230                    sorted: *sorted,
231                };
232                proto::data_type::Kind::Map(Box::new(map))
233            }
234            DataType::Struct(children) => {
235                let children = children.into_iter().map(|f| f.into_proto()).collect();
236                proto::data_type::Kind::Struct(proto::data_type::Struct { children })
237            }
238            other => unimplemented!("unsupported data type {other:?}"),
239        };
240
241        proto::DataType { kind: Some(kind) }
242    }
243
244    fn from_proto(proto: proto::DataType) -> Result<Self, TryFromProtoError> {
245        let data_type = proto
246            .kind
247            .ok_or_else(|| TryFromProtoError::missing_field("kind"))?;
248        let data_type = match data_type {
249            proto::data_type::Kind::Null(()) => DataType::Null,
250            proto::data_type::Kind::Boolean(()) => DataType::Boolean,
251            proto::data_type::Kind::Uint8(()) => DataType::UInt8,
252            proto::data_type::Kind::Uint16(()) => DataType::UInt16,
253            proto::data_type::Kind::Uint32(()) => DataType::UInt32,
254            proto::data_type::Kind::Uint64(()) => DataType::UInt64,
255            proto::data_type::Kind::Int8(()) => DataType::Int8,
256            proto::data_type::Kind::Int16(()) => DataType::Int16,
257            proto::data_type::Kind::Int32(()) => DataType::Int32,
258            proto::data_type::Kind::Int64(()) => DataType::Int64,
259            proto::data_type::Kind::Float32(()) => DataType::Float32,
260            proto::data_type::Kind::Float64(()) => DataType::Float64,
261            proto::data_type::Kind::String(()) => DataType::Utf8,
262            proto::data_type::Kind::Binary(()) => DataType::Binary,
263            proto::data_type::Kind::FixedBinary(size) => DataType::FixedSizeBinary(size),
264            proto::data_type::Kind::List(inner) => DataType::List(Arc::new((*inner).into_rust()?)),
265            proto::data_type::Kind::Map(inner) => {
266                let value = inner
267                    .value
268                    .ok_or_else(|| TryFromProtoError::missing_field("map.value"))?;
269                DataType::Map(Arc::new((*value).into_rust()?), inner.sorted)
270            }
271            proto::data_type::Kind::Struct(inner) => {
272                let children: Vec<Field> = inner
273                    .children
274                    .into_iter()
275                    .map(|c| c.into_rust())
276                    .collect::<Result<_, _>>()?;
277                DataType::Struct(Fields::from(children))
278            }
279        };
280
281        Ok(data_type)
282    }
283}
284
285impl RustType<proto::Field> for arrow::datatypes::Field {
286    fn into_proto(&self) -> proto::Field {
287        proto::Field {
288            name: self.name().clone(),
289            nullable: self.is_nullable(),
290            data_type: Some(Box::new(self.data_type().into_proto())),
291        }
292    }
293
294    fn from_proto(proto: proto::Field) -> Result<Self, TryFromProtoError> {
295        let proto::Field {
296            name,
297            nullable,
298            data_type,
299        } = proto;
300        let data_type =
301            data_type.ok_or_else(|| TryFromProtoError::missing_field("field.data_type"))?;
302        let data_type = (*data_type).into_rust()?;
303
304        Ok(Field::new(name, data_type, nullable))
305    }
306}
307
308impl RustType<proto::Buffer> for arrow::buffer::Buffer {
309    fn into_proto(&self) -> proto::Buffer {
310        // Wrapping since arrow's buffer doesn't implement AsRef, though the deref impl exists.
311        #[repr(transparent)]
312        struct BufferWrapper(arrow::buffer::Buffer);
313        impl AsRef<[u8]> for BufferWrapper {
314            fn as_ref(&self) -> &[u8] {
315                &*self.0
316            }
317        }
318        proto::Buffer {
319            data: bytes::Bytes::from_owner(BufferWrapper(self.clone())),
320        }
321    }
322
323    fn from_proto(proto: proto::Buffer) -> Result<Self, TryFromProtoError> {
324        Ok(arrow::buffer::Buffer::from(proto.data))
325    }
326}
327
328impl RustType<proto::BooleanBuffer> for arrow::buffer::BooleanBuffer {
329    fn into_proto(&self) -> proto::BooleanBuffer {
330        proto::BooleanBuffer {
331            buffer: Some(self.sliced().into_proto()),
332            length: u64::cast_from(self.len()),
333        }
334    }
335
336    fn from_proto(proto: proto::BooleanBuffer) -> Result<Self, TryFromProtoError> {
337        let proto::BooleanBuffer { buffer, length } = proto;
338        let buffer = buffer.into_rust_if_some("buffer")?;
339        Ok(BooleanBuffer::new(buffer, 0, usize::cast_from(length)))
340    }
341}
342
343/// Wraps a single arrow array, downcasted to a specific type.
344#[derive(Clone)]
345pub enum ArrayOrd {
346    /// Wraps a `NullArray`.
347    Null(NullArray),
348    /// Wraps a `Bool` array.
349    Bool(BooleanArray),
350    /// Wraps a `Int8` array.
351    Int8(Int8Array),
352    /// Wraps a `Int16` array.
353    Int16(Int16Array),
354    /// Wraps a `Int32` array.
355    Int32(Int32Array),
356    /// Wraps a `Int64` array.
357    Int64(Int64Array),
358    /// Wraps a `UInt8` array.
359    UInt8(UInt8Array),
360    /// Wraps a `UInt16` array.
361    UInt16(UInt16Array),
362    /// Wraps a `UInt32` array.
363    UInt32(UInt32Array),
364    /// Wraps a `UInt64` array.
365    UInt64(UInt64Array),
366    /// Wraps a `Float32` array.
367    Float32(Float32Array),
368    /// Wraps a `Float64` array.
369    Float64(Float64Array),
370    /// Wraps a `String` array.
371    String(StringArray),
372    /// Wraps a `Binary` array.
373    Binary(BinaryArray),
374    /// Wraps a `FixedSizeBinary` array.
375    FixedSizeBinary(FixedSizeBinaryArray),
376    /// Wraps a `List` array.
377    List(Option<NullBuffer>, OffsetBuffer<i32>, Box<ArrayOrd>),
378    /// Wraps a `Struct` array.
379    Struct(Option<NullBuffer>, Vec<ArrayOrd>),
380}
381
382impl ArrayOrd {
383    /// Downcast the provided array to a specific type in our enum.
384    pub fn new(array: &dyn Array) -> Self {
385        match array.data_type() {
386            DataType::Null => ArrayOrd::Null(NullArray::from(array.to_data())),
387            DataType::Boolean => ArrayOrd::Bool(array.as_boolean().clone()),
388            DataType::Int8 => ArrayOrd::Int8(array.as_primitive().clone()),
389            DataType::Int16 => ArrayOrd::Int16(array.as_primitive().clone()),
390            DataType::Int32 => ArrayOrd::Int32(array.as_primitive().clone()),
391            DataType::Int64 => ArrayOrd::Int64(array.as_primitive().clone()),
392            DataType::UInt8 => ArrayOrd::UInt8(array.as_primitive().clone()),
393            DataType::UInt16 => ArrayOrd::UInt16(array.as_primitive().clone()),
394            DataType::UInt32 => ArrayOrd::UInt32(array.as_primitive().clone()),
395            DataType::UInt64 => ArrayOrd::UInt64(array.as_primitive().clone()),
396            DataType::Float32 => ArrayOrd::Float32(array.as_primitive().clone()),
397            DataType::Float64 => ArrayOrd::Float64(array.as_primitive().clone()),
398            DataType::Binary => ArrayOrd::Binary(array.as_binary().clone()),
399            DataType::Utf8 => ArrayOrd::String(array.as_string().clone()),
400            DataType::FixedSizeBinary(_) => {
401                ArrayOrd::FixedSizeBinary(array.as_fixed_size_binary().clone())
402            }
403            DataType::List(_) => {
404                let list_array = array.as_list();
405                ArrayOrd::List(
406                    list_array.nulls().cloned(),
407                    list_array.offsets().clone(),
408                    Box::new(ArrayOrd::new(list_array.values())),
409                )
410            }
411            DataType::Struct(_) => {
412                let struct_array = array.as_struct();
413                let nulls = array.nulls().cloned();
414                let columns: Vec<_> = struct_array
415                    .columns()
416                    .iter()
417                    .map(|a| ArrayOrd::new(a))
418                    .collect();
419                ArrayOrd::Struct(nulls, columns)
420            }
421            data_type => unimplemented!("array type {data_type:?} not yet supported"),
422        }
423    }
424
425    /// Returns the rough amount of space required for the data in this array in bytes.
426    /// (Not counting nulls, dictionary encoding, or other space optimizations.)
427    pub fn goodbytes(&self) -> usize {
428        match self {
429            ArrayOrd::Null(_) => 0,
430            // This is, strictly speaking, wrong - but consistent with `ArrayIdx::goodbytes`,
431            // which counts one byte per bool.
432            ArrayOrd::Bool(b) => b.len(),
433            ArrayOrd::Int8(a) => a.values().inner().len(),
434            ArrayOrd::Int16(a) => a.values().inner().len(),
435            ArrayOrd::Int32(a) => a.values().inner().len(),
436            ArrayOrd::Int64(a) => a.values().inner().len(),
437            ArrayOrd::UInt8(a) => a.values().inner().len(),
438            ArrayOrd::UInt16(a) => a.values().inner().len(),
439            ArrayOrd::UInt32(a) => a.values().inner().len(),
440            ArrayOrd::UInt64(a) => a.values().inner().len(),
441            ArrayOrd::Float32(a) => a.values().inner().len(),
442            ArrayOrd::Float64(a) => a.values().inner().len(),
443            ArrayOrd::String(a) => a.values().len(),
444            ArrayOrd::Binary(a) => a.values().len(),
445            ArrayOrd::FixedSizeBinary(a) => a.values().len(),
446            ArrayOrd::List(_, _, nested) => nested.goodbytes(),
447            ArrayOrd::Struct(_, nested) => nested.iter().map(|a| a.goodbytes()).sum(),
448        }
449    }
450
451    /// Return a struct representing the value at a particular index in this array.
452    pub fn at(&self, idx: usize) -> ArrayIdx<'_> {
453        ArrayIdx { idx, array: self }
454    }
455}
456
457impl Debug for ArrayOrd {
458    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
459        struct DebugType<'a>(&'a ArrayOrd);
460
461        impl Debug for DebugType<'_> {
462            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
463                match self.0 {
464                    ArrayOrd::Null(_) => write!(f, "Null"),
465                    ArrayOrd::Bool(_) => write!(f, "Bool"),
466                    ArrayOrd::Int8(_) => write!(f, "Int8"),
467                    ArrayOrd::Int16(_) => write!(f, "Int16"),
468                    ArrayOrd::Int32(_) => write!(f, "Int32"),
469                    ArrayOrd::Int64(_) => write!(f, "Int64"),
470                    ArrayOrd::UInt8(_) => write!(f, "UInt8"),
471                    ArrayOrd::UInt16(_) => write!(f, "UInt16"),
472                    ArrayOrd::UInt32(_) => write!(f, "UInt32"),
473                    ArrayOrd::UInt64(_) => write!(f, "UInt64"),
474                    ArrayOrd::Float32(_) => write!(f, "Float32"),
475                    ArrayOrd::Float64(_) => write!(f, "Float64"),
476                    ArrayOrd::String(_) => write!(f, "String"),
477                    ArrayOrd::Binary(_) => write!(f, "Binary"),
478                    ArrayOrd::FixedSizeBinary(a) => f
479                        .debug_tuple("FixedSizeBinary")
480                        .field(&a.value_length())
481                        .finish(),
482                    ArrayOrd::List(_, _, nested) => f.debug_tuple("List").field(&*nested).finish(),
483                    ArrayOrd::Struct(_, fields) => {
484                        let mut tuple = f.debug_tuple("Struct");
485                        for field in fields {
486                            tuple.field(field);
487                        }
488                        tuple.finish()
489                    }
490                }
491            }
492        }
493
494        f.debug_struct("ArrayOrd")
495            .field("type", &DebugType(self))
496            .field("goodbytes", &self.goodbytes())
497            .finish()
498    }
499}
500
501/// A struct representing a particular entry in a particular array. Most useful for its `Ord`
502/// implementation, which can compare entire rows across similarly-typed arrays.
503///
504/// It is an error to compare indices from arrays with different types, with one exception:
505/// it is valid to compare two `StructArray`s, one of which is a prefix of the other...
506/// in which case we'll compare the values on that subset of the fields, and the shorter
507/// of the two structs will compare less if they're otherwise equal.
508#[derive(Clone, Copy, Debug)]
509pub struct ArrayIdx<'a> {
510    /// An index into a particular array.
511    pub idx: usize,
512    /// The particular array.
513    pub array: &'a ArrayOrd,
514}
515
516impl Display for ArrayIdx<'_> {
517    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
518        match self.array {
519            ArrayOrd::Null(_) => write!(f, "null"),
520            ArrayOrd::Bool(a) => write!(f, "{}", a.value(self.idx)),
521            ArrayOrd::Int8(a) => write!(f, "{}", a.value(self.idx)),
522            ArrayOrd::Int16(a) => write!(f, "{}", a.value(self.idx)),
523            ArrayOrd::Int32(a) => write!(f, "{}", a.value(self.idx)),
524            ArrayOrd::Int64(a) => write!(f, "{}", a.value(self.idx)),
525            ArrayOrd::UInt8(a) => write!(f, "{}", a.value(self.idx)),
526            ArrayOrd::UInt16(a) => write!(f, "{}", a.value(self.idx)),
527            ArrayOrd::UInt32(a) => write!(f, "{}", a.value(self.idx)),
528            ArrayOrd::UInt64(a) => write!(f, "{}", a.value(self.idx)),
529            ArrayOrd::Float32(a) => write!(f, "{}", a.value(self.idx)),
530            ArrayOrd::Float64(a) => write!(f, "{}", a.value(self.idx)),
531            ArrayOrd::String(a) => write!(f, "{}", a.value(self.idx)),
532            ArrayOrd::Binary(a) => {
533                for byte in a.value(self.idx) {
534                    write!(f, "{:02x}", byte)?;
535                }
536                Ok(())
537            }
538            ArrayOrd::FixedSizeBinary(a) => {
539                for byte in a.value(self.idx) {
540                    write!(f, "{:02x}", byte)?;
541                }
542                Ok(())
543            }
544            ArrayOrd::List(_, offsets, nested) => {
545                write!(
546                    f,
547                    "[{}]",
548                    mz_ore::str::separated(", ", list_range(offsets, nested, self.idx))
549                )
550            }
551            ArrayOrd::Struct(_, nested) => write!(
552                f,
553                "{{{}}}",
554                mz_ore::str::separated(", ", nested.iter().map(|f| f.at(self.idx)))
555            ),
556        }
557    }
558}
559
560#[inline]
561fn list_range<'a>(
562    offsets: &OffsetBuffer<i32>,
563    values: &'a ArrayOrd,
564    idx: usize,
565) -> impl Iterator<Item = ArrayIdx<'a>> + Clone {
566    let offsets = offsets.inner();
567    let from = offsets[idx].as_usize();
568    let to = offsets[idx + 1].as_usize();
569    (from..to).map(|i| values.at(i))
570}
571
572impl<'a> ArrayIdx<'a> {
573    /// Returns the rough amount of space required for this entry in bytes.
574    /// (Not counting nulls, dictionary encoding, or other space optimizations.)
575    pub fn goodbytes(&self) -> usize {
576        match self.array {
577            ArrayOrd::Null(_) => 0,
578            ArrayOrd::Bool(_) => size_of::<bool>(),
579            ArrayOrd::Int8(_) => size_of::<i8>(),
580            ArrayOrd::Int16(_) => size_of::<i16>(),
581            ArrayOrd::Int32(_) => size_of::<i32>(),
582            ArrayOrd::Int64(_) => size_of::<i64>(),
583            ArrayOrd::UInt8(_) => size_of::<u8>(),
584            ArrayOrd::UInt16(_) => size_of::<u16>(),
585            ArrayOrd::UInt32(_) => size_of::<u32>(),
586            ArrayOrd::UInt64(_) => size_of::<u64>(),
587            ArrayOrd::Float32(_) => size_of::<f32>(),
588            ArrayOrd::Float64(_) => size_of::<f64>(),
589            ArrayOrd::String(a) => a.value(self.idx).len(),
590            ArrayOrd::Binary(a) => a.value(self.idx).len(),
591            ArrayOrd::FixedSizeBinary(a) => a.value_length().as_usize(),
592            ArrayOrd::List(_, offsets, nested) => {
593                // Range over the list, summing up the bytes for each entry.
594                list_range(offsets, nested, self.idx)
595                    .map(|a| a.goodbytes())
596                    .sum()
597            }
598            ArrayOrd::Struct(_, nested) => nested.iter().map(|a| a.at(self.idx).goodbytes()).sum(),
599        }
600    }
601}
602
603impl<'a> Ord for ArrayIdx<'a> {
604    fn cmp(&self, other: &Self) -> Ordering {
605        #[inline]
606        fn is_null(buffer: &Option<NullBuffer>, idx: usize) -> bool {
607            buffer.as_ref().map_or(false, |b| b.is_null(idx))
608        }
609        #[inline]
610        fn cmp<A: ArrayAccessor>(
611            left: A,
612            left_idx: usize,
613            right: A,
614            right_idx: usize,
615            cmp: fn(&A::Item, &A::Item) -> Ordering,
616        ) -> Ordering {
617            // NB: nulls sort last, conveniently matching psql / mz_repr
618            match (left.is_null(left_idx), right.is_null(right_idx)) {
619                (false, true) => Ordering::Less,
620                (true, true) => Ordering::Equal,
621                (true, false) => Ordering::Greater,
622                (false, false) => cmp(&left.value(left_idx), &right.value(right_idx)),
623            }
624        }
625        match (&self.array, &other.array) {
626            (ArrayOrd::Null(s), ArrayOrd::Null(o)) => {
627                debug_assert!(
628                    self.idx < s.len() && other.idx < o.len(),
629                    "null array indices in bounds"
630                );
631                Ordering::Equal
632            }
633            // For arrays with "simple" value types, we fetch and compare the underlying values directly.
634            (ArrayOrd::Bool(s), ArrayOrd::Bool(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
635            (ArrayOrd::Int8(s), ArrayOrd::Int8(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
636            (ArrayOrd::Int16(s), ArrayOrd::Int16(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
637            (ArrayOrd::Int32(s), ArrayOrd::Int32(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
638            (ArrayOrd::Int64(s), ArrayOrd::Int64(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
639            (ArrayOrd::UInt8(s), ArrayOrd::UInt8(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
640            (ArrayOrd::UInt16(s), ArrayOrd::UInt16(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
641            (ArrayOrd::UInt32(s), ArrayOrd::UInt32(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
642            (ArrayOrd::UInt64(s), ArrayOrd::UInt64(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
643            (ArrayOrd::Float32(s), ArrayOrd::Float32(o)) => {
644                cmp(s, self.idx, o, other.idx, f32::total_cmp)
645            }
646            (ArrayOrd::Float64(s), ArrayOrd::Float64(o)) => {
647                cmp(s, self.idx, o, other.idx, f64::total_cmp)
648            }
649            (ArrayOrd::String(s), ArrayOrd::String(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
650            (ArrayOrd::Binary(s), ArrayOrd::Binary(o)) => cmp(s, self.idx, o, other.idx, Ord::cmp),
651            (ArrayOrd::FixedSizeBinary(s), ArrayOrd::FixedSizeBinary(o)) => {
652                cmp(s, self.idx, o, other.idx, Ord::cmp)
653            }
654            // For lists, we generate an iterator for each side that ranges over the correct
655            // indices into the value buffer, then compare them lexicographically.
656            (
657                ArrayOrd::List(s_nulls, s_offset, s_values),
658                ArrayOrd::List(o_nulls, o_offset, o_values),
659            ) => match (is_null(s_nulls, self.idx), is_null(o_nulls, other.idx)) {
660                (false, true) => Ordering::Less,
661                (true, true) => Ordering::Equal,
662                (true, false) => Ordering::Greater,
663                (false, false) => list_range(s_offset, s_values, self.idx)
664                    .cmp(list_range(o_offset, o_values, other.idx)),
665            },
666            // For structs, we iterate over the same index in each field for each input,
667            // comparing them lexicographically in order.
668            (ArrayOrd::Struct(s_nulls, s_cols), ArrayOrd::Struct(o_nulls, o_cols)) => {
669                match (is_null(s_nulls, self.idx), is_null(o_nulls, other.idx)) {
670                    (false, true) => Ordering::Less,
671                    (true, true) => Ordering::Equal,
672                    (true, false) => Ordering::Greater,
673                    (false, false) => {
674                        let s = s_cols.iter().map(|array| array.at(self.idx));
675                        let o = o_cols.iter().map(|array| array.at(other.idx));
676                        s.cmp(o)
677                    }
678                }
679            }
680            (a, b) => panic!("array types did not match! {a:?} vs. {b:?}",),
681        }
682    }
683}
684
685impl<'a> PartialOrd for ArrayIdx<'a> {
686    fn partial_cmp(&self, other: &ArrayIdx) -> Option<Ordering> {
687        Some(self.cmp(other))
688    }
689}
690
691impl<'a> PartialEq for ArrayIdx<'a> {
692    fn eq(&self, other: &ArrayIdx) -> bool {
693        self.cmp(other) == Ordering::Equal
694    }
695}
696
697impl<'a> Eq for ArrayIdx<'a> {}
698
699/// An array with precisely one entry, for use as a lower bound.
700#[derive(Debug, Clone)]
701pub struct ArrayBound {
702    raw: ArrayRef,
703    ord: ArrayOrd,
704    index: usize,
705}
706
707impl PartialEq for ArrayBound {
708    fn eq(&self, other: &Self) -> bool {
709        self.get().eq(&other.get())
710    }
711}
712
713impl Eq for ArrayBound {}
714
715impl ArrayBound {
716    /// Create a new `ArrayBound` for this array, with the bound at the provided index.
717    pub fn new(array: ArrayRef, index: usize) -> Self {
718        Self {
719            ord: ArrayOrd::new(array.as_ref()),
720            raw: array,
721            index,
722        }
723    }
724
725    /// Get the value of the bound.
726    pub fn get(&self) -> ArrayIdx<'_> {
727        self.ord.at(self.index)
728    }
729
730    /// Convert to an array-data proto, respecting a maximum data size. The resulting proto will
731    /// decode to a single-row array, such that `ArrayBound::new(decoded, 0).get() <= self.get()`,
732    /// which makes it suitable as a lower bound.
733    pub fn to_proto_lower(&self, max_len: usize) -> Option<ProtoArrayData> {
734        // Use `take` instead of slice to make sure we encode just the relevant row to proto,
735        // instead of some larger buffer with an offset.
736        let indices = UInt64Array::from_value(u64::usize_as(self.index), 1);
737        let taken = arrow::compute::take(self.raw.as_ref(), &indices, None).ok()?;
738        let array_data = taken.into_data();
739
740        let mut proto = array_data.into_proto();
741        let original_len = proto.encoded_len();
742        if original_len <= max_len {
743            return Some(proto);
744        }
745
746        let mut data_type = proto.data_type.take()?;
747        maybe_trim_proto(&mut data_type, &mut proto, max_len);
748        proto.data_type = Some(data_type);
749
750        if cfg!(debug_assertions) {
751            let array: ArrayData = proto
752                .clone()
753                .into_rust()
754                .expect("trimmed array data can still be decoded");
755            assert_eq!(array.len(), 1);
756            let new_bound = Self::new(make_array(array), 0);
757            assert!(
758                new_bound.get() <= self.get(),
759                "trimmed bound should be comparable to and no larger than the original data"
760            )
761        }
762
763        if proto.encoded_len() <= max_len {
764            Some(proto)
765        } else {
766            None
767        }
768    }
769}
770
771/// Makes a best effort to shrink the proto while preserving the ordering.
772/// (The proto might not be smaller after this method is called, but it should always
773/// be a valid lower bound.)
774///
775/// Note that we pass in the data type and the array data separately, since we only keep
776/// type info at the top level. If a caller does have a top-level `ArrayData` instance,
777/// they should take that type and pass it in separately.
778fn maybe_trim_proto(data_type: &mut proto::DataType, body: &mut ProtoArrayData, max_len: usize) {
779    assert!(body.data_type.is_none(), "expected separate data type");
780    // TODO: consider adding cases for strings and byte arrays
781    let encoded_len = data_type.encoded_len() + body.encoded_len();
782    match &mut data_type.kind {
783        Some(data_type::Kind::Struct(data_type::Struct { children: fields })) => {
784            // Pop off fields one by one, keeping an estimate of the encoded length.
785            let mut struct_len = encoded_len;
786            while struct_len > max_len {
787                let Some(mut child) = body.children.pop() else {
788                    break;
789                };
790                let Some(mut field) = fields.pop() else { break };
791
792                struct_len -= field.encoded_len() + child.encoded_len();
793                if let Some(remaining_len) = max_len.checked_sub(struct_len) {
794                    // We're under budget after removing this field! See if we can
795                    // shrink it to fit, but exit the loop regardless.
796                    let Some(field_type) = field.data_type.as_mut() else {
797                        break;
798                    };
799                    maybe_trim_proto(field_type, &mut child, remaining_len);
800                    if field.encoded_len() + child.encoded_len() <= remaining_len {
801                        fields.push(field);
802                        body.children.push(child);
803                    }
804                    break;
805                }
806            }
807        }
808        _ => {}
809    };
810}
811
812#[cfg(test)]
813mod tests {
814    use crate::arrow::{ArrayBound, ArrayOrd};
815    use arrow::array::{
816        ArrayRef, AsArray, BooleanArray, StringArray, StructArray, UInt64Array, make_array,
817    };
818    use arrow::datatypes::{DataType, Field, Fields};
819    use mz_ore::assert_none;
820    use mz_proto::ProtoType;
821    use std::sync::Arc;
822
823    #[mz_ore::test]
824    fn from_proto_child_count_mismatch_is_error() {
825        // A `ProtoArrayData` whose child count disagrees with its declared data
826        // type must decode to an error, not panic via `zip_eq`. Regression for
827        // the array_data_proto_roundtrip cargo-fuzz finding.
828        use prost::Message;
829        let bytes: &[u8] = &[
830            0x0a, 0x0a, 0x18, 0x32, 0x22, 0x00, 0x18, 0x0a, 0x18, 0x32, 0x22, 0x00, 0x2a, 0x00,
831            0x2a, 0x00, 0xe0, 0x32, 0x24,
832        ];
833        let proto =
834            crate::arrow::ProtoArrayData::decode(bytes).expect("crash input decodes as a proto");
835        let result: Result<arrow::array::ArrayData, _> = proto.into_rust();
836        assert!(
837            result.is_err(),
838            "child-count mismatch must be a decode error"
839        );
840    }
841
842    #[mz_ore::test]
843    fn trim_proto() {
844        let nested_fields: Fields = vec![Field::new("a", DataType::UInt64, true)].into();
845        let array: ArrayRef = Arc::new(StructArray::new(
846            vec![
847                Field::new("a", DataType::UInt64, true),
848                Field::new("b", DataType::Utf8, true),
849                Field::new_struct("c", nested_fields.clone(), true),
850            ]
851            .into(),
852            vec![
853                Arc::new(UInt64Array::from_iter_values([1])),
854                Arc::new(StringArray::from_iter_values(["large".repeat(50)])),
855                Arc::new(StructArray::new_null(nested_fields, 1)),
856            ],
857            None,
858        ));
859        let bound = ArrayBound::new(array, 0);
860
861        assert_none!(bound.to_proto_lower(0));
862        assert_none!(bound.to_proto_lower(1));
863
864        let proto = bound
865            .to_proto_lower(100)
866            .expect("can fit something in less than 100 bytes");
867        let array = make_array(proto.into_rust().expect("valid proto"));
868        assert_eq!(
869            array.as_struct().column_names().as_slice(),
870            &["a"],
871            "only the first column should fit"
872        );
873
874        let proto = bound
875            .to_proto_lower(1000)
876            .expect("can fit everything in less than 1000 bytes");
877        let array = make_array(proto.into_rust().expect("valid proto"));
878        assert_eq!(
879            array.as_struct().column_names().as_slice(),
880            &["a", "b", "c"],
881            "all columns should fit"
882        )
883    }
884
885    #[mz_ore::test]
886    fn struct_ord() {
887        let prefix = StructArray::new(
888            vec![Field::new("a", DataType::UInt64, true)].into(),
889            vec![Arc::new(UInt64Array::from_iter_values([1, 3, 5]))],
890            None,
891        );
892        let full = StructArray::new(
893            vec![
894                Field::new("a", DataType::UInt64, true),
895                Field::new("b", DataType::Utf8, true),
896            ]
897            .into(),
898            vec![
899                Arc::new(UInt64Array::from_iter_values([2, 3, 4])),
900                Arc::new(StringArray::from_iter_values(["a", "b", "c"])),
901            ],
902            None,
903        );
904        let prefix_ord = ArrayOrd::new(&prefix);
905        let full_ord = ArrayOrd::new(&full);
906
907        // Comparison works as normal over the shared columns... but when those columns are identical,
908        // the shorter struct is always smaller.
909        assert!(prefix_ord.at(0) < full_ord.at(0), "(1) < (2, 'a')");
910        assert!(prefix_ord.at(1) < full_ord.at(1), "(3) < (3, 'b')");
911        assert!(prefix_ord.at(2) > full_ord.at(2), "(5) < (4, 'c')");
912    }
913
914    #[mz_ore::test]
915    #[should_panic(expected = "array types did not match")]
916    fn struct_ord_incompat() {
917        // This test is descriptive, not prescriptive: we declare it is an error to compare
918        // structs like this, but not what the result of comparing them is.
919        let string = StructArray::new(
920            vec![
921                Field::new("a", DataType::UInt64, true),
922                Field::new("b", DataType::Utf8, true),
923            ]
924            .into(),
925            vec![
926                Arc::new(UInt64Array::from_iter_values([1])),
927                Arc::new(StringArray::from_iter_values(["a"])),
928            ],
929            None,
930        );
931        let boolean = StructArray::new(
932            vec![
933                Field::new("a", DataType::UInt64, true),
934                Field::new("b", DataType::Boolean, true),
935            ]
936            .into(),
937            vec![
938                Arc::new(UInt64Array::from_iter_values([1])),
939                Arc::new(BooleanArray::from_iter([Some(true)])),
940            ],
941            None,
942        );
943        let string_ord = ArrayOrd::new(&string);
944        let bool_ord = ArrayOrd::new(&boolean);
945
946        // Despite the matching first column, this will panic with a type mismatch.
947        assert!(string_ord.at(0) < bool_ord.at(0));
948    }
949}
mz_persist_types/arrow.rs

mz_persist_types/
arrow.rs
