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