Skip to main content

mz_arrow_util/
reader.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//! Reader for [`arrow`] data that outputs [`Row`]s.
11
12use std::sync::Arc;
13
14use anyhow::Context;
15use arrow::array::{
16    Array, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array, Decimal128Array,
17    Decimal256Array, FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array, Int8Array,
18    Int16Array, Int32Array, Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray,
19    IntervalYearMonthArray, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray,
20    MapArray, StringArray, StringViewArray, StructArray, Time32MillisecondArray, Time32SecondArray,
21    Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray,
22    TimestampNanosecondArray, TimestampSecondArray, UInt8Array, UInt16Array, UInt32Array,
23    UInt64Array,
24};
25use arrow::buffer::{NullBuffer, OffsetBuffer};
26use arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
27use chrono::{DateTime, NaiveTime};
28use dec::OrderedDecimal;
29use itertools::Itertools;
30use mz_ore::cast::CastFrom;
31use mz_repr::adt::array::ArrayDimension;
32use mz_repr::adt::date::Date;
33use mz_repr::adt::interval::Interval;
34use mz_repr::adt::jsonb::JsonbPacker;
35use mz_repr::adt::numeric::{Numeric, rescale};
36use mz_repr::adt::range::{Range, RangeLowerBound, RangeUpperBound};
37use mz_repr::adt::timestamp::CheckedTimestamp;
38use mz_repr::{Datum, RelationDesc, Row, RowPacker, SharedRow, SqlScalarType};
39use ordered_float::OrderedFloat;
40use uuid::Uuid;
41
42use crate::mask_nulls;
43
44/// Type that can read out of an [`arrow::array::StructArray`] and into a [`Row`], given a
45/// [`RelationDesc`].
46///
47/// The inverse of a [`crate::builder::ArrowBuilder`].
48///
49/// Note: When creating an [`ArrowReader`] we perform a "one-time downcast" of the children Arrays
50/// from the [`StructArray`], into `enum ColReader`s. This is a much more verbose approach than the
51/// alternative of downcasting from a `dyn arrow::array::Array` every time we read a [`Row`], but
52/// it is _much_ more performant.
53pub struct ArrowReader {
54    len: usize,
55    readers: Vec<ColReader>,
56}
57
58impl ArrowReader {
59    /// Create an [`ArrowReader`] validating that the provided [`RelationDesc`] and [`StructArray`]
60    /// have a matching schema.
61    ///
62    /// The [`RelationDesc`] and [`StructArray`] need to uphold the following to be a valid pair:
63    ///
64    /// * Same number of columns.
65    /// * Columns of all the same name.
66    /// * Columns of compatible types.
67    ///
68    /// TODO(cf2): Relax some of these restrictions by allowing users to map column names, omit
69    /// columns, perform some lightweight casting, and matching not on column name but column
70    /// position.
71    /// TODO(cf2): Allow specifying an optional `arrow::Schema` for extra metadata.
72    pub fn new(desc: &RelationDesc, array: StructArray) -> Result<Self, anyhow::Error> {
73        let inner_columns = array.columns();
74        let desc_columns = desc.typ().columns();
75
76        if inner_columns.len() != desc_columns.len() {
77            return Err(anyhow::anyhow!(
78                "wrong number of columns {} vs {}",
79                inner_columns.len(),
80                desc_columns.len()
81            ));
82        }
83
84        let mut readers = Vec::with_capacity(desc_columns.len());
85        for (col_name, col_type) in desc.iter() {
86            let column = array
87                .column_by_name(col_name)
88                .ok_or_else(|| anyhow::anyhow!("'{col_name}' not found"))?;
89            let reader = scalar_type_and_array_to_reader(&col_type.scalar_type, Arc::clone(column))
90                .context(col_name.clone())?;
91
92            readers.push(reader);
93        }
94
95        Ok(ArrowReader {
96            len: array.len(),
97            readers,
98        })
99    }
100
101    /// Read the value at `idx` into the provided `Row`.
102    pub fn read(&self, idx: usize, row: &mut Row) -> Result<(), anyhow::Error> {
103        let mut packer = row.packer();
104        for reader in &self.readers {
105            reader.read(idx, &mut packer).context(idx)?;
106        }
107        Ok(())
108    }
109
110    /// Read all of the values in this [`ArrowReader`] into `rows`.
111    pub fn read_all(&self, rows: &mut Vec<Row>) -> Result<usize, anyhow::Error> {
112        for idx in 0..self.len {
113            let mut row = Row::default();
114            self.read(idx, &mut row).context(idx)?;
115            rows.push(row);
116        }
117        Ok(self.len)
118    }
119}
120
121fn scalar_type_and_array_to_reader(
122    scalar_type: &SqlScalarType,
123    array: Arc<dyn Array>,
124) -> Result<ColReader, anyhow::Error> {
125    fn downcast_array<T: arrow::array::Array + Clone + 'static>(array: Arc<dyn Array>) -> T {
126        array
127            .as_any()
128            .downcast_ref::<T>()
129            .expect("checked DataType")
130            .clone()
131    }
132
133    match (scalar_type, array.data_type()) {
134        (SqlScalarType::Bool, DataType::Boolean) => {
135            Ok(ColReader::Boolean(downcast_array::<BooleanArray>(array)))
136        }
137        (SqlScalarType::Int16 | SqlScalarType::Int32 | SqlScalarType::Int64, DataType::Int8) => {
138            let array = downcast_array::<Int8Array>(array);
139            let cast: fn(i8) -> Datum<'static> = match scalar_type {
140                SqlScalarType::Int16 => |x| Datum::Int16(i16::cast_from(x)),
141                SqlScalarType::Int32 => |x| Datum::Int32(i32::cast_from(x)),
142                SqlScalarType::Int64 => |x| Datum::Int64(i64::cast_from(x)),
143                _ => unreachable!("checked above"),
144            };
145            Ok(ColReader::Int8 { array, cast })
146        }
147        (SqlScalarType::Int16, DataType::Int16) => {
148            Ok(ColReader::Int16(downcast_array::<Int16Array>(array)))
149        }
150        (SqlScalarType::Int32, DataType::Int32) => {
151            Ok(ColReader::Int32(downcast_array::<Int32Array>(array)))
152        }
153        (SqlScalarType::Int64, DataType::Int64) => {
154            Ok(ColReader::Int64(downcast_array::<Int64Array>(array)))
155        }
156        (
157            SqlScalarType::UInt16 | SqlScalarType::UInt32 | SqlScalarType::UInt64,
158            DataType::UInt8,
159        ) => {
160            let array = downcast_array::<UInt8Array>(array);
161            let cast: fn(u8) -> Datum<'static> = match scalar_type {
162                SqlScalarType::UInt16 => |x| Datum::UInt16(u16::cast_from(x)),
163                SqlScalarType::UInt32 => |x| Datum::UInt32(u32::cast_from(x)),
164                SqlScalarType::UInt64 => |x| Datum::UInt64(u64::cast_from(x)),
165                _ => unreachable!("checked above"),
166            };
167            Ok(ColReader::UInt8 { array, cast })
168        }
169        (SqlScalarType::UInt16, DataType::UInt16) => {
170            Ok(ColReader::UInt16(downcast_array::<UInt16Array>(array)))
171        }
172        // `oid` shares `Datum::UInt32`'s representation, so it decodes identically.
173        (SqlScalarType::UInt32 | SqlScalarType::Oid, DataType::UInt32) => {
174            Ok(ColReader::UInt32(downcast_array::<UInt32Array>(array)))
175        }
176        (SqlScalarType::UInt64, DataType::UInt64) => {
177            Ok(ColReader::UInt64(downcast_array::<UInt64Array>(array)))
178        }
179        (SqlScalarType::MzTimestamp, DataType::UInt64) => {
180            Ok(ColReader::MzTimestamp(downcast_array::<UInt64Array>(array)))
181        }
182        (SqlScalarType::Float32 | SqlScalarType::Float64, DataType::Float16) => {
183            let array = downcast_array::<Float16Array>(array);
184            let cast: fn(half::f16) -> Datum<'static> = match scalar_type {
185                SqlScalarType::Float32 => |x| Datum::Float32(OrderedFloat::from(x.to_f32())),
186                SqlScalarType::Float64 => |x| Datum::Float64(OrderedFloat::from(x.to_f64())),
187                _ => unreachable!("checked above"),
188            };
189            Ok(ColReader::Float16 { array, cast })
190        }
191        (SqlScalarType::Float32, DataType::Float32) => {
192            Ok(ColReader::Float32(downcast_array::<Float32Array>(array)))
193        }
194        (SqlScalarType::Float64, DataType::Float64) => {
195            Ok(ColReader::Float64(downcast_array::<Float64Array>(array)))
196        }
197        (SqlScalarType::Numeric { max_scale }, DataType::Decimal128(precision, scale)) => {
198            use num_traits::Pow;
199
200            let base = Numeric::from(10);
201            let scale = Numeric::from(*scale);
202            let scale_factor = base.pow(scale);
203
204            let precision = usize::cast_from(*precision);
205            // Don't use the context here, but make sure the precision is valid.
206            let mut ctx = dec::Context::<Numeric>::default();
207            ctx.set_precision(precision).map_err(|e| {
208                anyhow::anyhow!("invalid precision from Decimal128, {precision}, {e}")
209            })?;
210
211            let array = downcast_array::<Decimal128Array>(array);
212
213            Ok(ColReader::Decimal128 {
214                array,
215                scale_factor,
216                precision,
217                destination_max_scale: (*max_scale).map(|s| s.into_u8()),
218            })
219        }
220        (SqlScalarType::Numeric { max_scale }, DataType::Decimal256(precision, scale)) => {
221            use num_traits::Pow;
222
223            let base = Numeric::from(10);
224            let scale = Numeric::from(*scale);
225            let scale_factor = base.pow(scale);
226
227            let precision = usize::cast_from(*precision);
228            // Don't use the context here, but make sure the precision is valid.
229            let mut ctx = dec::Context::<Numeric>::default();
230            ctx.set_precision(precision).map_err(|e| {
231                anyhow::anyhow!("invalid precision from Decimal256, {precision}, {e}")
232            })?;
233
234            let array = downcast_array::<Decimal256Array>(array);
235
236            Ok(ColReader::Decimal256 {
237                array,
238                scale_factor,
239                precision,
240                destination_max_scale: (*max_scale).map(|s| s.into_u8()),
241            })
242        }
243        (SqlScalarType::Bytes, DataType::Binary) => {
244            Ok(ColReader::Binary(downcast_array::<BinaryArray>(array)))
245        }
246        (SqlScalarType::Bytes, DataType::LargeBinary) => {
247            let array = downcast_array::<LargeBinaryArray>(array);
248            Ok(ColReader::LargeBinary(array))
249        }
250        (SqlScalarType::Bytes, DataType::FixedSizeBinary(_)) => {
251            let array = downcast_array::<FixedSizeBinaryArray>(array);
252            Ok(ColReader::FixedSizeBinary(array))
253        }
254        (SqlScalarType::Bytes, DataType::BinaryView) => {
255            let array = downcast_array::<BinaryViewArray>(array);
256            Ok(ColReader::BinaryView(array))
257        }
258        (
259            SqlScalarType::Uuid,
260            DataType::Binary
261            | DataType::BinaryView
262            | DataType::LargeBinary
263            | DataType::FixedSizeBinary(_),
264        ) => {
265            let reader = scalar_type_and_array_to_reader(&SqlScalarType::Bytes, array)
266                .context("uuid reader")?;
267            Ok(ColReader::Uuid(Box::new(reader)))
268        }
269        // `char` and `varchar` store `Datum::String`, so they decode as strings.
270        // The writer emits them as Utf8/LargeUtf8 depending on length.
271        (
272            SqlScalarType::String | SqlScalarType::Char { .. } | SqlScalarType::VarChar { .. },
273            DataType::Utf8,
274        ) => Ok(ColReader::String(downcast_array::<StringArray>(array))),
275        (
276            SqlScalarType::String | SqlScalarType::Char { .. } | SqlScalarType::VarChar { .. },
277            DataType::LargeUtf8,
278        ) => {
279            let array = downcast_array::<LargeStringArray>(array);
280            Ok(ColReader::LargeString(array))
281        }
282        (
283            SqlScalarType::String | SqlScalarType::Char { .. } | SqlScalarType::VarChar { .. },
284            DataType::Utf8View,
285        ) => {
286            let array = downcast_array::<StringViewArray>(array);
287            Ok(ColReader::StringView(array))
288        }
289        (SqlScalarType::Jsonb, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View) => {
290            let reader = scalar_type_and_array_to_reader(&SqlScalarType::String, array)
291                .context("json reader")?;
292            Ok(ColReader::Jsonb(Box::new(reader)))
293        }
294        (SqlScalarType::Timestamp { .. }, DataType::Timestamp(TimeUnit::Second, None)) => {
295            let array = downcast_array::<TimestampSecondArray>(array);
296            Ok(ColReader::TimestampSecond(array))
297        }
298        (SqlScalarType::Timestamp { .. }, DataType::Timestamp(TimeUnit::Millisecond, None)) => {
299            let array = downcast_array::<TimestampMillisecondArray>(array);
300            Ok(ColReader::TimestampMillisecond(array))
301        }
302        (SqlScalarType::Timestamp { .. }, DataType::Timestamp(TimeUnit::Microsecond, None)) => {
303            let array = downcast_array::<TimestampMicrosecondArray>(array);
304            Ok(ColReader::TimestampMicrosecond(array))
305        }
306        (SqlScalarType::Timestamp { .. }, DataType::Timestamp(TimeUnit::Nanosecond, None)) => {
307            let array = downcast_array::<TimestampNanosecondArray>(array);
308            Ok(ColReader::TimestampNanosecond(array))
309        }
310        // A tz-aware timestamp array stores UTC-normalized instants, so the tz
311        // string is metadata we can ignore. The writer always emits microseconds.
312        (
313            SqlScalarType::TimestampTz { .. },
314            DataType::Timestamp(TimeUnit::Microsecond, Some(_)),
315        ) => {
316            let array = downcast_array::<TimestampMicrosecondArray>(array);
317            Ok(ColReader::TimestampTzMicrosecond(array))
318        }
319        (SqlScalarType::Date, DataType::Date32) => {
320            let array = downcast_array::<Date32Array>(array);
321            Ok(ColReader::Date32(array))
322        }
323        (SqlScalarType::Date, DataType::Date64) => {
324            let array = downcast_array::<Date64Array>(array);
325            Ok(ColReader::Date64(array))
326        }
327        (SqlScalarType::Time, DataType::Time32(TimeUnit::Second)) => {
328            let array = downcast_array::<Time32SecondArray>(array);
329            Ok(ColReader::Time32Seconds(array))
330        }
331        (SqlScalarType::Time, DataType::Time32(TimeUnit::Millisecond)) => {
332            let array = downcast_array::<Time32MillisecondArray>(array);
333            Ok(ColReader::Time32Milliseconds(array))
334        }
335        (SqlScalarType::Time, DataType::Time64(TimeUnit::Microsecond)) => {
336            let array = downcast_array::<Time64MicrosecondArray>(array);
337            Ok(ColReader::Time64Microseconds(array))
338        }
339        (
340            SqlScalarType::List {
341                element_type,
342                custom_id: _,
343            },
344            DataType::List(_),
345        ) => {
346            let array = downcast_array::<ListArray>(array);
347            let inner_decoder =
348                scalar_type_and_array_to_reader(element_type, Arc::clone(array.values()))
349                    .context("list")?;
350            Ok(ColReader::List {
351                offsets: array.offsets().clone(),
352                values: Box::new(inner_decoder),
353                nulls: array.nulls().cloned(),
354            })
355        }
356        (
357            SqlScalarType::List {
358                element_type,
359                custom_id: _,
360            },
361            DataType::LargeList(_),
362        ) => {
363            let array = downcast_array::<LargeListArray>(array);
364            let inner_decoder =
365                scalar_type_and_array_to_reader(element_type, Arc::clone(array.values()))
366                    .context("large list")?;
367            Ok(ColReader::LargeList {
368                offsets: array.offsets().clone(),
369                values: Box::new(inner_decoder),
370                nulls: array.nulls().cloned(),
371            })
372        }
373        (SqlScalarType::Array(element_type), DataType::Struct(_)) => {
374            // The builder encodes an array as a struct of an `items` list (the
375            // flat, row-major elements) and a `dimensions` count. Reverse that.
376            let struct_array = downcast_array::<StructArray>(array);
377
378            let items = struct_array
379                .column_by_name("items")
380                .ok_or_else(|| anyhow::anyhow!("array struct missing 'items' field"))?;
381            let items = items
382                .as_any()
383                .downcast_ref::<ListArray>()
384                .ok_or_else(|| anyhow::anyhow!("array 'items' field is not a List"))?;
385            let values = scalar_type_and_array_to_reader(element_type, Arc::clone(items.values()))
386                .context("array items")?;
387
388            let dims_col = struct_array
389                .column_by_name("dimensions")
390                .ok_or_else(|| anyhow::anyhow!("array struct missing 'dimensions' field"))?;
391            // The builder writes `dimensions` as UInt8. An Iceberg round-trip
392            // widens it to Int32 (Iceberg has no narrow integer types), so
393            // accept both.
394            let dims = match dims_col.data_type() {
395                DataType::UInt8 => {
396                    ArrayDims::UInt8(downcast_array::<UInt8Array>(Arc::clone(dims_col)))
397                }
398                DataType::Int32 => {
399                    ArrayDims::Int32(downcast_array::<Int32Array>(Arc::clone(dims_col)))
400                }
401                other => anyhow::bail!("unsupported array 'dimensions' type: {other:?}"),
402            };
403
404            Ok(ColReader::Array {
405                offsets: items.offsets().clone(),
406                values: Box::new(values),
407                dims,
408                nulls: struct_array.nulls().cloned(),
409            })
410        }
411        (
412            SqlScalarType::Record {
413                fields,
414                custom_id: _,
415            },
416            DataType::Struct(_),
417        ) => {
418            let record_array = downcast_array::<StructArray>(array);
419            let null_mask = record_array.nulls();
420
421            let mut decoders = Vec::with_capacity(fields.len());
422            for (name, typ) in fields.iter() {
423                let inner_array = record_array
424                    .column_by_name(name)
425                    .ok_or_else(|| anyhow::anyhow!("missing name '{name}'"))?;
426                let inner_array = mask_nulls(inner_array, null_mask);
427                let inner_decoder = scalar_type_and_array_to_reader(&typ.scalar_type, inner_array)
428                    .context(name.clone())?;
429
430                decoders.push(Box::new(inner_decoder));
431            }
432
433            Ok(ColReader::Record {
434                fields: decoders,
435                nulls: null_mask.cloned(),
436            })
437        }
438        (
439            SqlScalarType::Map {
440                value_type,
441                custom_id: _,
442            },
443            DataType::Map(_, _),
444        ) => {
445            let map_array = downcast_array::<MapArray>(array);
446            let keys = map_array
447                .keys()
448                .as_any()
449                .downcast_ref::<StringArray>()
450                .expect("map keys should be Utf8 strings")
451                .clone();
452            let values_reader =
453                scalar_type_and_array_to_reader(value_type, Arc::clone(map_array.values()))
454                    .context("map values")?;
455            Ok(ColReader::Map {
456                offsets: map_array.offsets().clone(),
457                keys,
458                values: Box::new(values_reader),
459                nulls: map_array.nulls().cloned(),
460            })
461        }
462        (SqlScalarType::Range { element_type }, DataType::Struct(_)) => {
463            let struct_array = downcast_array::<StructArray>(array);
464            let lower_array = struct_array
465                .column_by_name("lower")
466                .ok_or_else(|| anyhow::anyhow!("range struct missing 'lower' field"))?;
467            let upper_array = struct_array
468                .column_by_name("upper")
469                .ok_or_else(|| anyhow::anyhow!("range struct missing 'upper' field"))?;
470            let empty_array = struct_array
471                .column_by_name("empty")
472                .ok_or_else(|| anyhow::anyhow!("range struct missing 'empty' field"))?;
473            let lower_inclusive_array = struct_array
474                .column_by_name("lower_inclusive")
475                .ok_or_else(|| anyhow::anyhow!("range struct missing 'lower_inclusive' field"))?;
476            let upper_inclusive_array = struct_array
477                .column_by_name("upper_inclusive")
478                .ok_or_else(|| anyhow::anyhow!("range struct missing 'upper_inclusive' field"))?;
479
480            let lower_reader =
481                scalar_type_and_array_to_reader(element_type, Arc::clone(lower_array))
482                    .context("range lower")?;
483            let upper_reader =
484                scalar_type_and_array_to_reader(element_type, Arc::clone(upper_array))
485                    .context("range upper")?;
486
487            let empty = downcast_array::<BooleanArray>(Arc::clone(empty_array));
488            let lower_inclusive = downcast_array::<BooleanArray>(Arc::clone(lower_inclusive_array));
489            let upper_inclusive = downcast_array::<BooleanArray>(Arc::clone(upper_inclusive_array));
490
491            let lower_nulls = lower_array.nulls().cloned();
492            let upper_nulls = upper_array.nulls().cloned();
493
494            Ok(ColReader::Range {
495                lower: Box::new(lower_reader),
496                lower_nulls,
497                upper: Box::new(upper_reader),
498                upper_nulls,
499                lower_inclusive,
500                upper_inclusive,
501                empty,
502                nulls: struct_array.nulls().cloned(),
503            })
504        }
505        (SqlScalarType::Interval, DataType::Interval(IntervalUnit::YearMonth)) => {
506            Ok(ColReader::IntervalYearMonth(downcast_array::<
507                IntervalYearMonthArray,
508            >(array)))
509        }
510        (SqlScalarType::Interval, DataType::Interval(IntervalUnit::DayTime)) => {
511            Ok(ColReader::IntervalDayTime(downcast_array::<
512                IntervalDayTimeArray,
513            >(array)))
514        }
515        (SqlScalarType::Interval, DataType::Interval(IntervalUnit::MonthDayNano)) => {
516            Ok(ColReader::IntervalMonthDayNano(downcast_array::<
517                IntervalMonthDayNanoArray,
518            >(array)))
519        }
520        other => anyhow::bail!("unsupported: {other:?}"),
521    }
522}
523
524/// The `dimensions` field of an encoded array. The builder writes it as
525/// [`UInt8`](DataType::UInt8); an Iceberg round-trip widens it to
526/// [`Int32`](DataType::Int32).
527enum ArrayDims {
528    UInt8(UInt8Array),
529    Int32(Int32Array),
530}
531
532impl ArrayDims {
533    /// The number of dimensions of the array at `idx`.
534    fn ndims(&self, idx: usize) -> Result<u8, anyhow::Error> {
535        match self {
536            ArrayDims::UInt8(array) => Ok(array.value(idx)),
537            ArrayDims::Int32(array) => {
538                u8::try_from(array.value(idx)).context("array dimension count out of range")
539            }
540        }
541    }
542}
543
544/// A "downcasted" version of [`arrow::array::Array`] that supports reading [`Datum`]s.
545///
546/// Note: While this is fairly verbose, one-time "downcasting" to an enum is _much_ more performant
547/// than downcasting every time we read a [`Datum`].
548enum ColReader {
549    Boolean(arrow::array::BooleanArray),
550
551    Int8 {
552        array: arrow::array::Int8Array,
553        cast: fn(i8) -> Datum<'static>,
554    },
555    Int16(arrow::array::Int16Array),
556    Int32(arrow::array::Int32Array),
557    Int64(arrow::array::Int64Array),
558
559    UInt8 {
560        array: arrow::array::UInt8Array,
561        cast: fn(u8) -> Datum<'static>,
562    },
563    UInt16(arrow::array::UInt16Array),
564    UInt32(arrow::array::UInt32Array),
565    UInt64(arrow::array::UInt64Array),
566
567    Float16 {
568        array: arrow::array::Float16Array,
569        cast: fn(half::f16) -> Datum<'static>,
570    },
571    Float32(arrow::array::Float32Array),
572    Float64(arrow::array::Float64Array),
573
574    Decimal128 {
575        array: Decimal128Array,
576        scale_factor: Numeric,
577        precision: usize,
578        destination_max_scale: Option<u8>,
579    },
580    Decimal256 {
581        array: Decimal256Array,
582        scale_factor: Numeric,
583        precision: usize,
584        destination_max_scale: Option<u8>,
585    },
586
587    Binary(arrow::array::BinaryArray),
588    LargeBinary(arrow::array::LargeBinaryArray),
589    FixedSizeBinary(arrow::array::FixedSizeBinaryArray),
590    BinaryView(arrow::array::BinaryViewArray),
591    Uuid(Box<ColReader>),
592
593    String(arrow::array::StringArray),
594    LargeString(arrow::array::LargeStringArray),
595    StringView(arrow::array::StringViewArray),
596    Jsonb(Box<ColReader>),
597
598    TimestampSecond(arrow::array::TimestampSecondArray),
599    TimestampMillisecond(arrow::array::TimestampMillisecondArray),
600    TimestampMicrosecond(arrow::array::TimestampMicrosecondArray),
601    TimestampNanosecond(arrow::array::TimestampNanosecondArray),
602    TimestampTzMicrosecond(arrow::array::TimestampMicrosecondArray),
603
604    MzTimestamp(arrow::array::UInt64Array),
605
606    Date32(Date32Array),
607    Date64(Date64Array),
608
609    Time32Seconds(Time32SecondArray),
610    Time32Milliseconds(arrow::array::Time32MillisecondArray),
611    Time64Microseconds(arrow::array::Time64MicrosecondArray),
612
613    List {
614        offsets: OffsetBuffer<i32>,
615        values: Box<ColReader>,
616        nulls: Option<NullBuffer>,
617    },
618    LargeList {
619        offsets: OffsetBuffer<i64>,
620        values: Box<ColReader>,
621        nulls: Option<NullBuffer>,
622    },
623
624    Array {
625        offsets: OffsetBuffer<i32>,
626        values: Box<ColReader>,
627        dims: ArrayDims,
628        nulls: Option<NullBuffer>,
629    },
630
631    Record {
632        fields: Vec<Box<ColReader>>,
633        nulls: Option<NullBuffer>,
634    },
635
636    Map {
637        offsets: OffsetBuffer<i32>,
638        keys: StringArray,
639        values: Box<ColReader>,
640        nulls: Option<NullBuffer>,
641    },
642
643    Range {
644        lower: Box<ColReader>,
645        lower_nulls: Option<NullBuffer>,
646        upper: Box<ColReader>,
647        upper_nulls: Option<NullBuffer>,
648        lower_inclusive: BooleanArray,
649        upper_inclusive: BooleanArray,
650        empty: BooleanArray,
651        nulls: Option<NullBuffer>,
652    },
653
654    IntervalYearMonth(IntervalYearMonthArray),
655    IntervalDayTime(IntervalDayTimeArray),
656    IntervalMonthDayNano(IntervalMonthDayNanoArray),
657}
658
659impl ColReader {
660    fn read(&self, idx: usize, packer: &mut RowPacker) -> Result<(), anyhow::Error> {
661        let datum = match self {
662            ColReader::Boolean(array) => array
663                .is_valid(idx)
664                .then(|| array.value(idx))
665                .map(|x| if x { Datum::True } else { Datum::False }),
666            ColReader::Int8 { array, cast } => {
667                array.is_valid(idx).then(|| array.value(idx)).map(cast)
668            }
669            ColReader::Int16(array) => array
670                .is_valid(idx)
671                .then(|| array.value(idx))
672                .map(Datum::Int16),
673            ColReader::Int32(array) => array
674                .is_valid(idx)
675                .then(|| array.value(idx))
676                .map(Datum::Int32),
677            ColReader::Int64(array) => array
678                .is_valid(idx)
679                .then(|| array.value(idx))
680                .map(Datum::Int64),
681            ColReader::UInt8 { array, cast } => {
682                array.is_valid(idx).then(|| array.value(idx)).map(cast)
683            }
684            ColReader::UInt16(array) => array
685                .is_valid(idx)
686                .then(|| array.value(idx))
687                .map(Datum::UInt16),
688            ColReader::UInt32(array) => array
689                .is_valid(idx)
690                .then(|| array.value(idx))
691                .map(Datum::UInt32),
692            ColReader::UInt64(array) => array
693                .is_valid(idx)
694                .then(|| array.value(idx))
695                .map(Datum::UInt64),
696            ColReader::Float16 { array, cast } => {
697                array.is_valid(idx).then(|| array.value(idx)).map(cast)
698            }
699            ColReader::Float32(array) => array
700                .is_valid(idx)
701                .then(|| array.value(idx))
702                .map(|x| Datum::Float32(OrderedFloat(x))),
703            ColReader::Float64(array) => array
704                .is_valid(idx)
705                .then(|| array.value(idx))
706                .map(|x| Datum::Float64(OrderedFloat(x))),
707            ColReader::Decimal128 {
708                array,
709                scale_factor,
710                precision,
711                destination_max_scale,
712            } => array
713                .is_valid(idx)
714                .then(|| array.value(idx))
715                .map(|x| {
716                    // Create a Numeric from our i128 with precision.
717                    let mut ctx = dec::Context::<Numeric>::default();
718                    ctx.set_precision(*precision).expect("checked before");
719                    let mut num = ctx.from_i128(x);
720
721                    // Scale the number.
722                    ctx.div(&mut num, scale_factor);
723
724                    if let Some(destination_max_scale) = destination_max_scale {
725                        rescale(&mut num, *destination_max_scale)?;
726                    }
727
728                    Ok::<_, anyhow::Error>(Datum::Numeric(OrderedDecimal(num)))
729                })
730                .transpose()?,
731            ColReader::Decimal256 {
732                array,
733                scale_factor,
734                precision,
735                destination_max_scale,
736            } => array
737                .is_valid(idx)
738                .then(|| array.value(idx))
739                .map(|x| {
740                    let s = x.to_string();
741
742                    // Parse a i256 from it's String representation.
743                    //
744                    // TODO(cf3): See if we can add support for 256-bit numbers to the `dec` crate.
745                    let mut ctx = dec::Context::<Numeric>::default();
746                    ctx.set_precision(*precision).expect("checked before");
747                    let mut num = ctx
748                        .parse(s)
749                        .map_err(|e| anyhow::anyhow!("decimal out of range: {e}"))?;
750
751                    // Scale the number.
752                    ctx.div(&mut num, scale_factor);
753
754                    if let Some(destination_max_scale) = destination_max_scale {
755                        rescale(&mut num, *destination_max_scale)?;
756                    }
757
758                    Ok::<_, anyhow::Error>(Datum::Numeric(OrderedDecimal(num)))
759                })
760                .transpose()?,
761            ColReader::Binary(array) => array
762                .is_valid(idx)
763                .then(|| array.value(idx))
764                .map(Datum::Bytes),
765            ColReader::LargeBinary(array) => array
766                .is_valid(idx)
767                .then(|| array.value(idx))
768                .map(Datum::Bytes),
769            ColReader::FixedSizeBinary(array) => array
770                .is_valid(idx)
771                .then(|| array.value(idx))
772                .map(Datum::Bytes),
773            ColReader::BinaryView(array) => array
774                .is_valid(idx)
775                .then(|| array.value(idx))
776                .map(Datum::Bytes),
777            ColReader::Uuid(reader) => {
778                // First read a binary value into a temp row, and later parse that as UUID into our
779                // actual Row Packer.
780                let mut temp_row = SharedRow::get();
781                reader.read(idx, &mut temp_row.packer()).context("uuid")?;
782                let slice = match temp_row.unpack_first() {
783                    Datum::Bytes(slice) => slice,
784                    Datum::Null => {
785                        packer.push(Datum::Null);
786                        return Ok(());
787                    }
788                    other => anyhow::bail!("expected String, found {other:?}"),
789                };
790
791                let uuid = Uuid::from_slice(slice).context("parsing uuid")?;
792                Some(Datum::Uuid(uuid))
793            }
794            ColReader::String(array) => array
795                .is_valid(idx)
796                .then(|| array.value(idx))
797                .map(Datum::String),
798            ColReader::LargeString(array) => array
799                .is_valid(idx)
800                .then(|| array.value(idx))
801                .map(Datum::String),
802            ColReader::StringView(array) => array
803                .is_valid(idx)
804                .then(|| array.value(idx))
805                .map(Datum::String),
806            ColReader::Jsonb(reader) => {
807                // First read a string value into a temp row, and later parse that as JSON into our
808                // actual Row Packer.
809                let mut temp_row = SharedRow::get();
810                reader.read(idx, &mut temp_row.packer()).context("jsonb")?;
811                let value = match temp_row.unpack_first() {
812                    Datum::String(value) => value,
813                    Datum::Null => {
814                        packer.push(Datum::Null);
815                        return Ok(());
816                    }
817                    other => anyhow::bail!("expected String, found {other:?}"),
818                };
819
820                JsonbPacker::new(packer)
821                    .pack_str(value)
822                    .context("roundtrip json")?;
823
824                // Return early because we've already packed the necessasry Datums.
825                return Ok(());
826            }
827            ColReader::TimestampSecond(array) => array
828                .is_valid(idx)
829                .then(|| array.value(idx))
830                .map(|secs| {
831                    let dt = DateTime::from_timestamp(secs, 0)
832                        .ok_or_else(|| anyhow::anyhow!("invalid timestamp seconds {secs}"))?;
833                    let dt = CheckedTimestamp::from_timestamplike(dt.naive_utc())
834                        .context("TimestampSeconds")?;
835                    Ok::<_, anyhow::Error>(Datum::Timestamp(dt))
836                })
837                .transpose()?,
838            ColReader::TimestampMillisecond(array) => array
839                .is_valid(idx)
840                .then(|| array.value(idx))
841                .map(|millis| {
842                    let dt = DateTime::from_timestamp_millis(millis).ok_or_else(|| {
843                        anyhow::anyhow!("invalid timestamp milliseconds {millis}")
844                    })?;
845                    let dt = CheckedTimestamp::from_timestamplike(dt.naive_utc())
846                        .context("TimestampMillis")?;
847                    Ok::<_, anyhow::Error>(Datum::Timestamp(dt))
848                })
849                .transpose()?,
850            ColReader::TimestampMicrosecond(array) => array
851                .is_valid(idx)
852                .then(|| array.value(idx))
853                .map(|micros| {
854                    let dt = DateTime::from_timestamp_micros(micros).ok_or_else(|| {
855                        anyhow::anyhow!("invalid timestamp microseconds {micros}")
856                    })?;
857                    let dt = CheckedTimestamp::from_timestamplike(dt.naive_utc())
858                        .context("TimestampMicros")?;
859                    Ok::<_, anyhow::Error>(Datum::Timestamp(dt))
860                })
861                .transpose()?,
862            ColReader::TimestampNanosecond(array) => array
863                .is_valid(idx)
864                .then(|| array.value(idx))
865                .map(|nanos| {
866                    let dt = DateTime::from_timestamp_nanos(nanos);
867                    let dt = CheckedTimestamp::from_timestamplike(dt.naive_utc())
868                        .context("TimestampNanos")?;
869                    Ok::<_, anyhow::Error>(Datum::Timestamp(dt))
870                })
871                .transpose()?,
872            ColReader::TimestampTzMicrosecond(array) => array
873                .is_valid(idx)
874                .then(|| array.value(idx))
875                .map(|micros| {
876                    let dt = DateTime::from_timestamp_micros(micros).ok_or_else(|| {
877                        anyhow::anyhow!("invalid timestamptz microseconds {micros}")
878                    })?;
879                    let dt =
880                        CheckedTimestamp::from_timestamplike(dt).context("TimestampTzMicros")?;
881                    Ok::<_, anyhow::Error>(Datum::TimestampTz(dt))
882                })
883                .transpose()?,
884            ColReader::MzTimestamp(array) => array
885                .is_valid(idx)
886                .then(|| array.value(idx))
887                .map(|v| Datum::MzTimestamp(mz_repr::Timestamp::from(v))),
888            ColReader::Date32(array) => array
889                .is_valid(idx)
890                .then(|| array.value(idx))
891                .map(|unix_days| {
892                    let date = Date::from_unix_epoch(unix_days).context("date32")?;
893                    Ok::<_, anyhow::Error>(Datum::Date(date))
894                })
895                .transpose()?,
896            ColReader::Date64(array) => array
897                .is_valid(idx)
898                .then(|| array.value(idx))
899                .map(|unix_millis| {
900                    let date = DateTime::from_timestamp_millis(unix_millis)
901                        .ok_or_else(|| anyhow::anyhow!("invalid Date64 {unix_millis}"))?;
902                    let unix_epoch = DateTime::from_timestamp(0, 0)
903                        .expect("UNIX epoch")
904                        .date_naive();
905                    let delta = date.date_naive().signed_duration_since(unix_epoch);
906                    let days: i32 = delta.num_days().try_into().context("date64")?;
907                    let date = Date::from_unix_epoch(days).context("date64")?;
908                    Ok::<_, anyhow::Error>(Datum::Date(date))
909                })
910                .transpose()?,
911            ColReader::Time32Seconds(array) => array
912                .is_valid(idx)
913                .then(|| array.value(idx))
914                .map(|secs| {
915                    let usecs: u32 = secs.try_into().context("time32 seconds")?;
916                    let time = NaiveTime::from_num_seconds_from_midnight_opt(usecs, 0)
917                        .ok_or_else(|| anyhow::anyhow!("invalid Time32 Seconds {secs}"))?;
918                    Ok::<_, anyhow::Error>(Datum::Time(time))
919                })
920                .transpose()?,
921            ColReader::Time32Milliseconds(array) => array
922                .is_valid(idx)
923                .then(|| array.value(idx))
924                .map(|millis| {
925                    let umillis: u32 = millis.try_into().context("time32 milliseconds")?;
926                    let usecs = umillis / 1000;
927                    let unanos = (umillis % 1000).saturating_mul(1_000_000);
928                    let time = NaiveTime::from_num_seconds_from_midnight_opt(usecs, unanos)
929                        .ok_or_else(|| anyhow::anyhow!("invalid Time32 Milliseconds {umillis}"))?;
930                    Ok::<_, anyhow::Error>(Datum::Time(time))
931                })
932                .transpose()?,
933            ColReader::Time64Microseconds(array) => array
934                .is_valid(idx)
935                .then(|| array.value(idx))
936                .map(|micros| {
937                    // Inverse of the builder's `secs * 1_000_000 + nanos / 1_000`.
938                    let secs: u32 = (micros / 1_000_000).try_into().context("time64 seconds")?;
939                    let nanos: u32 = (micros % 1_000_000)
940                        .try_into()
941                        .map(|us: u32| us.saturating_mul(1_000))
942                        .context("time64 microseconds")?;
943                    let time = NaiveTime::from_num_seconds_from_midnight_opt(secs, nanos)
944                        .ok_or_else(|| anyhow::anyhow!("invalid Time64 Microseconds {micros}"))?;
945                    Ok::<_, anyhow::Error>(Datum::Time(time))
946                })
947                .transpose()?,
948            ColReader::List {
949                offsets,
950                values,
951                nulls,
952            } => {
953                let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
954                if !is_valid {
955                    packer.push(Datum::Null);
956                    return Ok(());
957                }
958
959                let start: usize = offsets[idx].try_into().context("list start offset")?;
960                let end: usize = offsets[idx + 1].try_into().context("list end offset")?;
961
962                packer
963                    .push_list_with(|packer| {
964                        for idx in start..end {
965                            values.read(idx, packer)?;
966                        }
967                        Ok::<_, anyhow::Error>(())
968                    })
969                    .context("pack list")?;
970
971                // Return early because we've already packed the necessasry Datums.
972                return Ok(());
973            }
974            ColReader::LargeList {
975                offsets,
976                values,
977                nulls,
978            } => {
979                let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
980                if !is_valid {
981                    packer.push(Datum::Null);
982                    return Ok(());
983                }
984
985                let start: usize = offsets[idx].try_into().context("list start offset")?;
986                let end: usize = offsets[idx + 1].try_into().context("list end offset")?;
987
988                packer
989                    .push_list_with(|packer| {
990                        for idx in start..end {
991                            values.read(idx, packer)?;
992                        }
993                        Ok::<_, anyhow::Error>(())
994                    })
995                    .context("pack list")?;
996
997                // Return early because we've already packed the necessasry Datums.
998                return Ok(());
999            }
1000            ColReader::Array {
1001                offsets,
1002                values,
1003                dims,
1004                nulls,
1005            } => {
1006                let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
1007                if !is_valid {
1008                    packer.push(Datum::Null);
1009                    return Ok(());
1010                }
1011
1012                let start: usize = offsets[idx].try_into().context("array start offset")?;
1013                let end: usize = offsets[idx + 1].try_into().context("array end offset")?;
1014                let nelements = end - start;
1015
1016                // The encoding stores only the dimension count and the flat,
1017                // row-major elements (see the builder), so per-dimension extents
1018                // are recoverable only for 0- and 1-dimensional arrays. A
1019                // higher-dimensional array's extents cannot be reconstructed, so
1020                // reject it rather than guess a shape.
1021                let ndims = dims.ndims(idx)?;
1022                if ndims > 1 {
1023                    anyhow::bail!(
1024                        "cannot decode {ndims}-dimensional array from parquet: the encoding \
1025                         records only the dimension count, not per-dimension extents"
1026                    );
1027                }
1028                let one_dim = [ArrayDimension {
1029                    lower_bound: 1,
1030                    length: nelements,
1031                }];
1032                let array_dims: &[ArrayDimension] = if ndims == 0 { &[] } else { &one_dim };
1033
1034                // SAFETY: the closure returns exactly the number of elements it
1035                // pushes (`end - start`).
1036                unsafe {
1037                    packer.push_array_with_unchecked(array_dims, |packer| {
1038                        for idx in start..end {
1039                            values.read(idx, packer)?;
1040                        }
1041                        Ok::<_, anyhow::Error>(end - start)
1042                    })
1043                }
1044                .context("pack array")?;
1045
1046                // Return early because we've already packed the necessasry Datums.
1047                return Ok(());
1048            }
1049            ColReader::Record { fields, nulls } => {
1050                let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
1051                if !is_valid {
1052                    packer.push(Datum::Null);
1053                    return Ok(());
1054                }
1055
1056                packer
1057                    .push_list_with(|packer| {
1058                        for field in fields {
1059                            field.read(idx, packer)?;
1060                        }
1061                        Ok::<_, anyhow::Error>(())
1062                    })
1063                    .context("pack record")?;
1064
1065                // Return early because we've already packed the necessasry Datums.
1066                return Ok(());
1067            }
1068            ColReader::Map {
1069                offsets,
1070                keys,
1071                values,
1072                nulls,
1073            } => {
1074                let is_non_null = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
1075                if !is_non_null {
1076                    packer.push(Datum::Null);
1077                    return Ok(());
1078                }
1079
1080                let start: usize = offsets[idx].try_into().context("map start offset")?;
1081                let end: usize = offsets[idx + 1].try_into().context("map end offset")?;
1082
1083                // Arrow's MapArray doesn't guarantee that keys are in sorted order, but Materialize's
1084                // Datum::Map does, so we need to sort the keys here before packing them, or else
1085                // many assumptions will break.
1086                let mut kv_sorted = (start..end)
1087                    .map(|i| (keys.value(i), i))
1088                    .sorted_by_key(|(k, _)| *k)
1089                    .peekable();
1090
1091                packer
1092                    .push_dict_with(|packer| {
1093                        while let Some((key, i)) = kv_sorted.next() {
1094                            // Parquet docs state that if there are duplicate keys, the last value
1095                            // should be used, so skip duplicates here.
1096                            //
1097                            // sorted_by_key is a stable sort, so entries with duplicate keys will
1098                            // maintain their original order, and we can pick the last one here.
1099                            if let Some((next_key, _)) = kv_sorted.peek() {
1100                                if key == *next_key {
1101                                    continue;
1102                                }
1103                            }
1104                            packer.push(Datum::String(key));
1105                            values.read(i, packer)?;
1106                        }
1107                        Ok::<_, anyhow::Error>(())
1108                    })
1109                    .context("pack map")?;
1110
1111                // Return early because we've already packed the necessary Datums.
1112                return Ok(());
1113            }
1114            ColReader::Range {
1115                lower,
1116                lower_nulls,
1117                upper,
1118                upper_nulls,
1119                lower_inclusive,
1120                upper_inclusive,
1121                empty,
1122                nulls,
1123            } => {
1124                let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
1125                if !is_valid {
1126                    packer.push(Datum::Null);
1127                    return Ok(());
1128                }
1129
1130                if empty.value(idx) {
1131                    packer.push(Datum::Range(Range { inner: None }));
1132                    return Ok(());
1133                }
1134
1135                let lower_is_infinite = lower_nulls
1136                    .as_ref()
1137                    .map(|n| !n.is_valid(idx))
1138                    .unwrap_or(false);
1139                let upper_is_infinite = upper_nulls
1140                    .as_ref()
1141                    .map(|n| !n.is_valid(idx))
1142                    .unwrap_or(false);
1143
1144                // Read finite bounds into owned Rows that live for the rest of
1145                // this block, so the Datums we borrow out of them stay valid
1146                // for the `push_range` call below.
1147                let lower_row = if lower_is_infinite {
1148                    None
1149                } else {
1150                    let mut temp = SharedRow::get();
1151                    lower.read(idx, &mut temp.packer())?;
1152                    Some(temp.clone())
1153                };
1154                let upper_row = if upper_is_infinite {
1155                    None
1156                } else {
1157                    let mut temp = SharedRow::get();
1158                    upper.read(idx, &mut temp.packer())?;
1159                    Some(temp.clone())
1160                };
1161
1162                let lower_bound = RangeLowerBound {
1163                    inclusive: lower_inclusive.value(idx),
1164                    bound: lower_row.as_ref().map(|row| row.unpack_first()),
1165                };
1166                let upper_bound = RangeUpperBound {
1167                    inclusive: upper_inclusive.value(idx),
1168                    bound: upper_row.as_ref().map(|row| row.unpack_first()),
1169                };
1170
1171                // Use `push_range` (not `push_range_with`) so the range is
1172                // canonicalized before being packed. Parquet files authored by
1173                // external engines may encode discrete ranges in non-canonical
1174                // form (e.g. `[1,10]` for int4range, which MZ stores as
1175                // `[1,11)`); without canonicalization those rows would not
1176                // compare or hash equal to MZ-constructed values.
1177                packer
1178                    .push_range(Range::new(Some((lower_bound, upper_bound))))
1179                    .context("pack range")?;
1180
1181                return Ok(());
1182            }
1183            ColReader::IntervalYearMonth(array) => array
1184                .is_valid(idx)
1185                .then(|| array.value(idx))
1186                .map(|months| Datum::Interval(Interval::new(months, 0, 0))),
1187            ColReader::IntervalDayTime(array) => {
1188                array.is_valid(idx).then(|| array.value(idx)).map(|v| {
1189                    let micros = i64::from(v.milliseconds) * 1_000;
1190                    Datum::Interval(Interval::new(0, v.days, micros))
1191                })
1192            }
1193            ColReader::IntervalMonthDayNano(array) => {
1194                array.is_valid(idx).then(|| array.value(idx)).map(|v| {
1195                    let micros = v.nanoseconds / 1_000;
1196                    Datum::Interval(Interval::new(v.months, v.days, micros))
1197                })
1198            }
1199        };
1200
1201        match datum {
1202            Some(d) => packer.push(d),
1203            None => packer.push(Datum::Null),
1204        }
1205
1206        Ok(())
1207    }
1208}
1209
1210#[cfg(test)]
1211mod tests {
1212    use arrow::datatypes::Field;
1213    use mz_ore::collections::CollectionExt;
1214
1215    use super::*;
1216
1217    #[mz_ore::test]
1218    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1219    fn smoketest_reader() {
1220        let desc = RelationDesc::builder()
1221            .with_column("bool", SqlScalarType::Bool.nullable(true))
1222            .with_column("int4", SqlScalarType::Int32.nullable(true))
1223            .with_column("uint8", SqlScalarType::UInt64.nullable(true))
1224            .with_column("float32", SqlScalarType::Float32.nullable(true))
1225            .with_column("string", SqlScalarType::String.nullable(true))
1226            .with_column("bytes", SqlScalarType::Bytes.nullable(true))
1227            .with_column("uuid", SqlScalarType::Uuid.nullable(true))
1228            .with_column("json", SqlScalarType::Jsonb.nullable(true))
1229            .with_column(
1230                "list",
1231                SqlScalarType::List {
1232                    element_type: Box::new(SqlScalarType::UInt32),
1233                    custom_id: None,
1234                }
1235                .nullable(true),
1236            )
1237            .finish();
1238
1239        let mut og_row = Row::default();
1240        let mut packer = og_row.packer();
1241
1242        packer.extend([
1243            Datum::True,
1244            Datum::Int32(42),
1245            Datum::UInt64(10000),
1246            Datum::Float32(OrderedFloat::from(-1.1f32)),
1247            Datum::String("hello world"),
1248            Datum::Bytes(b"1010101"),
1249            Datum::Uuid(uuid::Uuid::new_v4()),
1250        ]);
1251        JsonbPacker::new(&mut packer)
1252            .pack_serde_json(
1253                serde_json::json!({"code": 200, "email": "space_monkey@materialize.com"}),
1254            )
1255            .expect("failed to pack JSON");
1256        packer.push_list([Datum::UInt32(200), Datum::UInt32(300)]);
1257
1258        let null_row = Row::pack(vec![Datum::Null; 9]);
1259
1260        // Encode our data with our ArrowBuilder.
1261        let mut builder = crate::builder::ArrowBuilder::new(&desc, 2, 46).unwrap();
1262        builder.add_row(&og_row).unwrap();
1263        builder.add_row(&null_row).unwrap();
1264        let record_batch = builder.to_record_batch().unwrap();
1265
1266        // Decode our data!
1267        let reader =
1268            ArrowReader::new(&desc, arrow::array::StructArray::from(record_batch)).unwrap();
1269        let mut rnd_row = Row::default();
1270
1271        reader.read(0, &mut rnd_row).unwrap();
1272        assert_eq!(&og_row, &rnd_row);
1273
1274        // Create a packer to clear the row alloc.
1275        rnd_row.packer();
1276
1277        reader.read(1, &mut rnd_row).unwrap();
1278        assert_eq!(&null_row, &rnd_row);
1279    }
1280
1281    /// Regression for SS-341: `COPY TO PARQUET` can write these scalar types,
1282    /// so `COPY FROM PARQUET` must read them back. Each must survive a
1283    /// builder -> reader round-trip.
1284    #[mz_ore::test]
1285    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1286    fn smoketest_extra_scalar_types() {
1287        let desc = RelationDesc::builder()
1288            .with_column("oid", SqlScalarType::Oid.nullable(true))
1289            .with_column("time", SqlScalarType::Time.nullable(true))
1290            .with_column(
1291                "timestamptz",
1292                SqlScalarType::TimestampTz { precision: None }.nullable(true),
1293            )
1294            .with_column("char", SqlScalarType::Char { length: None }.nullable(true))
1295            .with_column(
1296                "varchar",
1297                SqlScalarType::VarChar { max_length: None }.nullable(true),
1298            )
1299            .with_column("mz_timestamp", SqlScalarType::MzTimestamp.nullable(true))
1300            .finish();
1301
1302        let tstz = CheckedTimestamp::from_timestamplike(
1303            DateTime::from_timestamp(1_600_000_000, 0).expect("valid timestamp"),
1304        )
1305        .expect("valid CheckedTimestamp");
1306
1307        let og_row = Row::pack(vec![
1308            Datum::UInt32(42),
1309            // Sub-second component at microsecond resolution: the writer stores
1310            // Time64 microseconds, so nanoseconds would not round-trip.
1311            Datum::Time(NaiveTime::from_hms_micro_opt(12, 34, 56, 789_012).unwrap()),
1312            Datum::TimestampTz(tstz),
1313            Datum::String("abc"),
1314            Datum::String("hello world"),
1315            Datum::MzTimestamp(mz_repr::Timestamp::from(123_456_u64)),
1316        ]);
1317        let null_row = Row::pack(vec![Datum::Null; 6]);
1318
1319        let mut builder = crate::builder::ArrowBuilder::new(&desc, 2, 64).unwrap();
1320        builder.add_row(&og_row).unwrap();
1321        builder.add_row(&null_row).unwrap();
1322        let record_batch = builder.to_record_batch().unwrap();
1323
1324        let reader = ArrowReader::new(&desc, StructArray::from(record_batch)).unwrap();
1325
1326        let mut got = Row::default();
1327        reader.read(0, &mut got).unwrap();
1328        assert_eq!(&og_row, &got, "values did not round-trip");
1329
1330        got.packer();
1331        reader.read(1, &mut got).unwrap();
1332        assert_eq!(&null_row, &got, "NULLs did not round-trip");
1333    }
1334
1335    /// Regression: an array column must survive a builder -> reader round-trip.
1336    /// The builder encodes an array as a struct of `{items, dimensions}`; the
1337    /// reader must reverse that back into a `Datum::Array`.
1338    #[mz_ore::test]
1339    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1340    fn smoketest_array() {
1341        let desc = RelationDesc::builder()
1342            .with_column(
1343                "arr",
1344                SqlScalarType::Array(Box::new(SqlScalarType::Int32)).nullable(true),
1345            )
1346            .finish();
1347
1348        let mut row_1d = Row::default();
1349        row_1d
1350            .packer()
1351            .try_push_array(
1352                &[ArrayDimension {
1353                    lower_bound: 1,
1354                    length: 3,
1355                }],
1356                [Datum::Int32(1), Datum::Null, Datum::Int32(3)],
1357            )
1358            .unwrap();
1359
1360        let mut row_empty = Row::default();
1361        row_empty
1362            .packer()
1363            .try_push_array(&[], std::iter::empty::<Datum>())
1364            .unwrap();
1365
1366        let row_null = Row::pack(vec![Datum::Null]);
1367
1368        // Encode with the builder, decode with the reader.
1369        let mut builder = crate::builder::ArrowBuilder::new(&desc, 3, 128).unwrap();
1370        builder.add_row(&row_1d).unwrap();
1371        builder.add_row(&row_empty).unwrap();
1372        builder.add_row(&row_null).unwrap();
1373        let record_batch = builder.to_record_batch().unwrap();
1374
1375        let reader = ArrowReader::new(&desc, StructArray::from(record_batch)).unwrap();
1376        let mut got = Row::default();
1377
1378        reader.read(0, &mut got).unwrap();
1379        assert_eq!(got, row_1d, "1-D array did not round-trip");
1380
1381        got.packer();
1382        reader.read(1, &mut got).unwrap();
1383        assert_eq!(got, row_empty, "empty array did not round-trip");
1384
1385        got.packer();
1386        reader.read(2, &mut got).unwrap();
1387        assert_eq!(got, row_null, "NULL array did not round-trip");
1388    }
1389
1390    #[mz_ore::test]
1391    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1392    fn smoketest_decimal128() {
1393        let desc = RelationDesc::builder()
1394            .with_column(
1395                "a",
1396                SqlScalarType::Numeric { max_scale: None }.nullable(true),
1397            )
1398            .finish();
1399
1400        let mut dec128 = arrow::array::Decimal128Builder::new();
1401        dec128 = dec128.with_precision_and_scale(12, 3).unwrap();
1402
1403        // 1.234
1404        dec128.append_value(1234);
1405        dec128.append_null();
1406        // 100000000.009
1407        dec128.append_value(100000000009);
1408
1409        let dec128 = dec128.finish();
1410        #[allow(clippy::as_conversions)]
1411        let batch = StructArray::from(vec![(
1412            Arc::new(Field::new("a", dec128.data_type().clone(), true)),
1413            Arc::new(dec128) as arrow::array::ArrayRef,
1414        )]);
1415
1416        // Decode our data!
1417        let reader = ArrowReader::new(&desc, batch).unwrap();
1418        let mut rnd_row = Row::default();
1419
1420        reader.read(0, &mut rnd_row).unwrap();
1421        let num = rnd_row.into_element().unwrap_numeric();
1422        assert_eq!(num.0, Numeric::from(1.234f64));
1423
1424        // Create a packer to clear the row alloc.
1425        rnd_row.packer();
1426
1427        reader.read(1, &mut rnd_row).unwrap();
1428        let num = rnd_row.into_element();
1429        assert_eq!(num, Datum::Null);
1430
1431        // Create a packer to clear the row alloc.
1432        rnd_row.packer();
1433
1434        reader.read(2, &mut rnd_row).unwrap();
1435        let num = rnd_row.into_element().unwrap_numeric();
1436        assert_eq!(num.0, Numeric::from(100000000.009f64));
1437    }
1438
1439    #[mz_ore::test]
1440    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1441    fn smoketest_decimal256() {
1442        let desc = RelationDesc::builder()
1443            .with_column(
1444                "a",
1445                SqlScalarType::Numeric { max_scale: None }.nullable(true),
1446            )
1447            .finish();
1448
1449        let mut dec256 = arrow::array::Decimal256Builder::new();
1450        dec256 = dec256.with_precision_and_scale(12, 3).unwrap();
1451
1452        // 1.234
1453        dec256.append_value(arrow::datatypes::i256::from(1234));
1454        dec256.append_null();
1455        // 100000000.009
1456        dec256.append_value(arrow::datatypes::i256::from(100000000009i64));
1457
1458        let dec256 = dec256.finish();
1459        #[allow(clippy::as_conversions)]
1460        let batch = StructArray::from(vec![(
1461            Arc::new(Field::new("a", dec256.data_type().clone(), true)),
1462            Arc::new(dec256) as arrow::array::ArrayRef,
1463        )]);
1464
1465        // Decode our data!
1466        let reader = ArrowReader::new(&desc, batch).unwrap();
1467        let mut rnd_row = Row::default();
1468
1469        reader.read(0, &mut rnd_row).unwrap();
1470        let num = rnd_row.into_element().unwrap_numeric();
1471        assert_eq!(num.0, Numeric::from(1.234f64));
1472
1473        // Create a packer to clear the row alloc.
1474        rnd_row.packer();
1475
1476        reader.read(1, &mut rnd_row).unwrap();
1477        let num = rnd_row.into_element();
1478        assert_eq!(num, Datum::Null);
1479
1480        // Create a packer to clear the row alloc.
1481        rnd_row.packer();
1482
1483        reader.read(2, &mut rnd_row).unwrap();
1484        let num = rnd_row.into_element().unwrap_numeric();
1485        assert_eq!(num.0, Numeric::from(100000000.009f64));
1486    }
1487
1488    /// Regression test for SS-193: when the destination column declares a
1489    /// `max_scale`, the reader should round the decoded value to that scale
1490    /// rather than preserving the source file's scale.
1491    #[mz_ore::test]
1492    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1493    fn decimal_applies_destination_max_scale() {
1494        use mz_repr::adt::numeric::NumericMaxScale;
1495
1496        // Destination column is numeric(10, 2): scale 2.
1497        let desc = RelationDesc::builder()
1498            .with_column(
1499                "a",
1500                SqlScalarType::Numeric {
1501                    max_scale: Some(NumericMaxScale::try_from(2_i64).unwrap()),
1502                }
1503                .nullable(true),
1504            )
1505            .finish();
1506
1507        let expected = Numeric::from(10.45f64);
1508
1509        // The source files carry scale-3 values (10.447) that don't round
1510        // evenly to the destination's scale 2.
1511        let mut dec128 = arrow::array::Decimal128Builder::new();
1512        dec128 = dec128.with_precision_and_scale(12, 3).unwrap();
1513        dec128.append_value(10447);
1514        let dec128 = dec128.finish();
1515        #[allow(clippy::as_conversions)]
1516        let batch128 = StructArray::from(vec![(
1517            Arc::new(Field::new("a", dec128.data_type().clone(), true)),
1518            Arc::new(dec128) as arrow::array::ArrayRef,
1519        )]);
1520
1521        let reader = ArrowReader::new(&desc, batch128).unwrap();
1522        let mut rnd_row = Row::default();
1523        reader.read(0, &mut rnd_row).unwrap();
1524        let num = rnd_row.into_element().unwrap_numeric();
1525        assert_eq!(num.0, expected, "Decimal128 did not round to max_scale");
1526
1527        let mut dec256 = arrow::array::Decimal256Builder::new();
1528        dec256 = dec256.with_precision_and_scale(12, 3).unwrap();
1529        dec256.append_value(arrow::datatypes::i256::from(10447));
1530        let dec256 = dec256.finish();
1531        #[allow(clippy::as_conversions)]
1532        let batch256 = StructArray::from(vec![(
1533            Arc::new(Field::new("a", dec256.data_type().clone(), true)),
1534            Arc::new(dec256) as arrow::array::ArrayRef,
1535        )]);
1536
1537        let reader = ArrowReader::new(&desc, batch256).unwrap();
1538        let mut rnd_row = Row::default();
1539        reader.read(0, &mut rnd_row).unwrap();
1540        let num = rnd_row.into_element().unwrap_numeric();
1541        assert_eq!(num.0, expected, "Decimal256 did not round to max_scale");
1542    }
1543
1544    #[mz_ore::test]
1545    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1546    fn decimal_rescale_to_higher_scale_overflows() {
1547        use mz_repr::adt::numeric::NumericMaxScale;
1548
1549        let value: i128 = 10_i128.pow(33) - 1;
1550
1551        let build_batch = || {
1552            let mut dec128 = arrow::array::Decimal128Builder::new();
1553            dec128 = dec128.with_precision_and_scale(38, 0).unwrap();
1554            dec128.append_value(value);
1555            let dec128 = dec128.finish();
1556            #[allow(clippy::as_conversions)]
1557            StructArray::from(vec![(
1558                Arc::new(Field::new("a", dec128.data_type().clone(), true)),
1559                Arc::new(dec128) as arrow::array::ArrayRef,
1560            )])
1561        };
1562
1563        let desc_with_scale = |scale: i64| {
1564            RelationDesc::builder()
1565                .with_column(
1566                    "a",
1567                    SqlScalarType::Numeric {
1568                        max_scale: Some(NumericMaxScale::try_from(scale).unwrap()),
1569                    }
1570                    .nullable(true),
1571                )
1572                .finish()
1573        };
1574
1575        // Test working case with supportable scale then broken case.
1576        let reader = ArrowReader::new(&desc_with_scale(2), build_batch()).unwrap();
1577        let mut row = Row::default();
1578        reader
1579            .read(0, &mut row)
1580            .expect("value must decode at destination scale 2");
1581        let num = row.into_element().unwrap_numeric();
1582        let mut expected = Numeric::try_from(value).unwrap();
1583        rescale(&mut expected, 2).unwrap();
1584        assert_eq!(num.0, expected, "value did not rescale to scale 2");
1585
1586        let reader = ArrowReader::new(&desc_with_scale(8), build_batch()).unwrap();
1587        let mut row = Row::default();
1588        let err = reader
1589            .read(0, &mut row)
1590            .expect_err("value must overflow at destination scale 8");
1591
1592        assert!(
1593            format!("{err:#}").contains("exceed maximum precision"),
1594            "unexpected error: {err:#}",
1595        );
1596    }
1597
1598    /// Regression test for database-issues#11330: when a Parquet file authored
1599    /// by an external engine encodes a discrete range in non-canonical form
1600    /// (e.g. `[1,10]` for `int4range`), the reader must canonicalize it to MZ's
1601    /// internal form (`[1,11)`). Otherwise rows ingested via `COPY FROM PARQUET`
1602    /// don't compare or hash equal to logically-identical rows constructed
1603    /// inside MZ.
1604    #[mz_ore::test]
1605    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decContextDefault` on OS `linux`
1606    fn range_canonicalizes_noncanonical_input() {
1607        use arrow::array::ArrayRef;
1608        use arrow::datatypes::DataType;
1609        use mz_repr::adt::range::{Range, RangeLowerBound, RangeUpperBound};
1610
1611        let desc = RelationDesc::builder()
1612            .with_column(
1613                "r",
1614                SqlScalarType::Range {
1615                    element_type: Box::new(SqlScalarType::Int32),
1616                }
1617                .nullable(true),
1618            )
1619            .finish();
1620
1621        // Build a range StructArray containing two non-canonical encodings:
1622        //   row 0: `[1,10]`  -> canonicalizes to `[1,11)`
1623        //   row 1: `(5,15)`  -> canonicalizes to `[6,15)`
1624        let lower = Int32Array::from(vec![Some(1), Some(5)]);
1625        let upper = Int32Array::from(vec![Some(10), Some(15)]);
1626        let lower_inclusive = BooleanArray::from(vec![true, false]);
1627        let upper_inclusive = BooleanArray::from(vec![true, false]);
1628        let empty = BooleanArray::from(vec![false, false]);
1629
1630        #[allow(clippy::as_conversions)]
1631        let range_fields: Vec<(Arc<Field>, ArrayRef)> = vec![
1632            (
1633                Arc::new(Field::new("lower", DataType::Int32, true)),
1634                Arc::new(lower) as ArrayRef,
1635            ),
1636            (
1637                Arc::new(Field::new("upper", DataType::Int32, true)),
1638                Arc::new(upper) as ArrayRef,
1639            ),
1640            (
1641                Arc::new(Field::new("lower_inclusive", DataType::Boolean, false)),
1642                Arc::new(lower_inclusive) as ArrayRef,
1643            ),
1644            (
1645                Arc::new(Field::new("upper_inclusive", DataType::Boolean, false)),
1646                Arc::new(upper_inclusive) as ArrayRef,
1647            ),
1648            (
1649                Arc::new(Field::new("empty", DataType::Boolean, false)),
1650                Arc::new(empty) as ArrayRef,
1651            ),
1652        ];
1653        let range_struct = StructArray::from(range_fields);
1654
1655        #[allow(clippy::as_conversions)]
1656        let batch = StructArray::from(vec![(
1657            Arc::new(Field::new("r", range_struct.data_type().clone(), true)),
1658            Arc::new(range_struct) as ArrayRef,
1659        )]);
1660
1661        let reader = ArrowReader::new(&desc, batch).unwrap();
1662
1663        // Row 0: `[1,10]` -> `[1,11)`
1664        let mut got = Row::default();
1665        reader.read(0, &mut got).unwrap();
1666        let mut want = Row::default();
1667        want.packer()
1668            .push_range(Range::new(Some((
1669                RangeLowerBound {
1670                    inclusive: true,
1671                    bound: Some(Datum::Int32(1)),
1672                },
1673                RangeUpperBound {
1674                    inclusive: true,
1675                    bound: Some(Datum::Int32(10)),
1676                },
1677            ))))
1678            .unwrap();
1679        assert_eq!(got, want, "row 0: [1,10] should canonicalize to [1,11)");
1680
1681        // Row 1: `(5,15)` -> `[6,15)`
1682        let mut got = Row::default();
1683        reader.read(1, &mut got).unwrap();
1684        let mut want = Row::default();
1685        want.packer()
1686            .push_range(Range::new(Some((
1687                RangeLowerBound {
1688                    inclusive: false,
1689                    bound: Some(Datum::Int32(5)),
1690                },
1691                RangeUpperBound {
1692                    inclusive: false,
1693                    bound: Some(Datum::Int32(15)),
1694                },
1695            ))))
1696            .unwrap();
1697        assert_eq!(got, want, "row 1: (5,15) should canonicalize to [6,15)");
1698    }
1699}