1use 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
44pub struct ArrowReader {
54 len: usize,
55 readers: Vec<ColReader>,
56}
57
58impl ArrowReader {
59 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 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 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 (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 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 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 (
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 (
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 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 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
524enum ArrayDims {
528 UInt8(UInt8Array),
529 Int32(Int32Array),
530}
531
532impl ArrayDims {
533 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
544enum 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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)] 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 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 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 rnd_row.packer();
1276
1277 reader.read(1, &mut rnd_row).unwrap();
1278 assert_eq!(&null_row, &rnd_row);
1279 }
1280
1281 #[mz_ore::test]
1285 #[cfg_attr(miri, ignore)] 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 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 #[mz_ore::test]
1339 #[cfg_attr(miri, ignore)] 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 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)] 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 dec128.append_value(1234);
1405 dec128.append_null();
1406 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 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 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 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)] 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 dec256.append_value(arrow::datatypes::i256::from(1234));
1454 dec256.append_null();
1455 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 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 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 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 #[mz_ore::test]
1492 #[cfg_attr(miri, ignore)] fn decimal_applies_destination_max_scale() {
1494 use mz_repr::adt::numeric::NumericMaxScale;
1495
1496 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 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)] 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 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 #[mz_ore::test]
1605 #[cfg_attr(miri, ignore)] 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 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 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 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}