Skip to main content

mz_repr/
row.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//! In-memory `Tag`-based encoding of a tuple of `Datum`s.
11//!
12//! See `doc/developer/row-encoding.md` for the size limits this encoding and
13//! its datum types impose.
14
15use std::borrow::Borrow;
16use std::cell::{Cell, RefCell};
17use std::cmp::Ordering;
18use std::convert::{TryFrom, TryInto};
19use std::fmt::{self, Debug};
20use std::hash::{Hash, Hasher};
21use std::marker::PhantomData;
22use std::mem::{size_of, transmute};
23use std::ops::Deref;
24use std::str;
25
26use chrono::{DateTime, Datelike, NaiveDate, NaiveDateTime, NaiveTime, Timelike, Utc};
27use compact_bytes::CompactBytes;
28use mz_ore::cast::{CastFrom, ReinterpretCast};
29use mz_ore::soft_assert_no_log;
30use mz_ore::vec::Vector;
31use mz_persist_types::Codec64;
32use num_enum::{IntoPrimitive, TryFromPrimitive};
33use ordered_float::OrderedFloat;
34#[cfg(any(test, feature = "proptest"))]
35use proptest::prelude::*;
36#[cfg(any(test, feature = "proptest"))]
37use proptest::strategy::{BoxedStrategy, Strategy};
38use serde::{Deserialize, Serialize};
39use uuid::Uuid;
40
41use crate::adt::array::{
42    Array, ArrayDimension, ArrayDimensions, InvalidArrayError, MAX_ARRAY_DIMENSIONS,
43};
44use crate::adt::date::Date;
45use crate::adt::interval::Interval;
46use crate::adt::mz_acl_item::{AclItem, MzAclItem};
47use crate::adt::numeric;
48use crate::adt::numeric::Numeric;
49use crate::adt::range::{
50    self, InvalidRangeError, Range, RangeBound, RangeInner, RangeLowerBound, RangeUpperBound,
51};
52use crate::adt::timestamp::CheckedTimestamp;
53#[cfg(any(test, feature = "proptest"))]
54use crate::scalar::arb_datum;
55use crate::scalar::{DatumKind, SqlScalarType};
56use crate::{Datum, RelationDesc, Timestamp};
57
58pub(crate) mod encode;
59pub mod iter;
60
61include!(concat!(env!("OUT_DIR"), "/mz_repr.row.rs"));
62
63/// A packed representation for `Datum`s.
64///
65/// `Datum` is easy to work with but very space inefficient. A `Datum::Int32(42)`
66/// is laid out in memory like this:
67///
68///   tag: 3
69///   padding: 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
70///   data: 0 0 0 42
71///   padding: 0 0 0 0 0 0 0 0 0 0 0 0
72///
73/// For a total of 32 bytes! The second set of padding is needed in case we were
74/// to write a 16-byte datum into this location. The first set of padding is
75/// needed to align that hypothetical decimal to a 16 bytes boundary.
76///
77/// A `Row` stores zero or more `Datum`s without any padding. We avoid the need
78/// for the first set of padding by only providing access to the `Datum`s via
79/// calls to `ptr::read_unaligned`, which on modern x86 is barely penalized. We
80/// avoid the need for the second set of padding by not providing mutable access
81/// to the `Datum`. Instead, `Row` is append-only.
82///
83/// A `Row` can be built from a collection of `Datum`s using `Row::pack`, but it
84/// is more efficient to use `Row::pack_slice` so that a right-sized allocation
85/// can be created. If that is not possible, consider using the row buffer
86/// pattern: allocate one row, pack into it, and then call [`Row::clone`] to
87/// receive a copy of that row, leaving behind the original allocation to pack
88/// future rows.
89///
90/// Creating a row via [`Row::pack_slice`]:
91///
92/// ```
93/// # use mz_repr::{Row, Datum};
94/// let row = Row::pack_slice(&[Datum::Int32(0), Datum::Int32(1), Datum::Int32(2)]);
95/// assert_eq!(row.unpack(), vec![Datum::Int32(0), Datum::Int32(1), Datum::Int32(2)])
96/// ```
97///
98/// `Row`s can be unpacked by iterating over them:
99///
100/// ```
101/// # use mz_repr::{Row, Datum};
102/// let row = Row::pack_slice(&[Datum::Int32(0), Datum::Int32(1), Datum::Int32(2)]);
103/// assert_eq!(row.iter().nth(1).unwrap(), Datum::Int32(1));
104/// ```
105///
106/// If you want random access to the `Datum`s in a `Row`, use `Row::unpack` to create a `Vec<Datum>`
107/// ```
108/// # use mz_repr::{Row, Datum};
109/// let row = Row::pack_slice(&[Datum::Int32(0), Datum::Int32(1), Datum::Int32(2)]);
110/// let datums = row.unpack();
111/// assert_eq!(datums[1], Datum::Int32(1));
112/// ```
113///
114/// # Performance
115///
116/// Rows are dynamically sized, but up to a fixed size their data is stored in-line.
117/// It is best to re-use a `Row` across multiple `Row` creation calls, as this
118/// avoids the allocations involved in `Row::new()`.
119#[derive(Default, Eq, PartialEq, Serialize, Deserialize)]
120pub struct Row {
121    data: CompactBytes,
122}
123
124impl Row {
125    const SIZE: usize = CompactBytes::MAX_INLINE;
126
127    /// A variant of `Row::from_proto` that allows for reuse of internal allocs
128    /// and validates the decoding against a provided [`RelationDesc`].
129    pub fn decode_from_proto(
130        &mut self,
131        proto: &ProtoRow,
132        desc: &RelationDesc,
133    ) -> Result<(), String> {
134        let mut packer = self.packer();
135        for (col_idx, _, _) in desc.iter_all() {
136            let d = match proto.datums.get(col_idx.to_raw()) {
137                Some(x) => x,
138                None => {
139                    packer.push(Datum::Null);
140                    continue;
141                }
142            };
143            packer.try_push_proto(d)?;
144        }
145
146        Ok(())
147    }
148
149    /// Allocate an empty `Row` with a pre-allocated capacity.
150    #[inline]
151    pub fn with_capacity(cap: usize) -> Self {
152        Self {
153            data: CompactBytes::with_capacity(cap),
154        }
155    }
156
157    /// Create an empty `Row`.
158    #[inline]
159    pub const fn empty() -> Self {
160        Self {
161            data: CompactBytes::empty(),
162        }
163    }
164
165    /// Creates a new row from supplied bytes.
166    ///
167    /// # Safety
168    ///
169    /// This method relies on `data` being an appropriate row encoding, and can
170    /// result in unsafety if this is not the case.
171    pub unsafe fn from_bytes_unchecked(data: &[u8]) -> Self {
172        Row {
173            data: CompactBytes::new(data),
174        }
175    }
176
177    /// Constructs a [`RowPacker`] that will pack datums into this row's
178    /// allocation.
179    ///
180    /// This method clears the existing contents of the row, but retains the
181    /// allocation.
182    pub fn packer(&mut self) -> RowPacker<'_> {
183        self.clear();
184        RowPacker { row: self }
185    }
186
187    /// Take some `Datum`s and pack them into a `Row`.
188    ///
189    /// This method builds a `Row` by repeatedly increasing the backing
190    /// allocation. If the contents of the iterator are known ahead of
191    /// time, consider [`Row::with_capacity`] to right-size the allocation
192    /// first, and then [`RowPacker::extend`] to populate it with `Datum`s.
193    /// This avoids the repeated allocation resizing and copying.
194    pub fn pack<'a, I, D>(iter: I) -> Row
195    where
196        I: IntoIterator<Item = D>,
197        D: Borrow<Datum<'a>>,
198    {
199        let mut row = Row::default();
200        row.packer().extend(iter);
201        row
202    }
203
204    /// Use `self` to pack `iter`, and then clone the result.
205    ///
206    /// This is a convenience method meant to reduce boilerplate around row
207    /// formation.
208    pub fn pack_using<'a, I, D>(&mut self, iter: I) -> Row
209    where
210        I: IntoIterator<Item = D>,
211        D: Borrow<Datum<'a>>,
212    {
213        self.packer().extend(iter);
214        self.clone()
215    }
216
217    /// Like [`Row::pack`], but the provided iterator is allowed to produce an
218    /// error, in which case the packing operation is aborted and the error
219    /// returned.
220    pub fn try_pack<'a, I, D, E>(iter: I) -> Result<Row, E>
221    where
222        I: IntoIterator<Item = Result<D, E>>,
223        D: Borrow<Datum<'a>>,
224    {
225        let mut row = Row::default();
226        row.packer().try_extend(iter)?;
227        Ok(row)
228    }
229
230    /// Pack a slice of `Datum`s into a `Row`.
231    ///
232    /// This method has the advantage over `pack` that it can determine the required
233    /// allocation before packing the elements, ensuring only one allocation and no
234    /// redundant copies required.
235    pub fn pack_slice<'a>(slice: &[Datum<'a>]) -> Row {
236        // Pre-allocate the needed number of bytes.
237        let mut row = Row::with_capacity(datums_size(slice.iter()));
238        row.packer().extend(slice.iter());
239        row
240    }
241
242    /// Returns the total amount of bytes used by this row.
243    pub fn byte_len(&self) -> usize {
244        let heap_size = if self.data.spilled() {
245            self.data.len()
246        } else {
247            0
248        };
249        let inline_size = std::mem::size_of::<Self>();
250        inline_size.saturating_add(heap_size)
251    }
252
253    /// The length of the encoded row in bytes. Does not include the size of the `Row` struct itself.
254    pub fn data_len(&self) -> usize {
255        self.data.len()
256    }
257
258    /// Returns the total capacity in bytes used by this row.
259    pub fn byte_capacity(&self) -> usize {
260        self.data.capacity()
261    }
262
263    /// Extracts a Row slice containing the entire [`Row`].
264    #[inline]
265    pub fn as_row_ref(&self) -> &RowRef {
266        // SAFETY: `Row` contains valid row data, by construction.
267        unsafe { RowRef::from_slice(self.data.as_slice()) }
268    }
269
270    /// Clear the contents of the [`Row`], leaving any allocation in place.
271    #[inline]
272    fn clear(&mut self) {
273        self.data.clear();
274    }
275}
276
277impl Borrow<RowRef> for Row {
278    #[inline]
279    fn borrow(&self) -> &RowRef {
280        self.as_row_ref()
281    }
282}
283
284impl AsRef<RowRef> for Row {
285    #[inline]
286    fn as_ref(&self) -> &RowRef {
287        self.as_row_ref()
288    }
289}
290
291impl Deref for Row {
292    type Target = RowRef;
293
294    #[inline]
295    fn deref(&self) -> &Self::Target {
296        self.as_row_ref()
297    }
298}
299
300// Nothing depends on Row being exactly 24, we just want to add visibility to the size.
301static_assertions::const_assert_eq!(std::mem::size_of::<Row>(), 24);
302
303impl Clone for Row {
304    fn clone(&self) -> Self {
305        Row {
306            data: self.data.clone(),
307        }
308    }
309
310    fn clone_from(&mut self, source: &Self) {
311        self.data.clone_from(&source.data);
312    }
313}
314
315// Row's `Hash` implementation defers to `RowRef` to ensure they hash equivalently.
316impl std::hash::Hash for Row {
317    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
318        self.as_row_ref().hash(state)
319    }
320}
321
322#[cfg(any(test, feature = "proptest"))]
323impl Arbitrary for Row {
324    type Parameters = prop::collection::SizeRange;
325    type Strategy = BoxedStrategy<Row>;
326
327    fn arbitrary_with(size: Self::Parameters) -> Self::Strategy {
328        prop::collection::vec(arb_datum(true), size)
329            .prop_map(|items| {
330                let mut row = Row::default();
331                let mut packer = row.packer();
332                for item in items.iter() {
333                    let datum: Datum<'_> = item.into();
334                    packer.push(datum);
335                }
336                row
337            })
338            .boxed()
339    }
340}
341
342impl PartialOrd for Row {
343    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
344        Some(self.cmp(other))
345    }
346}
347
348impl Ord for Row {
349    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
350        self.as_ref().cmp(other.as_ref())
351    }
352}
353
354/// A [`Row`] that serializes as protobuf-encoded [`ProtoRow`] bytes.
355///
356/// `Row`'s own serde impl emits the raw bytes of the in-memory `Tag` based
357/// datum encoding, which is free to change between releases. Use this wrapper
358/// instead wherever a row is serialized into a durable, cross-version format,
359/// such as the stable LIR plan format. `ProtoRow` already carries the needed
360/// backward compatibility obligation: it is persist's storage codec for
361/// `SourceData`, and `row.proto` is covered by the buf breaking lint.
362#[derive(
363    Clone,
364    Default,
365    Eq,
366    PartialEq,
367    Ord,
368    PartialOrd,
369    Hash,
370    Serialize,
371    Deserialize
372)]
373pub struct StableRow(#[serde(with = "stable_row_proto")] pub Row);
374
375impl From<Row> for StableRow {
376    fn from(row: Row) -> Self {
377        StableRow(row)
378    }
379}
380
381impl Deref for StableRow {
382    type Target = Row;
383
384    fn deref(&self) -> &Row {
385        &self.0
386    }
387}
388
389impl Debug for StableRow {
390    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
391        self.0.fmt(f)
392    }
393}
394
395mod stable_row_proto {
396    use mz_proto::RustType;
397    use prost::Message;
398    use serde::de::Error;
399    use serde::{Deserialize, Deserializer, Serializer};
400
401    use crate::row::{ProtoRow, Row};
402
403    pub fn serialize<S: Serializer>(row: &Row, serializer: S) -> Result<S::Ok, S::Error> {
404        serializer.serialize_bytes(&row.into_proto().encode_to_vec())
405    }
406
407    pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Row, D::Error> {
408        let bytes = serde_bytes::ByteBuf::deserialize(deserializer)?;
409        let proto = ProtoRow::decode(bytes.as_slice()).map_err(D::Error::custom)?;
410        Row::from_proto(proto).map_err(D::Error::custom)
411    }
412}
413
414#[allow(missing_debug_implementations)]
415mod columnation {
416    use columnation::{Columnation, Region};
417    use mz_ore::region::LgAllocRegion;
418
419    use crate::Row;
420
421    /// Region allocation for `Row` data.
422    ///
423    /// Content bytes are stored in stable contiguous memory locations,
424    /// and then a `Row` referencing them is falsified.
425    pub struct RowStack {
426        region: LgAllocRegion<u8>,
427    }
428
429    impl RowStack {
430        const LIMIT: usize = 2 << 20;
431    }
432
433    // Implement `Default` manually to specify a region allocation limit.
434    impl Default for RowStack {
435        fn default() -> Self {
436            Self {
437                // Limit the region size to 2MiB.
438                region: LgAllocRegion::with_limit(Self::LIMIT),
439            }
440        }
441    }
442
443    impl Columnation for Row {
444        type InnerRegion = RowStack;
445    }
446
447    impl Region for RowStack {
448        type Item = Row;
449        #[inline]
450        fn clear(&mut self) {
451            self.region.clear();
452        }
453        #[inline(always)]
454        unsafe fn copy(&mut self, item: &Row) -> Row {
455            if item.data.spilled() {
456                let bytes = self.region.copy_slice(&item.data[..]);
457                Row {
458                    data: compact_bytes::CompactBytes::from_raw_parts(
459                        bytes.as_mut_ptr(),
460                        item.data.len(),
461                        item.data.capacity(),
462                    ),
463                }
464            } else {
465                item.clone()
466            }
467        }
468
469        fn reserve_items<'a, I>(&mut self, items: I)
470        where
471            Self: 'a,
472            I: Iterator<Item = &'a Self::Item> + Clone,
473        {
474            let size = items
475                .filter(|row| row.data.spilled())
476                .map(|row| row.data.len())
477                .sum();
478            let size = std::cmp::min(size, Self::LIMIT);
479            self.region.reserve(size);
480        }
481
482        fn reserve_regions<'a, I>(&mut self, regions: I)
483        where
484            Self: 'a,
485            I: Iterator<Item = &'a Self> + Clone,
486        {
487            let size = regions.map(|r| r.region.len()).sum();
488            let size = std::cmp::min(size, Self::LIMIT);
489            self.region.reserve(size);
490        }
491
492        fn heap_size(&self, callback: impl FnMut(usize, usize)) {
493            self.region.heap_size(callback)
494        }
495    }
496}
497
498mod columnar {
499    use columnar::common::PushIndexAs;
500    use columnar::{
501        AsBytes, Borrow, Clear, Columnar, Container, FromBytes, Index, IndexAs, Len, Push,
502    };
503    use mz_ore::cast::CastFrom;
504    use std::ops::Range;
505
506    use crate::{Row, RowRef};
507
508    #[derive(
509        Copy,
510        Clone,
511        Debug,
512        Default,
513        PartialEq,
514        serde::Serialize,
515        serde::Deserialize
516    )]
517    pub struct Rows<BC = Vec<u64>, VC = Vec<u8>> {
518        /// Bounds container; provides indexed access to offsets.
519        bounds: BC,
520        /// Values container; provides slice access to bytes.
521        values: VC,
522    }
523
524    impl Columnar for Row {
525        #[inline(always)]
526        fn copy_from(&mut self, other: columnar::Ref<'_, Self>) {
527            self.clear();
528            self.data.extend_from_slice(other.data());
529        }
530        #[inline(always)]
531        fn into_owned(other: columnar::Ref<'_, Self>) -> Self {
532            other.to_owned()
533        }
534        type Container = Rows;
535        #[inline(always)]
536        fn reborrow<'b, 'a: 'b>(thing: columnar::Ref<'a, Self>) -> columnar::Ref<'b, Self>
537        where
538            Self: 'a,
539        {
540            thing
541        }
542    }
543
544    impl<BC: PushIndexAs<u64>> Borrow for Rows<BC, Vec<u8>> {
545        type Ref<'a> = &'a RowRef;
546        type Borrowed<'a>
547            = Rows<BC::Borrowed<'a>, &'a [u8]>
548        where
549            Self: 'a;
550        #[inline(always)]
551        fn borrow<'a>(&'a self) -> Self::Borrowed<'a> {
552            Rows {
553                bounds: self.bounds.borrow(),
554                values: self.values.borrow(),
555            }
556        }
557        #[inline(always)]
558        fn reborrow<'c, 'a: 'c>(item: Self::Borrowed<'a>) -> Self::Borrowed<'c>
559        where
560            Self: 'a,
561        {
562            Rows {
563                bounds: BC::reborrow(item.bounds),
564                values: item.values,
565            }
566        }
567
568        fn reborrow_ref<'b, 'a: 'b>(item: Self::Ref<'a>) -> Self::Ref<'b>
569        where
570            Self: 'a,
571        {
572            item
573        }
574    }
575
576    impl<BC: PushIndexAs<u64>> Container for Rows<BC, Vec<u8>> {
577        fn extend_from_self(&mut self, other: Self::Borrowed<'_>, range: Range<usize>) {
578            if !range.is_empty() {
579                // Imported bounds will be relative to this starting offset.
580                let values_len: u64 = self.values.len().try_into().expect("must fit");
581
582                // Push all bytes that we can, all at once.
583                let other_lower = if range.start == 0 {
584                    0
585                } else {
586                    other.bounds.index_as(range.start - 1)
587                };
588                let other_upper = other.bounds.index_as(range.end - 1);
589                self.values.extend_from_self(
590                    other.values,
591                    usize::try_from(other_lower).expect("must fit")
592                        ..usize::try_from(other_upper).expect("must fit"),
593                );
594
595                // Each bound needs to be shifted by `values_len - other_lower`.
596                if values_len == other_lower {
597                    self.bounds.extend_from_self(other.bounds, range);
598                } else {
599                    for index in range {
600                        let shifted = other.bounds.index_as(index) - other_lower + values_len;
601                        self.bounds.push(&shifted)
602                    }
603                }
604            }
605        }
606        fn reserve_for<'a, I>(&mut self, selves: I)
607        where
608            Self: 'a,
609            I: Iterator<Item = Self::Borrowed<'a>> + Clone,
610        {
611            self.bounds.reserve_for(selves.clone().map(|r| r.bounds));
612            self.values.reserve_for(selves.map(|r| r.values));
613        }
614    }
615
616    impl<'a, BC: AsBytes<'a>, VC: AsBytes<'a>> AsBytes<'a> for Rows<BC, VC> {
617        const SLICE_COUNT: usize = BC::SLICE_COUNT + VC::SLICE_COUNT;
618        #[inline(always)]
619        fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
620            mz_ore::soft_assert_no_log!(index < Self::SLICE_COUNT);
621            if index < BC::SLICE_COUNT {
622                self.bounds.get_byte_slice(index)
623            } else {
624                self.values.get_byte_slice(index - BC::SLICE_COUNT)
625            }
626        }
627    }
628    impl<'a, BC: FromBytes<'a>, VC: FromBytes<'a>> FromBytes<'a> for Rows<BC, VC> {
629        const SLICE_COUNT: usize = BC::SLICE_COUNT + VC::SLICE_COUNT;
630        #[inline(always)]
631        fn from_bytes(bytes: &mut impl Iterator<Item = &'a [u8]>) -> Self {
632            Self {
633                bounds: FromBytes::from_bytes(bytes),
634                values: FromBytes::from_bytes(bytes),
635            }
636        }
637    }
638
639    impl<BC: Len, VC> Len for Rows<BC, VC> {
640        #[inline(always)]
641        fn len(&self) -> usize {
642            self.bounds.len()
643        }
644    }
645
646    impl<'a, BC: Len + IndexAs<u64>> Index for Rows<BC, &'a [u8]> {
647        type Ref = &'a RowRef;
648        #[inline(always)]
649        fn get(&self, index: usize) -> Self::Ref {
650            let lower = if index == 0 {
651                0
652            } else {
653                self.bounds.index_as(index - 1)
654            };
655            let upper = self.bounds.index_as(index);
656            let lower = usize::cast_from(lower);
657            let upper = usize::cast_from(upper);
658            // SAFETY: self.values contains only valid row data, and self.metadata delimits only ranges
659            // that correspond to the original rows.
660            unsafe { RowRef::from_slice(&self.values[lower..upper]) }
661        }
662    }
663    impl<'a, BC: Len + IndexAs<u64>> Index for &'a Rows<BC, Vec<u8>> {
664        type Ref = &'a RowRef;
665        #[inline(always)]
666        fn get(&self, index: usize) -> Self::Ref {
667            let lower = if index == 0 {
668                0
669            } else {
670                self.bounds.index_as(index - 1)
671            };
672            let upper = self.bounds.index_as(index);
673            let lower = usize::cast_from(lower);
674            let upper = usize::cast_from(upper);
675            // SAFETY: self.values contains only valid row data, and self.metadata delimits only ranges
676            // that correspond to the original rows.
677            unsafe { RowRef::from_slice(&self.values[lower..upper]) }
678        }
679    }
680
681    impl<BC: Push<u64>> Push<&Row> for Rows<BC> {
682        #[inline(always)]
683        fn push(&mut self, item: &Row) {
684            self.values.extend_from_slice(item.data.as_slice());
685            self.bounds.push(u64::cast_from(self.values.len()));
686        }
687    }
688    impl<BC: for<'a> Push<&'a u64>> Push<&RowRef> for Rows<BC> {
689        #[inline(always)]
690        fn push(&mut self, item: &RowRef) {
691            self.values.extend_from_slice(item.data());
692            self.bounds.push(&u64::cast_from(self.values.len()));
693        }
694    }
695    impl<BC: Clear, VC: Clear> Clear for Rows<BC, VC> {
696        #[inline(always)]
697        fn clear(&mut self) {
698            self.bounds.clear();
699            self.values.clear();
700        }
701    }
702}
703
704/// A contiguous slice of bytes that are row data.
705///
706/// A [`RowRef`] is to [`Row`] as [`prim@str`] is to [`String`].
707#[derive(PartialEq, Eq, Hash)]
708#[repr(transparent)]
709pub struct RowRef([u8]);
710
711impl RowRef {
712    /// Create a [`RowRef`] from a slice of data.
713    ///
714    /// # Safety
715    ///
716    /// We do not check that the provided slice is valid [`Row`] data; the caller is required to
717    /// ensure this.
718    pub unsafe fn from_slice(row: &[u8]) -> &RowRef {
719        #[allow(clippy::as_conversions)]
720        let ptr = row as *const [u8] as *const RowRef;
721        // SAFETY: We know `ptr` is non-null and aligned because it came from a &[u8].
722        unsafe { &*ptr }
723    }
724
725    /// Unpack `self` into a `Vec<Datum>` for efficient random access.
726    pub fn unpack(&self) -> Vec<Datum<'_>> {
727        // It's usually cheaper to unpack twice to figure out the right length than it is to grow the vec as we go
728        let len = self.iter().count();
729        let mut vec = Vec::with_capacity(len);
730        vec.extend(self.iter());
731        vec
732    }
733
734    /// Return the first [`Datum`] in `self`
735    ///
736    /// Panics if the [`RowRef`] is empty.
737    pub fn unpack_first(&self) -> Datum<'_> {
738        self.iter().next().unwrap()
739    }
740
741    /// Iterate the [`Datum`] elements of the [`RowRef`].
742    pub fn iter(&self) -> DatumListIter<'_> {
743        DatumListIter { data: &self.0 }
744    }
745
746    /// Return the byte length of this [`RowRef`].
747    pub fn byte_len(&self) -> usize {
748        self.0.len()
749    }
750
751    /// For debugging only.
752    pub fn data(&self) -> &[u8] {
753        &self.0
754    }
755
756    /// True iff there is no data in this [`RowRef`].
757    pub fn is_empty(&self) -> bool {
758        self.0.is_empty()
759    }
760}
761
762impl ToOwned for RowRef {
763    type Owned = Row;
764
765    fn to_owned(&self) -> Self::Owned {
766        // SAFETY: RowRef has the invariant that the wrapped data must be a valid Row encoding.
767        unsafe { Row::from_bytes_unchecked(&self.0) }
768    }
769}
770
771impl<'a> IntoIterator for &'a RowRef {
772    type Item = Datum<'a>;
773    type IntoIter = DatumListIter<'a>;
774
775    fn into_iter(self) -> DatumListIter<'a> {
776        DatumListIter { data: &self.0 }
777    }
778}
779
780/// These implementations order first by length, and then by slice contents.
781/// This allows many comparisons to complete without dereferencing memory.
782/// Warning: These order by the u8 array representation, and NOT by Datum::cmp.
783impl PartialOrd for RowRef {
784    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
785        Some(self.cmp(other))
786    }
787}
788
789impl Ord for RowRef {
790    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
791        match self.0.len().cmp(&other.0.len()) {
792            std::cmp::Ordering::Less => std::cmp::Ordering::Less,
793            std::cmp::Ordering::Greater => std::cmp::Ordering::Greater,
794            std::cmp::Ordering::Equal => self.0.cmp(&other.0),
795        }
796    }
797}
798
799impl fmt::Debug for RowRef {
800    /// Debug representation using the internal datums
801    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
802        f.write_str("RowRef{")?;
803        f.debug_list().entries(&*self).finish()?;
804        f.write_str("}")
805    }
806}
807
808/// Packs datums into a [`Row`].
809///
810/// Creating a `RowPacker` via [`Row::packer`] starts a packing operation on the
811/// row. A packing operation always starts from scratch: the existing contents
812/// of the underlying row are cleared.
813///
814/// To complete a packing operation, drop the `RowPacker`.
815#[derive(Debug)]
816pub struct RowPacker<'a> {
817    row: &'a mut Row,
818}
819
820/// Infallible conversion from a [`Datum`] to a typed value.
821///
822/// Used by [`DatumList::typed_iter`] to yield elements as `T` rather than
823/// raw `Datum`s. At runtime, `T` is always `Datum<'a>`, so the conversion
824/// is identity.
825///
826/// See `doc/developer/design/20260311_sqlfunc_generic.md` for the design
827/// behind the generic type parameter and type erasure.
828///
829/// This trait is sealed and cannot be implemented outside of this crate.
830pub trait FromDatum<'a>:
831    Sized + PartialEq + std::borrow::Borrow<Datum<'a>> + sealed::Sealed
832{
833    fn from_datum(datum: Datum<'a>) -> Self;
834}
835
836mod sealed {
837    use crate::Datum;
838
839    pub trait Sealed {}
840    impl<'a> Sealed for Datum<'a> {}
841}
842
843impl<'a> FromDatum<'a> for Datum<'a> {
844    #[inline]
845    fn from_datum(datum: Datum<'a>) -> Self {
846        datum
847    }
848}
849
850#[derive(Debug, Clone)]
851pub struct DatumListIter<'a> {
852    data: &'a [u8],
853}
854
855#[derive(Debug, Clone)]
856pub struct DatumListTypedIter<'a, T> {
857    inner: DatumListIter<'a>,
858    _phantom: PhantomData<fn() -> T>,
859}
860
861#[derive(Debug, Clone)]
862pub struct DatumDictIter<'a> {
863    data: &'a [u8],
864    prev_key: Option<&'a str>,
865}
866
867#[derive(Debug, Clone)]
868pub struct DatumDictTypedIter<'a, T> {
869    inner: DatumDictIter<'a>,
870    _phantom: PhantomData<fn() -> T>,
871}
872
873/// `RowArena` is used to hold on to temporary `Row`s for functions like `eval` that need to create complex `Datum`s but don't have a `Row` to put them in yet.
874#[derive(Debug)]
875pub struct RowArena {
876    // A stack of byte regions, used as a bump allocator. Bytes handed to
877    // `push_bytes` are *copied* into the active (last) region and a reference
878    // into that region is returned.
879    //
880    // The invariant that keeps returned references valid for the arena's
881    // lifetime is that a region is never reallocated once it holds data: when
882    // the active region lacks spare capacity for a push we allocate a *new*,
883    // larger region rather than growing the current one (which would move its
884    // bytes and dangle outstanding references). The outer `Vec` may itself
885    // reallocate as regions are added, but that only moves the `Vec<u8>`
886    // headers, not the heap buffers they own, so references remain valid.
887    //
888    // `clear` retains only the largest region (emptied) to right-size the arena
889    // for reuse; reusing one region across `clear` cycles makes a steady-state
890    // workload (e.g. decoding rows one at a time) allocation-free.
891    inner: RefCell<Vec<Vec<u8>>>,
892    // A single recycled scratch buffer backing `RowArena::writer`. A writer takes ownership of this
893    // buffer (or allocates a fresh one if absent), builds into it, and on drop returns it here for
894    // the next writer to reuse — so building values incrementally does not allocate per use once the
895    // buffer reaches its high-water mark. Holding `Option` (rather than the buffer directly) means
896    // `writer` borrows this cell only transiently, to take and return the buffer, never across the
897    // writer's lifetime. That keeps nested writers sound: a writer obtained while another is live
898    // finds the slot empty and allocates its own buffer instead of double-borrowing.
899    scratch: RefCell<Option<Vec<u8>>>,
900    // Optional ceiling on the bytes this arena will hold, and a running total of what it holds.
901    // `None` is unbounded, which is what every arena in a dataflow must keep using: a budget that
902    // can change mid-run would make dataflow evaluation non-deterministic (see
903    // [`RowArena::with_budget`]). A budget is for evaluating a user-authored expression in a shared
904    // process, where that expression's memory use must be bounded (see `mz_adapter::webhook`).
905    //
906    // NOTE: exceeding the budget does not make a push fail. The pushes are infallible, and a
907    // refused push would hand back a truncated value, i.e. a corrupt datum. The budget is instead a
908    // *reported* condition: `over_budget` is polled by whoever is able to return an error, which
909    // for scalar expressions is the evaluator between calls.
910    budget: Option<usize>,
911    allocated: Cell<usize>,
912}
913
914// DatumList and DatumDict defined here rather than near Datum because we need private access to the unsafe data field
915
916/// A sequence of Datums
917///
918/// The type parameter `T` represents the element type of the list. It is a
919/// phantom parameter that carries no runtime data — the actual elements are
920/// stored as serialized bytes and `T` is not enforced at runtime. It is up
921/// to the caller to ensure `T` matches the actual element type. The default
922/// `T = Datum<'a>` means existing code that writes `DatumList<'a>` continues
923/// to work unchanged.
924///
925/// See `doc/developer/design/20260311_sqlfunc_generic.md` for the design
926/// behind the generic type parameter.
927pub struct DatumList<'a, T = Datum<'a>> {
928    /// Points at the serialized datums
929    data: &'a [u8],
930    _phantom: PhantomData<fn() -> T>,
931}
932
933impl<'a, T> DatumList<'a, T> {
934    /// Private constructor. All `DatumList` values should be created through
935    /// this function to keep the `PhantomData` bookkeeping in one place.
936    pub(crate) fn new(data: &'a [u8]) -> Self {
937        DatumList {
938            data,
939            _phantom: PhantomData,
940        }
941    }
942}
943
944impl<'a, T> Clone for DatumList<'a, T> {
945    fn clone(&self) -> Self {
946        *self
947    }
948}
949
950impl<'a, T> Copy for DatumList<'a, T> {}
951
952impl<'a, T> Debug for DatumList<'a, T> {
953    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
954        f.debug_list().entries(self.iter()).finish()
955    }
956}
957
958impl<'a, T> PartialEq for DatumList<'a, T> {
959    #[inline(always)]
960    fn eq(&self, other: &DatumList<'a, T>) -> bool {
961        self.iter().eq(other.iter())
962    }
963}
964
965impl<'a, T> Eq for DatumList<'a, T> {}
966
967impl<'a, T> Hash for DatumList<'a, T> {
968    #[inline(always)]
969    fn hash<H: Hasher>(&self, state: &mut H) {
970        for d in self.iter() {
971            d.hash(state);
972        }
973    }
974}
975
976impl<T> Ord for DatumList<'_, T> {
977    #[inline(always)]
978    fn cmp(&self, other: &DatumList<'_, T>) -> Ordering {
979        // Grow the stack: lists can be arbitrarily deeply nested (e.g. jsonb).
980        mz_ore::stack::maybe_grow(|| self.iter().cmp(other.iter()))
981    }
982}
983
984impl<T> PartialOrd for DatumList<'_, T> {
985    #[inline(always)]
986    fn partial_cmp(&self, other: &DatumList<'_, T>) -> Option<Ordering> {
987        Some(self.cmp(other))
988    }
989}
990
991/// A mapping from string keys to Datums
992///
993/// The type parameter `T` represents the value type of the map. It is a
994/// phantom parameter — the actual values are stored as serialized bytes and
995/// `T` is not enforced at runtime. It is up to the caller to ensure `T`
996/// matches the actual value type. The default `T = Datum<'a>` means existing
997/// code that writes `DatumMap<'a>` continues to work unchanged.
998///
999/// See `doc/developer/design/20260311_sqlfunc_generic.md` for the design
1000/// behind the generic type parameter.
1001pub struct DatumMap<'a, T = Datum<'a>> {
1002    /// Points at the serialized datums, which should be sorted in key order
1003    data: &'a [u8],
1004    _phantom: PhantomData<fn() -> T>,
1005}
1006
1007impl<'a, T> DatumMap<'a, T> {
1008    /// Private constructor. All `DatumMap` values should be created through
1009    /// this function to keep the `PhantomData` bookkeeping in one place.
1010    pub(crate) fn new(data: &'a [u8]) -> Self {
1011        DatumMap {
1012            data,
1013            _phantom: PhantomData,
1014        }
1015    }
1016}
1017
1018impl<'a, T> Clone for DatumMap<'a, T> {
1019    fn clone(&self) -> Self {
1020        *self
1021    }
1022}
1023
1024impl<'a, T> Copy for DatumMap<'a, T> {}
1025
1026impl<'a, T> PartialEq for DatumMap<'a, T> {
1027    #[inline(always)]
1028    fn eq(&self, other: &DatumMap<'a, T>) -> bool {
1029        self.iter().eq(other.iter())
1030    }
1031}
1032
1033impl<'a, T> Eq for DatumMap<'a, T> {}
1034
1035impl<'a, T> Hash for DatumMap<'a, T> {
1036    #[inline(always)]
1037    fn hash<H: Hasher>(&self, state: &mut H) {
1038        for (k, v) in self.iter() {
1039            k.hash(state);
1040            v.hash(state);
1041        }
1042    }
1043}
1044
1045impl<'a, T> Ord for DatumMap<'a, T> {
1046    #[inline(always)]
1047    fn cmp(&self, other: &DatumMap<'a, T>) -> Ordering {
1048        // Grow the stack: maps can be arbitrarily deeply nested (e.g. jsonb).
1049        mz_ore::stack::maybe_grow(|| self.iter().cmp(other.iter()))
1050    }
1051}
1052
1053impl<'a, T> PartialOrd for DatumMap<'a, T> {
1054    #[inline(always)]
1055    fn partial_cmp(&self, other: &DatumMap<'a, T>) -> Option<Ordering> {
1056        Some(self.cmp(other))
1057    }
1058}
1059
1060impl<'a> crate::scalar::SqlContainerType for DatumList<'a, Datum<'a>> {
1061    fn unwrap_element_type(container: &SqlScalarType) -> &SqlScalarType {
1062        container.unwrap_list_element_type()
1063    }
1064    fn wrap_element_type(element: SqlScalarType) -> SqlScalarType {
1065        SqlScalarType::List {
1066            element_type: Box::new(element),
1067            custom_id: None,
1068        }
1069    }
1070}
1071
1072impl<'a> crate::scalar::SqlContainerType for DatumMap<'a, Datum<'a>> {
1073    fn unwrap_element_type(container: &SqlScalarType) -> &SqlScalarType {
1074        container.unwrap_map_value_type()
1075    }
1076    fn wrap_element_type(element: SqlScalarType) -> SqlScalarType {
1077        SqlScalarType::Map {
1078            value_type: Box::new(element),
1079            custom_id: None,
1080        }
1081    }
1082}
1083
1084/// Represents a single `Datum`, appropriate to be nested inside other
1085/// `Datum`s.
1086#[derive(Clone, Copy, Eq, PartialEq, Hash)]
1087pub struct DatumNested<'a> {
1088    val: &'a [u8],
1089}
1090
1091impl<'a> std::fmt::Display for DatumNested<'a> {
1092    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1093        std::fmt::Display::fmt(&self.datum(), f)
1094    }
1095}
1096
1097impl<'a> std::fmt::Debug for DatumNested<'a> {
1098    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1099        f.debug_struct("DatumNested")
1100            .field("val", &self.datum())
1101            .finish()
1102    }
1103}
1104
1105impl<'a> DatumNested<'a> {
1106    // Figure out which bytes `read_datum` returns (e.g. including the tag),
1107    // and then store a reference to those bytes, so we can "replay" this same
1108    // call later on without storing the datum itself.
1109    pub fn extract(data: &mut &'a [u8]) -> DatumNested<'a> {
1110        let prev = *data;
1111        let _ = unsafe { read_datum(data) };
1112        DatumNested {
1113            val: &prev[..(prev.len() - data.len())],
1114        }
1115    }
1116
1117    /// Returns the datum `self` contains.
1118    pub fn datum(&self) -> Datum<'a> {
1119        let mut temp = self.val;
1120        unsafe { read_datum(&mut temp) }
1121    }
1122}
1123
1124impl<'a> Ord for DatumNested<'a> {
1125    fn cmp(&self, other: &Self) -> Ordering {
1126        // Grow the stack: this recurses once per level of nested list/map values.
1127        mz_ore::stack::maybe_grow(|| self.datum().cmp(&other.datum()))
1128    }
1129}
1130
1131impl<'a> PartialOrd for DatumNested<'a> {
1132    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1133        Some(self.cmp(other))
1134    }
1135}
1136
1137// Prefer adding new tags to the end of the enum. Certain behavior, like row ordering and EXPLAIN
1138// PHYSICAL PLAN, rely on the ordering of this enum. Neither of these are breaking changes, but
1139// it's annoying when they change.
1140#[derive(Debug, Clone, Copy, PartialEq, Eq, IntoPrimitive, TryFromPrimitive)]
1141#[repr(u8)]
1142enum Tag {
1143    Null,
1144    False,
1145    True,
1146    Float32,
1147    Float64,
1148    Date,
1149    Time,
1150    Timestamp,
1151    TimestampTz,
1152    Interval,
1153    // The length-prefixed tags: for each of the three kinds, four tags ordered by the width of
1154    // the length prefix in front of the payload, so a tag's distance from its kind's first tag
1155    // selects the width. Each kind has its own `read_datum` arm naming the `Datum` constructor,
1156    // since the kind never varies within a column; the width does, so it stays inside the arm.
1157    BytesTiny,
1158    BytesShort,
1159    BytesLong,
1160    BytesHuge,
1161    StringTiny,
1162    StringShort,
1163    StringLong,
1164    StringHuge,
1165    ListTiny,
1166    ListShort,
1167    ListLong,
1168    ListHuge,
1169    Uuid,
1170    Array,
1171    Dict,
1172    JsonNull,
1173    Dummy,
1174    Numeric,
1175    MzTimestamp,
1176    Range,
1177    MzAclItem,
1178    AclItem,
1179    // Everything except leap seconds and times beyond the range of
1180    // i64 nanoseconds. (Note that Materialize does not support leap
1181    // seconds, but this module does).
1182    CheapTimestamp,
1183    // Everything except leap seconds and times beyond the range of
1184    // i64 nanoseconds. (Note that Materialize does not support leap
1185    // seconds, but this module does).
1186    CheapTimestampTz,
1187    // The next several tags are for variable-length signed integer encoding.
1188    // The basic idea is that `NonNegativeIntN_K` is used to encode a datum of type
1189    // IntN whose actual value is positive or zero and fits in K bits, and similarly for
1190    // NegativeIntN_K with negative values.
1191    //
1192    // The order of these tags matters, because we want to be able to choose the
1193    // tag for a given datum quickly, with arithmetic, rather than slowly, with a
1194    // stack of `if` statements.
1195    //
1196    // Separate tags for non-negative and negative numbers are used to avoid having to
1197    // waste one bit in the actual data space to encode the sign.
1198    // A signed family alternates non-negative and negative at each payload width, so a tag
1199    // splits into both by arithmetic: the width is its distance from the family's first tag
1200    // shifted right by one, and the sign is that distance's low bit. Keeping the two signs
1201    // apart would make the width a subtraction from one of two bases, chosen by a compare.
1202    //
1203    // The family ends at the width where the payload is the whole integer. There the sign is
1204    // already the payload's top bit, so one tag serves both and the alternation stops, which is
1205    // why the fixed-width tag sits there rather than in a family of its own.
1206    NonNegativeInt16_0, // i.e., 0
1207    NegativeInt16_0,    // i.e., -1
1208    NonNegativeInt16_8,
1209    NegativeInt16_8,
1210    Int16,
1211
1212    NonNegativeInt32_0,
1213    NegativeInt32_0,
1214    NonNegativeInt32_8,
1215    NegativeInt32_8,
1216    NonNegativeInt32_16,
1217    NegativeInt32_16,
1218    NonNegativeInt32_24,
1219    NegativeInt32_24,
1220    Int32,
1221
1222    NonNegativeInt64_0,
1223    NegativeInt64_0,
1224    NonNegativeInt64_8,
1225    NegativeInt64_8,
1226    NonNegativeInt64_16,
1227    NegativeInt64_16,
1228    NonNegativeInt64_24,
1229    NegativeInt64_24,
1230    NonNegativeInt64_32,
1231    NegativeInt64_32,
1232    NonNegativeInt64_40,
1233    NegativeInt64_40,
1234    NonNegativeInt64_48,
1235    NegativeInt64_48,
1236    NonNegativeInt64_56,
1237    NegativeInt64_56,
1238    Int64,
1239
1240    // These are like the ones above, but for unsigned types. The situation is slightly simpler
1241    // as we don't have negatives, so the width is the whole distance from the family's first
1242    // tag, and the fixed-width tag is again the widest member.
1243    UInt8_0, // i.e., 0
1244    UInt8,
1245
1246    UInt16_0,
1247    UInt16_8,
1248    UInt16,
1249
1250    UInt32_0,
1251    UInt32_8,
1252    UInt32_16,
1253    UInt32_24,
1254    UInt32,
1255
1256    UInt64_0,
1257    UInt64_8,
1258    UInt64_16,
1259    UInt64_24,
1260    UInt64_32,
1261    UInt64_40,
1262    UInt64_48,
1263    UInt64_56,
1264    UInt64,
1265}
1266
1267impl Tag {
1268    /// The tag's discriminant, usable in const context.
1269    #[allow(clippy::as_conversions)]
1270    const fn byte(self) -> u8 {
1271        self as u8
1272    }
1273}
1274
1275/// Assert that the listed tags are consecutive, in the order given.
1276///
1277/// Every family below is addressed by arithmetic rather than by name: `push_datum` writes
1278/// `first + n` for a payload of `n` bytes, and `read_signed_varint`/`read_unsigned_varint`
1279/// invert that by subtraction. A variant inserted into the middle of a family, or two members
1280/// swapped, silently changes how many bytes a tag claims, which corrupts the datum rather than
1281/// failing to compile. These assertions turn that into a compile error at the definition.
1282macro_rules! assert_consecutive {
1283    ($($tag:ident),+ $(,)?) => {
1284        const _: () = {
1285            let tags: &[u8] = &[$(Tag::$tag.byte()),+];
1286            let mut i = 1;
1287            while i < tags.len() {
1288                assert!(
1289                    tags[i] == tags[i - 1] + 1,
1290                    concat!("tags are not consecutive: ", stringify!($($tag),+))
1291                );
1292                i += 1;
1293            }
1294        };
1295    };
1296}
1297
1298assert_consecutive!(BytesTiny, BytesShort, BytesLong, BytesHuge);
1299assert_consecutive!(StringTiny, StringShort, StringLong, StringHuge);
1300assert_consecutive!(ListTiny, ListShort, ListLong, ListHuge);
1301assert_consecutive!(
1302    NonNegativeInt16_0,
1303    NegativeInt16_0,
1304    NonNegativeInt16_8,
1305    NegativeInt16_8,
1306    Int16,
1307);
1308assert_consecutive!(
1309    NonNegativeInt32_0,
1310    NegativeInt32_0,
1311    NonNegativeInt32_8,
1312    NegativeInt32_8,
1313    NonNegativeInt32_16,
1314    NegativeInt32_16,
1315    NonNegativeInt32_24,
1316    NegativeInt32_24,
1317    Int32,
1318);
1319assert_consecutive!(
1320    NonNegativeInt64_0,
1321    NegativeInt64_0,
1322    NonNegativeInt64_8,
1323    NegativeInt64_8,
1324    NonNegativeInt64_16,
1325    NegativeInt64_16,
1326    NonNegativeInt64_24,
1327    NegativeInt64_24,
1328    NonNegativeInt64_32,
1329    NegativeInt64_32,
1330    NonNegativeInt64_40,
1331    NegativeInt64_40,
1332    NonNegativeInt64_48,
1333    NegativeInt64_48,
1334    NonNegativeInt64_56,
1335    NegativeInt64_56,
1336    Int64,
1337);
1338assert_consecutive!(UInt8_0, UInt8,);
1339assert_consecutive!(UInt16_0, UInt16_8, UInt16,);
1340assert_consecutive!(UInt32_0, UInt32_8, UInt32_16, UInt32_24, UInt32,);
1341assert_consecutive!(
1342    UInt64_0, UInt64_8, UInt64_16, UInt64_24, UInt64_32, UInt64_40, UInt64_48, UInt64_56, UInt64,
1343);
1344
1345// --------------------------------------------------------------------------------
1346// reading data
1347
1348/// Read a byte slice starting at byte `offset`.
1349///
1350/// Updates `offset` to point to the first byte after the end of the read region.
1351fn read_untagged_bytes<'a>(data: &mut &'a [u8]) -> &'a [u8] {
1352    let len = u64::from_le_bytes(read_byte_array(data));
1353    let len = usize::cast_from(len);
1354    let (bytes, next) = data.split_at(len);
1355    *data = next;
1356    bytes
1357}
1358
1359/// Read a byte slice preceded by its length, the width of the length prefix given by the tag's
1360/// distance from `first`, its kind's `*Tiny` tag.
1361///
1362/// Each arm reads the prefix at a constant width, which is a plain load; deriving the width and
1363/// reading that many bytes measured twice as slow on long payloads.
1364///
1365/// # Safety
1366///
1367/// The contents are whatever `push_lengthed_bytes` wrote, so a caller may treat them as UTF-8
1368/// only for a `String` tag.
1369#[inline(always)]
1370fn read_lengthed_bytes<'a>(data: &mut &'a [u8], tag: Tag, first: Tag) -> &'a [u8] {
1371    let len = match u8::from(tag).wrapping_sub(u8::from(first)) {
1372        0 => usize::from(read_byte(data)),
1373        1 => usize::from(u16::from_le_bytes(read_byte_array(data))),
1374        2 => usize::cast_from(u32::from_le_bytes(read_byte_array(data))),
1375        _ => usize::cast_from(u64::from_le_bytes(read_byte_array(data))),
1376    };
1377    let (bytes, next) = data.split_at(len);
1378    *data = next;
1379    bytes
1380}
1381
1382#[inline(always)]
1383fn read_byte(data: &mut &[u8]) -> u8 {
1384    let byte = data[0];
1385    *data = &data[1..];
1386    byte
1387}
1388
1389/// The payload of a variable-length integer whose datum ends within eight bytes of the end of
1390/// `data`, so the wide load in [`read_varint_word`] would run off the end.
1391///
1392/// Out of line and cold: this is at most the tail of a row, or of a nested list or map.
1393#[cold]
1394#[inline(never)]
1395fn read_varint_word_tail(data: &[u8], len: usize) -> u64 {
1396    #[inline(always)]
1397    fn ext<const L: usize>(data: &[u8]) -> u64 {
1398        let mut raw = [0; 8];
1399        raw[..L].copy_from_slice(&data[..L]);
1400        u64::from_le_bytes(raw)
1401    }
1402    // Each arm reads a constant width, which is a load rather than a `memcpy`. This path serves
1403    // the last datum or two of every row, which at low arity is a real share of all datums.
1404    match len {
1405        0 => 0,
1406        1 => u64::from(data[0]),
1407        2 => ext::<2>(data),
1408        3 => ext::<3>(data),
1409        4 => ext::<4>(data),
1410        5 => ext::<5>(data),
1411        6 => ext::<6>(data),
1412        7 => ext::<7>(data),
1413        // An eight-byte payload needs eight bytes past the tag, which this path's caller found
1414        // wanting, so a valid row cannot reach here.
1415        _ => panic!("payload runs past the end of the row"),
1416    }
1417}
1418
1419/// Read the `len` payload bytes of a variable-length integer, returning them in the low
1420/// `len * 8` bits of a word. Bits above that are zero or belong to whatever follows in the row,
1421/// so a caller must mask them off.
1422///
1423/// Loading a fixed eight bytes is what keeps `len` out of the load, and so off a branch. Those
1424/// bytes exist except at the very end of the buffer.
1425#[inline(always)]
1426fn read_varint_word(data: &mut &[u8], len: usize) -> u64 {
1427    let word = match data.first_chunk::<8>() {
1428        Some(chunk) => u64::from_le_bytes(*chunk),
1429        None => read_varint_word_tail(data, len),
1430    };
1431    *data = &data[len..];
1432    word
1433}
1434
1435/// Mask covering the low `len` bytes of a word.
1436///
1437/// `len` reaches eight for a 64-bit value that needs every byte, where a shift of 64 would be
1438/// undefined, so that case saturates instead.
1439#[inline(always)]
1440fn payload_mask(len: usize) -> u64 {
1441    if len >= 8 {
1442        u64::MAX
1443    } else {
1444        (1u64 << (len * 8)) - 1
1445    }
1446}
1447
1448/// The low `N` bytes of `word`, little-endian.
1449#[inline(always)]
1450fn truncate<const N: usize>(word: u64) -> [u8; N] {
1451    word.to_le_bytes()[..N].try_into().expect("N <= 8")
1452}
1453
1454/// Read the payload of a variable-length integer of either sign, extended to `N` bytes.
1455///
1456/// A signed family alternates the two signs at each payload width, so the tag's distance from
1457/// `first` holds the width above its low bit and the sign in it. Both fall out by shifting and
1458/// masking, with no compare and nothing to dispatch on: a column's tag varies with the magnitude
1459/// and sign of every value, so any branch on it is one the predictor cannot learn.
1460///
1461/// # Correctness
1462///
1463/// `tag` must belong to the family starting at `first`, and `data` must hold its payload.
1464#[inline(always)]
1465fn read_signed_varint<const N: usize>(data: &mut &[u8], tag: Tag, first: Tag) -> [u8; N] {
1466    let delta = u8::from(tag).wrapping_sub(u8::from(first));
1467    let len = usize::from(delta >> 1);
1468    // All ones for a negative value, so the bytes above the payload sign-extend, and zero
1469    // otherwise. A `len` of zero leaves the whole word filled, which is the -1 the encoder means.
1470    let negative = delta & 1 == 1;
1471    read_varint_payload(data, len, negative)
1472}
1473
1474/// Read a `len` byte payload and extend it to `N` bytes, with ones above it when `negative`.
1475///
1476/// `len` must not exceed `N`, since `truncate` keeps only the low `N` bytes and a wider payload
1477/// would lose its top ones.
1478#[inline(always)]
1479fn read_varint_payload<const N: usize>(data: &mut &[u8], len: usize, negative: bool) -> [u8; N] {
1480    let mask = payload_mask(len);
1481    let fill = 0u64.wrapping_sub(u64::from(negative));
1482    truncate((read_varint_word(data, len) & mask) | (fill & !mask))
1483}
1484
1485/// Read the payload of an unsigned variable-length integer, zero-extended to `N` bytes.
1486///
1487/// As [`read_signed_varint`], without the sign.
1488///
1489/// # Correctness
1490///
1491/// `tag` must belong to the family starting at `first`, and `data` must hold its payload.
1492#[inline(always)]
1493fn read_unsigned_varint<const N: usize>(data: &mut &[u8], tag: Tag, first: Tag) -> [u8; N] {
1494    let len = usize::from(u8::from(tag).wrapping_sub(u8::from(first)));
1495    read_varint_payload(data, len, false)
1496}
1497
1498#[inline(always)]
1499pub(super) fn read_byte_array<const N: usize>(data: &mut &[u8]) -> [u8; N] {
1500    let (prev, next) = data.split_first_chunk().unwrap();
1501    *data = next;
1502    *prev
1503}
1504
1505pub(super) fn read_date(data: &mut &[u8]) -> Date {
1506    let days = i32::from_le_bytes(read_byte_array(data));
1507    Date::from_pg_epoch(days).expect("unexpected date")
1508}
1509
1510pub(super) fn read_naive_date(data: &mut &[u8]) -> NaiveDate {
1511    let year = i32::from_le_bytes(read_byte_array(data));
1512    let ordinal = u32::from_le_bytes(read_byte_array(data));
1513    NaiveDate::from_yo_opt(year, ordinal).unwrap()
1514}
1515
1516pub(super) fn read_time(data: &mut &[u8]) -> NaiveTime {
1517    let secs = u32::from_le_bytes(read_byte_array(data));
1518    let nanos = u32::from_le_bytes(read_byte_array(data));
1519    NaiveTime::from_num_seconds_from_midnight_opt(secs, nanos).unwrap()
1520}
1521
1522/// Read a datum starting at byte `offset`.
1523///
1524/// Updates `offset` to point to the first byte after the end of the read region.
1525///
1526/// # Safety
1527///
1528/// This function is safe if a `Datum` was previously written at this offset by `push_datum`.
1529/// Otherwise it could return invalid values, which is Undefined Behavior.
1530pub unsafe fn read_datum<'a>(data: &mut &'a [u8]) -> Datum<'a> {
1531    let tag = Tag::try_from_primitive(read_byte(data)).expect("unknown row tag");
1532    match tag {
1533        Tag::Null => Datum::Null,
1534        Tag::False => Datum::False,
1535        Tag::True => Datum::True,
1536        Tag::NonNegativeInt16_0
1537        | Tag::NegativeInt16_0
1538        | Tag::NonNegativeInt16_8
1539        | Tag::NegativeInt16_8
1540        | Tag::Int16 => Datum::Int16(i16::from_le_bytes(read_signed_varint(
1541            data,
1542            tag,
1543            Tag::NonNegativeInt16_0,
1544        ))),
1545        Tag::NonNegativeInt32_0
1546        | Tag::NegativeInt32_0
1547        | Tag::NonNegativeInt32_8
1548        | Tag::NegativeInt32_8
1549        | Tag::NonNegativeInt32_16
1550        | Tag::NegativeInt32_16
1551        | Tag::NonNegativeInt32_24
1552        | Tag::NegativeInt32_24
1553        | Tag::Int32 => Datum::Int32(i32::from_le_bytes(read_signed_varint(
1554            data,
1555            tag,
1556            Tag::NonNegativeInt32_0,
1557        ))),
1558        Tag::NonNegativeInt64_0
1559        | Tag::NegativeInt64_0
1560        | Tag::NonNegativeInt64_8
1561        | Tag::NegativeInt64_8
1562        | Tag::NonNegativeInt64_16
1563        | Tag::NegativeInt64_16
1564        | Tag::NonNegativeInt64_24
1565        | Tag::NegativeInt64_24
1566        | Tag::NonNegativeInt64_32
1567        | Tag::NegativeInt64_32
1568        | Tag::NonNegativeInt64_40
1569        | Tag::NegativeInt64_40
1570        | Tag::NonNegativeInt64_48
1571        | Tag::NegativeInt64_48
1572        | Tag::NonNegativeInt64_56
1573        | Tag::NegativeInt64_56
1574        | Tag::Int64 => Datum::Int64(i64::from_le_bytes(read_signed_varint(
1575            data,
1576            tag,
1577            Tag::NonNegativeInt64_0,
1578        ))),
1579        Tag::UInt8_0 | Tag::UInt8 => Datum::UInt8(u8::from_le_bytes(read_unsigned_varint(
1580            data,
1581            tag,
1582            Tag::UInt8_0,
1583        ))),
1584        Tag::UInt16_0 | Tag::UInt16_8 | Tag::UInt16 => Datum::UInt16(u16::from_le_bytes(
1585            read_unsigned_varint(data, tag, Tag::UInt16_0),
1586        )),
1587        Tag::UInt32_0 | Tag::UInt32_8 | Tag::UInt32_16 | Tag::UInt32_24 | Tag::UInt32 => {
1588            Datum::UInt32(u32::from_le_bytes(read_unsigned_varint(
1589                data,
1590                tag,
1591                Tag::UInt32_0,
1592            )))
1593        }
1594        Tag::UInt64_0
1595        | Tag::UInt64_8
1596        | Tag::UInt64_16
1597        | Tag::UInt64_24
1598        | Tag::UInt64_32
1599        | Tag::UInt64_40
1600        | Tag::UInt64_48
1601        | Tag::UInt64_56
1602        | Tag::UInt64 => Datum::UInt64(u64::from_le_bytes(read_unsigned_varint(
1603            data,
1604            tag,
1605            Tag::UInt64_0,
1606        ))),
1607
1608        Tag::Float32 => {
1609            let f = f32::from_bits(u32::from_le_bytes(read_byte_array(data)));
1610            Datum::Float32(OrderedFloat::from(f))
1611        }
1612        Tag::Float64 => {
1613            let f = f64::from_bits(u64::from_le_bytes(read_byte_array(data)));
1614            Datum::Float64(OrderedFloat::from(f))
1615        }
1616        Tag::Date => Datum::Date(read_date(data)),
1617        Tag::Time => Datum::Time(read_time(data)),
1618        Tag::CheapTimestamp => {
1619            let ts = i64::from_le_bytes(read_byte_array(data));
1620            let secs = ts.div_euclid(1_000_000_000);
1621            let nsecs: u32 = ts.rem_euclid(1_000_000_000).try_into().unwrap();
1622            let ndt = DateTime::from_timestamp(secs, nsecs)
1623                .expect("We only write round-trippable timestamps")
1624                .naive_utc();
1625            Datum::Timestamp(
1626                CheckedTimestamp::from_timestamplike(ndt).expect("unexpected timestamp"),
1627            )
1628        }
1629        Tag::CheapTimestampTz => {
1630            let ts = i64::from_le_bytes(read_byte_array(data));
1631            let secs = ts.div_euclid(1_000_000_000);
1632            let nsecs: u32 = ts.rem_euclid(1_000_000_000).try_into().unwrap();
1633            let dt = DateTime::from_timestamp(secs, nsecs)
1634                .expect("We only write round-trippable timestamps");
1635            Datum::TimestampTz(
1636                CheckedTimestamp::from_timestamplike(dt).expect("unexpected timestamp"),
1637            )
1638        }
1639        Tag::Timestamp => {
1640            let date = read_naive_date(data);
1641            let time = read_time(data);
1642            Datum::Timestamp(
1643                CheckedTimestamp::from_timestamplike(date.and_time(time))
1644                    .expect("unexpected timestamp"),
1645            )
1646        }
1647        Tag::TimestampTz => {
1648            let date = read_naive_date(data);
1649            let time = read_time(data);
1650            Datum::TimestampTz(
1651                CheckedTimestamp::from_timestamplike(DateTime::from_naive_utc_and_offset(
1652                    date.and_time(time),
1653                    Utc,
1654                ))
1655                .expect("unexpected timestamptz"),
1656            )
1657        }
1658        Tag::Interval => {
1659            let months = i32::from_le_bytes(read_byte_array(data));
1660            let days = i32::from_le_bytes(read_byte_array(data));
1661            let micros = i64::from_le_bytes(read_byte_array(data));
1662            Datum::Interval(Interval {
1663                months,
1664                days,
1665                micros,
1666            })
1667        }
1668        Tag::BytesTiny | Tag::BytesShort | Tag::BytesLong | Tag::BytesHuge => {
1669            Datum::Bytes(read_lengthed_bytes(data, tag, Tag::BytesTiny))
1670        }
1671        Tag::StringTiny | Tag::StringShort | Tag::StringLong | Tag::StringHuge => {
1672            // SAFETY: the bytes were written from a `str` under a `String` tag.
1673            Datum::String(str::from_utf8_unchecked(read_lengthed_bytes(
1674                data,
1675                tag,
1676                Tag::StringTiny,
1677            )))
1678        }
1679        Tag::ListTiny | Tag::ListShort | Tag::ListLong | Tag::ListHuge => Datum::List(
1680            DatumList::new(read_lengthed_bytes(data, tag, Tag::ListTiny)),
1681        ),
1682        Tag::Uuid => Datum::Uuid(Uuid::from_bytes(read_byte_array(data))),
1683        Tag::Array => {
1684            // See the comment in `Row::push_array` for details on the encoding
1685            // of arrays.
1686            let ndims = read_byte(data);
1687            let dims_size = usize::from(ndims) * size_of::<u64>() * 2;
1688            let (dims, next) = data.split_at(dims_size);
1689            *data = next;
1690            let bytes = read_untagged_bytes(data);
1691            Datum::Array(Array {
1692                dims: ArrayDimensions { data: dims },
1693                elements: DatumList::new(bytes),
1694            })
1695        }
1696        Tag::Dict => {
1697            let bytes = read_untagged_bytes(data);
1698            Datum::Map(DatumMap::new(bytes))
1699        }
1700        Tag::JsonNull => Datum::JsonNull,
1701        Tag::Dummy => Datum::Dummy,
1702        Tag::Numeric => {
1703            let digits = read_byte(data).into();
1704            let exponent = i8::reinterpret_cast(read_byte(data));
1705            let bits = read_byte(data);
1706
1707            let lsu_u16_len = Numeric::digits_to_lsu_elements_len(digits);
1708            let lsu_u8_len = lsu_u16_len * 2;
1709            let (lsu_u8, next) = data.split_at(lsu_u8_len);
1710            *data = next;
1711
1712            // TODO: if we refactor the decimal library to accept the owned
1713            // array as a parameter to `from_raw_parts` below, we could likely
1714            // avoid a copy because it is exactly the value we want
1715            let mut lsu = [0; numeric::NUMERIC_DATUM_WIDTH_USIZE];
1716            for (i, c) in lsu_u8.chunks(2).enumerate() {
1717                lsu[i] = u16::from_le_bytes(c.try_into().unwrap());
1718            }
1719
1720            let d = Numeric::from_raw_parts(digits, exponent.into(), bits, lsu);
1721            Datum::from(d)
1722        }
1723        Tag::MzTimestamp => {
1724            let t = Timestamp::decode(read_byte_array(data));
1725            Datum::MzTimestamp(t)
1726        }
1727        Tag::Range => {
1728            // See notes on `push_range_with` for details about encoding.
1729            let flag_byte = read_byte(data);
1730            let flags = range::InternalFlags::from_bits(flag_byte)
1731                .expect("range flags must be encoded validly");
1732
1733            if flags.contains(range::InternalFlags::EMPTY) {
1734                assert!(
1735                    flags == range::InternalFlags::EMPTY,
1736                    "empty ranges contain only RANGE_EMPTY flag"
1737                );
1738
1739                return Datum::Range(Range { inner: None });
1740            }
1741
1742            let lower_bound = if flags.contains(range::InternalFlags::LB_INFINITE) {
1743                None
1744            } else {
1745                Some(DatumNested::extract(data))
1746            };
1747
1748            let lower = RangeBound {
1749                inclusive: flags.contains(range::InternalFlags::LB_INCLUSIVE),
1750                bound: lower_bound,
1751            };
1752
1753            let upper_bound = if flags.contains(range::InternalFlags::UB_INFINITE) {
1754                None
1755            } else {
1756                Some(DatumNested::extract(data))
1757            };
1758
1759            let upper = RangeBound {
1760                inclusive: flags.contains(range::InternalFlags::UB_INCLUSIVE),
1761                bound: upper_bound,
1762            };
1763
1764            Datum::Range(Range {
1765                inner: Some(RangeInner { lower, upper }),
1766            })
1767        }
1768        Tag::MzAclItem => {
1769            const N: usize = MzAclItem::binary_size();
1770            let mz_acl_item =
1771                MzAclItem::decode_binary(&read_byte_array::<N>(data)).expect("invalid mz_aclitem");
1772            Datum::MzAclItem(mz_acl_item)
1773        }
1774        Tag::AclItem => {
1775            const N: usize = AclItem::binary_size();
1776            let acl_item =
1777                AclItem::decode_binary(&read_byte_array::<N>(data)).expect("invalid aclitem");
1778            Datum::AclItem(acl_item)
1779        }
1780    }
1781}
1782
1783// --------------------------------------------------------------------------------
1784// writing data
1785
1786fn push_untagged_bytes<D>(data: &mut D, bytes: &[u8])
1787where
1788    D: Vector<u8>,
1789{
1790    let len = u64::cast_from(bytes.len());
1791    data.extend_from_slice(&len.to_le_bytes());
1792    data.extend_from_slice(bytes);
1793}
1794
1795fn push_lengthed_bytes<D>(data: &mut D, bytes: &[u8], tag: Tag)
1796where
1797    D: Vector<u8>,
1798{
1799    match tag {
1800        Tag::BytesTiny | Tag::StringTiny | Tag::ListTiny => {
1801            let len = bytes.len().to_le_bytes();
1802            data.push(len[0]);
1803        }
1804        Tag::BytesShort | Tag::StringShort | Tag::ListShort => {
1805            let len = bytes.len().to_le_bytes();
1806            data.extend_from_slice(&len[0..2]);
1807        }
1808        Tag::BytesLong | Tag::StringLong | Tag::ListLong => {
1809            let len = bytes.len().to_le_bytes();
1810            data.extend_from_slice(&len[0..4]);
1811        }
1812        Tag::BytesHuge | Tag::StringHuge | Tag::ListHuge => {
1813            let len = bytes.len().to_le_bytes();
1814            data.extend_from_slice(&len);
1815        }
1816        _ => unreachable!(),
1817    }
1818    data.extend_from_slice(bytes);
1819}
1820
1821pub(super) fn date_to_array(date: Date) -> [u8; size_of::<i32>()] {
1822    i32::to_le_bytes(date.pg_epoch_days())
1823}
1824
1825fn push_date<D>(data: &mut D, date: Date)
1826where
1827    D: Vector<u8>,
1828{
1829    data.extend_from_slice(&date_to_array(date));
1830}
1831
1832pub(super) fn naive_date_to_arrays(
1833    date: NaiveDate,
1834) -> ([u8; size_of::<i32>()], [u8; size_of::<u32>()]) {
1835    (
1836        i32::to_le_bytes(date.year()),
1837        u32::to_le_bytes(date.ordinal()),
1838    )
1839}
1840
1841fn push_naive_date<D>(data: &mut D, date: NaiveDate)
1842where
1843    D: Vector<u8>,
1844{
1845    let (ds1, ds2) = naive_date_to_arrays(date);
1846    data.extend_from_slice(&ds1);
1847    data.extend_from_slice(&ds2);
1848}
1849
1850pub(super) fn time_to_arrays(time: NaiveTime) -> ([u8; size_of::<u32>()], [u8; size_of::<u32>()]) {
1851    (
1852        u32::to_le_bytes(time.num_seconds_from_midnight()),
1853        u32::to_le_bytes(time.nanosecond()),
1854    )
1855}
1856
1857fn push_time<D>(data: &mut D, time: NaiveTime)
1858where
1859    D: Vector<u8>,
1860{
1861    let (ts1, ts2) = time_to_arrays(time);
1862    data.extend_from_slice(&ts1);
1863    data.extend_from_slice(&ts2);
1864}
1865
1866/// Returns an i64 representing a `NaiveDateTime`, if
1867/// said i64 can be round-tripped back to a `NaiveDateTime`.
1868///
1869/// The only exotic NDTs for which this can't happen are those that
1870/// are hundreds of years in the future or past, or those that
1871/// represent a leap second. (Note that Materialize does not support
1872/// leap seconds, but this module does).
1873// This function is inspired by `NaiveDateTime::timestamp_nanos`,
1874// with extra checking.
1875fn checked_timestamp_nanos(dt: NaiveDateTime) -> Option<i64> {
1876    let subsec_nanos = dt.and_utc().timestamp_subsec_nanos();
1877    if subsec_nanos >= 1_000_000_000 {
1878        return None;
1879    }
1880    let as_ns = dt.and_utc().timestamp().checked_mul(1_000_000_000)?;
1881    as_ns.checked_add(i64::from(subsec_nanos))
1882}
1883
1884// This function is extremely hot, so
1885// we just use `as` to avoid the overhead of
1886// `try_into` followed by `unwrap`.
1887// `leading_ones` and `leading_zeros`
1888// can never return values greater than 64, so the conversion is safe.
1889#[inline(always)]
1890#[allow(clippy::as_conversions)]
1891fn min_bytes_signed<T>(i: T) -> u8
1892where
1893    T: Into<i64>,
1894{
1895    let i: i64 = i.into();
1896
1897    // To fit in n bytes, we require that
1898    // everything but the leading sign bits fits in n*8
1899    // bits.
1900    let n_sign_bits = if i.is_negative() {
1901        i.leading_ones() as u8
1902    } else {
1903        i.leading_zeros() as u8
1904    };
1905
1906    (64 - n_sign_bits + 7) / 8
1907}
1908
1909// In principle we could just use `min_bytes_signed`, rather than
1910// having a separate function here, as long as we made that one take
1911// `T: Into<i128>` instead of 64. But LLVM doesn't seem smart enough
1912// to realize that that function is the same as the current version,
1913// and generates worse code.
1914//
1915// Justification for `as` is the same as in `min_bytes_signed`.
1916#[inline(always)]
1917#[allow(clippy::as_conversions)]
1918fn min_bytes_unsigned<T>(i: T) -> u8
1919where
1920    T: Into<u64>,
1921{
1922    let i: u64 = i.into();
1923
1924    let n_sign_bits = i.leading_zeros() as u8;
1925
1926    (64 - n_sign_bits + 7) / 8
1927}
1928
1929const TINY: usize = 1 << 8;
1930const SHORT: usize = 1 << 16;
1931const LONG: usize = 1 << 32;
1932
1933fn push_datum<D>(data: &mut D, datum: Datum)
1934where
1935    D: Vector<u8>,
1936{
1937    match datum {
1938        Datum::Null => data.push(Tag::Null.into()),
1939        Datum::False => data.push(Tag::False.into()),
1940        Datum::True => data.push(Tag::True.into()),
1941        Datum::Int16(i) => {
1942            // The family alternates the signs at each width, so the width is two tags apart and
1943            // the sign is the low bit. The clamp folds both signs onto the widest tag, where the
1944            // payload carries its own sign. See `read_signed_varint`, which takes this apart.
1945            const WIDEST_DELTA: u8 = Tag::Int16.byte() - Tag::NonNegativeInt16_0.byte();
1946            let mbs = min_bytes_signed(i);
1947            let delta = ((mbs << 1) + u8::from(i.is_negative())).min(WIDEST_DELTA);
1948            let tag = u8::from(Tag::NonNegativeInt16_0) + delta;
1949
1950            data.push(tag);
1951            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbs)]);
1952        }
1953        Datum::Int32(i) => {
1954            // The family alternates the signs at each width, so the width is two tags apart and
1955            // the sign is the low bit. The clamp folds both signs onto the widest tag, where the
1956            // payload carries its own sign. See `read_signed_varint`, which takes this apart.
1957            const WIDEST_DELTA: u8 = Tag::Int32.byte() - Tag::NonNegativeInt32_0.byte();
1958            let mbs = min_bytes_signed(i);
1959            let delta = ((mbs << 1) + u8::from(i.is_negative())).min(WIDEST_DELTA);
1960            let tag = u8::from(Tag::NonNegativeInt32_0) + delta;
1961
1962            data.push(tag);
1963            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbs)]);
1964        }
1965        Datum::Int64(i) => {
1966            // The family alternates the signs at each width, so the width is two tags apart and
1967            // the sign is the low bit. The clamp folds both signs onto the widest tag, where the
1968            // payload carries its own sign. See `read_signed_varint`, which takes this apart.
1969            const WIDEST_DELTA: u8 = Tag::Int64.byte() - Tag::NonNegativeInt64_0.byte();
1970            let mbs = min_bytes_signed(i);
1971            let delta = ((mbs << 1) + u8::from(i.is_negative())).min(WIDEST_DELTA);
1972            let tag = u8::from(Tag::NonNegativeInt64_0) + delta;
1973
1974            data.push(tag);
1975            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbs)]);
1976        }
1977        Datum::UInt8(i) => {
1978            let mbu = min_bytes_unsigned(i);
1979            let tag = u8::from(Tag::UInt8_0) + mbu;
1980            data.push(tag);
1981            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbu)]);
1982        }
1983        Datum::UInt16(i) => {
1984            let mbu = min_bytes_unsigned(i);
1985            let tag = u8::from(Tag::UInt16_0) + mbu;
1986            data.push(tag);
1987            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbu)]);
1988        }
1989        Datum::UInt32(i) => {
1990            let mbu = min_bytes_unsigned(i);
1991            let tag = u8::from(Tag::UInt32_0) + mbu;
1992            data.push(tag);
1993            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbu)]);
1994        }
1995        Datum::UInt64(i) => {
1996            let mbu = min_bytes_unsigned(i);
1997            let tag = u8::from(Tag::UInt64_0) + mbu;
1998            data.push(tag);
1999            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbu)]);
2000        }
2001        Datum::Float32(f) => {
2002            data.push(Tag::Float32.into());
2003            data.extend_from_slice(&f.to_bits().to_le_bytes());
2004        }
2005        Datum::Float64(f) => {
2006            data.push(Tag::Float64.into());
2007            data.extend_from_slice(&f.to_bits().to_le_bytes());
2008        }
2009        Datum::Date(d) => {
2010            data.push(Tag::Date.into());
2011            push_date(data, d);
2012        }
2013        Datum::Time(t) => {
2014            data.push(Tag::Time.into());
2015            push_time(data, t);
2016        }
2017        Datum::Timestamp(t) => {
2018            let datetime = t.to_naive();
2019            if let Some(nanos) = checked_timestamp_nanos(datetime) {
2020                data.push(Tag::CheapTimestamp.into());
2021                data.extend_from_slice(&nanos.to_le_bytes());
2022            } else {
2023                data.push(Tag::Timestamp.into());
2024                push_naive_date(data, datetime.date());
2025                push_time(data, datetime.time());
2026            }
2027        }
2028        Datum::TimestampTz(t) => {
2029            let datetime = t.to_naive();
2030            if let Some(nanos) = checked_timestamp_nanos(datetime) {
2031                data.push(Tag::CheapTimestampTz.into());
2032                data.extend_from_slice(&nanos.to_le_bytes());
2033            } else {
2034                data.push(Tag::TimestampTz.into());
2035                push_naive_date(data, datetime.date());
2036                push_time(data, datetime.time());
2037            }
2038        }
2039        Datum::Interval(i) => {
2040            data.push(Tag::Interval.into());
2041            data.extend_from_slice(&i.months.to_le_bytes());
2042            data.extend_from_slice(&i.days.to_le_bytes());
2043            data.extend_from_slice(&i.micros.to_le_bytes());
2044        }
2045        Datum::Bytes(bytes) => {
2046            let tag = match bytes.len() {
2047                0..TINY => Tag::BytesTiny,
2048                TINY..SHORT => Tag::BytesShort,
2049                SHORT..LONG => Tag::BytesLong,
2050                _ => Tag::BytesHuge,
2051            };
2052            data.push(tag.into());
2053            push_lengthed_bytes(data, bytes, tag);
2054        }
2055        Datum::String(string) => {
2056            let tag = match string.len() {
2057                0..TINY => Tag::StringTiny,
2058                TINY..SHORT => Tag::StringShort,
2059                SHORT..LONG => Tag::StringLong,
2060                _ => Tag::StringHuge,
2061            };
2062            data.push(tag.into());
2063            push_lengthed_bytes(data, string.as_bytes(), tag);
2064        }
2065        Datum::List(list) => {
2066            let tag = match list.data.len() {
2067                0..TINY => Tag::ListTiny,
2068                TINY..SHORT => Tag::ListShort,
2069                SHORT..LONG => Tag::ListLong,
2070                _ => Tag::ListHuge,
2071            };
2072            data.push(tag.into());
2073            push_lengthed_bytes(data, list.data, tag);
2074        }
2075        Datum::Uuid(u) => {
2076            data.push(Tag::Uuid.into());
2077            data.extend_from_slice(u.as_bytes());
2078        }
2079        Datum::Array(array) => {
2080            // See the comment in `Row::push_array` for details on the encoding
2081            // of arrays.
2082            data.push(Tag::Array.into());
2083            data.push(array.dims.ndims());
2084            data.extend_from_slice(array.dims.data);
2085            push_untagged_bytes(data, array.elements.data);
2086        }
2087        Datum::Map(dict) => {
2088            data.push(Tag::Dict.into());
2089            push_untagged_bytes(data, dict.data);
2090        }
2091        Datum::JsonNull => data.push(Tag::JsonNull.into()),
2092        Datum::MzTimestamp(t) => {
2093            data.push(Tag::MzTimestamp.into());
2094            data.extend_from_slice(&t.encode());
2095        }
2096        Datum::Dummy => data.push(Tag::Dummy.into()),
2097        Datum::Numeric(mut n) => {
2098            // Pseudo-canonical representation of decimal values with
2099            // insignificant zeroes trimmed. This compresses the number further
2100            // than `Numeric::trim` by removing all zeroes, and not only those in
2101            // the fractional component.
2102            numeric::cx_datum().reduce(&mut n.0);
2103            let (digits, exponent, bits, lsu) = n.0.to_raw_parts();
2104            data.push(Tag::Numeric.into());
2105            data.push(u8::try_from(digits).expect("digits to fit within u8; should not exceed 39"));
2106            data.push(
2107                i8::try_from(exponent)
2108                    .expect("exponent to fit within i8; should not exceed +/- 39")
2109                    .to_le_bytes()[0],
2110            );
2111            data.push(bits);
2112
2113            let lsu = &lsu[..Numeric::digits_to_lsu_elements_len(digits)];
2114
2115            // Little endian machines can take the lsu directly from u16 to u8.
2116            if cfg!(target_endian = "little") {
2117                // SAFETY: `lsu` (returned by `coefficient_units()`) is a `&[u16]`, so
2118                // each element can safely be transmuted into two `u8`s.
2119                let (prefix, lsu_bytes, suffix) = unsafe { lsu.align_to::<u8>() };
2120                // The `u8` aligned version of the `lsu` should have twice as many
2121                // elements as we expect for the `u16` version.
2122                soft_assert_no_log!(
2123                    lsu_bytes.len() == Numeric::digits_to_lsu_elements_len(digits) * 2,
2124                    "u8 version of numeric LSU contained the wrong number of elements; expected {}, but got {}",
2125                    Numeric::digits_to_lsu_elements_len(digits) * 2,
2126                    lsu_bytes.len()
2127                );
2128                // There should be no unaligned elements in the prefix or suffix.
2129                soft_assert_no_log!(prefix.is_empty() && suffix.is_empty());
2130                data.extend_from_slice(lsu_bytes);
2131            } else {
2132                for u in lsu {
2133                    data.extend_from_slice(&u.to_le_bytes());
2134                }
2135            }
2136        }
2137        Datum::Range(range) => {
2138            // See notes on `push_range_with` for details about encoding.
2139            data.push(Tag::Range.into());
2140            data.push(range.internal_flag_bits());
2141
2142            if let Some(RangeInner { lower, upper }) = range.inner {
2143                for bound in [lower.bound, upper.bound] {
2144                    if let Some(bound) = bound {
2145                        match bound.datum() {
2146                            Datum::Null => panic!("cannot push Datum::Null into range"),
2147                            d => push_datum::<D>(data, d),
2148                        }
2149                    }
2150                }
2151            }
2152        }
2153        Datum::MzAclItem(mz_acl_item) => {
2154            data.push(Tag::MzAclItem.into());
2155            data.extend_from_slice(&mz_acl_item.encode_binary());
2156        }
2157        Datum::AclItem(acl_item) => {
2158            data.push(Tag::AclItem.into());
2159            data.extend_from_slice(&acl_item.encode_binary());
2160        }
2161    }
2162}
2163
2164/// Return the number of bytes these Datums would use if packed as a Row.
2165pub fn row_size<'a, I>(a: I) -> usize
2166where
2167    I: IntoIterator<Item = Datum<'a>>,
2168{
2169    // Using datums_size instead of a.data().len() here is safer because it will
2170    // return the size of the datums if they were packed into a Row. Although
2171    // a.data().len() happens to give the correct answer (and is faster), data()
2172    // is documented as for debugging only.
2173    let sz = datums_size::<_, _>(a);
2174    let size_of_row = std::mem::size_of::<Row>();
2175    // The Row struct attempts to inline data until it can't fit in the
2176    // preallocated size. Otherwise it spills to heap, and uses the Row to point
2177    // to that.
2178    if sz > Row::SIZE {
2179        sz + size_of_row
2180    } else {
2181        size_of_row
2182    }
2183}
2184
2185/// Number of bytes required by the datum.
2186/// This is used to optimistically pre-allocate buffers for packing rows.
2187pub fn datum_size(datum: &Datum) -> usize {
2188    match datum {
2189        Datum::Null => 1,
2190        Datum::False => 1,
2191        Datum::True => 1,
2192        Datum::Int16(i) => 1 + usize::from(min_bytes_signed(*i)),
2193        Datum::Int32(i) => 1 + usize::from(min_bytes_signed(*i)),
2194        Datum::Int64(i) => 1 + usize::from(min_bytes_signed(*i)),
2195        Datum::UInt8(i) => 1 + usize::from(min_bytes_unsigned(*i)),
2196        Datum::UInt16(i) => 1 + usize::from(min_bytes_unsigned(*i)),
2197        Datum::UInt32(i) => 1 + usize::from(min_bytes_unsigned(*i)),
2198        Datum::UInt64(i) => 1 + usize::from(min_bytes_unsigned(*i)),
2199        Datum::Float32(_) => 1 + size_of::<f32>(),
2200        Datum::Float64(_) => 1 + size_of::<f64>(),
2201        Datum::Date(_) => 1 + size_of::<i32>(),
2202        Datum::Time(_) => 1 + 8,
2203        Datum::Timestamp(t) => {
2204            1 + if checked_timestamp_nanos(t.to_naive()).is_some() {
2205                8
2206            } else {
2207                16
2208            }
2209        }
2210        Datum::TimestampTz(t) => {
2211            1 + if checked_timestamp_nanos(t.naive_utc()).is_some() {
2212                8
2213            } else {
2214                16
2215            }
2216        }
2217        Datum::Interval(_) => 1 + size_of::<i32>() + size_of::<i32>() + size_of::<i64>(),
2218        Datum::Bytes(bytes) => {
2219            // We use a variable length representation of slice length.
2220            let bytes_for_length = match bytes.len() {
2221                0..TINY => 1,
2222                TINY..SHORT => 2,
2223                SHORT..LONG => 4,
2224                _ => 8,
2225            };
2226            1 + bytes_for_length + bytes.len()
2227        }
2228        Datum::String(string) => {
2229            // We use a variable length representation of slice length.
2230            let bytes_for_length = match string.len() {
2231                0..TINY => 1,
2232                TINY..SHORT => 2,
2233                SHORT..LONG => 4,
2234                _ => 8,
2235            };
2236            1 + bytes_for_length + string.len()
2237        }
2238        Datum::Uuid(_) => 1 + size_of::<uuid::Bytes>(),
2239        Datum::Array(array) => {
2240            1 + size_of::<u8>()
2241                + array.dims.data.len()
2242                + size_of::<u64>()
2243                + array.elements.data.len()
2244        }
2245        Datum::List(list) => 1 + size_of::<u64>() + list.data.len(),
2246        Datum::Map(dict) => 1 + size_of::<u64>() + dict.data.len(),
2247        Datum::JsonNull => 1,
2248        Datum::MzTimestamp(_) => 1 + size_of::<Timestamp>(),
2249        Datum::Dummy => 1,
2250        Datum::Numeric(d) => {
2251            let mut d = d.0.clone();
2252            // Values must be reduced to determine appropriate number of
2253            // coefficient units.
2254            numeric::cx_datum().reduce(&mut d);
2255            // 4 = 1 bit each for tag, digits, exponent, bits
2256            4 + (d.coefficient_units().len() * 2)
2257        }
2258        Datum::Range(Range { inner }) => {
2259            // Tag + flags
2260            2 + match inner {
2261                None => 0,
2262                Some(RangeInner { lower, upper }) => [lower.bound, upper.bound]
2263                    .iter()
2264                    .map(|bound| match bound {
2265                        None => 0,
2266                        Some(bound) => bound.val.len(),
2267                    })
2268                    .sum(),
2269            }
2270        }
2271        Datum::MzAclItem(_) => 1 + MzAclItem::binary_size(),
2272        Datum::AclItem(_) => 1 + AclItem::binary_size(),
2273    }
2274}
2275
2276/// Number of bytes required by a sequence of datums.
2277///
2278/// This method can be used to right-size the allocation for a `Row`
2279/// before calling [`RowPacker::extend`].
2280pub fn datums_size<'a, I, D>(iter: I) -> usize
2281where
2282    I: IntoIterator<Item = D>,
2283    D: Borrow<Datum<'a>>,
2284{
2285    iter.into_iter().map(|d| datum_size(d.borrow())).sum()
2286}
2287
2288/// Number of bytes required by a list of datums. This computes the size that would be required if
2289/// the given datums were packed into a list.
2290///
2291/// This is used to optimistically pre-allocate buffers for packing rows.
2292pub fn datum_list_size<'a, I, D>(iter: I) -> usize
2293where
2294    I: IntoIterator<Item = D>,
2295    D: Borrow<Datum<'a>>,
2296{
2297    1 + size_of::<u64>() + datums_size(iter)
2298}
2299
2300impl RowPacker<'_> {
2301    /// Constructs a row packer that will pack additional datums into the
2302    /// provided row.
2303    ///
2304    /// This function is intentionally somewhat inconvenient to call. You
2305    /// usually want to call [`Row::packer`] instead to start packing from
2306    /// scratch.
2307    pub fn for_existing_row(row: &mut Row) -> RowPacker<'_> {
2308        RowPacker { row }
2309    }
2310
2311    /// Extend an existing `Row` with a `Datum`.
2312    #[inline]
2313    pub fn push<'a, D>(&mut self, datum: D)
2314    where
2315        D: Borrow<Datum<'a>>,
2316    {
2317        push_datum(&mut self.row.data, *datum.borrow());
2318    }
2319
2320    /// Extend an existing `Row` with additional `Datum`s.
2321    #[inline]
2322    pub fn extend<'a, I, D>(&mut self, iter: I)
2323    where
2324        I: IntoIterator<Item = D>,
2325        D: Borrow<Datum<'a>>,
2326    {
2327        for datum in iter {
2328            push_datum(&mut self.row.data, *datum.borrow())
2329        }
2330    }
2331
2332    /// Extend an existing `Row` with additional `Datum`s.
2333    ///
2334    /// In the case the iterator produces an error, the pushing of
2335    /// datums in terminated and the error returned. The `Row` will
2336    /// be incomplete, but it will be safe to read datums from it.
2337    #[inline]
2338    pub fn try_extend<'a, I, E, D>(&mut self, iter: I) -> Result<(), E>
2339    where
2340        I: IntoIterator<Item = Result<D, E>>,
2341        D: Borrow<Datum<'a>>,
2342    {
2343        for datum in iter {
2344            push_datum(&mut self.row.data, *datum?.borrow());
2345        }
2346        Ok(())
2347    }
2348
2349    /// Appends the datums of an entire `Row`.
2350    pub fn extend_by_row(&mut self, row: &Row) {
2351        self.row.data.extend_from_slice(row.data.as_slice());
2352    }
2353
2354    /// Appends the datums of an entire `Row`.
2355    pub fn extend_by_row_ref(&mut self, row: &RowRef) {
2356        self.row.data.extend_from_slice(row.data());
2357    }
2358
2359    /// Appends the slice of data representing an entire `Row`. The data is not validated.
2360    ///
2361    /// # Safety
2362    ///
2363    /// The requirements from [`Row::from_bytes_unchecked`] apply here, too:
2364    /// This method relies on `data` being an appropriate row encoding, and can
2365    /// result in unsafety if this is not the case.
2366    #[inline]
2367    pub unsafe fn extend_by_slice_unchecked(&mut self, data: &[u8]) {
2368        self.row.data.extend_from_slice(data)
2369    }
2370
2371    /// Pushes a [`DatumList`] that is built from a closure.
2372    ///
2373    /// The supplied closure will be invoked once with a `Row` that can be used
2374    /// to populate the list. It is valid to call any method on the
2375    /// [`RowPacker`] except for [`RowPacker::clear`], [`RowPacker::truncate`],
2376    /// or [`RowPacker::truncate_datums`].
2377    ///
2378    /// Returns the value returned by the closure, if any.
2379    ///
2380    /// ```
2381    /// # use mz_repr::{Row, Datum};
2382    /// let mut row = Row::default();
2383    /// row.packer().push_list_with(|row| {
2384    ///     row.push(Datum::String("age"));
2385    ///     row.push(Datum::Int64(42));
2386    /// });
2387    /// assert_eq!(
2388    ///     row.unpack_first().unwrap_list().iter().collect::<Vec<_>>(),
2389    ///     vec![Datum::String("age"), Datum::Int64(42)],
2390    /// );
2391    /// ```
2392    #[inline]
2393    pub fn push_list_with<F, R>(&mut self, f: F) -> R
2394    where
2395        F: FnOnce(&mut RowPacker) -> R,
2396    {
2397        // First, assume that the list will fit in 255 bytes, and thus the length will fit in
2398        // 1 byte. If not, we'll fix it up later.
2399        let start = self.row.data.len();
2400        self.row.data.push(Tag::ListTiny.into());
2401        // Write a dummy len, will fix it up later.
2402        self.row.data.push(0);
2403
2404        let out = f(self);
2405
2406        // The `- 1 - 1` is for the tag and the len.
2407        let len = self.row.data.len() - start - 1 - 1;
2408        // We now know the real len.
2409        if len < TINY {
2410            // If the len fits in 1 byte, we just need to fix up the len.
2411            self.row.data[start + 1] = len.to_le_bytes()[0];
2412        } else {
2413            // Note: We move this code path into its own function, so that the common case can be
2414            // inlined.
2415            long_list(&mut self.row.data, start, len);
2416        }
2417
2418        /// 1. Fix up the tag.
2419        /// 2. Move the actual data a bit (for which we also need to make room at the end).
2420        /// 3. Fix up the len.
2421        /// `data`: The row's backing data.
2422        /// `start`: where `push_list_with` started writing in `data`.
2423        /// `len`: the length of the data, excluding the tag and the length.
2424        #[cold]
2425        fn long_list(data: &mut CompactBytes, start: usize, len: usize) {
2426            // `len_len`: the length of the length. (Possible values are: 2, 4, 8. 1 is handled
2427            // elsewhere.) The other parameters are the same as for `long_list`.
2428            let long_list_inner = |data: &mut CompactBytes, len_len| {
2429                // We'll need memory for the new, bigger length, so make the `CompactBytes` bigger.
2430                // The `- 1` is because the old length was 1 byte.
2431                const ZEROS: [u8; 8] = [0; 8];
2432                data.extend_from_slice(&ZEROS[0..len_len - 1]);
2433                // Move the data to the end of the `CompactBytes`, to make space for the new length.
2434                // Originally, it started after the 1-byte tag and the 1-byte length, now it will
2435                // start after the 1-byte tag and the len_len-byte length.
2436                //
2437                // Note that this is the only operation in `long_list` whose cost is proportional
2438                // to `len`. Since `len` is at least 256 here, the other operations' cost are
2439                // negligible. `copy_within` is a memmove, which is probably a fair bit faster per
2440                // Datum than a Datum encoding in the `f` closure.
2441                data.copy_within(start + 1 + 1..start + 1 + 1 + len, start + 1 + len_len);
2442                // Write the new length.
2443                data[start + 1..start + 1 + len_len]
2444                    .copy_from_slice(&len.to_le_bytes()[0..len_len]);
2445            };
2446            match len {
2447                0..TINY => {
2448                    unreachable!()
2449                }
2450                TINY..SHORT => {
2451                    data[start] = Tag::ListShort.into();
2452                    long_list_inner(data, 2);
2453                }
2454                SHORT..LONG => {
2455                    data[start] = Tag::ListLong.into();
2456                    long_list_inner(data, 4);
2457                }
2458                _ => {
2459                    data[start] = Tag::ListHuge.into();
2460                    long_list_inner(data, 8);
2461                }
2462            };
2463        }
2464
2465        out
2466    }
2467
2468    /// Pushes a [`DatumMap`] that is built from a closure.
2469    ///
2470    /// The supplied closure will be invoked once with a `Row` that can be used
2471    /// to populate the dict.
2472    ///
2473    /// The closure **must** alternate pushing string keys and arbitrary values,
2474    /// otherwise reading the dict will cause a panic.
2475    ///
2476    /// The closure **must** push keys in ascending order, otherwise equality
2477    /// checks on the resulting `Row` may be wrong and reading the dict IN DEBUG
2478    /// MODE will cause a panic.
2479    ///
2480    /// The closure **must not** call [`RowPacker::clear`],
2481    /// [`RowPacker::truncate`], or [`RowPacker::truncate_datums`].
2482    ///
2483    /// # Example
2484    ///
2485    /// ```
2486    /// # use mz_repr::{Row, Datum};
2487    /// let mut row = Row::default();
2488    /// row.packer().push_dict_with(|row| {
2489    ///
2490    ///     // key
2491    ///     row.push(Datum::String("age"));
2492    ///     // value
2493    ///     row.push(Datum::Int64(42));
2494    ///
2495    ///     // key
2496    ///     row.push(Datum::String("name"));
2497    ///     // value
2498    ///     row.push(Datum::String("bob"));
2499    /// });
2500    /// assert_eq!(
2501    ///     row.unpack_first().unwrap_map().iter().collect::<Vec<_>>(),
2502    ///     vec![("age", Datum::Int64(42)), ("name", Datum::String("bob"))]
2503    /// );
2504    /// ```
2505    pub fn push_dict_with<F, R>(&mut self, f: F) -> R
2506    where
2507        F: FnOnce(&mut RowPacker) -> R,
2508    {
2509        self.row.data.push(Tag::Dict.into());
2510        let start = self.row.data.len();
2511        // write a dummy len, will fix it up later
2512        self.row.data.extend_from_slice(&[0; size_of::<u64>()]);
2513
2514        let res = f(self);
2515
2516        let len = u64::cast_from(self.row.data.len() - start - size_of::<u64>());
2517        // fix up the len
2518        self.row.data[start..start + size_of::<u64>()].copy_from_slice(&len.to_le_bytes());
2519
2520        res
2521    }
2522
2523    /// Like [`RowPacker::push_dict_with`], but accepts a fallible closure.
2524    pub fn try_push_dict_with<F, E>(&mut self, f: F) -> Result<(), E>
2525    where
2526        F: FnOnce(&mut RowPacker) -> Result<(), E>,
2527    {
2528        self.push_dict_with(f)
2529    }
2530
2531    /// Convenience function to construct an array from an iter of `Datum`s.
2532    ///
2533    /// Returns an error if the number of elements in `iter` does not match
2534    /// the cardinality of the array as described by `dims`, or if the
2535    /// number of dimensions exceeds [`MAX_ARRAY_DIMENSIONS`]. If an error
2536    /// occurs, the packer's state will be unchanged.
2537    pub fn try_push_array<'a, I, D>(
2538        &mut self,
2539        dims: &[ArrayDimension],
2540        iter: I,
2541    ) -> Result<(), InvalidArrayError>
2542    where
2543        I: IntoIterator<Item = D>,
2544        D: Borrow<Datum<'a>>,
2545    {
2546        // SAFETY: The function returns the exact number of elements pushed into the array.
2547        unsafe {
2548            self.push_array_with_unchecked(dims, |packer| {
2549                let mut nelements = 0;
2550                for datum in iter {
2551                    packer.push(datum);
2552                    nelements += 1;
2553                }
2554                Ok::<_, InvalidArrayError>(nelements)
2555            })
2556        }
2557    }
2558
2559    /// Like [`RowPacker::try_push_array`], but accepts a fallible iterator of
2560    /// elements.
2561    pub fn try_push_array_fallible<'a, I, D, E>(
2562        &mut self,
2563        dims: &[ArrayDimension],
2564        iter: I,
2565    ) -> Result<Result<(), E>, InvalidArrayError>
2566    where
2567        I: IntoIterator<Item = Result<D, E>>,
2568        D: Borrow<Datum<'a>>,
2569    {
2570        enum Error<E> {
2571            Usage(InvalidArrayError),
2572            Inner(E),
2573        }
2574
2575        impl<E> From<InvalidArrayError> for Error<E> {
2576            fn from(e: InvalidArrayError) -> Self {
2577                Self::Usage(e)
2578            }
2579        }
2580
2581        // SAFETY: The function returns the exact number of elements pushed into the array.
2582        let result = unsafe {
2583            self.push_array_with_unchecked(dims, |packer| {
2584                let mut nelements = 0;
2585                for datum in iter {
2586                    packer.push(datum.map_err(Error::Inner)?);
2587                    nelements += 1;
2588                }
2589                Ok(nelements)
2590            })
2591        };
2592        match result {
2593            Ok(()) => Ok(Ok(())),
2594            Err(Error::Usage(e)) => Err(e),
2595            Err(Error::Inner(e)) => Ok(Err(e)),
2596        }
2597    }
2598
2599    /// Convenience function to construct an array from a function. The function must return the
2600    /// number of elements it pushed into the array. It is undefined behavior if the function returns
2601    /// a number different to the number of elements it pushed.
2602    ///
2603    /// Returns an error if the number of elements pushed by `f` does not match
2604    /// the cardinality of the array as described by `dims`, or if the
2605    /// number of dimensions exceeds [`MAX_ARRAY_DIMENSIONS`], or if `f` errors. If an error
2606    /// occurs, the packer's state will be unchanged.
2607    pub unsafe fn push_array_with_unchecked<F, E>(
2608        &mut self,
2609        dims: &[ArrayDimension],
2610        f: F,
2611    ) -> Result<(), E>
2612    where
2613        F: FnOnce(&mut RowPacker) -> Result<usize, E>,
2614        E: From<InvalidArrayError>,
2615    {
2616        // Arrays are encoded as follows.
2617        //
2618        // u8    ndims
2619        // u64   dim_0 lower bound
2620        // u64   dim_0 length
2621        // ...
2622        // u64   dim_n lower bound
2623        // u64   dim_n length
2624        // u64   element data size in bytes
2625        // u8    element data, where elements are encoded in row-major order
2626
2627        if dims.len() > usize::from(MAX_ARRAY_DIMENSIONS) {
2628            return Err(InvalidArrayError::TooManyDimensions(dims.len()).into());
2629        }
2630
2631        let start = self.row.data.len();
2632        self.row.data.push(Tag::Array.into());
2633
2634        // Write dimension information.
2635        self.row
2636            .data
2637            .push(dims.len().try_into().expect("ndims verified to fit in u8"));
2638        for dim in dims {
2639            self.row
2640                .data
2641                .extend_from_slice(&i64::cast_from(dim.lower_bound).to_le_bytes());
2642            self.row
2643                .data
2644                .extend_from_slice(&u64::cast_from(dim.length).to_le_bytes());
2645        }
2646
2647        // Write elements.
2648        let off = self.row.data.len();
2649        self.row.data.extend_from_slice(&[0; size_of::<u64>()]);
2650        let nelements = match f(self) {
2651            Ok(nelements) => nelements,
2652            Err(e) => {
2653                self.row.data.truncate(start);
2654                return Err(e);
2655            }
2656        };
2657        let len = u64::cast_from(self.row.data.len() - off - size_of::<u64>());
2658        self.row.data[off..off + size_of::<u64>()].copy_from_slice(&len.to_le_bytes());
2659
2660        // Check that the number of elements written matches the dimension
2661        // information.
2662        let cardinality = match dims {
2663            [] => 0,
2664            // Saturate the product: a cardinality that overflows `usize` is
2665            // impossibly large (no array can hold that many elements), so it can
2666            // never equal the actual `nelements` and the check below rejects it as
2667            // `WrongCardinality`. A plain `product()` would panic under overflow
2668            // checks (debug/fuzz) and silently wrap in release — and a wrapped
2669            // value could even spuriously match `nelements`, accepting a corrupt
2670            // array (e.g. dims claiming `[2^32, 2^32]` wrap to 0 elements).
2671            dims => dims
2672                .iter()
2673                .map(|d| d.length)
2674                .fold(1usize, usize::saturating_mul),
2675        };
2676        if nelements != cardinality {
2677            self.row.data.truncate(start);
2678            return Err(InvalidArrayError::WrongCardinality {
2679                actual: nelements,
2680                expected: cardinality,
2681            }
2682            .into());
2683        }
2684
2685        Ok(())
2686    }
2687
2688    /// Pushes an [`Array`] that is built from a closure.
2689    ///
2690    /// __WARNING__: This is fairly "sharp" tool that is easy to get wrong. You
2691    /// should prefer [`RowPacker::try_push_array`] when possible.
2692    ///
2693    /// Returns an error if the number of elements pushed does not match
2694    /// the cardinality of the array as described by `dims`, or if the
2695    /// number of dimensions exceeds [`MAX_ARRAY_DIMENSIONS`]. If an error
2696    /// occurs, the packer's state will be unchanged.
2697    pub fn push_array_with_row_major<F, I>(
2698        &mut self,
2699        dims: I,
2700        f: F,
2701    ) -> Result<(), InvalidArrayError>
2702    where
2703        I: IntoIterator<Item = ArrayDimension>,
2704        F: FnOnce(&mut RowPacker) -> usize,
2705    {
2706        let start = self.row.data.len();
2707        self.row.data.push(Tag::Array.into());
2708
2709        // Write dummy dimension length for now, we'll fix it up.
2710        let dims_start = self.row.data.len();
2711        self.row.data.push(42);
2712
2713        let mut num_dims: u8 = 0;
2714        let mut cardinality: usize = 1;
2715        for dim in dims {
2716            num_dims += 1;
2717            // Saturate: an overflowing cardinality is impossibly large and is
2718            // rejected by the `nelements` check below. See the matching note in
2719            // `push_array_with_unchecked`.
2720            cardinality = cardinality.saturating_mul(dim.length);
2721
2722            self.row
2723                .data
2724                .extend_from_slice(&i64::cast_from(dim.lower_bound).to_le_bytes());
2725            self.row
2726                .data
2727                .extend_from_slice(&u64::cast_from(dim.length).to_le_bytes());
2728        }
2729
2730        if num_dims > MAX_ARRAY_DIMENSIONS {
2731            // Reset the packer state so we don't have invalid data.
2732            self.row.data.truncate(start);
2733            return Err(InvalidArrayError::TooManyDimensions(usize::from(num_dims)));
2734        }
2735        // Fix up our dimension length.
2736        self.row.data[dims_start..dims_start + size_of::<u8>()]
2737            .copy_from_slice(&num_dims.to_le_bytes());
2738
2739        // Write elements.
2740        let off = self.row.data.len();
2741        self.row.data.extend_from_slice(&[0; size_of::<u64>()]);
2742
2743        let nelements = f(self);
2744
2745        let len = u64::cast_from(self.row.data.len() - off - size_of::<u64>());
2746        self.row.data[off..off + size_of::<u64>()].copy_from_slice(&len.to_le_bytes());
2747
2748        // Check that the number of elements written matches the dimension
2749        // information.
2750        let cardinality = match num_dims {
2751            0 => 0,
2752            _ => cardinality,
2753        };
2754        if nelements != cardinality {
2755            self.row.data.truncate(start);
2756            return Err(InvalidArrayError::WrongCardinality {
2757                actual: nelements,
2758                expected: cardinality,
2759            });
2760        }
2761
2762        Ok(())
2763    }
2764
2765    /// Convenience function to push a `DatumList` from an iter of `Datum`s
2766    ///
2767    /// See [`RowPacker::push_dict_with`] if you need to be able to handle errors
2768    pub fn push_list<'a, I, D>(&mut self, iter: I)
2769    where
2770        I: IntoIterator<Item = D>,
2771        D: Borrow<Datum<'a>>,
2772    {
2773        self.push_list_with(|packer| {
2774            for elem in iter {
2775                packer.push(*elem.borrow())
2776            }
2777        });
2778    }
2779
2780    /// Convenience function to push a `DatumMap` from an iter of `(&str, Datum)` pairs
2781    pub fn push_dict<'a, I, D>(&mut self, iter: I)
2782    where
2783        I: IntoIterator<Item = (&'a str, D)>,
2784        D: Borrow<Datum<'a>>,
2785    {
2786        self.push_dict_with(|packer| {
2787            for (k, v) in iter {
2788                packer.push(Datum::String(k));
2789                packer.push(*v.borrow())
2790            }
2791        })
2792    }
2793
2794    /// Pushes a `Datum::Range` derived from the `Range<Datum<'a>`.
2795    ///
2796    /// # Panics
2797    /// - If lower and upper express finite values and they are datums of
2798    ///   different types.
2799    /// - If lower or upper express finite values and are equal to
2800    ///   `Datum::Null`. To handle `Datum::Null` properly, use
2801    ///   [`RangeBound::new`].
2802    ///
2803    /// # Notes
2804    /// - This function canonicalizes the range before pushing it to the row.
2805    /// - Prefer this function over `push_range_with` because of its
2806    ///   canonicaliztion.
2807    /// - Prefer creating [`RangeBound`]s using [`RangeBound::new`], which
2808    ///   handles `Datum::Null` in a SQL-friendly way.
2809    pub fn push_range<'a>(&mut self, mut range: Range<Datum<'a>>) -> Result<(), InvalidRangeError> {
2810        range.canonicalize()?;
2811        match range.inner {
2812            None => {
2813                self.row.data.push(Tag::Range.into());
2814                // Untagged bytes only contains the `RANGE_EMPTY` flag value.
2815                self.row.data.push(range::InternalFlags::EMPTY.bits());
2816                Ok(())
2817            }
2818            Some(inner) => self.push_range_with(
2819                RangeLowerBound {
2820                    inclusive: inner.lower.inclusive,
2821                    bound: inner
2822                        .lower
2823                        .bound
2824                        .map(|value| move |row: &mut RowPacker| Ok(row.push(value))),
2825                },
2826                RangeUpperBound {
2827                    inclusive: inner.upper.inclusive,
2828                    bound: inner
2829                        .upper
2830                        .bound
2831                        .map(|value| move |row: &mut RowPacker| Ok(row.push(value))),
2832                },
2833            ),
2834        }
2835    }
2836
2837    /// Pushes a `DatumRange` built from the specified arguments.
2838    ///
2839    /// # Warning
2840    /// Unlike `push_range`, `push_range_with` _does not_ canonicalize its
2841    /// inputs. Consequentially, this means it's possible to generate ranges
2842    /// that will not reflect the proper ordering and equality.
2843    ///
2844    /// # Panics
2845    /// - If lower or upper expresses a finite value and does not push exactly
2846    ///   one value into the `RowPacker`.
2847    /// - If lower and upper express finite values and they are datums of
2848    ///   different types.
2849    /// - If lower or upper express finite values and push `Datum::Null`.
2850    ///
2851    /// # Notes
2852    /// - Prefer `push_range_with` over this function. This function should be
2853    ///   used only when you are not pushing `Datum`s to the inner row.
2854    /// - Range encoding is `[<flag bytes>,<lower>?,<upper>?]`, where `lower`
2855    ///   and `upper` are optional, contingent on the flag value expressing an
2856    ///   empty range (where neither will be present) or infinite bounds (where
2857    ///   each infinite bound will be absent).
2858    /// - To push an emtpy range, use `push_range` using `Range { inner: None }`.
2859    pub fn push_range_with<L, U, E>(
2860        &mut self,
2861        lower: RangeLowerBound<L>,
2862        upper: RangeUpperBound<U>,
2863    ) -> Result<(), E>
2864    where
2865        L: FnOnce(&mut RowPacker) -> Result<(), E>,
2866        U: FnOnce(&mut RowPacker) -> Result<(), E>,
2867        E: From<InvalidRangeError>,
2868    {
2869        let start = self.row.data.len();
2870        self.row.data.push(Tag::Range.into());
2871
2872        let mut flags = range::InternalFlags::empty();
2873
2874        flags.set(range::InternalFlags::LB_INFINITE, lower.bound.is_none());
2875        flags.set(range::InternalFlags::UB_INFINITE, upper.bound.is_none());
2876        flags.set(range::InternalFlags::LB_INCLUSIVE, lower.inclusive);
2877        flags.set(range::InternalFlags::UB_INCLUSIVE, upper.inclusive);
2878
2879        let mut expected_datums = 0;
2880
2881        self.row.data.push(flags.bits());
2882
2883        let datum_check = self.row.data.len();
2884
2885        if let Some(value) = lower.bound {
2886            let start = self.row.data.len();
2887            value(self)?;
2888            assert!(
2889                start < self.row.data.len(),
2890                "finite values must each push exactly one value; expected 1 but got 0"
2891            );
2892            expected_datums += 1;
2893        }
2894
2895        if let Some(value) = upper.bound {
2896            let start = self.row.data.len();
2897            value(self)?;
2898            assert!(
2899                start < self.row.data.len(),
2900                "finite values must each push exactly one value; expected 1 but got 0"
2901            );
2902            expected_datums += 1;
2903        }
2904
2905        // Validate the invariants that 0, 1, or 2 elements were pushed, none are Null,
2906        // and if two are pushed then the second is not less than the first. Panic in
2907        // some cases and error in others.
2908        let mut actual_datums = 0;
2909        let mut seen = None;
2910        let mut dataz = &self.row.data[datum_check..];
2911        while !dataz.is_empty() {
2912            let d = unsafe { read_datum(&mut dataz) };
2913            // These checks only fail when decoding untrusted/corrupted bytes;
2914            // valid callers always push consistent, non-null bounds. Return an
2915            // error rather than asserting so a crafted proto doesn't panic.
2916            if d == Datum::Null {
2917                self.row.data.truncate(start);
2918                return Err(InvalidRangeError::InvalidRangeData.into());
2919            }
2920
2921            match seen {
2922                None => seen = Some(d),
2923                Some(seen) => {
2924                    let seen_kind = DatumKind::from(seen);
2925                    let d_kind = DatumKind::from(d);
2926                    if seen_kind != d_kind {
2927                        self.row.data.truncate(start);
2928                        return Err(InvalidRangeError::InvalidRangeData.into());
2929                    }
2930
2931                    if seen > d {
2932                        self.row.data.truncate(start);
2933                        return Err(InvalidRangeError::MisorderedRangeBounds.into());
2934                    }
2935                }
2936            }
2937            actual_datums += 1;
2938        }
2939
2940        if actual_datums != expected_datums {
2941            self.row.data.truncate(start);
2942            return Err(InvalidRangeError::InvalidRangeData.into());
2943        }
2944
2945        Ok(())
2946    }
2947
2948    /// Clears the contents of the packer without de-allocating its backing memory.
2949    pub fn clear(&mut self) {
2950        self.row.data.clear();
2951    }
2952
2953    /// Truncates the underlying storage to the specified byte position.
2954    ///
2955    /// # Safety
2956    ///
2957    /// `pos` MUST specify a byte offset that lies on a datum boundary.
2958    /// If `pos` specifies a byte offset that is *within* a datum, the row
2959    /// packer will produce an invalid row, the unpacking of which may
2960    /// trigger undefined behavior!
2961    ///
2962    /// To find the byte offset of a datum boundary, inspect the packer's
2963    /// byte length by calling `packer.data().len()` after pushing the desired
2964    /// number of datums onto the packer.
2965    pub unsafe fn truncate(&mut self, pos: usize) {
2966        self.row.data.truncate(pos)
2967    }
2968
2969    /// Truncates the underlying row to contain at most the first `n` datums.
2970    pub fn truncate_datums(&mut self, n: usize) {
2971        let prev_len = self.row.data.len();
2972        let mut iter = self.row.iter();
2973        for _ in iter.by_ref().take(n) {}
2974        let next_len = iter.data.len();
2975        // SAFETY: iterator offsets always lie on a datum boundary.
2976        unsafe { self.truncate(prev_len - next_len) }
2977    }
2978
2979    /// Returns the total amount of bytes used by the underlying row.
2980    pub fn byte_len(&self) -> usize {
2981        self.row.byte_len()
2982    }
2983}
2984
2985impl<'a> IntoIterator for &'a Row {
2986    type Item = Datum<'a>;
2987    type IntoIter = DatumListIter<'a>;
2988    fn into_iter(self) -> DatumListIter<'a> {
2989        self.iter()
2990    }
2991}
2992
2993impl fmt::Debug for Row {
2994    /// Debug representation using the internal datums
2995    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
2996        f.write_str("Row{")?;
2997        f.debug_list().entries(self.iter()).finish()?;
2998        f.write_str("}")
2999    }
3000}
3001
3002impl fmt::Display for Row {
3003    /// Display representation using the internal datums
3004    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
3005        f.write_str("(")?;
3006        for (i, datum) in self.iter().enumerate() {
3007            if i != 0 {
3008                f.write_str(", ")?;
3009            }
3010            write!(f, "{}", datum)?;
3011        }
3012        f.write_str(")")
3013    }
3014}
3015
3016impl<'a, T> DatumList<'a, T> {
3017    pub fn iter(&self) -> DatumListIter<'a> {
3018        DatumListIter { data: self.data }
3019    }
3020
3021    /// Iterate elements as typed `T` values rather than raw `Datum`s.
3022    ///
3023    /// Each datum is decoded and converted via [`FromDatum`]. Since generic
3024    /// type parameters in `#[sqlfunc]` are erased to `Datum<'a>` before code
3025    /// generation, this is monomorphized to an identity conversion at runtime.
3026    pub fn typed_iter(&self) -> DatumListTypedIter<'a, T>
3027    where
3028        T: FromDatum<'a>,
3029    {
3030        DatumListTypedIter {
3031            inner: self.iter(),
3032            _phantom: PhantomData,
3033        }
3034    }
3035
3036    /// For debugging only
3037    pub fn data(&self) -> &'a [u8] {
3038        self.data
3039    }
3040}
3041
3042impl<T> DatumList<'static, T> {
3043    pub fn empty() -> Self {
3044        DatumList::new(&[])
3045    }
3046}
3047
3048impl<'a> IntoIterator for DatumList<'a> {
3049    type Item = Datum<'a>;
3050    type IntoIter = DatumListIter<'a>;
3051    fn into_iter(self) -> DatumListIter<'a> {
3052        self.iter()
3053    }
3054}
3055
3056impl<'a> Iterator for DatumListIter<'a> {
3057    type Item = Datum<'a>;
3058    fn next(&mut self) -> Option<Self::Item> {
3059        if self.data.is_empty() {
3060            None
3061        } else {
3062            Some(unsafe { read_datum(&mut self.data) })
3063        }
3064    }
3065}
3066
3067impl<'a, T: FromDatum<'a>> Iterator for DatumListTypedIter<'a, T> {
3068    type Item = T;
3069    fn next(&mut self) -> Option<Self::Item> {
3070        self.inner.next().map(T::from_datum)
3071    }
3072}
3073
3074impl<'a, T> DatumMap<'a, T> {
3075    pub fn iter(&self) -> DatumDictIter<'a> {
3076        DatumDictIter {
3077            data: self.data,
3078            prev_key: None,
3079        }
3080    }
3081
3082    /// Iterate entries as `(&str, T)` pairs rather than `(&str, Datum)`.
3083    ///
3084    /// Each value datum is converted via [`FromDatum`]. Since generic type
3085    /// parameters in `#[sqlfunc]` are erased to `Datum<'a>` before code
3086    /// generation, this is monomorphized to an identity conversion at runtime.
3087    pub fn typed_iter(&self) -> DatumDictTypedIter<'a, T>
3088    where
3089        T: FromDatum<'a>,
3090    {
3091        DatumDictTypedIter {
3092            inner: self.iter(),
3093            _phantom: PhantomData,
3094        }
3095    }
3096
3097    /// For debugging only
3098    pub fn data(&self) -> &'a [u8] {
3099        self.data
3100    }
3101}
3102
3103impl<T> DatumMap<'static, T> {
3104    pub fn empty() -> Self {
3105        DatumMap::new(&[])
3106    }
3107}
3108
3109impl<'a, T> Debug for DatumMap<'a, T> {
3110    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3111        f.debug_map().entries(self.iter()).finish()
3112    }
3113}
3114
3115impl<'a> IntoIterator for &'a DatumMap<'a> {
3116    type Item = (&'a str, Datum<'a>);
3117    type IntoIter = DatumDictIter<'a>;
3118    fn into_iter(self) -> DatumDictIter<'a> {
3119        self.iter()
3120    }
3121}
3122
3123impl<'a> Iterator for DatumDictIter<'a> {
3124    type Item = (&'a str, Datum<'a>);
3125    fn next(&mut self) -> Option<Self::Item> {
3126        if self.data.is_empty() {
3127            None
3128        } else {
3129            let key_tag =
3130                Tag::try_from_primitive(read_byte(&mut self.data)).expect("unknown row tag");
3131            assert!(
3132                key_tag == Tag::StringTiny
3133                    || key_tag == Tag::StringShort
3134                    || key_tag == Tag::StringLong
3135                    || key_tag == Tag::StringHuge,
3136                "Dict keys must be strings, got {:?}",
3137                key_tag
3138            );
3139            let bytes = read_lengthed_bytes(&mut self.data, key_tag, Tag::StringTiny);
3140            // SAFETY: the bytes were written from a `str` under a `String` tag.
3141            let key = unsafe { str::from_utf8_unchecked(bytes) };
3142            let val = unsafe { read_datum(&mut self.data) };
3143
3144            // Gate the `prev_key` bookkeeping on the same flag as the assert it feeds, so builds
3145            // with soft assertions off pay nothing for it.
3146            if mz_ore::assert::soft_assertions_enabled() {
3147                if let Some(prev_key) = self.prev_key {
3148                    mz_ore::soft_assert_no_log!(
3149                        prev_key < key,
3150                        "Dict keys must be unique and given in ascending order: {} came before {}",
3151                        prev_key,
3152                        key
3153                    );
3154                }
3155                self.prev_key = Some(key);
3156            }
3157
3158            Some((key, val))
3159        }
3160    }
3161}
3162
3163impl<'a, T: FromDatum<'a>> Iterator for DatumDictTypedIter<'a, T> {
3164    type Item = (&'a str, T);
3165    fn next(&mut self) -> Option<Self::Item> {
3166        self.inner.next().map(|(k, v)| (k, T::from_datum(v)))
3167    }
3168}
3169
3170impl RowArena {
3171    pub fn new() -> Self {
3172        RowArena {
3173            inner: RefCell::new(vec![]),
3174            scratch: RefCell::new(None),
3175            budget: None,
3176            allocated: Cell::new(0),
3177        }
3178    }
3179
3180    /// Creates a `RowArena` that reports itself [`RowArena::over_budget`] once it holds more than
3181    /// `budget` bytes.
3182    ///
3183    /// The budget is advisory to the arena itself: pushes still succeed, because handing back a
3184    /// truncated value would corrupt the datum. It is the caller's job to poll `over_budget` at a
3185    /// point where it can fail, so the bytes an arena actually reaches is `budget` plus whatever the
3186    /// operation in flight at the time added.
3187    ///
3188    /// NOTE: a budget bounds a single ad-hoc evaluation in a shared process (see
3189    /// `mz_adapter::webhook`). It must not be given to an arena that feeds a compute dataflow. A
3190    /// dataflow re-evaluates the same expression against the same input and must return the same
3191    /// result every time. Whether an evaluation is over budget depends on what else the arena has
3192    /// accumulated, and the webhook budget is a runtime dyncfg, so a budgeted dataflow arena would
3193    /// make the result depend on when it ran. Differential then turns a changed-but-not-retracted
3194    /// result into non-accumulating diffs that corrupt the collection. A dataflow that ever needs a
3195    /// budget must fix it for the lifetime of a cluster replica.
3196    pub fn with_budget(budget: usize) -> Self {
3197        RowArena {
3198            budget: Some(budget),
3199            ..RowArena::new()
3200        }
3201    }
3202
3203    /// Bytes this arena currently holds.
3204    pub fn allocated_bytes(&self) -> usize {
3205        self.allocated.get()
3206    }
3207
3208    /// Whether this arena holds more than its budget. Always false without one.
3209    pub fn over_budget(&self) -> bool {
3210        self.budget
3211            .is_some_and(|budget| self.allocated.get() > budget)
3212    }
3213
3214    /// Bytes this arena can still take before it is [`RowArena::over_budget`], or `usize::MAX`
3215    /// without a budget.
3216    ///
3217    /// Intended for an operation that can predict its own size and would rather fail than build a
3218    /// value it is about to be told is too big.
3219    pub fn budget_remaining(&self) -> usize {
3220        match self.budget {
3221            None => usize::MAX,
3222            Some(budget) => budget.saturating_sub(self.allocated.get()),
3223        }
3224    }
3225
3226    /// Creates a `RowArena` with an initial region sized to hold `capacity` bytes, to avoid
3227    /// reallocations as the first datums are created in the arena.
3228    pub fn with_capacity(capacity: usize) -> Self {
3229        let mut inner = Vec::new();
3230        if capacity > 0 {
3231            inner.push(Vec::with_capacity(capacity));
3232        }
3233        RowArena {
3234            inner: RefCell::new(inner),
3235            ..RowArena::new()
3236        }
3237    }
3238
3239    /// Ensures the active region can hold at least `additional` more bytes without allocating a
3240    /// new region. Call this when you expect to push roughly `additional` bytes next.
3241    pub fn reserve(&self, additional: usize) {
3242        if additional == 0 {
3243            return;
3244        }
3245        let mut inner = self.inner.borrow_mut();
3246        match inner.last_mut() {
3247            // The active region is empty, so nothing references it yet and it is safe to grow it
3248            // in place (a reallocation cannot dangle a live reference).
3249            Some(active) if active.is_empty() => {
3250                if active.capacity() < additional {
3251                    active.reserve_exact(additional);
3252                }
3253            }
3254            // The active region holds live data; we cannot grow it without moving those bytes, so
3255            // stage a fresh region. Size it like `push_bytes` does (at least double the current
3256            // region) so a sequence of small `reserve`s still yields at most log-many regions
3257            // rather than many small ones.
3258            Some(active) => {
3259                let new_cap = std::cmp::max(additional, active.capacity().saturating_mul(2));
3260                inner.push(Vec::with_capacity(new_cap));
3261            }
3262            None => inner.push(Vec::with_capacity(additional)),
3263        }
3264    }
3265
3266    /// Copies `bytes` into the arena and returns a reference valid for its lifetime.
3267    ///
3268    /// Accepts anything that derefs to `[u8]` (e.g. `Vec<u8>`, `&[u8]`); the bytes are copied, so
3269    /// the caller's allocation is not retained.
3270    #[allow(clippy::transmute_ptr_to_ptr)]
3271    pub fn push_bytes<'a, B: Deref<Target = [u8]>>(&'a self, bytes: B) -> &'a [u8] {
3272        let bytes: &[u8] = &bytes;
3273        let need = bytes.len();
3274        if need == 0 {
3275            return &[];
3276        }
3277        let mut inner = self.inner.borrow_mut();
3278
3279        // Find or create a region with spare capacity for `need` bytes, never growing a region
3280        // that already holds data (see the type-level comment for why this preserves references).
3281        let has_room = inner
3282            .last()
3283            .map_or(false, |region| region.capacity() - region.len() >= need);
3284        if !has_room {
3285            let last_cap = inner.last().map_or(0, |region| region.capacity());
3286            let new_cap = std::cmp::max(need, last_cap.saturating_mul(2));
3287            inner.push(Vec::with_capacity(new_cap));
3288        }
3289
3290        let region = inner.last_mut().expect("region present");
3291        let start = region.len();
3292        region.extend_from_slice(bytes);
3293        self.allocated.set(self.allocated.get() + need);
3294        let copied = &region[start..];
3295        unsafe {
3296            // This is safe because:
3297            //   * `copied` references bytes inside `region`'s heap buffer, which we just sized to
3298            //     fit without reallocating; that buffer is never resized again while it holds data
3299            //     (we allocate a new region instead), so the reference stays valid.
3300            //   * The buffer lives as long as the arena: regions are only dropped by `clear`/`drop`,
3301            //     both of which take `&mut`/ownership, so no `&'a self`-tied reference can outlive
3302            //     them.
3303            //   * Pushing further regions may reallocate `self.inner`, but that moves only the
3304            //     `Vec<u8>` headers, not the heap buffers they own.
3305            transmute::<&[u8], &'a [u8]>(copied)
3306        }
3307    }
3308
3309    /// Moves `bytes` into the arena and returns a reference valid for its lifetime.
3310    ///
3311    /// Prefer this to [`RowArena::push_bytes`] whenever the bytes are already owned: a value large
3312    /// enough that it would get a region to itself has its allocation adopted as that region, rather
3313    /// than a fresh region being allocated and copied into, which for a large value halves the peak.
3314    /// Smaller values are copied, so the arena keeps bump allocating.
3315    pub fn push_owned_bytes<'a>(&'a self, bytes: Vec<u8>) -> &'a [u8] {
3316        /// Never adopt below this, however empty the arena. `last_cap` alone would let every value
3317        /// on a fresh or small arena look big enough, and then each gets an exactly-sized region
3318        /// with no headroom: one region and one `Vec<u8>` header per value.
3319        const MIN_ADOPT_BYTES: usize = 4 * 1024;
3320
3321        let need = bytes.len();
3322        if need == 0 {
3323            return &[];
3324        }
3325
3326        let mut inner = self.inner.borrow_mut();
3327        // Adopt only when `push_bytes` would have given these bytes a dedicated, `need`-sized region
3328        // anyway, i.e. when `need` exceeds the `last_cap * 2` it would otherwise allocate. There's
3329        // no headroom to lose, so adoption saves a copy for free. Below that we copy, because
3330        // `push_bytes` grows a region *with* headroom that later values reuse. Adopting there would
3331        // defeat the bump allocator: adoption leaves an empty region (capacity 0) on top, so nothing
3332        // would ever grow a region with headroom again.
3333        let last_cap = inner.last().map_or(0, |region| region.capacity());
3334        let adopt = need > std::cmp::max(MIN_ADOPT_BYTES, last_cap.saturating_mul(2));
3335        if !adopt {
3336            drop(inner);
3337            return self.push_bytes(&bytes[..]);
3338        }
3339
3340        // `push_bytes` would allocate a fresh region here and copy into it, so adopt the caller's
3341        // allocation as that region. Sound for the same reasons as `push_bytes`: the reference
3342        // points into a heap buffer the arena now owns for `'a`, and the buffer is never resized
3343        // while it holds data.
3344        //
3345        // Inserted *below* the active region rather than appended, because `push_bytes` sizes a new
3346        // region as twice the last one's capacity: leaving a large adopted buffer on top would make
3347        // the next push allocate twice its size.
3348        self.allocated.set(self.allocated.get() + need);
3349        let idx = inner.len().saturating_sub(1);
3350        inner.insert(idx, bytes);
3351        if inner.len() == 1 {
3352            // There was no active region to insert below, so keep an empty one on top for the same
3353            // reason. `Vec::new` does not allocate.
3354            inner.push(Vec::new());
3355        }
3356        let adopted = &inner[idx][..];
3357        unsafe { transmute::<&[u8], &'a [u8]>(adopted) }
3358    }
3359
3360    /// Moves `string` into the arena and returns a reference valid for its lifetime.
3361    pub fn push_string<'a>(&'a self, string: String) -> &'a str {
3362        let copied = self.push_owned_bytes(string.into_bytes());
3363        unsafe {
3364            // This is safe because we just moved in the bytes of a valid `String`.
3365            std::str::from_utf8_unchecked(copied)
3366        }
3367    }
3368
3369    /// Returns a growable, writeable byte buffer for assembling a value incrementally.
3370    ///
3371    /// Write into it with [`RowArenaBuf::push`], [`RowArenaBuf::extend_from_slice`], or
3372    /// [`std::io::Write`], then call [`RowArenaBuf::finish`] to copy the result into the arena and
3373    /// obtain a reference valid for the arena's lifetime. The backing buffer is a single scratch
3374    /// allocation reused across writers, so this lets a producer that builds bytes piecewise (e.g.
3375    /// decoding a row) avoid managing its own scratch.
3376    ///
3377    /// Nested writers are sound but not free: a writer obtained while another is still live can't
3378    /// reuse the (in-use) scratch, so it allocates its own buffer. Steady-state, non-nested use
3379    /// stays allocation-free.
3380    pub fn writer(&self) -> RowArenaBuf<'_> {
3381        // Take the recycled buffer if one is available, else allocate a fresh one. The cell is
3382        // borrowed only for this `take`, never for the writer's lifetime, so a nested `writer` call
3383        // doesn't double-borrow: it simply finds the slot empty and allocates its own buffer.
3384        let mut buf = self.scratch.borrow_mut().take().unwrap_or_default();
3385        buf.clear();
3386        RowArenaBuf { arena: self, buf }
3387    }
3388
3389    /// Take ownership of `row` for the lifetime of the arena, returning a
3390    /// reference to the first datum in the row.
3391    ///
3392    /// If we had an owned datum type, this method would be much clearer, and
3393    /// would be called `push_owned_datum`.
3394    pub fn push_unary_row<'a>(&'a self, row: Row) -> Datum<'a> {
3395        let copied = self.push_bytes(row.data());
3396        unsafe {
3397            // This is safe because `copied` is a valid encoding of a single datum (we just packed
3398            // it into `row`), backed by the arena for the lifetime `'a`. Copying the bytes also
3399            // sidesteps the `Row`'s inline (`SmallVec`) storage entirely.
3400            let datum = read_datum(&mut &copied[..]);
3401            transmute::<Datum<'_>, Datum<'a>>(datum)
3402        }
3403    }
3404
3405    /// Equivalent to `push_unary_row` but returns a `DatumNested` rather than a
3406    /// `Datum`.
3407    fn push_unary_row_datum_nested<'a>(&'a self, row: Row) -> DatumNested<'a> {
3408        let copied = self.push_bytes(row.data());
3409        unsafe {
3410            // Safe for the same reasons as `push_unary_row`.
3411            let nested = DatumNested::extract(&mut &copied[..]);
3412            transmute::<DatumNested<'_>, DatumNested<'a>>(nested)
3413        }
3414    }
3415
3416    /// Convenience function to make a new `Row` containing a single datum, and
3417    /// take ownership of it for the lifetime of the arena
3418    ///
3419    /// ```
3420    /// # use mz_repr::{RowArena, Datum};
3421    /// let arena = RowArena::new();
3422    /// let datum = arena.make_datum(|packer| {
3423    ///   packer.push_list(&[Datum::String("hello"), Datum::String("world")]);
3424    /// });
3425    /// assert_eq!(datum.unwrap_list().iter().collect::<Vec<_>>(), vec![Datum::String("hello"), Datum::String("world")]);
3426    /// ```
3427    pub fn make_datum<'a, F>(&'a self, f: F) -> Datum<'a>
3428    where
3429        F: FnOnce(&mut RowPacker),
3430    {
3431        let mut row = Row::default();
3432        f(&mut row.packer());
3433        self.push_unary_row(row)
3434    }
3435
3436    /// Convenience function to build a list datum from an iterator of typed
3437    /// elements and return it as a `DatumList<'a, T>`.
3438    ///
3439    /// By accepting an iterator of `T: Borrow<Datum>` instead of a raw
3440    /// `RowPacker` closure, this guarantees that only elements of type `T`
3441    /// are pushed.
3442    pub fn make_datum_list<'a, T: std::borrow::Borrow<Datum<'a>>>(
3443        &'a self,
3444        iter: impl IntoIterator<Item = T>,
3445    ) -> DatumList<'a, T> {
3446        let datum = self.make_datum(|packer| {
3447            packer.push_list_with(|packer| {
3448                for elem in iter {
3449                    packer.push(*elem.borrow());
3450                }
3451            });
3452        });
3453        DatumList::new(datum.unwrap_list().data())
3454    }
3455
3456    /// Convenience function identical to `make_datum` but instead returns a
3457    /// `DatumNested`.
3458    pub fn make_datum_nested<'a, F>(&'a self, f: F) -> DatumNested<'a>
3459    where
3460        F: FnOnce(&mut RowPacker),
3461    {
3462        let mut row = Row::default();
3463        f(&mut row.packer());
3464        self.push_unary_row_datum_nested(row)
3465    }
3466
3467    /// Like [`RowArena::make_datum`], but the provided closure can return an error.
3468    pub fn try_make_datum<'a, F, E>(&'a self, f: F) -> Result<Datum<'a>, E>
3469    where
3470        F: FnOnce(&mut RowPacker) -> Result<(), E>,
3471    {
3472        let mut row = Row::default();
3473        f(&mut row.packer())?;
3474        Ok(self.push_unary_row(row))
3475    }
3476
3477    /// Clear the contents of the arena.
3478    ///
3479    /// Retains the single largest region (emptied) so the arena can be reused without
3480    /// reallocating; a workload that clears between uses of similar size becomes allocation-free.
3481    pub fn clear(&mut self) {
3482        let inner = self.inner.get_mut();
3483        // Keep only the largest-capacity region, reset to empty, and drop the rest. Because region
3484        // capacities only ever grow (each new region at least doubles the previous), the largest is
3485        // normally the last; we scan for it defensively, which is cheap given log-many regions.
3486        if let Some(largest) = (0..inner.len()).max_by_key(|&i| inner[i].capacity()) {
3487            inner.swap(0, largest);
3488            inner.truncate(1);
3489            inner[0].clear();
3490        }
3491        self.allocated.set(0);
3492    }
3493}
3494
3495impl Default for RowArena {
3496    fn default() -> RowArena {
3497        RowArena::new()
3498    }
3499}
3500
3501/// A growable, writeable byte buffer that builds a value into a [`RowArena`].
3502///
3503/// Obtained from [`RowArena::writer`]. Behaves like a writeable byte slice (push/extend bytes,
3504/// read back as `&[u8]`); [`RowArenaBuf::finish`] copies the assembled bytes into the arena and
3505/// returns a reference valid for the arena's lifetime. The buffer is owned for the writer's
3506/// lifetime and, on drop, returned to the arena to be reused by the next writer.
3507#[derive(Debug)]
3508pub struct RowArenaBuf<'a> {
3509    arena: &'a RowArena,
3510    buf: Vec<u8>,
3511}
3512
3513impl<'a> RowArenaBuf<'a> {
3514    /// Appends a single byte.
3515    pub fn push(&mut self, byte: u8) {
3516        self.buf.push(byte);
3517    }
3518
3519    /// Appends a slice of bytes.
3520    pub fn extend_from_slice(&mut self, bytes: &[u8]) {
3521        self.buf.extend_from_slice(bytes);
3522    }
3523
3524    /// The bytes written so far.
3525    pub fn as_slice(&self) -> &[u8] {
3526        &self.buf
3527    }
3528
3529    /// The number of bytes written so far.
3530    pub fn len(&self) -> usize {
3531        self.buf.len()
3532    }
3533
3534    /// Whether no bytes have been written.
3535    pub fn is_empty(&self) -> bool {
3536        self.buf.is_empty()
3537    }
3538
3539    /// Copies the written bytes into the arena, returning a reference valid for its lifetime.
3540    pub fn finish(self) -> &'a [u8] {
3541        // `self` is dropped at the end of this call, returning `buf` to the arena for reuse; the
3542        // returned reference points into a committed region, not `buf`, so it stays valid.
3543        self.arena.push_bytes(self.buf.as_slice())
3544    }
3545
3546    /// Like [`RowArenaBuf::finish`], but returns the bytes as a `&str`.
3547    ///
3548    /// Intended for buffers written via [`std::fmt::Write`] (e.g. `write!`), whose contents are
3549    /// valid UTF-8. Panics if the bytes are not valid UTF-8.
3550    pub fn finish_str(self) -> &'a str {
3551        let bytes = self.arena.push_bytes(self.buf.as_slice());
3552        std::str::from_utf8(bytes).expect("RowArenaBuf::finish_str on non-UTF-8 contents")
3553    }
3554}
3555
3556impl<'a> Drop for RowArenaBuf<'a> {
3557    fn drop(&mut self) {
3558        // Return the buffer to the arena so the next writer can reuse its allocation. We keep only
3559        // one buffer: if the slot is already occupied — an outer writer is still live, or a nested
3560        // writer beat us to it — we drop ours rather than growing an unbounded pool. The borrow is
3561        // transient and never overlaps a live writer's, so this can't double-borrow.
3562        let mut slot = self.arena.scratch.borrow_mut();
3563        if slot.is_none() {
3564            *slot = Some(std::mem::take(&mut self.buf));
3565        }
3566    }
3567}
3568
3569impl<'a> std::ops::Deref for RowArenaBuf<'a> {
3570    type Target = [u8];
3571    fn deref(&self) -> &[u8] {
3572        &self.buf
3573    }
3574}
3575
3576impl<'a> std::io::Write for RowArenaBuf<'a> {
3577    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
3578        self.buf.extend_from_slice(bytes);
3579        Ok(bytes.len())
3580    }
3581
3582    fn flush(&mut self) -> std::io::Result<()> {
3583        Ok(())
3584    }
3585}
3586
3587impl<'a> std::fmt::Write for RowArenaBuf<'a> {
3588    fn write_str(&mut self, s: &str) -> std::fmt::Result {
3589        self.buf.extend_from_slice(s.as_bytes());
3590        Ok(())
3591    }
3592}
3593
3594/// A thread-local row, which can be borrowed and returned.
3595/// # Example
3596///
3597/// Use this type instead of creating a new row:
3598/// ```
3599/// use mz_repr::SharedRow;
3600///
3601/// let mut row_builder = SharedRow::get();
3602/// ```
3603///
3604/// This allows us to reuse an existing row allocation instead of creating a new one or retaining
3605/// an allocation locally. Additionally, we can observe the size of the local row in a central
3606/// place and potentially reallocate to reduce memory needs.
3607///
3608/// # Panic
3609///
3610/// [`SharedRow::get`] panics when trying to obtain multiple references to the shared row.
3611#[derive(Debug)]
3612pub struct SharedRow(Row);
3613
3614impl SharedRow {
3615    thread_local! {
3616        /// A thread-local slot containing a shared Row that can be temporarily used by a function.
3617        /// There can be at most one active user of this Row, which is tracked by the state of the
3618        /// `Option<_>` wrapper. When it is `Some(..)`, the row is available for using. When it
3619        /// is `None`, it is not, and the constructor will panic if a thread attempts to use it.
3620        static SHARED_ROW: Cell<Option<Row>> = const { Cell::new(Some(Row::empty())) }
3621    }
3622
3623    /// Get the shared row.
3624    ///
3625    /// The row's contents are cleared before returning it.
3626    ///
3627    /// # Panic
3628    ///
3629    /// Panics when the row is already borrowed elsewhere.
3630    pub fn get() -> Self {
3631        let mut row = Self::SHARED_ROW
3632            .take()
3633            .expect("attempted to borrow already borrowed SharedRow");
3634        // Clear row
3635        row.packer();
3636        Self(row)
3637    }
3638
3639    /// Gets the shared row and uses it to pack `iter`.
3640    pub fn pack<'a, I, D>(iter: I) -> Row
3641    where
3642        I: IntoIterator<Item = D>,
3643        D: Borrow<Datum<'a>>,
3644    {
3645        let mut row_builder = Self::get();
3646        let mut row_packer = row_builder.packer();
3647        row_packer.extend(iter);
3648        row_builder.clone()
3649    }
3650}
3651
3652impl std::ops::Deref for SharedRow {
3653    type Target = Row;
3654
3655    fn deref(&self) -> &Self::Target {
3656        &self.0
3657    }
3658}
3659
3660impl std::ops::DerefMut for SharedRow {
3661    fn deref_mut(&mut self) -> &mut Self::Target {
3662        &mut self.0
3663    }
3664}
3665
3666impl Drop for SharedRow {
3667    fn drop(&mut self) {
3668        // Take the Row allocation from this instance and put it back in the thread local slot for
3669        // the next user. The Row in `self` is replaced with an empty Row which does not allocate.
3670        Self::SHARED_ROW.set(Some(std::mem::take(&mut self.0)))
3671    }
3672}
3673
3674#[cfg(test)]
3675mod tests {
3676    use std::cmp::Ordering;
3677    use std::collections::hash_map::DefaultHasher;
3678    use std::hash::{Hash, Hasher};
3679
3680    use chrono::{DateTime, NaiveDate};
3681    use itertools::Itertools;
3682    use mz_ore::{assert_err, assert_none};
3683    use ordered_float::OrderedFloat;
3684
3685    use crate::SqlScalarType;
3686
3687    use super::*;
3688
3689    // StableRow's wire format is proto bytes, not the in-memory datum
3690    // encoding, so rows of every column type must roundtrip exactly through
3691    // both a self-describing format (JSON) and a compact binary one
3692    // (bincode). Equality on Row compares the packed in-memory bytes, so
3693    // this also catches any datum normalization sneaking into the
3694    // Row -> ProtoRow -> Row conversion.
3695    proptest! {
3696        #![proptest_config(ProptestConfig::with_cases(1000))]
3697
3698        #[mz_ore::test]
3699        #[cfg_attr(miri, ignore)] // too slow, and decNumber uses FFI
3700        fn stable_row_serde_roundtrip(
3701            stable in crate::relation::arb_relation_desc(1..8)
3702                .prop_flat_map(|desc| crate::relation::arb_row_for_relation(&desc))
3703                .prop_map(StableRow)
3704        ) {
3705            let json = serde_json::to_string(&stable).expect("serializes to JSON");
3706            let from_json: StableRow =
3707                serde_json::from_str(&json).expect("deserializes from JSON");
3708            prop_assert_eq!(&stable, &from_json);
3709
3710            let bytes = bincode::serialize(&stable).expect("serializes to bincode");
3711            let from_bincode: StableRow =
3712                bincode::deserialize(&bytes).expect("deserializes from bincode");
3713            prop_assert_eq!(&stable, &from_bincode);
3714        }
3715    }
3716
3717    /// Every width a variable-length integer can take, at both ends of its range and either
3718    /// side of each byte boundary, plus the values whose payload is empty.
3719    fn varint_edge_cases() -> Vec<Datum<'static>> {
3720        let mut datums = vec![
3721            Datum::Int16(0),
3722            Datum::Int16(-1),
3723            Datum::Int16(i16::MIN),
3724            Datum::Int16(i16::MAX),
3725            Datum::Int32(0),
3726            Datum::Int32(-1),
3727            Datum::Int32(i32::MIN),
3728            Datum::Int32(i32::MAX),
3729            Datum::Int64(0),
3730            Datum::Int64(-1),
3731            Datum::Int64(i64::MIN),
3732            Datum::Int64(i64::MAX),
3733            Datum::UInt8(0),
3734            Datum::UInt8(u8::MAX),
3735            Datum::UInt16(0),
3736            Datum::UInt16(u16::MAX),
3737            Datum::UInt32(0),
3738            Datum::UInt32(u32::MAX),
3739            Datum::UInt64(0),
3740            Datum::UInt64(u64::MAX),
3741        ];
3742        // One below, at, and one above every point where the payload grows a byte.
3743        for bits in 1..64 {
3744            let boundary = 1u64 << bits;
3745            for delta in [-1i64, 0, 1] {
3746                let Some(v) = boundary.checked_add_signed(delta) else {
3747                    continue;
3748                };
3749                datums.push(Datum::UInt64(v));
3750                if let Ok(v) = u32::try_from(v) {
3751                    datums.push(Datum::UInt32(v));
3752                }
3753                if let Ok(v) = u16::try_from(v) {
3754                    datums.push(Datum::UInt16(v));
3755                }
3756                if let Ok(v) = u8::try_from(v) {
3757                    datums.push(Datum::UInt8(v));
3758                }
3759                let Ok(v) = i64::try_from(v) else {
3760                    continue;
3761                };
3762                datums.push(Datum::Int64(v));
3763                datums.push(Datum::Int64(-v));
3764                if let Ok(v) = i32::try_from(v) {
3765                    datums.push(Datum::Int32(v));
3766                    datums.push(Datum::Int32(-v));
3767                }
3768                if let Ok(v) = i16::try_from(v) {
3769                    datums.push(Datum::Int16(v));
3770                    datums.push(Datum::Int16(-v));
3771                }
3772            }
3773        }
3774        datums
3775    }
3776
3777    /// Both signed families of a width share one match arm, which recovers the payload width by
3778    /// subtracting the family's first tag. A value whose tag falls outside the family it is
3779    /// decoded as would read the wrong number of bytes, so check every boundary lands where the
3780    /// arithmetic expects.
3781    #[mz_ore::test]
3782    fn varint_tags_land_in_their_family() {
3783        for datum in varint_edge_cases() {
3784            let row = Row::pack_slice(&[datum]);
3785            let tag = Tag::try_from_primitive(row.data[0]).expect("valid tag");
3786            let (first, len, negative, widest) = match datum {
3787                Datum::Int16(i) => (Tag::NonNegativeInt16_0, min_bytes_signed(i), i < 0, 2),
3788                Datum::Int32(i) => (Tag::NonNegativeInt32_0, min_bytes_signed(i), i < 0, 4),
3789                Datum::Int64(i) => (Tag::NonNegativeInt64_0, min_bytes_signed(i), i < 0, 8),
3790                Datum::UInt8(u) => (Tag::UInt8_0, min_bytes_unsigned(u), false, 1),
3791                Datum::UInt16(u) => (Tag::UInt16_0, min_bytes_unsigned(u), false, 2),
3792                Datum::UInt32(u) => (Tag::UInt32_0, min_bytes_unsigned(u), false, 4),
3793                Datum::UInt64(u) => (Tag::UInt64_0, min_bytes_unsigned(u), false, 8),
3794                other => panic!("not a variable-length integer: {other:?}"),
3795            };
3796            let delta = u8::from(tag) - u8::from(first);
3797            let signed = matches!(datum, Datum::Int16(_) | Datum::Int32(_) | Datum::Int64(_));
3798            if signed {
3799                // Interleaved: the width sits above the low bit, the sign in it. At the widest
3800                // width the alternation stops, one tag serving both signs.
3801                assert_eq!(delta >> 1, len, "wrong payload width in tag for {datum:?}");
3802                let sign_bit = if len == widest { 0 } else { u8::from(negative) };
3803                assert_eq!(delta & 1, sign_bit, "wrong sign in tag for {datum:?}");
3804            } else {
3805                assert_eq!(delta, len, "wrong payload width in tag for {datum:?}");
3806            }
3807            assert_eq!(row.unpack_first(), datum, "did not round-trip: {datum:?}");
3808        }
3809    }
3810
3811    /// The wide load in `read_varint_word` reads eight bytes whatever the payload's width, so
3812    /// the datums at the end of a row, where those bytes do not exist, take the tail path.
3813    #[mz_ore::test]
3814    fn varint_at_end_of_row_reads_exact_width() {
3815        for datum in varint_edge_cases() {
3816            // Alone in a row, and behind padding long enough to push the fast path back in.
3817            for prefix in [None, Some(Datum::String("0123456789abcdef"))] {
3818                let row = match prefix {
3819                    Some(p) => Row::pack_slice(&[p, datum]),
3820                    None => Row::pack_slice(&[datum]),
3821                };
3822                assert_eq!(
3823                    row.iter().last(),
3824                    Some(datum),
3825                    "did not round-trip at end of row: {datum:?}"
3826                );
3827            }
3828            // And nested, where the list's slice ends before the row's data does.
3829            let mut row = Row::default();
3830            let mut packer = row.packer();
3831            packer.push_list_with(|packer| packer.push(datum));
3832            packer.push(Datum::Int64(1));
3833            let list = match row.unpack_first() {
3834                Datum::List(list) => list,
3835                other => panic!("expected a list, got {other:?}"),
3836            };
3837            assert_eq!(list.iter().next(), Some(datum), "did not round-trip nested");
3838        }
3839    }
3840
3841    // Regression: comparing deeply nested list values must not overflow the
3842    // stack (STACK-7). `Datum` ordering recurses once per nesting level.
3843    #[mz_ore::test]
3844    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
3845    fn cmp_deep_nested_list_does_not_overflow() {
3846        fn deep() -> Row {
3847            // `push_list` byte-copies the inner value, so building does not recurse.
3848            let mut row = Row::pack_slice(&[Datum::Int64(1)]);
3849            for _ in 0..50_000 {
3850                let mut next = Row::default();
3851                next.packer().push_list([row.unpack_first()]);
3852                row = next;
3853            }
3854            row
3855        }
3856        let a = deep();
3857        let b = deep();
3858        assert_eq!(a.unpack_first().cmp(&b.unpack_first()), Ordering::Equal);
3859    }
3860
3861    fn hash<T: Hash>(t: &T) -> u64 {
3862        let mut hasher = DefaultHasher::new();
3863        t.hash(&mut hasher);
3864        hasher.finish()
3865    }
3866
3867    #[mz_ore::test]
3868    fn test_assumptions() {
3869        assert_eq!(size_of::<Tag>(), 1);
3870        #[cfg(target_endian = "big")]
3871        {
3872            // if you want to run this on a big-endian cpu, we'll need big-endian versions of the serialization code
3873            assert!(false);
3874        }
3875    }
3876
3877    #[mz_ore::test]
3878    fn miri_test_arena() {
3879        let arena = RowArena::new();
3880
3881        assert_eq!(arena.push_string("".to_owned()), "");
3882        assert_eq!(arena.push_string("العَرَبِيَّة".to_owned()), "العَرَبِيَّة");
3883
3884        let empty: &[u8] = &[];
3885        assert_eq!(arena.push_bytes(vec![]), empty);
3886        assert_eq!(arena.push_bytes(vec![0, 2, 1, 255]), &[0, 2, 1, 255]);
3887
3888        let mut row = Row::default();
3889        let mut packer = row.packer();
3890        packer.push_dict_with(|row| {
3891            row.push(Datum::String("a"));
3892            row.push_list_with(|row| {
3893                row.push(Datum::String("one"));
3894                row.push(Datum::String("two"));
3895                row.push(Datum::String("three"));
3896            });
3897            row.push(Datum::String("b"));
3898            row.push(Datum::String("c"));
3899        });
3900        assert_eq!(arena.push_unary_row(row.clone()), row.unpack_first());
3901    }
3902
3903    #[mz_ore::test]
3904    fn miri_test_arena_growth_keeps_references() {
3905        // References returned by `push_bytes` must stay valid as later pushes allocate new
3906        // regions; this exercises the "never resize a region that holds data" invariant.
3907        let arena = RowArena::new();
3908        let chunks: Vec<Vec<u8>> = (0..128u16)
3909            .map(|i| vec![u8::try_from(i % 256).unwrap(); usize::from(i % 13) + 1])
3910            .collect();
3911        let refs: Vec<&[u8]> = chunks
3912            .iter()
3913            .map(|c| arena.push_bytes(c.as_slice()))
3914            .collect();
3915        for (i, r) in refs.iter().enumerate() {
3916            assert_eq!(*r, chunks[i].as_slice());
3917        }
3918    }
3919
3920    #[mz_ore::test]
3921    fn miri_test_arena_unary_row_at_offset() {
3922        // A row pushed after other bytes lands at a non-zero offset within a region; reading it
3923        // back must not depend on the row starting at offset zero or on any alignment.
3924        let arena = RowArena::new();
3925        arena.reserve(4096);
3926        let _pad = arena.push_bytes(vec![0xAB; 5]);
3927        let row = Row::pack_slice(&[Datum::String("hello"), Datum::Int64(42), Datum::True]);
3928        assert_eq!(arena.push_unary_row(row.clone()), row.unpack_first());
3929    }
3930
3931    #[mz_ore::test]
3932    fn miri_test_arena_clear_reuse() {
3933        // After `clear` the arena retains a region and remains usable across cycles.
3934        let mut arena = RowArena::new();
3935        for i in 0..100u8 {
3936            let _ = arena.push_bytes(vec![i; 16]);
3937        }
3938        arena.clear();
3939        assert_eq!(arena.push_bytes(vec![7u8; 8]), &[7u8; 8]);
3940        assert_eq!(arena.push_string("after clear".to_owned()), "after clear");
3941        arena.clear();
3942        let empty: &[u8] = &[];
3943        assert_eq!(arena.push_bytes(Vec::<u8>::new()), empty);
3944    }
3945
3946    #[mz_ore::test]
3947    fn miri_test_arena_adopts_owned_bytes_and_keeps_references() {
3948        // `push_owned_bytes` adopts a buffer too large for the active region instead of copying it,
3949        // which puts a region the arena never wrote into in the middle of the stack. References
3950        // handed out before and after that must all stay valid.
3951        let arena = RowArena::new();
3952        let before = arena.push_bytes(vec![1u8; 8]);
3953        let adopted = arena.push_owned_bytes(vec![2u8; 64 * 1024]);
3954        let after = arena.push_bytes(vec![3u8; 8]);
3955        // A small buffer fits the active region, so it is copied rather than given a region.
3956        let small = arena.push_owned_bytes(vec![4u8; 4]);
3957
3958        assert_eq!(before, &[1u8; 8]);
3959        assert_eq!(adopted, &vec![2u8; 64 * 1024][..]);
3960        assert_eq!(after, &[3u8; 8]);
3961        assert_eq!(small, &[4u8; 4]);
3962
3963        let empty: &[u8] = &[];
3964        assert_eq!(arena.push_owned_bytes(vec![]), empty);
3965    }
3966
3967    #[mz_ore::test]
3968    fn test_arena_owned_pushes_keep_bump_allocating() {
3969        // Adoption never *creates* a region with headroom: it inserts the caller's buffer, whose
3970        // capacity equals its length, below whatever is on top. Only `push_bytes` grows the arena
3971        // geometrically (`new_cap = max(need, last_cap * 2)`), so once the active region cannot fit
3972        // an incoming value it never can again and every later owned push adopts: one retained
3973        // allocation and one `Vec<u8>` header per value, rather than `O(log n)` regions. That is the
3974        // default path for every `String`- and `Vec<u8>`-returning scalar function, and the arenas
3975        // in the MFP and join paths outlive a single row, so the region list grows with the number
3976        // of string values in a batch. `RowArena::clear` scans every region, so it degrades too.
3977        const VALUES: usize = 500;
3978        const VALUE: &str = "0123456789";
3979
3980        let regions = |arena: &RowArena| arena.inner.borrow().len();
3981        let push_all = |arena: &RowArena, owned: bool| {
3982            for _ in 0..VALUES {
3983                match owned {
3984                    true => _ = arena.push_string(VALUE.to_string()),
3985                    false => _ = arena.push_bytes(VALUE.as_bytes()),
3986                }
3987            }
3988        };
3989
3990        // The bump allocator working as intended, as the baseline to hold the owned path to.
3991        let copied = RowArena::new();
3992        push_all(&copied, false);
3993
3994        let owned = RowArena::new();
3995        push_all(&owned, true);
3996
3997        // Seeding with ordinary copies first must not change the answer. The arena never recovers,
3998        // so this is not just the empty-arena case where the placeholder on top has capacity 0.
3999        let seeded = RowArena::new();
4000        let _ = seeded.push_bytes(VALUE.as_bytes());
4001        push_all(&seeded, true);
4002
4003        // Compare against the copy path rather than an absolute count, so this pins the property (a
4004        // run of small owned pushes still ends with a region that has headroom) and leaves the
4005        // adoption predicate to the fix.
4006        let (copied, owned, seeded) = (regions(&copied), regions(&owned), regions(&seeded));
4007        assert!(
4008            owned <= copied * 2 && seeded <= copied * 2,
4009            "{VALUES} owned pushes left {owned} regions on an empty arena and {seeded} on a seeded \
4010             one, against {copied} for the same bytes copied",
4011        );
4012    }
4013
4014    #[mz_ore::test]
4015    fn miri_test_arena_budget() {
4016        // Without a budget nothing is ever over it, however much is pushed.
4017        let arena = RowArena::new();
4018        let _ = arena.push_bytes(vec![0u8; 1024]);
4019        assert!(!arena.over_budget());
4020        assert_eq!(arena.budget_remaining(), usize::MAX);
4021
4022        let arena = RowArena::with_budget(100);
4023        assert!(!arena.over_budget());
4024        assert_eq!(arena.budget_remaining(), 100);
4025
4026        // Staying within the budget leaves it satisfied, and the remaining count tracks what a
4027        // caller that predicts its own size would consult.
4028        let _ = arena.push_bytes(vec![0u8; 60]);
4029        assert!(!arena.over_budget());
4030        assert_eq!(arena.budget_remaining(), 40);
4031        assert_eq!(arena.allocated_bytes(), 60);
4032
4033        // Crossing it reports, rather than refusing the push: a truncated push would corrupt the
4034        // datum, so the value is intact and it is the caller's job to fail.
4035        let pushed = arena.push_bytes(vec![7u8; 80]);
4036        assert_eq!(pushed, &[7u8; 80]);
4037        assert!(arena.over_budget());
4038        assert_eq!(arena.budget_remaining(), 0);
4039
4040        // An adopted buffer counts against the budget too, or adoption would be a way around it.
4041        // Large enough to actually be adopted rather than copied.
4042        let mut arena = RowArena::with_budget(100);
4043        let _ = arena.push_owned_bytes(vec![0u8; 8 * 1024]);
4044        assert!(arena.over_budget());
4045
4046        arena.clear();
4047        assert!(!arena.over_budget());
4048        assert_eq!(arena.allocated_bytes(), 0);
4049    }
4050
4051    #[mz_ore::test]
4052    fn miri_test_arena_writer() {
4053        use std::io::Write;
4054
4055        let arena = RowArena::new();
4056
4057        // Build a value incrementally and commit it.
4058        let mut w = arena.writer();
4059        let mut expected = Vec::new();
4060        for i in 0..1000u16 {
4061            let byte = u8::try_from(i % 256).unwrap();
4062            w.push(byte);
4063            expected.push(byte);
4064            w.extend_from_slice(&[byte, byte]);
4065            expected.extend_from_slice(&[byte, byte]);
4066        }
4067        assert_eq!(w.as_slice(), expected.as_slice());
4068        assert_eq!(w.len(), expected.len());
4069        let first = w.finish();
4070        assert_eq!(first, expected.as_slice());
4071
4072        // A second writer reuses the scratch; its result is independent of the first, which stays
4073        // valid because `finish` copied it into the arena.
4074        let mut w2 = arena.writer();
4075        write!(w2, "hello").unwrap();
4076        let second = w2.finish();
4077        assert_eq!(second, b"hello");
4078        assert_eq!(first, expected.as_slice());
4079
4080        // An empty writer commits to an empty slice.
4081        let empty: &[u8] = &[];
4082        assert_eq!(arena.writer().finish(), empty);
4083
4084        // Abandoning a writer without finishing is fine; the next writer starts empty.
4085        {
4086            let mut w3 = arena.writer();
4087            w3.extend_from_slice(b"discarded");
4088        }
4089        assert_eq!(arena.writer().as_slice(), empty);
4090    }
4091
4092    #[mz_ore::test]
4093    fn miri_test_arena_writer_nested() {
4094        // Reentrancy: a writer obtained while another is still live must not panic (no `RefCell`
4095        // double-borrow) and must not disturb the outer writer. The nested writer just gets its own
4096        // buffer; the outer one keeps building independently.
4097        let arena = RowArena::new();
4098
4099        let mut outer = arena.writer();
4100        outer.extend_from_slice(b"outer-before-");
4101
4102        // Take a second writer while `outer` is still live -- the case that double-borrowed before.
4103        let inner_bytes = {
4104            let mut inner = arena.writer();
4105            inner.extend_from_slice(b"inner");
4106            // The outer writer is unaffected by the nested one.
4107            assert_eq!(outer.as_slice(), b"outer-before-");
4108            inner.finish()
4109        };
4110        assert_eq!(inner_bytes, b"inner");
4111
4112        // `outer` is intact and still writable after the nested writer committed.
4113        outer.extend_from_slice(b"after");
4114        let outer_bytes = outer.finish();
4115        assert_eq!(outer_bytes, b"outer-before-after");
4116        // Both committed slices stay valid and independent.
4117        assert_eq!(inner_bytes, b"inner");
4118
4119        // Once all writers have dropped, the recycled buffer is reusable (and cleared on acquire).
4120        let mut again = arena.writer();
4121        again.extend_from_slice(b"reused");
4122        assert_eq!(again.finish(), b"reused");
4123    }
4124
4125    #[mz_ore::test]
4126    fn miri_test_arena_writer_fmt() {
4127        use std::fmt::Write;
4128
4129        // Format text into the writer (e.g. building a cast-to-string result) and commit as `&str`.
4130        let arena = RowArena::new();
4131        let mut w = arena.writer();
4132        for i in 0..5 {
4133            write!(w, "{i},").unwrap();
4134        }
4135        assert_eq!(w.finish_str(), "0,1,2,3,4,");
4136    }
4137
4138    #[mz_ore::test]
4139    fn miri_test_round_trip() {
4140        fn round_trip(datums: Vec<Datum>) {
4141            let row = Row::pack(datums.clone());
4142
4143            // When run under miri this catches undefined bytes written to data
4144            // eg by calling push_copy! on a type which contains undefined padding values
4145            println!("{:?}", row.data());
4146
4147            let datums2 = row.iter().collect::<Vec<_>>();
4148            let datums3 = row.unpack();
4149            assert_eq!(datums, datums2);
4150            assert_eq!(datums, datums3);
4151        }
4152
4153        round_trip(vec![]);
4154        round_trip(
4155            SqlScalarType::enumerate()
4156                .iter()
4157                .flat_map(|r#type| r#type.interesting_datums())
4158                .collect(),
4159        );
4160        round_trip(vec![
4161            Datum::Null,
4162            Datum::Null,
4163            Datum::False,
4164            Datum::True,
4165            Datum::Int16(-21),
4166            Datum::Int32(-42),
4167            Datum::Int64(-2_147_483_648 - 42),
4168            Datum::UInt8(0),
4169            Datum::UInt8(1),
4170            Datum::UInt16(0),
4171            Datum::UInt16(1),
4172            Datum::UInt16(1 << 8),
4173            Datum::UInt32(0),
4174            Datum::UInt32(1),
4175            Datum::UInt32(1 << 8),
4176            Datum::UInt32(1 << 16),
4177            Datum::UInt32(1 << 24),
4178            Datum::UInt64(0),
4179            Datum::UInt64(1),
4180            Datum::UInt64(1 << 8),
4181            Datum::UInt64(1 << 16),
4182            Datum::UInt64(1 << 24),
4183            Datum::UInt64(1 << 32),
4184            Datum::UInt64(1 << 40),
4185            Datum::UInt64(1 << 48),
4186            Datum::UInt64(1 << 56),
4187            Datum::Float32(OrderedFloat::from(-42.12)),
4188            Datum::Float64(OrderedFloat::from(-2_147_483_648.0 - 42.12)),
4189            Datum::Date(Date::from_pg_epoch(365 * 45 + 21).unwrap()),
4190            Datum::Timestamp(
4191                CheckedTimestamp::from_timestamplike(
4192                    NaiveDate::from_isoywd_opt(2019, 30, chrono::Weekday::Wed)
4193                        .unwrap()
4194                        .and_hms_opt(14, 32, 11)
4195                        .unwrap(),
4196                )
4197                .unwrap(),
4198            ),
4199            Datum::TimestampTz(
4200                CheckedTimestamp::from_timestamplike(DateTime::from_timestamp(61, 0).unwrap())
4201                    .unwrap(),
4202            ),
4203            Datum::Interval(Interval {
4204                months: 312,
4205                ..Default::default()
4206            }),
4207            Datum::Interval(Interval::new(0, 0, 1_012_312)),
4208            Datum::Bytes(&[]),
4209            Datum::Bytes(&[0, 2, 1, 255]),
4210            Datum::String(""),
4211            Datum::String("العَرَبِيَّة"),
4212        ]);
4213    }
4214
4215    #[mz_ore::test]
4216    fn test_array() {
4217        // Construct an array using `Row::push_array` and verify that it unpacks
4218        // correctly.
4219        const DIM: ArrayDimension = ArrayDimension {
4220            lower_bound: 2,
4221            length: 2,
4222        };
4223        let mut row = Row::default();
4224        let mut packer = row.packer();
4225        packer
4226            .try_push_array(&[DIM], vec![Datum::Int32(1), Datum::Int32(2)])
4227            .unwrap();
4228        let arr1 = row.unpack_first().unwrap_array();
4229        assert_eq!(arr1.dims().into_iter().collect::<Vec<_>>(), vec![DIM]);
4230        assert_eq!(
4231            arr1.elements().into_iter().collect::<Vec<_>>(),
4232            vec![Datum::Int32(1), Datum::Int32(2)]
4233        );
4234
4235        // Pack a previously-constructed `Datum::Array` and verify that it
4236        // unpacks correctly.
4237        let row = Row::pack_slice(&[Datum::Array(arr1)]);
4238        let arr2 = row.unpack_first().unwrap_array();
4239        assert_eq!(arr1, arr2);
4240    }
4241
4242    #[mz_ore::test]
4243    fn test_multidimensional_array() {
4244        let datums = vec![
4245            Datum::Int32(1),
4246            Datum::Int32(2),
4247            Datum::Int32(3),
4248            Datum::Int32(4),
4249            Datum::Int32(5),
4250            Datum::Int32(6),
4251            Datum::Int32(7),
4252            Datum::Int32(8),
4253        ];
4254
4255        let mut row = Row::default();
4256        let mut packer = row.packer();
4257        packer
4258            .try_push_array(
4259                &[
4260                    ArrayDimension {
4261                        lower_bound: 1,
4262                        length: 1,
4263                    },
4264                    ArrayDimension {
4265                        lower_bound: 1,
4266                        length: 4,
4267                    },
4268                    ArrayDimension {
4269                        lower_bound: 1,
4270                        length: 2,
4271                    },
4272                ],
4273                &datums,
4274            )
4275            .unwrap();
4276        let array = row.unpack_first().unwrap_array();
4277        assert_eq!(array.elements().into_iter().collect::<Vec<_>>(), datums);
4278    }
4279
4280    #[mz_ore::test]
4281    fn test_array_max_dimensions() {
4282        let mut row = Row::default();
4283        let max_dims = usize::from(MAX_ARRAY_DIMENSIONS);
4284
4285        // An array with one too many dimensions should be rejected.
4286        let res = row.packer().try_push_array(
4287            &vec![
4288                ArrayDimension {
4289                    lower_bound: 1,
4290                    length: 1
4291                };
4292                max_dims + 1
4293            ],
4294            vec![Datum::Int32(4)],
4295        );
4296        assert_eq!(res, Err(InvalidArrayError::TooManyDimensions(max_dims + 1)));
4297        assert!(row.data.is_empty());
4298
4299        // An array with exactly the maximum allowable dimensions should be
4300        // accepted.
4301        row.packer()
4302            .try_push_array(
4303                &vec![
4304                    ArrayDimension {
4305                        lower_bound: 1,
4306                        length: 1
4307                    };
4308                    max_dims
4309                ],
4310                vec![Datum::Int32(4)],
4311            )
4312            .unwrap();
4313    }
4314
4315    #[mz_ore::test]
4316    fn test_array_wrong_cardinality() {
4317        let mut row = Row::default();
4318        let res = row.packer().try_push_array(
4319            &[
4320                ArrayDimension {
4321                    lower_bound: 1,
4322                    length: 2,
4323                },
4324                ArrayDimension {
4325                    lower_bound: 1,
4326                    length: 3,
4327                },
4328            ],
4329            vec![Datum::Int32(1), Datum::Int32(2)],
4330        );
4331        assert_eq!(
4332            res,
4333            Err(InvalidArrayError::WrongCardinality {
4334                actual: 2,
4335                expected: 6,
4336            })
4337        );
4338        assert!(row.data.is_empty());
4339    }
4340
4341    #[mz_ore::test]
4342    fn test_array_cardinality_overflow() {
4343        // Dimension lengths whose product overflows `usize` must be rejected as
4344        // a `WrongCardinality` error, not panic (under overflow checks) or wrap
4345        // (in release, which could spuriously accept a corrupt array). The
4346        // product saturates to `usize::MAX`, which no real element count matches.
4347        let mut row = Row::default();
4348        let res = row.packer().try_push_array(
4349            &[
4350                ArrayDimension {
4351                    lower_bound: 1,
4352                    length: usize::MAX,
4353                },
4354                ArrayDimension {
4355                    lower_bound: 1,
4356                    length: 2,
4357                },
4358            ],
4359            vec![Datum::Int32(1), Datum::Int32(2)],
4360        );
4361        assert_eq!(
4362            res,
4363            Err(InvalidArrayError::WrongCardinality {
4364                actual: 2,
4365                expected: usize::MAX,
4366            })
4367        );
4368        assert!(row.data.is_empty());
4369    }
4370
4371    #[mz_ore::test]
4372    fn test_nesting() {
4373        let mut row = Row::default();
4374        row.packer().push_dict_with(|row| {
4375            row.push(Datum::String("favourites"));
4376            row.push_list_with(|row| {
4377                row.push(Datum::String("ice cream"));
4378                row.push(Datum::String("oreos"));
4379                row.push(Datum::String("cheesecake"));
4380            });
4381            row.push(Datum::String("name"));
4382            row.push(Datum::String("bob"));
4383        });
4384
4385        let mut iter = row.unpack_first().unwrap_map().iter();
4386
4387        let (k, v) = iter.next().unwrap();
4388        assert_eq!(k, "favourites");
4389        assert_eq!(
4390            v.unwrap_list().iter().collect::<Vec<_>>(),
4391            vec![
4392                Datum::String("ice cream"),
4393                Datum::String("oreos"),
4394                Datum::String("cheesecake"),
4395            ]
4396        );
4397
4398        let (k, v) = iter.next().unwrap();
4399        assert_eq!(k, "name");
4400        assert_eq!(v, Datum::String("bob"));
4401    }
4402
4403    #[mz_ore::test]
4404    fn test_dict_errors() -> Result<(), Box<dyn std::error::Error>> {
4405        let pack = |ok| {
4406            let mut row = Row::default();
4407            row.packer().push_dict_with(|row| {
4408                if ok {
4409                    row.push(Datum::String("key"));
4410                    row.push(Datum::Int32(42));
4411                    Ok(7)
4412                } else {
4413                    Err("fail")
4414                }
4415            })?;
4416            Ok(row)
4417        };
4418
4419        assert_eq!(pack(false), Err("fail"));
4420
4421        let row = pack(true)?;
4422        let mut dict = row.unpack_first().unwrap_map().iter();
4423        assert_eq!(dict.next(), Some(("key", Datum::Int32(42))));
4424        assert_eq!(dict.next(), None);
4425
4426        Ok(())
4427    }
4428
4429    #[mz_ore::test]
4430    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decNumberFromInt32` on OS `linux`
4431    fn test_datum_sizes() {
4432        let arena = RowArena::new();
4433
4434        // Test the claims about various datum sizes.
4435        let values_of_interest = vec![
4436            Datum::Null,
4437            Datum::False,
4438            Datum::Int16(0),
4439            Datum::Int32(0),
4440            Datum::Int64(0),
4441            Datum::UInt8(0),
4442            Datum::UInt8(1),
4443            Datum::UInt16(0),
4444            Datum::UInt16(1),
4445            Datum::UInt16(1 << 8),
4446            Datum::UInt32(0),
4447            Datum::UInt32(1),
4448            Datum::UInt32(1 << 8),
4449            Datum::UInt32(1 << 16),
4450            Datum::UInt32(1 << 24),
4451            Datum::UInt64(0),
4452            Datum::UInt64(1),
4453            Datum::UInt64(1 << 8),
4454            Datum::UInt64(1 << 16),
4455            Datum::UInt64(1 << 24),
4456            Datum::UInt64(1 << 32),
4457            Datum::UInt64(1 << 40),
4458            Datum::UInt64(1 << 48),
4459            Datum::UInt64(1 << 56),
4460            Datum::Float32(OrderedFloat(0.0)),
4461            Datum::Float64(OrderedFloat(0.0)),
4462            Datum::from(numeric::Numeric::from(0)),
4463            Datum::from(numeric::Numeric::from(1000)),
4464            Datum::from(numeric::Numeric::from(9999)),
4465            Datum::Date(
4466                NaiveDate::from_ymd_opt(1, 1, 1)
4467                    .unwrap()
4468                    .try_into()
4469                    .unwrap(),
4470            ),
4471            Datum::Timestamp(
4472                CheckedTimestamp::from_timestamplike(
4473                    DateTime::from_timestamp(0, 0).unwrap().naive_utc(),
4474                )
4475                .unwrap(),
4476            ),
4477            Datum::TimestampTz(
4478                CheckedTimestamp::from_timestamplike(DateTime::from_timestamp(0, 0).unwrap())
4479                    .unwrap(),
4480            ),
4481            Datum::Interval(Interval::default()),
4482            Datum::Bytes(&[]),
4483            Datum::String(""),
4484            Datum::JsonNull,
4485            Datum::Range(Range { inner: None }),
4486            arena.make_datum(|packer| {
4487                packer
4488                    .push_range(Range::new(Some((
4489                        RangeLowerBound::new(Datum::Int32(-1), true),
4490                        RangeUpperBound::new(Datum::Int32(1), true),
4491                    ))))
4492                    .unwrap();
4493            }),
4494        ];
4495        for value in values_of_interest {
4496            if datum_size(&value) != Row::pack_slice(&[value]).data.len() {
4497                panic!("Disparity in claimed size for {:?}", value);
4498            }
4499        }
4500    }
4501
4502    #[mz_ore::test]
4503    fn test_range_errors() {
4504        fn test_range_errors_inner<'a>(
4505            datums: Vec<Vec<Datum<'a>>>,
4506        ) -> Result<(), InvalidRangeError> {
4507            let mut row = Row::default();
4508            let row_len = row.byte_len();
4509            let mut packer = row.packer();
4510            let r = packer.push_range_with(
4511                RangeLowerBound {
4512                    inclusive: true,
4513                    bound: Some(|row: &mut RowPacker| {
4514                        for d in &datums[0] {
4515                            row.push(d);
4516                        }
4517                        Ok(())
4518                    }),
4519                },
4520                RangeUpperBound {
4521                    inclusive: true,
4522                    bound: Some(|row: &mut RowPacker| {
4523                        for d in &datums[1] {
4524                            row.push(d);
4525                        }
4526                        Ok(())
4527                    }),
4528                },
4529            );
4530
4531            assert_eq!(row_len, row.byte_len());
4532
4533            r
4534        }
4535
4536        // A finite bound whose closure pushes zero values violates the
4537        // `push_range_with` caller contract and still panics. This is
4538        // unreachable when decoding a `ProtoRow`: each decoded bound pushes
4539        // exactly one datum (or fails), so only an in-process caller can hit it.
4540        for panicking_case in [
4541            vec![vec![Datum::Int32(1)], vec![]],
4542            vec![vec![Datum::Int32(1), Datum::Int32(2)], vec![]],
4543        ] {
4544            #[allow(clippy::disallowed_methods)] // not using enhanced panic handler in tests
4545            let result = std::panic::catch_unwind(|| test_range_errors_inner(panicking_case));
4546            assert_err!(result);
4547        }
4548
4549        // Inconsistent bound counts, mismatched datum kinds, and Null bounds are
4550        // all reachable from a crafted/corrupted `ProtoRow`, so they return an
4551        // error instead of panicking.
4552        for error_case in [
4553            vec![
4554                vec![Datum::Int32(1), Datum::Int32(2)],
4555                vec![Datum::Int32(3)],
4556            ],
4557            vec![
4558                vec![Datum::Int32(1)],
4559                vec![Datum::Int32(2), Datum::Int32(3)],
4560            ],
4561            vec![vec![Datum::Int32(1)], vec![Datum::UInt16(2)]],
4562            vec![vec![Datum::Null], vec![Datum::Int32(2)]],
4563            vec![vec![Datum::Int32(1)], vec![Datum::Null]],
4564        ] {
4565            assert_eq!(
4566                test_range_errors_inner(error_case),
4567                Err(InvalidRangeError::InvalidRangeData)
4568            );
4569        }
4570
4571        let e = test_range_errors_inner(vec![vec![Datum::Int32(2)], vec![Datum::Int32(1)]]);
4572        assert_eq!(e, Err(InvalidRangeError::MisorderedRangeBounds));
4573    }
4574
4575    /// Lists have a variable-length encoding for their lengths. We test each case here.
4576    #[mz_ore::test]
4577    #[cfg_attr(miri, ignore)] // slow
4578    fn test_list_encoding() {
4579        fn test_list_encoding_inner(len: usize) {
4580            let list_elem = |i: usize| {
4581                if i % 2 == 0 {
4582                    Datum::False
4583                } else {
4584                    Datum::True
4585                }
4586            };
4587            let mut row = Row::default();
4588            {
4589                // Push some stuff.
4590                let mut packer = row.packer();
4591                packer.push(Datum::String("start"));
4592                packer.push_list_with(|packer| {
4593                    for i in 0..len {
4594                        packer.push(list_elem(i));
4595                    }
4596                });
4597                packer.push(Datum::String("end"));
4598            }
4599            // Check that we read back exactly what we pushed.
4600            let mut row_it = row.iter();
4601            assert_eq!(row_it.next().unwrap(), Datum::String("start"));
4602            match row_it.next().unwrap() {
4603                Datum::List(list) => {
4604                    let mut list_it = list.iter();
4605                    for i in 0..len {
4606                        assert_eq!(list_it.next().unwrap(), list_elem(i));
4607                    }
4608                    assert_none!(list_it.next());
4609                }
4610                _ => panic!("expected Datum::List"),
4611            }
4612            assert_eq!(row_it.next().unwrap(), Datum::String("end"));
4613            assert_none!(row_it.next());
4614        }
4615
4616        test_list_encoding_inner(0);
4617        test_list_encoding_inner(1);
4618        test_list_encoding_inner(10);
4619        test_list_encoding_inner(TINY - 1); // tiny
4620        test_list_encoding_inner(TINY + 1); // short
4621        test_list_encoding_inner(SHORT + 1); // long
4622
4623        // The biggest one takes 40 s on my laptop, probably not worth it.
4624        //test_list_encoding_inner(LONG + 1); // huge
4625    }
4626
4627    /// Demonstrates that DatumList's Eq (bytewise) and Ord (datum-by-datum) are now consistent.
4628    /// A list containing -0.0 and one containing +0.0 have different byte representations
4629    /// (IEEE 754 distinguishes them), originally Eq says they are not equal. But after
4630    /// using the new Datum::cmp, Eq says they are equal, which matches what Ord
4631    /// compares via iter().cmp(other.iter()), and them as equal.
4632    #[mz_ore::test]
4633    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
4634    fn test_datum_list_eq_ord_consistency() {
4635        // Build list containing +0.0
4636        let mut row_pos = Row::default();
4637        row_pos.packer().push_list_with(|p| {
4638            p.push(Datum::Float64(OrderedFloat::from(0.0)));
4639        });
4640        let list_pos = row_pos.unpack_first().unwrap_list();
4641
4642        // Build list containing -0.0 (distinct bit pattern from +0.0)
4643        let mut row_neg = Row::default();
4644        row_neg.packer().push_list_with(|p| {
4645            p.push(Datum::Float64(OrderedFloat::from(-0.0)));
4646        });
4647        let list_neg = row_neg.unpack_first().unwrap_list();
4648
4649        // Eq is bytewise: different encodings => not equal
4650        // This was a bug in the past, so we test it.
4651        assert_eq!(
4652            list_pos, list_neg,
4653            "Eq should see different encodings as equal"
4654        );
4655
4656        // Ord is datum-by-datum: -0.0 and +0.0 compare equal as Datums
4657        assert_eq!(
4658            list_pos.cmp(&list_neg),
4659            Ordering::Equal,
4660            "Ord (datum-by-datum) should see -0.0 and +0.0 as equal"
4661        );
4662    }
4663
4664    /// Demonstrates that DatumMap's derived Eq (bytewise) can make maps with equal keys and
4665    /// values compare equal when values have different encodings (e.g. -0.0 vs +0.0).
4666    #[mz_ore::test]
4667    fn test_datum_map_eq_bytewise_consistency() {
4668        // Build map {"k": +0.0}
4669        let mut row_pos = Row::default();
4670        row_pos.packer().push_dict_with(|p| {
4671            p.push(Datum::String("k"));
4672            p.push(Datum::Float64(OrderedFloat::from(0.0)));
4673        });
4674        let map_pos = row_pos.unpack_first().unwrap_map();
4675
4676        // Build map {"k": -0.0}
4677        let mut row_neg = Row::default();
4678        row_neg.packer().push_dict_with(|p| {
4679            p.push(Datum::String("k"));
4680            p.push(Datum::Float64(OrderedFloat::from(-0.0)));
4681        });
4682        let map_neg = row_neg.unpack_first().unwrap_map();
4683
4684        // Same keys and semantically equal values, but Eq (bytewise) says not equal
4685        assert_eq!(
4686            map_pos, map_neg,
4687            "DatumMap Eq is semantic; -0.0 and +0.0 have different encodings but are equal"
4688        );
4689        // Verify they have the same logical content
4690        let entries_pos: Vec<_> = map_pos.iter().collect();
4691        let entries_neg: Vec<_> = map_neg.iter().collect();
4692        assert_eq!(entries_pos.len(), entries_neg.len());
4693        for ((k1, v1), (k2, v2)) in entries_pos.iter().zip_eq(entries_neg.iter()) {
4694            assert_eq!(k1, k2);
4695            assert_eq!(
4696                v1, v2,
4697                "Datum-level comparison treats -0.0 and +0.0 as equal"
4698            );
4699        }
4700    }
4701
4702    /// Hash must agree with Eq: equal lists must have the same hash.
4703    #[mz_ore::test]
4704    fn test_datum_list_hash_consistency() {
4705        // Equal lists (including -0.0 vs +0.0) must hash the same
4706        let mut row_pos = Row::default();
4707        row_pos.packer().push_list_with(|p| {
4708            p.push(Datum::Float64(OrderedFloat::from(0.0)));
4709        });
4710        let list_pos = row_pos.unpack_first().unwrap_list();
4711
4712        let mut row_neg = Row::default();
4713        row_neg.packer().push_list_with(|p| {
4714            p.push(Datum::Float64(OrderedFloat::from(-0.0)));
4715        });
4716        let list_neg = row_neg.unpack_first().unwrap_list();
4717
4718        assert_eq!(list_pos, list_neg);
4719        assert_eq!(
4720            hash(&list_pos),
4721            hash(&list_neg),
4722            "equal lists must have same hash"
4723        );
4724
4725        // Unequal lists should have different hashes (with asymptotic probability 1)
4726        let mut row_a = Row::default();
4727        row_a.packer().push_list_with(|p| {
4728            p.push(Datum::Int32(1));
4729            p.push(Datum::Int32(2));
4730        });
4731        let list_a = row_a.unpack_first().unwrap_list();
4732
4733        let mut row_b = Row::default();
4734        row_b.packer().push_list_with(|p| {
4735            p.push(Datum::Int32(1));
4736            p.push(Datum::Int32(3));
4737        });
4738        let list_b = row_b.unpack_first().unwrap_list();
4739
4740        assert_ne!(list_a, list_b);
4741        assert_ne!(
4742            hash(&list_a),
4743            hash(&list_b),
4744            "unequal lists must have different hashes"
4745        );
4746    }
4747
4748    /// Ord/PartialOrd for DatumList: less, equal, greater.
4749    #[mz_ore::test]
4750    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
4751    fn test_datum_list_ordering() {
4752        let mut row_12 = Row::default();
4753        row_12.packer().push_list_with(|p| {
4754            p.push(Datum::Int32(1));
4755            p.push(Datum::Int32(2));
4756        });
4757        let list_12 = row_12.unpack_first().unwrap_list();
4758
4759        let mut row_13 = Row::default();
4760        row_13.packer().push_list_with(|p| {
4761            p.push(Datum::Int32(1));
4762            p.push(Datum::Int32(3));
4763        });
4764        let list_13 = row_13.unpack_first().unwrap_list();
4765
4766        let mut row_123 = Row::default();
4767        row_123.packer().push_list_with(|p| {
4768            p.push(Datum::Int32(1));
4769            p.push(Datum::Int32(2));
4770            p.push(Datum::Int32(3));
4771        });
4772        let list_123 = row_123.unpack_first().unwrap_list();
4773
4774        // [1, 2] < [1, 3] due to the second element being different
4775        assert_eq!(list_12.cmp(&list_13), Ordering::Less);
4776        assert_eq!(list_13.cmp(&list_12), Ordering::Greater);
4777        assert_eq!(list_12.cmp(&list_12), Ordering::Equal);
4778        // shorter prefix compares less
4779        assert_eq!(list_12.cmp(&list_123), Ordering::Less);
4780    }
4781
4782    /// Hash must agree with Eq: equal maps must have the same hash.
4783    #[mz_ore::test]
4784    fn test_datum_map_hash_consistency() {
4785        let mut row_pos = Row::default();
4786        row_pos.packer().push_dict_with(|p| {
4787            p.push(Datum::String("x"));
4788            p.push(Datum::Float64(OrderedFloat::from(0.0)));
4789        });
4790        let map_pos = row_pos.unpack_first().unwrap_map();
4791
4792        let mut row_neg = Row::default();
4793        row_neg.packer().push_dict_with(|p| {
4794            p.push(Datum::String("x"));
4795            p.push(Datum::Float64(OrderedFloat::from(-0.0)));
4796        });
4797        let map_neg = row_neg.unpack_first().unwrap_map();
4798
4799        assert_eq!(map_pos, map_neg);
4800        assert_eq!(
4801            hash(&map_pos),
4802            hash(&map_neg),
4803            "equal maps must have same hash"
4804        );
4805
4806        let mut row_a = Row::default();
4807        row_a.packer().push_dict_with(|p| {
4808            p.push(Datum::String("a"));
4809            p.push(Datum::Int32(1));
4810        });
4811        let map_a = row_a.unpack_first().unwrap_map();
4812
4813        let mut row_b = Row::default();
4814        row_b.packer().push_dict_with(|p| {
4815            p.push(Datum::String("a"));
4816            p.push(Datum::Int32(2));
4817        });
4818        let map_b = row_b.unpack_first().unwrap_map();
4819
4820        assert_ne!(map_a, map_b);
4821        assert_ne!(
4822            hash(&map_a),
4823            hash(&map_b),
4824            "unequal maps must have different hashes"
4825        );
4826    }
4827
4828    /// Ord/PartialOrd for DatumMap: less, equal, greater (by key then value).
4829    #[mz_ore::test]
4830    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
4831    fn test_datum_map_ordering() {
4832        let mut row_a1 = Row::default();
4833        row_a1.packer().push_dict_with(|p| {
4834            p.push(Datum::String("a"));
4835            p.push(Datum::Int32(1));
4836        });
4837        let map_a1 = row_a1.unpack_first().unwrap_map();
4838
4839        let mut row_a2 = Row::default();
4840        row_a2.packer().push_dict_with(|p| {
4841            p.push(Datum::String("a"));
4842            p.push(Datum::Int32(2));
4843        });
4844        let map_a2 = row_a2.unpack_first().unwrap_map();
4845
4846        let mut row_b1 = Row::default();
4847        row_b1.packer().push_dict_with(|p| {
4848            p.push(Datum::String("b"));
4849            p.push(Datum::Int32(1));
4850        });
4851        let map_b1 = row_b1.unpack_first().unwrap_map();
4852
4853        assert_eq!(map_a1.cmp(&map_a2), Ordering::Less);
4854        assert_eq!(map_a2.cmp(&map_a1), Ordering::Greater);
4855        assert_eq!(map_a1.cmp(&map_a1), Ordering::Equal);
4856        assert_eq!(map_a1.cmp(&map_b1), Ordering::Less); // "a" < "b"
4857    }
4858
4859    /// Datum puts Null last in the enum so that nulls sort last (PostgreSQL default).
4860    /// This ordering is used when comparing DatumList/DatumMap (e.g. jsonb_agg tiebreaker).
4861    #[mz_ore::test]
4862    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
4863    fn test_datum_list_and_map_null_sorts_last() {
4864        // DatumList: [1] < [null] so non-null sorts before null
4865        let mut row_list_1 = Row::default();
4866        row_list_1
4867            .packer()
4868            .push_list_with(|p| p.push(Datum::Int32(1)));
4869        let list_1 = row_list_1.unpack_first().unwrap_list();
4870
4871        let mut row_list_null = Row::default();
4872        row_list_null
4873            .packer()
4874            .push_list_with(|p| p.push(Datum::Null));
4875        let list_null = row_list_null.unpack_first().unwrap_list();
4876
4877        assert_eq!(list_1.cmp(&list_null), Ordering::Less);
4878        assert_eq!(list_null.cmp(&list_1), Ordering::Greater);
4879
4880        // DatumMap: {"k": 1} < {"k": null} so non-null sorts before null (same as jsonb_agg)
4881        let mut row_map_1 = Row::default();
4882        row_map_1.packer().push_dict_with(|p| {
4883            p.push(Datum::String("k"));
4884            p.push(Datum::Int32(1));
4885        });
4886        let map_1 = row_map_1.unpack_first().unwrap_map();
4887
4888        let mut row_map_null = Row::default();
4889        row_map_null.packer().push_dict_with(|p| {
4890            p.push(Datum::String("k"));
4891            p.push(Datum::Null);
4892        });
4893        let map_null = row_map_null.unpack_first().unwrap_map();
4894
4895        assert_eq!(map_1.cmp(&map_null), Ordering::Less);
4896        assert_eq!(map_null.cmp(&map_1), Ordering::Greater);
4897    }
4898}