1use std::fmt::Debug;
15use std::ops::AddAssign;
16use std::sync::Arc;
17
18use anyhow::anyhow;
19use arrow::array::{
20 Array, ArrayBuilder, ArrayRef, BinaryArray, BinaryBuilder, BooleanArray, BooleanBufferBuilder,
21 BooleanBuilder, FixedSizeBinaryArray, FixedSizeBinaryBuilder, Float32Array, Float32Builder,
22 Float64Array, Float64Builder, Int16Array, Int16Builder, Int32Array, Int32Builder, Int64Array,
23 Int64Builder, ListArray, ListBuilder, MapArray, StringArray, StringBuilder, StructArray,
24 UInt8Array, UInt8Builder, UInt16Array, UInt16Builder, UInt32Array, UInt32Builder, UInt64Array,
25 UInt64Builder, make_array,
26};
27use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
28use arrow::datatypes::{DataType, Field, Fields, ToByteSlice};
29use bytes::{BufMut, Bytes};
30use chrono::Timelike;
31use dec::{Context, Decimal, OrderedDecimal};
32use itertools::{EitherOrBoth, Itertools};
33use mz_ore::assert_none;
34use mz_ore::cast::CastFrom;
35use mz_persist_types::Codec;
36use mz_persist_types::arrow::ArrayOrd;
37use mz_persist_types::columnar::{ColumnDecoder, ColumnEncoder, FixedSizeCodec, Schema};
38use mz_persist_types::stats::{
39 ColumnNullStats, ColumnStatKinds, ColumnarStats, ColumnarStatsBuilder, FixedSizeBytesStatsKind,
40 OptionStats, PrimitiveStats, StructStats,
41};
42use mz_proto::chrono::ProtoNaiveTime;
43use mz_proto::{ProtoType, RustType, TryFromProtoError};
44use prost::Message;
45use uuid::Uuid;
46
47use crate::adt::array::{ArrayDimension, PackedArrayDimension};
48use crate::adt::date::Date;
49use crate::adt::datetime::PackedNaiveTime;
50use crate::adt::interval::PackedInterval;
51use crate::adt::jsonb::{JsonbPacker, JsonbRef};
52use crate::adt::mz_acl_item::{PackedAclItem, PackedMzAclItem};
53use crate::adt::numeric::{NUMERIC_DATUM_MAX_PRECISION, Numeric, PackedNumeric};
54use crate::adt::range::{Range, RangeInner, RangeLowerBound, RangeUpperBound};
55use crate::adt::timestamp::{CheckedTimestamp, PackedNaiveDateTime};
56use crate::row::proto_datum::DatumType;
57use crate::row::{
58 ProtoArray, ProtoArrayDimension, ProtoDatum, ProtoDatumOther, ProtoDict, ProtoDictElement,
59 ProtoNumeric, ProtoRange, ProtoRangeInner, ProtoRow,
60};
61use crate::stats::{fixed_stats_from_column, numeric_stats_from_column, stats_for_json};
62use crate::{Datum, ProtoRelationDesc, RelationDesc, Row, RowPacker, SqlScalarType, Timestamp};
63
64#[allow(clippy::as_conversions)]
69mod fixed_binary_sizes {
70 use super::*;
71
72 pub const TIME_FIXED_BYTES: i32 = PackedNaiveTime::SIZE as i32;
73 pub const TIMESTAMP_FIXED_BYTES: i32 = PackedNaiveDateTime::SIZE as i32;
74 pub const INTERVAL_FIXED_BYTES: i32 = PackedInterval::SIZE as i32;
75 pub const ACL_ITEM_FIXED_BYTES: i32 = PackedAclItem::SIZE as i32;
76 pub const _MZ_ACL_ITEM_FIXED_BYTES: i32 = PackedMzAclItem::SIZE as i32;
77 pub const ARRAY_DIMENSION_FIXED_BYTES: i32 = PackedArrayDimension::SIZE as i32;
78
79 pub const UUID_FIXED_BYTES: i32 = 16;
80 static_assertions::const_assert_eq!(UUID_FIXED_BYTES as usize, std::mem::size_of::<Uuid>());
81}
82use fixed_binary_sizes::*;
83
84pub fn preserves_order(scalar_type: &SqlScalarType) -> bool {
88 match scalar_type {
89 SqlScalarType::Bool
91 | SqlScalarType::Int16
92 | SqlScalarType::Int32
93 | SqlScalarType::Int64
94 | SqlScalarType::UInt16
95 | SqlScalarType::UInt32
96 | SqlScalarType::UInt64
97 | SqlScalarType::Date
98 | SqlScalarType::Time
99 | SqlScalarType::Timestamp { .. }
100 | SqlScalarType::TimestampTz { .. }
101 | SqlScalarType::Interval
102 | SqlScalarType::Bytes
103 | SqlScalarType::String
104 | SqlScalarType::Uuid
105 | SqlScalarType::MzTimestamp
106 | SqlScalarType::MzAclItem
107 | SqlScalarType::AclItem => true,
108 SqlScalarType::Record { fields, .. } => fields
110 .iter()
111 .all(|(_, field_type)| preserves_order(&field_type.scalar_type)),
112 SqlScalarType::Float32 | SqlScalarType::Float64 => false,
115 SqlScalarType::Numeric { .. } => false,
118 SqlScalarType::PgLegacyChar
121 | SqlScalarType::PgLegacyName
122 | SqlScalarType::Char { .. }
123 | SqlScalarType::VarChar { .. }
124 | SqlScalarType::Jsonb
125 | SqlScalarType::Array(_)
126 | SqlScalarType::List { .. }
127 | SqlScalarType::Oid
128 | SqlScalarType::Map { .. }
129 | SqlScalarType::RegProc
130 | SqlScalarType::RegType
131 | SqlScalarType::RegClass
132 | SqlScalarType::Int2Vector
133 | SqlScalarType::Range { .. } => false,
134 }
135}
136
137#[derive(Debug)]
139struct DatumEncoder {
140 nullable: bool,
141 encoder: DatumColumnEncoder,
142}
143
144impl DatumEncoder {
145 fn goodbytes(&self) -> usize {
146 self.encoder.goodbytes()
147 }
148
149 fn push(&mut self, datum: Datum) {
150 assert!(
151 !datum.is_null() || self.nullable,
152 "tried pushing Null into non-nullable column"
153 );
154 self.encoder.push(datum);
155 }
156
157 fn push_invalid(&mut self) {
158 self.encoder.push_invalid();
159 }
160
161 fn finish(self) -> ArrayRef {
162 self.encoder.finish()
163 }
164}
165
166#[derive(Debug)]
172enum DatumColumnEncoder {
173 Bool(BooleanBuilder),
174 U8(UInt8Builder),
175 U16(UInt16Builder),
176 U32(UInt32Builder),
177 U64(UInt64Builder),
178 I16(Int16Builder),
179 I32(Int32Builder),
180 I64(Int64Builder),
181 F32(Float32Builder),
182 F64(Float64Builder),
183 Numeric {
184 binary_values: BinaryBuilder,
186 approx_values: Float64Builder,
188 numeric_context: Context<Numeric>,
190 },
191 Bytes(BinaryBuilder),
192 String(StringBuilder),
193 Date(Int32Builder),
194 Time(FixedSizeBinaryBuilder),
195 Timestamp(FixedSizeBinaryBuilder),
196 TimestampTz(FixedSizeBinaryBuilder),
197 MzTimestamp(UInt64Builder),
198 Interval(FixedSizeBinaryBuilder),
199 Uuid(FixedSizeBinaryBuilder),
200 AclItem(FixedSizeBinaryBuilder),
201 MzAclItem(BinaryBuilder),
202 Range(BinaryBuilder),
203 Jsonb {
210 offsets: Vec<i32>,
212 buf: Vec<u8>,
214 nulls: Option<BooleanBufferBuilder>,
216 },
217 Array {
218 dims: ListBuilder<FixedSizeBinaryBuilder>,
220 val_lengths: Vec<usize>,
222 vals: Box<DatumColumnEncoder>,
224 nulls: Option<BooleanBufferBuilder>,
226 },
227 List {
228 lengths: Vec<usize>,
230 values: Box<DatumColumnEncoder>,
232 nulls: Option<BooleanBufferBuilder>,
234 },
235 Map {
236 lengths: Vec<usize>,
238 keys: StringBuilder,
240 vals: Box<DatumColumnEncoder>,
242 nulls: Option<BooleanBufferBuilder>,
244 },
245 Record {
246 fields: Vec<DatumEncoder>,
248 nulls: Option<BooleanBufferBuilder>,
250 length: usize,
252 },
253 RecordEmpty(BooleanBuilder),
258}
259
260impl DatumColumnEncoder {
261 fn goodbytes(&self) -> usize {
262 match self {
263 DatumColumnEncoder::Bool(a) => a.len(),
264 DatumColumnEncoder::U8(a) => a.values_slice().to_byte_slice().len(),
265 DatumColumnEncoder::U16(a) => a.values_slice().to_byte_slice().len(),
266 DatumColumnEncoder::U32(a) => a.values_slice().to_byte_slice().len(),
267 DatumColumnEncoder::U64(a) => a.values_slice().to_byte_slice().len(),
268 DatumColumnEncoder::I16(a) => a.values_slice().to_byte_slice().len(),
269 DatumColumnEncoder::I32(a) => a.values_slice().to_byte_slice().len(),
270 DatumColumnEncoder::I64(a) => a.values_slice().to_byte_slice().len(),
271 DatumColumnEncoder::F32(a) => a.values_slice().to_byte_slice().len(),
272 DatumColumnEncoder::F64(a) => a.values_slice().to_byte_slice().len(),
273 DatumColumnEncoder::Numeric {
274 binary_values,
275 approx_values,
276 ..
277 } => {
278 binary_values.values_slice().len()
279 + approx_values.values_slice().to_byte_slice().len()
280 }
281 DatumColumnEncoder::Bytes(a) => a.values_slice().len(),
282 DatumColumnEncoder::String(a) => a.values_slice().len(),
283 DatumColumnEncoder::Date(a) => a.values_slice().to_byte_slice().len(),
284 DatumColumnEncoder::Time(a) => a.len() * PackedNaiveTime::SIZE,
285 DatumColumnEncoder::Timestamp(a) => a.len() * PackedNaiveDateTime::SIZE,
286 DatumColumnEncoder::TimestampTz(a) => a.len() * PackedNaiveDateTime::SIZE,
287 DatumColumnEncoder::MzTimestamp(a) => a.values_slice().to_byte_slice().len(),
288 DatumColumnEncoder::Interval(a) => a.len() * PackedInterval::SIZE,
289 DatumColumnEncoder::Uuid(a) => a.len() * size_of::<Uuid>(),
290 DatumColumnEncoder::AclItem(a) => a.len() * PackedAclItem::SIZE,
291 DatumColumnEncoder::MzAclItem(a) => a.values_slice().len(),
292 DatumColumnEncoder::Range(a) => a.values_slice().len(),
293 DatumColumnEncoder::Jsonb { buf, .. } => buf.len(),
294 DatumColumnEncoder::Array { dims, vals, .. } => {
295 dims.len() * PackedArrayDimension::SIZE + vals.goodbytes()
296 }
297 DatumColumnEncoder::List { values, .. } => values.goodbytes(),
298 DatumColumnEncoder::Map { keys, vals, .. } => {
299 keys.values_slice().len() + vals.goodbytes()
300 }
301 DatumColumnEncoder::Record { fields, .. } => fields.iter().map(|f| f.goodbytes()).sum(),
302 DatumColumnEncoder::RecordEmpty(a) => a.len(),
303 }
304 }
305
306 fn push<'e, 'd>(&'e mut self, datum: Datum<'d>) {
307 match (self, datum) {
308 (DatumColumnEncoder::Bool(bool_builder), Datum::True) => {
309 bool_builder.append_value(true)
310 }
311 (DatumColumnEncoder::Bool(bool_builder), Datum::False) => {
312 bool_builder.append_value(false)
313 }
314 (DatumColumnEncoder::U8(builder), Datum::UInt8(val)) => builder.append_value(val),
315 (DatumColumnEncoder::U16(builder), Datum::UInt16(val)) => builder.append_value(val),
316 (DatumColumnEncoder::U32(builder), Datum::UInt32(val)) => builder.append_value(val),
317 (DatumColumnEncoder::U64(builder), Datum::UInt64(val)) => builder.append_value(val),
318 (DatumColumnEncoder::I16(builder), Datum::Int16(val)) => builder.append_value(val),
319 (DatumColumnEncoder::I32(builder), Datum::Int32(val)) => builder.append_value(val),
320 (DatumColumnEncoder::I64(builder), Datum::Int64(val)) => builder.append_value(val),
321 (DatumColumnEncoder::F32(builder), Datum::Float32(val)) => builder.append_value(*val),
322 (DatumColumnEncoder::F64(builder), Datum::Float64(val)) => builder.append_value(*val),
323 (
324 DatumColumnEncoder::Numeric {
325 approx_values,
326 binary_values,
327 numeric_context,
328 },
329 Datum::Numeric(val),
330 ) => {
331 let float_approx = numeric_context.try_into_f64(val.0).unwrap_or_else(|_| {
332 numeric_context.clear_status();
333 if val.0.is_negative() {
334 f64::NEG_INFINITY
335 } else {
336 f64::INFINITY
337 }
338 });
339 let packed = PackedNumeric::from_value(val.0);
340
341 approx_values.append_value(float_approx);
342 binary_values.append_value(packed.as_bytes());
343 }
344 (DatumColumnEncoder::String(builder), Datum::String(val)) => builder.append_value(val),
345 (DatumColumnEncoder::Bytes(builder), Datum::Bytes(val)) => builder.append_value(val),
346 (DatumColumnEncoder::Date(builder), Datum::Date(val)) => {
347 builder.append_value(val.pg_epoch_days())
348 }
349 (DatumColumnEncoder::Time(builder), Datum::Time(val)) => {
350 let packed = PackedNaiveTime::from_value(val);
351 builder
352 .append_value(packed.as_bytes())
353 .expect("known correct size");
354 }
355 (DatumColumnEncoder::Timestamp(builder), Datum::Timestamp(val)) => {
356 let packed = PackedNaiveDateTime::from_value(val.to_naive());
357 builder
358 .append_value(packed.as_bytes())
359 .expect("known correct size");
360 }
361 (DatumColumnEncoder::TimestampTz(builder), Datum::TimestampTz(val)) => {
362 let packed = PackedNaiveDateTime::from_value(val.to_naive());
363 builder
364 .append_value(packed.as_bytes())
365 .expect("known correct size");
366 }
367 (DatumColumnEncoder::MzTimestamp(builder), Datum::MzTimestamp(val)) => {
368 builder.append_value(val.into());
369 }
370 (DatumColumnEncoder::Interval(builder), Datum::Interval(val)) => {
371 let packed = PackedInterval::from_value(val);
372 builder
373 .append_value(packed.as_bytes())
374 .expect("known correct size");
375 }
376 (DatumColumnEncoder::Uuid(builder), Datum::Uuid(val)) => builder
377 .append_value(val.as_bytes())
378 .expect("known correct size"),
379 (DatumColumnEncoder::AclItem(builder), Datum::AclItem(val)) => {
380 let packed = PackedAclItem::from_value(val);
381 builder
382 .append_value(packed.as_bytes())
383 .expect("known correct size");
384 }
385 (DatumColumnEncoder::MzAclItem(builder), Datum::MzAclItem(val)) => {
386 let packed = PackedMzAclItem::from_value(val);
387 builder.append_value(packed.as_bytes());
388 }
389 (DatumColumnEncoder::Range(builder), d @ Datum::Range(_)) => {
390 let proto = ProtoDatum::from(d);
391 let bytes = proto.encode_to_vec();
392 builder.append_value(&bytes);
393 }
394 (
395 DatumColumnEncoder::Jsonb {
396 offsets,
397 buf,
398 nulls,
399 },
400 d @ Datum::JsonNull
401 | d @ Datum::True
402 | d @ Datum::False
403 | d @ Datum::Numeric(_)
404 | d @ Datum::String(_)
405 | d @ Datum::List(_)
406 | d @ Datum::Map(_),
407 ) => {
408 let mut buf = buf;
410 let json = JsonbRef::from_datum(d);
411
412 json.to_writer(&mut buf)
414 .expect("failed to serialize Datum to jsonb");
415 let offset: i32 = buf.len().try_into().expect("wrote more than 4GB of JSON");
416 offsets.push(offset);
417
418 if let Some(nulls) = nulls {
419 nulls.append(true);
420 }
421 }
422 (
423 DatumColumnEncoder::Array {
424 dims,
425 val_lengths,
426 vals,
427 nulls,
428 },
429 Datum::Array(array),
430 ) => {
431 for dimension in array.dims() {
433 let packed = PackedArrayDimension::from_value(dimension);
434 dims.values()
435 .append_value(packed.as_bytes())
436 .expect("known correct size");
437 }
438 dims.append(true);
439
440 let mut count = 0;
442 for datum in array.elements() {
443 count += 1;
444 vals.push(datum);
445 }
446 val_lengths.push(count);
447
448 if let Some(nulls) = nulls {
449 nulls.append(true);
450 }
451 }
452 (
453 DatumColumnEncoder::List {
454 lengths,
455 values,
456 nulls,
457 },
458 Datum::List(list),
459 ) => {
460 let mut count = 0;
461 for datum in list {
462 count += 1;
463 values.push(datum);
464 }
465 lengths.push(count);
466
467 if let Some(nulls) = nulls {
468 nulls.append(true);
469 }
470 }
471 (
472 DatumColumnEncoder::Map {
473 lengths,
474 keys,
475 vals,
476 nulls,
477 },
478 Datum::Map(map),
479 ) => {
480 let mut count = 0;
481 for (key, datum) in &map {
482 count += 1;
483 keys.append_value(key);
484 vals.push(datum);
485 }
486 lengths.push(count);
487
488 if let Some(nulls) = nulls {
489 nulls.append(true);
490 }
491 }
492 (
493 DatumColumnEncoder::Record {
494 fields,
495 nulls,
496 length,
497 },
498 Datum::List(records),
499 ) => {
500 let mut count = 0;
501 for (datum, encoder) in records.into_iter().zip_eq(fields.iter_mut()) {
503 count += 1;
504 encoder.push(datum);
505 }
506 assert_eq!(count, fields.len());
507
508 length.add_assign(1);
509 if let Some(nulls) = nulls.as_mut() {
510 nulls.append(true);
511 }
512 }
513 (DatumColumnEncoder::RecordEmpty(builder), Datum::List(records)) => {
514 assert_none!(records.into_iter().next());
515 builder.append_value(true);
516 }
517 (encoder, Datum::Null) => encoder.push_invalid(),
518 (encoder, datum) => panic!("can't encode {datum:?} into {encoder:?}"),
519 }
520 }
521
522 fn push_invalid(&mut self) {
523 match self {
524 DatumColumnEncoder::Bool(builder) => builder.append_null(),
525 DatumColumnEncoder::U8(builder) => builder.append_null(),
526 DatumColumnEncoder::U16(builder) => builder.append_null(),
527 DatumColumnEncoder::U32(builder) => builder.append_null(),
528 DatumColumnEncoder::U64(builder) => builder.append_null(),
529 DatumColumnEncoder::I16(builder) => builder.append_null(),
530 DatumColumnEncoder::I32(builder) => builder.append_null(),
531 DatumColumnEncoder::I64(builder) => builder.append_null(),
532 DatumColumnEncoder::F32(builder) => builder.append_null(),
533 DatumColumnEncoder::F64(builder) => builder.append_null(),
534 DatumColumnEncoder::Numeric {
535 approx_values,
536 binary_values,
537 numeric_context: _,
538 } => {
539 approx_values.append_null();
540 binary_values.append_null();
541 }
542 DatumColumnEncoder::String(builder) => builder.append_null(),
543 DatumColumnEncoder::Bytes(builder) => builder.append_null(),
544 DatumColumnEncoder::Date(builder) => builder.append_null(),
545 DatumColumnEncoder::Time(builder) => builder.append_null(),
546 DatumColumnEncoder::Timestamp(builder) => builder.append_null(),
547 DatumColumnEncoder::TimestampTz(builder) => builder.append_null(),
548 DatumColumnEncoder::MzTimestamp(builder) => builder.append_null(),
549 DatumColumnEncoder::Interval(builder) => builder.append_null(),
550 DatumColumnEncoder::Uuid(builder) => builder.append_null(),
551 DatumColumnEncoder::AclItem(builder) => builder.append_null(),
552 DatumColumnEncoder::MzAclItem(builder) => builder.append_null(),
553 DatumColumnEncoder::Range(builder) => builder.append_null(),
554 DatumColumnEncoder::Jsonb {
555 offsets,
556 buf: _,
557 nulls,
558 } => {
559 let nulls = nulls.get_or_insert_with(|| {
560 let mut buf = BooleanBufferBuilder::new(offsets.len());
561 buf.append_n(offsets.len() - 1, true);
563 buf
564 });
565
566 offsets.push(offsets.last().copied().unwrap_or(0));
567 nulls.append(false);
568 }
569 DatumColumnEncoder::Array {
570 dims,
571 val_lengths,
572 vals: _,
573 nulls,
574 } => {
575 let nulls = nulls.get_or_insert_with(|| {
576 let mut buf = BooleanBufferBuilder::new(dims.len() + 1);
577 buf.append_n(dims.len(), true);
578 buf
579 });
580 dims.append_null();
581
582 val_lengths.push(0);
583 nulls.append(false);
584 }
585 DatumColumnEncoder::List {
586 lengths,
587 values: _,
588 nulls,
589 } => {
590 let nulls = nulls.get_or_insert_with(|| {
591 let mut buf = BooleanBufferBuilder::new(lengths.len() + 1);
592 buf.append_n(lengths.len(), true);
593 buf
594 });
595
596 lengths.push(0);
597 nulls.append(false);
598 }
599 DatumColumnEncoder::Map {
600 lengths,
601 keys: _,
602 vals: _,
603 nulls,
604 } => {
605 let nulls = nulls.get_or_insert_with(|| {
606 let mut buf = BooleanBufferBuilder::new(lengths.len() + 1);
607 buf.append_n(lengths.len(), true);
608 buf
609 });
610
611 lengths.push(0);
612 nulls.append(false);
613 }
614 DatumColumnEncoder::Record {
615 fields,
616 nulls,
617 length,
618 } => {
619 let nulls = nulls.get_or_insert_with(|| {
620 let mut buf = BooleanBufferBuilder::new(*length + 1);
621 buf.append_n(*length, true);
622 buf
623 });
624 nulls.append(false);
625 length.add_assign(1);
626
627 for field in fields {
628 field.push_invalid();
629 }
630 }
631 DatumColumnEncoder::RecordEmpty(builder) => builder.append_null(),
632 }
633 }
634
635 fn finish(self) -> ArrayRef {
636 match self {
637 DatumColumnEncoder::Bool(mut builder) => {
638 let array = builder.finish();
639 Arc::new(array)
640 }
641 DatumColumnEncoder::U8(mut builder) => {
642 let array = builder.finish();
643 Arc::new(array)
644 }
645 DatumColumnEncoder::U16(mut builder) => {
646 let array = builder.finish();
647 Arc::new(array)
648 }
649 DatumColumnEncoder::U32(mut builder) => {
650 let array = builder.finish();
651 Arc::new(array)
652 }
653 DatumColumnEncoder::U64(mut builder) => {
654 let array = builder.finish();
655 Arc::new(array)
656 }
657 DatumColumnEncoder::I16(mut builder) => {
658 let array = builder.finish();
659 Arc::new(array)
660 }
661 DatumColumnEncoder::I32(mut builder) => {
662 let array = builder.finish();
663 Arc::new(array)
664 }
665 DatumColumnEncoder::I64(mut builder) => {
666 let array = builder.finish();
667 Arc::new(array)
668 }
669 DatumColumnEncoder::F32(mut builder) => {
670 let array = builder.finish();
671 Arc::new(array)
672 }
673 DatumColumnEncoder::F64(mut builder) => {
674 let array = builder.finish();
675 Arc::new(array)
676 }
677 DatumColumnEncoder::Numeric {
678 mut approx_values,
679 mut binary_values,
680 numeric_context: _,
681 } => {
682 let approx_array = approx_values.finish();
683 let binary_array = binary_values.finish();
684
685 assert_eq!(approx_array.len(), binary_array.len());
686 debug_assert_eq!(approx_array.logical_nulls(), binary_array.logical_nulls());
688
689 let fields = Fields::from(vec![
690 Field::new("approx", approx_array.data_type().clone(), true),
691 Field::new("binary", binary_array.data_type().clone(), true),
692 ]);
693 let nulls = approx_array.logical_nulls();
694 let array = StructArray::new(
695 fields,
696 vec![Arc::new(approx_array), Arc::new(binary_array)],
697 nulls,
698 );
699 Arc::new(array)
700 }
701 DatumColumnEncoder::String(mut builder) => {
702 let array = builder.finish();
703 Arc::new(array)
704 }
705 DatumColumnEncoder::Bytes(mut builder) => {
706 let array = builder.finish();
707 Arc::new(array)
708 }
709 DatumColumnEncoder::Date(mut builder) => {
710 let array = builder.finish();
711 Arc::new(array)
712 }
713 DatumColumnEncoder::Time(mut builder) => {
714 let array = builder.finish();
715 Arc::new(array)
716 }
717 DatumColumnEncoder::Timestamp(mut builder) => {
718 let array = builder.finish();
719 Arc::new(array)
720 }
721 DatumColumnEncoder::TimestampTz(mut builder) => {
722 let array = builder.finish();
723 Arc::new(array)
724 }
725 DatumColumnEncoder::MzTimestamp(mut builder) => {
726 let array = builder.finish();
727 Arc::new(array)
728 }
729 DatumColumnEncoder::Interval(mut builder) => {
730 let array = builder.finish();
731 Arc::new(array)
732 }
733 DatumColumnEncoder::Uuid(mut builder) => {
734 let array = builder.finish();
735 Arc::new(array)
736 }
737 DatumColumnEncoder::AclItem(mut builder) => Arc::new(builder.finish()),
738 DatumColumnEncoder::MzAclItem(mut builder) => Arc::new(builder.finish()),
739 DatumColumnEncoder::Range(mut builder) => Arc::new(builder.finish()),
740 DatumColumnEncoder::Jsonb {
741 offsets,
742 buf,
743 mut nulls,
744 } => {
745 let values = Buffer::from_vec(buf);
746 let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets));
747 let nulls = nulls.as_mut().map(|n| NullBuffer::from(n.finish()));
748 let array = StringArray::new(offsets, values, nulls);
749 Arc::new(array)
750 }
751 DatumColumnEncoder::Array {
752 mut dims,
753 val_lengths,
754 vals,
755 mut nulls,
756 } => {
757 let nulls = nulls.as_mut().map(|n| NullBuffer::from(n.finish()));
758 let vals = vals.finish();
759
760 let field = Field::new_list_field(vals.data_type().clone(), true);
763 let val_offsets = OffsetBuffer::from_lengths(val_lengths);
764 let values =
765 ListArray::new(Arc::new(field), val_offsets, Arc::new(vals), nulls.clone());
766
767 let dims = dims.finish();
768 assert_eq!(values.len(), dims.len());
769
770 let fields = Fields::from(vec![
773 Field::new("dims", dims.data_type().clone(), true),
774 Field::new("vals", values.data_type().clone(), true),
775 ]);
776 let array = StructArray::new(fields, vec![Arc::new(dims), Arc::new(values)], nulls);
777
778 Arc::new(array)
779 }
780 DatumColumnEncoder::List {
781 lengths,
782 values,
783 mut nulls,
784 } => {
785 let values = values.finish();
786
787 let field = Field::new_list_field(values.data_type().clone(), true);
790 let offsets = OffsetBuffer::<i32>::from_lengths(lengths.iter().copied());
791 let nulls = nulls.as_mut().map(|n| NullBuffer::from(n.finish()));
792
793 let array = ListArray::new(Arc::new(field), offsets, values, nulls);
794 Arc::new(array)
795 }
796 DatumColumnEncoder::Map {
797 lengths,
798 mut keys,
799 vals,
800 mut nulls,
801 } => {
802 let keys = keys.finish();
803 let vals = vals.finish();
804
805 let offsets = OffsetBuffer::<i32>::from_lengths(lengths.iter().copied());
806 let nulls = nulls.as_mut().map(|n| NullBuffer::from(n.finish()));
807
808 assert_none!(keys.logical_nulls());
811 let key_field = Arc::new(Field::new("key", keys.data_type().clone(), false));
812 let val_field = Arc::new(Field::new("val", vals.data_type().clone(), true));
813 let fields = Fields::from(vec![Arc::clone(&key_field), Arc::clone(&val_field)]);
814 let entries = StructArray::new(fields, vec![Arc::new(keys), vals], None);
815
816 let field = Field::new("map_entries", entries.data_type().clone(), false);
819 let array = ListArray::new(Arc::new(field), offsets, Arc::new(entries), nulls);
820 Arc::new(array)
821 }
822 DatumColumnEncoder::Record {
823 fields,
824 mut nulls,
825 length: _,
826 } => {
827 let (fields, arrays): (Vec<_>, Vec<_>) = fields
828 .into_iter()
829 .enumerate()
830 .map(|(tag, encoder)| {
831 let nullable = true;
837 let array = encoder.finish();
838 let field =
839 Field::new(tag.to_string(), array.data_type().clone(), nullable);
840 (field, array)
841 })
842 .unzip();
843 let nulls = nulls.as_mut().map(|n| NullBuffer::from(n.finish()));
844
845 let array = StructArray::new(Fields::from(fields), arrays, nulls);
846 Arc::new(array)
847 }
848 DatumColumnEncoder::RecordEmpty(mut builder) => Arc::new(builder.finish()),
849 }
850 }
851}
852
853#[derive(Debug)]
858enum DatumColumnDecoder {
859 Bool(BooleanArray),
860 U8(UInt8Array),
861 U16(UInt16Array),
862 U32(UInt32Array),
863 U64(UInt64Array),
864 I16(Int16Array),
865 I32(Int32Array),
866 I64(Int64Array),
867 F32(Float32Array),
868 F64(Float64Array),
869 Numeric(BinaryArray),
870 Bytes(BinaryArray),
871 String(StringArray),
872 Date(Int32Array),
873 Time(FixedSizeBinaryArray),
874 Timestamp(FixedSizeBinaryArray),
875 TimestampTz(FixedSizeBinaryArray),
876 MzTimestamp(UInt64Array),
877 Interval(FixedSizeBinaryArray),
878 Uuid(FixedSizeBinaryArray),
879 Json(StringArray),
880 Array {
881 dim_offsets: OffsetBuffer<i32>,
882 dims: FixedSizeBinaryArray,
883 val_offsets: OffsetBuffer<i32>,
884 vals: Box<DatumColumnDecoder>,
885 nulls: Option<NullBuffer>,
886 },
887 List {
888 offsets: OffsetBuffer<i32>,
889 values: Box<DatumColumnDecoder>,
890 nulls: Option<NullBuffer>,
891 },
892 Map {
893 offsets: OffsetBuffer<i32>,
894 keys: StringArray,
895 vals: Box<DatumColumnDecoder>,
896 nulls: Option<NullBuffer>,
897 },
898 RecordEmpty(BooleanArray),
899 Record {
900 fields: Vec<Box<DatumColumnDecoder>>,
901 nulls: Option<NullBuffer>,
902 },
903 Range(BinaryArray),
904 MzAclItem(BinaryArray),
905 AclItem(FixedSizeBinaryArray),
906}
907
908impl DatumColumnDecoder {
909 fn get<'a>(&'a self, idx: usize, packer: &'a mut RowPacker) {
910 let datum = match self {
911 DatumColumnDecoder::Bool(array) => array
912 .is_valid(idx)
913 .then(|| array.value(idx))
914 .map(|x| if x { Datum::True } else { Datum::False }),
915 DatumColumnDecoder::U8(array) => array
916 .is_valid(idx)
917 .then(|| array.value(idx))
918 .map(Datum::UInt8),
919 DatumColumnDecoder::U16(array) => array
920 .is_valid(idx)
921 .then(|| array.value(idx))
922 .map(Datum::UInt16),
923 DatumColumnDecoder::U32(array) => array
924 .is_valid(idx)
925 .then(|| array.value(idx))
926 .map(Datum::UInt32),
927 DatumColumnDecoder::U64(array) => array
928 .is_valid(idx)
929 .then(|| array.value(idx))
930 .map(Datum::UInt64),
931 DatumColumnDecoder::I16(array) => array
932 .is_valid(idx)
933 .then(|| array.value(idx))
934 .map(Datum::Int16),
935 DatumColumnDecoder::I32(array) => array
936 .is_valid(idx)
937 .then(|| array.value(idx))
938 .map(Datum::Int32),
939 DatumColumnDecoder::I64(array) => array
940 .is_valid(idx)
941 .then(|| array.value(idx))
942 .map(Datum::Int64),
943 DatumColumnDecoder::F32(array) => array
944 .is_valid(idx)
945 .then(|| array.value(idx))
946 .map(|x| Datum::Float32(ordered_float::OrderedFloat(x))),
947 DatumColumnDecoder::F64(array) => array
948 .is_valid(idx)
949 .then(|| array.value(idx))
950 .map(|x| Datum::Float64(ordered_float::OrderedFloat(x))),
951 DatumColumnDecoder::Numeric(array) => array.is_valid(idx).then(|| {
952 let val = array.value(idx);
953 let val = PackedNumeric::from_bytes(val)
954 .expect("failed to roundtrip Numeric")
955 .into_value();
956 Datum::Numeric(OrderedDecimal(val))
957 }),
958 DatumColumnDecoder::String(array) => array
959 .is_valid(idx)
960 .then(|| array.value(idx))
961 .map(Datum::String),
962 DatumColumnDecoder::Bytes(array) => array
963 .is_valid(idx)
964 .then(|| array.value(idx))
965 .map(Datum::Bytes),
966 DatumColumnDecoder::Date(array) => {
967 array.is_valid(idx).then(|| array.value(idx)).map(|x| {
968 let date = Date::from_pg_epoch(x).expect("failed to roundtrip");
969 Datum::Date(date)
970 })
971 }
972 DatumColumnDecoder::Time(array) => {
973 array.is_valid(idx).then(|| array.value(idx)).map(|x| {
974 let packed = PackedNaiveTime::from_bytes(x).expect("failed to roundtrip time");
975 Datum::Time(packed.into_value())
976 })
977 }
978 DatumColumnDecoder::Timestamp(array) => {
979 array.is_valid(idx).then(|| array.value(idx)).map(|x| {
980 let packed = PackedNaiveDateTime::from_bytes(x)
981 .expect("failed to roundtrip PackedNaiveDateTime");
982 let timestamp = CheckedTimestamp::from_timestamplike(packed.into_value())
983 .expect("failed to roundtrip timestamp");
984 Datum::Timestamp(timestamp)
985 })
986 }
987 DatumColumnDecoder::TimestampTz(array) => {
988 array.is_valid(idx).then(|| array.value(idx)).map(|x| {
989 let packed = PackedNaiveDateTime::from_bytes(x)
990 .expect("failed to roundtrip PackedNaiveDateTime");
991 let timestamp =
992 CheckedTimestamp::from_timestamplike(packed.into_value().and_utc())
993 .expect("failed to roundtrip timestamp");
994 Datum::TimestampTz(timestamp)
995 })
996 }
997 DatumColumnDecoder::MzTimestamp(array) => array
998 .is_valid(idx)
999 .then(|| array.value(idx))
1000 .map(|x| Datum::MzTimestamp(Timestamp::from(x))),
1001 DatumColumnDecoder::Interval(array) => {
1002 array.is_valid(idx).then(|| array.value(idx)).map(|x| {
1003 let packed =
1004 PackedInterval::from_bytes(x).expect("failed to roundtrip interval");
1005 Datum::Interval(packed.into_value())
1006 })
1007 }
1008 DatumColumnDecoder::Uuid(array) => {
1009 array.is_valid(idx).then(|| array.value(idx)).map(|x| {
1010 let uuid = Uuid::from_slice(x).expect("failed to roundtrip uuid");
1011 Datum::Uuid(uuid)
1012 })
1013 }
1014 DatumColumnDecoder::AclItem(array) => {
1015 array.is_valid(idx).then(|| array.value(idx)).map(|x| {
1016 let packed =
1017 PackedAclItem::from_bytes(x).expect("failed to roundtrip MzAclItem");
1018 Datum::AclItem(packed.into_value())
1019 })
1020 }
1021 DatumColumnDecoder::MzAclItem(array) => {
1022 array.is_valid(idx).then(|| array.value(idx)).map(|x| {
1023 let packed =
1024 PackedMzAclItem::from_bytes(x).expect("failed to roundtrip MzAclItem");
1025 Datum::MzAclItem(packed.into_value())
1026 })
1027 }
1028 DatumColumnDecoder::Range(array) => {
1029 let Some(val) = array.is_valid(idx).then(|| array.value(idx)) else {
1030 packer.push(Datum::Null);
1031 return;
1032 };
1033
1034 let proto = ProtoDatum::decode(val).expect("failed to roundtrip Range");
1035 packer
1036 .try_push_proto(&proto)
1037 .expect("failed to pack ProtoRange");
1038
1039 return;
1041 }
1042 DatumColumnDecoder::Json(array) => {
1043 let Some(val) = array.is_valid(idx).then(|| array.value(idx)) else {
1044 packer.push(Datum::Null);
1045 return;
1046 };
1047 JsonbPacker::new(packer)
1048 .pack_str(val)
1049 .expect("failed to roundtrip JSON");
1050
1051 return;
1053 }
1054 DatumColumnDecoder::Array {
1055 dim_offsets,
1056 dims,
1057 val_offsets,
1058 vals,
1059 nulls,
1060 } => {
1061 let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
1062 if !is_valid {
1063 packer.push(Datum::Null);
1064 return;
1065 }
1066
1067 let start: usize = dim_offsets[idx]
1068 .try_into()
1069 .expect("unexpected negative offset");
1070 let end: usize = dim_offsets[idx + 1]
1071 .try_into()
1072 .expect("unexpected negative offset");
1073 let dimensions = (start..end).map(|idx| {
1074 PackedArrayDimension::from_bytes(dims.value(idx))
1075 .expect("failed to roundtrip ArrayDimension")
1076 .into_value()
1077 });
1078
1079 let start: usize = val_offsets[idx]
1080 .try_into()
1081 .expect("unexpected negative offset");
1082 let end: usize = val_offsets[idx + 1]
1083 .try_into()
1084 .expect("unexpected negative offset");
1085 packer
1086 .push_array_with_row_major(dimensions, |packer| {
1087 for x in start..end {
1088 vals.get(x, packer);
1089 }
1090 end - start
1092 })
1093 .expect("failed to pack Array");
1094
1095 return;
1097 }
1098 DatumColumnDecoder::List {
1099 offsets,
1100 values,
1101 nulls,
1102 } => {
1103 let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
1104 if !is_valid {
1105 packer.push(Datum::Null);
1106 return;
1107 }
1108
1109 let start: usize = offsets[idx].try_into().expect("unexpected negative offset");
1110 let end: usize = offsets[idx + 1]
1111 .try_into()
1112 .expect("unexpected negative offset");
1113
1114 packer.push_list_with(|packer| {
1115 for idx in start..end {
1116 values.get(idx, packer)
1117 }
1118 });
1119
1120 return;
1122 }
1123 DatumColumnDecoder::Map {
1124 offsets,
1125 keys,
1126 vals,
1127 nulls,
1128 } => {
1129 let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
1130 if !is_valid {
1131 packer.push(Datum::Null);
1132 return;
1133 }
1134
1135 let start: usize = offsets[idx].try_into().expect("unexpected negative offset");
1136 let end: usize = offsets[idx + 1]
1137 .try_into()
1138 .expect("unexpected negative offset");
1139
1140 packer.push_dict_with(|packer| {
1141 for idx in start..end {
1142 packer.push(Datum::String(keys.value(idx)));
1143 vals.get(idx, packer);
1144 }
1145 });
1146
1147 return;
1149 }
1150 DatumColumnDecoder::RecordEmpty(array) => array.is_valid(idx).then(Datum::empty_list),
1151 DatumColumnDecoder::Record { fields, nulls } => {
1152 let is_valid = nulls.as_ref().map(|n| n.is_valid(idx)).unwrap_or(true);
1153 if !is_valid {
1154 packer.push(Datum::Null);
1155 return;
1156 }
1157
1158 packer.push_list_with(|packer| {
1160 for field in fields {
1161 field.get(idx, packer);
1162 }
1163 });
1164
1165 return;
1167 }
1168 };
1169
1170 match datum {
1171 Some(d) => packer.push(d),
1172 None => packer.push(Datum::Null),
1173 }
1174 }
1175
1176 fn stats(&self) -> ColumnStatKinds {
1177 match self {
1178 DatumColumnDecoder::Bool(a) => PrimitiveStats::<bool>::from_column(a).into(),
1179 DatumColumnDecoder::U8(a) => PrimitiveStats::<u8>::from_column(a).into(),
1180 DatumColumnDecoder::U16(a) => PrimitiveStats::<u16>::from_column(a).into(),
1181 DatumColumnDecoder::U32(a) => PrimitiveStats::<u32>::from_column(a).into(),
1182 DatumColumnDecoder::U64(a) => PrimitiveStats::<u64>::from_column(a).into(),
1183 DatumColumnDecoder::I16(a) => PrimitiveStats::<i16>::from_column(a).into(),
1184 DatumColumnDecoder::I32(a) => PrimitiveStats::<i32>::from_column(a).into(),
1185 DatumColumnDecoder::I64(a) => PrimitiveStats::<i64>::from_column(a).into(),
1186 DatumColumnDecoder::F32(a) => PrimitiveStats::<f32>::from_column(a).into(),
1187 DatumColumnDecoder::F64(a) => PrimitiveStats::<f64>::from_column(a).into(),
1188 DatumColumnDecoder::Numeric(a) => numeric_stats_from_column(a),
1189 DatumColumnDecoder::String(a) => PrimitiveStats::<String>::from_column(a).into(),
1190 DatumColumnDecoder::Bytes(a) => PrimitiveStats::<Vec<u8>>::from_column(a).into(),
1191 DatumColumnDecoder::Date(a) => PrimitiveStats::<i32>::from_column(a).into(),
1192 DatumColumnDecoder::Time(a) => {
1193 fixed_stats_from_column(a, FixedSizeBytesStatsKind::PackedTime)
1194 }
1195 DatumColumnDecoder::Timestamp(a) => {
1196 fixed_stats_from_column(a, FixedSizeBytesStatsKind::PackedDateTime)
1197 }
1198 DatumColumnDecoder::TimestampTz(a) => {
1199 fixed_stats_from_column(a, FixedSizeBytesStatsKind::PackedDateTime)
1200 }
1201 DatumColumnDecoder::MzTimestamp(a) => PrimitiveStats::<u64>::from_column(a).into(),
1202 DatumColumnDecoder::Interval(a) => {
1203 fixed_stats_from_column(a, FixedSizeBytesStatsKind::PackedInterval)
1204 }
1205 DatumColumnDecoder::Uuid(a) => {
1206 fixed_stats_from_column(a, FixedSizeBytesStatsKind::Uuid)
1207 }
1208 DatumColumnDecoder::AclItem(_)
1209 | DatumColumnDecoder::MzAclItem(_)
1210 | DatumColumnDecoder::Range(_) => ColumnStatKinds::None,
1211 DatumColumnDecoder::Json(a) => stats_for_json(a.iter()).values,
1212 DatumColumnDecoder::Array { .. }
1213 | DatumColumnDecoder::List { .. }
1214 | DatumColumnDecoder::Map { .. }
1215 | DatumColumnDecoder::Record { .. }
1216 | DatumColumnDecoder::RecordEmpty(_) => ColumnStatKinds::None,
1217 }
1218 }
1219
1220 fn goodbytes(&self) -> usize {
1221 match self {
1222 DatumColumnDecoder::Bool(a) => ArrayOrd::Bool(a.clone()).goodbytes(),
1223 DatumColumnDecoder::U8(a) => ArrayOrd::UInt8(a.clone()).goodbytes(),
1224 DatumColumnDecoder::U16(a) => ArrayOrd::UInt16(a.clone()).goodbytes(),
1225 DatumColumnDecoder::U32(a) => ArrayOrd::UInt32(a.clone()).goodbytes(),
1226 DatumColumnDecoder::U64(a) => ArrayOrd::UInt64(a.clone()).goodbytes(),
1227 DatumColumnDecoder::I16(a) => ArrayOrd::Int16(a.clone()).goodbytes(),
1228 DatumColumnDecoder::I32(a) => ArrayOrd::Int32(a.clone()).goodbytes(),
1229 DatumColumnDecoder::I64(a) => ArrayOrd::Int64(a.clone()).goodbytes(),
1230 DatumColumnDecoder::F32(a) => ArrayOrd::Float32(a.clone()).goodbytes(),
1231 DatumColumnDecoder::F64(a) => ArrayOrd::Float64(a.clone()).goodbytes(),
1232 DatumColumnDecoder::Numeric(a) => ArrayOrd::Binary(a.clone()).goodbytes(),
1233 DatumColumnDecoder::String(a) => ArrayOrd::String(a.clone()).goodbytes(),
1234 DatumColumnDecoder::Bytes(a) => ArrayOrd::Binary(a.clone()).goodbytes(),
1235 DatumColumnDecoder::Date(a) => ArrayOrd::Int32(a.clone()).goodbytes(),
1236 DatumColumnDecoder::Time(a) => ArrayOrd::FixedSizeBinary(a.clone()).goodbytes(),
1237 DatumColumnDecoder::Timestamp(a) => ArrayOrd::FixedSizeBinary(a.clone()).goodbytes(),
1238 DatumColumnDecoder::TimestampTz(a) => ArrayOrd::FixedSizeBinary(a.clone()).goodbytes(),
1239 DatumColumnDecoder::MzTimestamp(a) => ArrayOrd::UInt64(a.clone()).goodbytes(),
1240 DatumColumnDecoder::Interval(a) => ArrayOrd::FixedSizeBinary(a.clone()).goodbytes(),
1241 DatumColumnDecoder::Uuid(a) => ArrayOrd::FixedSizeBinary(a.clone()).goodbytes(),
1242 DatumColumnDecoder::AclItem(a) => ArrayOrd::FixedSizeBinary(a.clone()).goodbytes(),
1243 DatumColumnDecoder::MzAclItem(a) => ArrayOrd::Binary(a.clone()).goodbytes(),
1244 DatumColumnDecoder::Range(a) => ArrayOrd::Binary(a.clone()).goodbytes(),
1245 DatumColumnDecoder::Json(a) => ArrayOrd::String(a.clone()).goodbytes(),
1246 DatumColumnDecoder::Array { dims, vals, .. } => {
1247 (dims.len() * PackedArrayDimension::SIZE) + vals.goodbytes()
1248 }
1249 DatumColumnDecoder::List { values, .. } => values.goodbytes(),
1250 DatumColumnDecoder::Map { keys, vals, .. } => {
1251 ArrayOrd::String(keys.clone()).goodbytes() + vals.goodbytes()
1252 }
1253 DatumColumnDecoder::Record { fields, .. } => fields.iter().map(|f| f.goodbytes()).sum(),
1254 DatumColumnDecoder::RecordEmpty(a) => ArrayOrd::Bool(a.clone()).goodbytes(),
1255 }
1256 }
1257}
1258
1259impl Schema<Row> for RelationDesc {
1260 type ArrowColumn = arrow::array::StructArray;
1261 type Statistics = OptionStats<StructStats>;
1262
1263 type Decoder = RowColumnarDecoder;
1264 type Encoder = RowColumnarEncoder;
1265
1266 fn decoder(&self, col: Self::ArrowColumn) -> Result<Self::Decoder, anyhow::Error> {
1267 RowColumnarDecoder::new(col, self)
1268 }
1269
1270 fn encoder(&self) -> Result<Self::Encoder, anyhow::Error> {
1271 RowColumnarEncoder::new(self)
1272 .ok_or_else(|| anyhow::anyhow!("Cannot encode a RelationDesc with no columns"))
1273 }
1274}
1275
1276#[derive(Debug)]
1278pub struct RowColumnarDecoder {
1279 len: usize,
1282 decoders: Vec<(Arc<str>, Option<usize>, DatumColumnDecoder)>,
1285 nullability: Option<NullBuffer>,
1288}
1289
1290fn mask_nulls(column: &ArrayRef, null_mask: Option<&NullBuffer>) -> ArrayRef {
1292 if null_mask.is_none() {
1293 Arc::clone(column)
1294 } else {
1295 let nulls = NullBuffer::union(null_mask, column.nulls());
1298 let data = column
1299 .to_data()
1300 .into_builder()
1301 .nulls(nulls)
1302 .build()
1303 .expect("changed only null mask");
1304 make_array(data)
1305 }
1306}
1307
1308impl RowColumnarDecoder {
1309 pub fn new(col: StructArray, desc: &RelationDesc) -> Result<Self, anyhow::Error> {
1314 let inner_columns = col.columns();
1315 let desc_columns = desc.typ().columns();
1316
1317 if desc_columns.len() > inner_columns.len() {
1318 anyhow::bail!(
1319 "provided array has too few columns! {desc_columns:?} > {inner_columns:?}"
1320 );
1321 }
1322
1323 let mut decoders = Vec::with_capacity(desc_columns.len());
1325
1326 let null_mask = col.nulls();
1327
1328 for (col_idx, col_name, col_type) in desc.iter_all() {
1330 let field_name = col_idx.to_stable_name();
1331 let column = col.column_by_name(&field_name).ok_or_else(|| {
1332 anyhow::anyhow!(
1333 "StructArray did not contain column name {field_name}, found {:?}",
1334 col.column_names()
1335 )
1336 })?;
1337 let column = mask_nulls(column, null_mask);
1338 let null_count = col_type.nullable.then(|| column.null_count());
1339 let decoder = array_to_decoder(&column, &col_type.scalar_type)?;
1340 decoders.push((col_name.as_str().into(), null_count, decoder));
1341 }
1342
1343 Ok(RowColumnarDecoder {
1344 len: col.len(),
1345 decoders,
1346 nullability: col.logical_nulls(),
1347 })
1348 }
1349
1350 pub fn null_count(&self) -> usize {
1354 self.nullability.as_ref().map_or(0, |n| n.null_count())
1355 }
1356}
1357
1358impl ColumnDecoder<Row> for RowColumnarDecoder {
1359 fn decode(&self, idx: usize, val: &mut Row) {
1360 let mut packer = val.packer();
1361
1362 for (_, _, decoder) in &self.decoders {
1363 decoder.get(idx, &mut packer);
1364 }
1365 }
1366
1367 fn is_null(&self, idx: usize) -> bool {
1368 let Some(nullability) = self.nullability.as_ref() else {
1369 return false;
1370 };
1371 nullability.is_null(idx)
1372 }
1373
1374 fn goodbytes(&self) -> usize {
1375 let decoders_size: usize = self
1376 .decoders
1377 .iter()
1378 .map(|(_name, _null_count, decoder)| decoder.goodbytes())
1379 .sum();
1380
1381 decoders_size
1382 + self
1383 .nullability
1384 .as_ref()
1385 .map(|nulls| nulls.inner().inner().len())
1386 .unwrap_or(0)
1387 }
1388
1389 fn stats(&self) -> StructStats {
1390 StructStats {
1391 len: self.len,
1392 cols: self
1393 .decoders
1394 .iter()
1395 .map(|(name, null_count, decoder)| {
1396 let name = name.to_string();
1397 let stats = ColumnarStats {
1398 nulls: null_count.map(|count| ColumnNullStats { count }),
1399 values: decoder.stats(),
1400 };
1401 (name, stats)
1402 })
1403 .collect(),
1404 }
1405 }
1406}
1407
1408#[derive(Debug)]
1410pub struct RowColumnarEncoder {
1411 encoders: Vec<DatumEncoder>,
1412 col_names: Vec<(usize, Arc<str>)>,
1414 nullability: BooleanBufferBuilder,
1416}
1417
1418impl RowColumnarEncoder {
1419 pub fn new(desc: &RelationDesc) -> Option<Self> {
1429 if desc.typ().columns().is_empty() {
1430 return None;
1431 }
1432
1433 let (col_names, encoders): (Vec<_>, Vec<_>) = desc
1434 .iter_all()
1435 .map(|(col_idx, col_name, col_type)| {
1436 let encoder = scalar_type_to_encoder(&col_type.scalar_type)
1437 .expect("failed to create encoder");
1438 let encoder = DatumEncoder {
1439 nullable: col_type.nullable,
1440 encoder,
1441 };
1442
1443 let name = (col_idx.to_raw(), col_name.as_str().into());
1454
1455 (name, encoder)
1456 })
1457 .unzip();
1458
1459 Some(RowColumnarEncoder {
1460 encoders,
1461 col_names,
1462 nullability: BooleanBufferBuilder::new(100),
1463 })
1464 }
1465}
1466
1467impl ColumnEncoder<Row> for RowColumnarEncoder {
1468 type FinishedColumn = StructArray;
1469
1470 fn goodbytes(&self) -> usize {
1471 self.encoders.iter().map(|e| e.goodbytes()).sum()
1472 }
1473
1474 fn append(&mut self, val: &Row) {
1475 let mut num_datums = 0;
1476 for (datum, encoder) in val.iter().zip_eq(self.encoders.iter_mut()) {
1477 encoder.push(datum);
1478 num_datums += 1;
1479 }
1480 assert_eq!(
1481 num_datums,
1482 self.encoders.len(),
1483 "tried to encode {val:?}, but only have {:?}",
1484 self.encoders
1485 );
1486
1487 self.nullability.append(true);
1488 }
1489
1490 fn append_null(&mut self) {
1491 for encoder in self.encoders.iter_mut() {
1492 encoder.push_invalid();
1493 }
1494 self.nullability.append(false);
1495 }
1496
1497 fn finish(self) -> Self::FinishedColumn {
1498 let RowColumnarEncoder {
1499 encoders,
1500 col_names,
1501 nullability,
1502 ..
1503 } = self;
1504
1505 let (arrays, fields): (Vec<_>, Vec<_>) = col_names
1506 .iter()
1507 .zip_eq(encoders)
1508 .map(|((col_idx, _col_name), encoder)| {
1509 let nullable = true;
1515 let array = encoder.finish();
1516 let field = Field::new(col_idx.to_string(), array.data_type().clone(), nullable);
1517
1518 (array, field)
1519 })
1520 .multiunzip();
1521
1522 let null_buffer = NullBuffer::from(BooleanBuffer::from(nullability));
1523
1524 let array = StructArray::new(Fields::from(fields), arrays, Some(null_buffer));
1525
1526 array
1527 }
1528}
1529
1530#[inline]
1535fn downcast_array<T: 'static>(array: &Arc<dyn Array>) -> Result<&T, anyhow::Error> {
1536 array
1537 .as_any()
1538 .downcast_ref::<T>()
1539 .ok_or_else(|| anyhow!("expected {}, found {array:?}", std::any::type_name::<T>()))
1540}
1541
1542fn array_to_decoder(
1547 array: &Arc<dyn Array>,
1548 col_ty: &SqlScalarType,
1549) -> Result<DatumColumnDecoder, anyhow::Error> {
1550 let decoder = match (array.data_type(), col_ty) {
1551 (DataType::Boolean, SqlScalarType::Bool) => {
1552 let array = downcast_array::<BooleanArray>(array)?;
1553 DatumColumnDecoder::Bool(array.clone())
1554 }
1555 (DataType::UInt8, SqlScalarType::PgLegacyChar) => {
1556 let array = downcast_array::<UInt8Array>(array)?;
1557 DatumColumnDecoder::U8(array.clone())
1558 }
1559 (DataType::UInt16, SqlScalarType::UInt16) => {
1560 let array = downcast_array::<UInt16Array>(array)?;
1561 DatumColumnDecoder::U16(array.clone())
1562 }
1563 (
1564 DataType::UInt32,
1565 SqlScalarType::UInt32
1566 | SqlScalarType::Oid
1567 | SqlScalarType::RegClass
1568 | SqlScalarType::RegProc
1569 | SqlScalarType::RegType,
1570 ) => {
1571 let array = downcast_array::<UInt32Array>(array)?;
1572 DatumColumnDecoder::U32(array.clone())
1573 }
1574 (DataType::UInt64, SqlScalarType::UInt64) => {
1575 let array = downcast_array::<UInt64Array>(array)?;
1576 DatumColumnDecoder::U64(array.clone())
1577 }
1578 (DataType::Int16, SqlScalarType::Int16) => {
1579 let array = downcast_array::<Int16Array>(array)?;
1580 DatumColumnDecoder::I16(array.clone())
1581 }
1582 (DataType::Int32, SqlScalarType::Int32) => {
1583 let array = downcast_array::<Int32Array>(array)?;
1584 DatumColumnDecoder::I32(array.clone())
1585 }
1586 (DataType::Int64, SqlScalarType::Int64) => {
1587 let array = downcast_array::<Int64Array>(array)?;
1588 DatumColumnDecoder::I64(array.clone())
1589 }
1590 (DataType::Float32, SqlScalarType::Float32) => {
1591 let array = downcast_array::<Float32Array>(array)?;
1592 DatumColumnDecoder::F32(array.clone())
1593 }
1594 (DataType::Float64, SqlScalarType::Float64) => {
1595 let array = downcast_array::<Float64Array>(array)?;
1596 DatumColumnDecoder::F64(array.clone())
1597 }
1598 (DataType::Struct(_), SqlScalarType::Numeric { .. }) => {
1599 let array = downcast_array::<StructArray>(array)?;
1600 let binary_values = array
1603 .column_by_name("binary")
1604 .expect("missing binary column");
1605
1606 let array = downcast_array::<BinaryArray>(binary_values)?;
1607 DatumColumnDecoder::Numeric(array.clone())
1608 }
1609 (
1610 DataType::Utf8,
1611 SqlScalarType::String
1612 | SqlScalarType::PgLegacyName
1613 | SqlScalarType::Char { .. }
1614 | SqlScalarType::VarChar { .. },
1615 ) => {
1616 let array = downcast_array::<StringArray>(array)?;
1617 DatumColumnDecoder::String(array.clone())
1618 }
1619 (DataType::Binary, SqlScalarType::Bytes) => {
1620 let array = downcast_array::<BinaryArray>(array)?;
1621 DatumColumnDecoder::Bytes(array.clone())
1622 }
1623 (DataType::Int32, SqlScalarType::Date) => {
1624 let array = downcast_array::<Int32Array>(array)?;
1625 DatumColumnDecoder::Date(array.clone())
1626 }
1627 (DataType::FixedSizeBinary(TIME_FIXED_BYTES), SqlScalarType::Time) => {
1628 let array = downcast_array::<FixedSizeBinaryArray>(array)?;
1629 DatumColumnDecoder::Time(array.clone())
1630 }
1631 (DataType::FixedSizeBinary(TIMESTAMP_FIXED_BYTES), SqlScalarType::Timestamp { .. }) => {
1632 let array = downcast_array::<FixedSizeBinaryArray>(array)?;
1633 DatumColumnDecoder::Timestamp(array.clone())
1634 }
1635 (DataType::FixedSizeBinary(TIMESTAMP_FIXED_BYTES), SqlScalarType::TimestampTz { .. }) => {
1636 let array = downcast_array::<FixedSizeBinaryArray>(array)?;
1637 DatumColumnDecoder::TimestampTz(array.clone())
1638 }
1639 (DataType::UInt64, SqlScalarType::MzTimestamp) => {
1640 let array = downcast_array::<UInt64Array>(array)?;
1641 DatumColumnDecoder::MzTimestamp(array.clone())
1642 }
1643 (DataType::FixedSizeBinary(INTERVAL_FIXED_BYTES), SqlScalarType::Interval) => {
1644 let array = downcast_array::<FixedSizeBinaryArray>(array)?;
1645 DatumColumnDecoder::Interval(array.clone())
1646 }
1647 (DataType::FixedSizeBinary(UUID_FIXED_BYTES), SqlScalarType::Uuid) => {
1648 let array = downcast_array::<FixedSizeBinaryArray>(array)?;
1649 DatumColumnDecoder::Uuid(array.clone())
1650 }
1651 (DataType::FixedSizeBinary(ACL_ITEM_FIXED_BYTES), SqlScalarType::AclItem) => {
1652 let array = downcast_array::<FixedSizeBinaryArray>(array)?;
1653 DatumColumnDecoder::AclItem(array.clone())
1654 }
1655 (DataType::Binary, SqlScalarType::MzAclItem) => {
1656 let array = downcast_array::<BinaryArray>(array)?;
1657 DatumColumnDecoder::MzAclItem(array.clone())
1658 }
1659 (DataType::Binary, SqlScalarType::Range { .. }) => {
1660 let array = downcast_array::<BinaryArray>(array)?;
1661 DatumColumnDecoder::Range(array.clone())
1662 }
1663 (DataType::Utf8, SqlScalarType::Jsonb) => {
1664 let array = downcast_array::<StringArray>(array)?;
1665 DatumColumnDecoder::Json(array.clone())
1666 }
1667 (DataType::Struct(_), s @ SqlScalarType::Array(_) | s @ SqlScalarType::Int2Vector) => {
1668 let element_type = match s {
1669 SqlScalarType::Array(inner) => inner,
1670 SqlScalarType::Int2Vector => &SqlScalarType::Int16,
1671 _ => unreachable!("checked above"),
1672 };
1673
1674 let array = downcast_array::<StructArray>(array)?;
1675 let nulls = array.nulls().cloned();
1676
1677 let dims = array
1678 .column_by_name("dims")
1679 .expect("missing dimensions column");
1680 let dims = downcast_array::<ListArray>(dims).cloned()?;
1681 let dim_offsets = dims.offsets().clone();
1682 let dims = downcast_array::<FixedSizeBinaryArray>(dims.values()).cloned()?;
1683
1684 let vals = array.column_by_name("vals").expect("missing values column");
1685 let vals = downcast_array::<ListArray>(vals)?;
1686 let val_offsets = vals.offsets().clone();
1687 let vals = array_to_decoder(vals.values(), element_type)?;
1688
1689 DatumColumnDecoder::Array {
1690 dim_offsets,
1691 dims,
1692 val_offsets,
1693 vals: Box::new(vals),
1694 nulls,
1695 }
1696 }
1697 (DataType::List(_), SqlScalarType::List { element_type, .. }) => {
1698 let array = downcast_array::<ListArray>(array)?;
1699 let inner_decoder = array_to_decoder(array.values(), &*element_type)?;
1700 DatumColumnDecoder::List {
1701 offsets: array.offsets().clone(),
1702 values: Box::new(inner_decoder),
1703 nulls: array.nulls().cloned(),
1704 }
1705 }
1706 (DataType::Map(_, true), SqlScalarType::Map { value_type, .. }) => {
1707 let array = downcast_array::<MapArray>(array)?;
1708 let keys = downcast_array::<StringArray>(array.keys())?;
1709 let vals = array_to_decoder(array.values(), value_type)?;
1710 DatumColumnDecoder::Map {
1711 offsets: array.offsets().clone(),
1712 keys: keys.clone(),
1713 vals: Box::new(vals),
1714 nulls: array.nulls().cloned(),
1715 }
1716 }
1717 (DataType::List(_), SqlScalarType::Map { value_type, .. }) => {
1718 let array: &ListArray = downcast_array(array)?;
1719 let entries: &StructArray = downcast_array(array.values())?;
1720 let [keys, values]: &[ArrayRef; 2] = entries.columns().try_into()?;
1721 let keys: &StringArray = downcast_array(keys)?;
1722 let vals: DatumColumnDecoder = array_to_decoder(values, value_type)?;
1723 DatumColumnDecoder::Map {
1724 offsets: array.offsets().clone(),
1725 keys: keys.clone(),
1726 vals: Box::new(vals),
1727 nulls: array.nulls().cloned(),
1728 }
1729 }
1730 (DataType::Boolean, SqlScalarType::Record { fields, .. }) if fields.is_empty() => {
1731 let empty_record_array = downcast_array::<BooleanArray>(array)?;
1732 DatumColumnDecoder::RecordEmpty(empty_record_array.clone())
1733 }
1734 (DataType::Struct(_), SqlScalarType::Record { fields, .. }) => {
1735 let record_array = downcast_array::<StructArray>(array)?;
1736 let null_mask = record_array.nulls();
1737 let mut decoders = Vec::with_capacity(fields.len());
1738 for (tag, (_name, col_type)) in fields.iter().enumerate() {
1739 let inner_array = record_array
1740 .column_by_name(&tag.to_string())
1741 .ok_or_else(|| anyhow::anyhow!("no column named '{tag}'"))?;
1742 let inner_array = mask_nulls(inner_array, null_mask);
1743 let inner_decoder = array_to_decoder(&inner_array, &col_type.scalar_type)?;
1744
1745 decoders.push(Box::new(inner_decoder));
1746 }
1747
1748 DatumColumnDecoder::Record {
1749 fields: decoders,
1750 nulls: record_array.nulls().cloned(),
1751 }
1752 }
1753 (x, y) => {
1754 let msg = format!("can't decode column of {x:?} for scalar type {y:?}");
1755 mz_ore::soft_panic_or_log!("{msg}");
1756 anyhow::bail!("{msg}");
1757 }
1758 };
1759
1760 Ok(decoder)
1761}
1762
1763fn scalar_type_to_encoder(col_ty: &SqlScalarType) -> Result<DatumColumnEncoder, anyhow::Error> {
1769 let encoder = match &col_ty {
1770 SqlScalarType::Bool => DatumColumnEncoder::Bool(BooleanBuilder::with_capacity(0)),
1771 SqlScalarType::PgLegacyChar => DatumColumnEncoder::U8(UInt8Builder::with_capacity(0)),
1772 SqlScalarType::UInt16 => DatumColumnEncoder::U16(UInt16Builder::with_capacity(0)),
1773 SqlScalarType::UInt32
1774 | SqlScalarType::Oid
1775 | SqlScalarType::RegClass
1776 | SqlScalarType::RegProc
1777 | SqlScalarType::RegType => DatumColumnEncoder::U32(UInt32Builder::with_capacity(0)),
1778 SqlScalarType::UInt64 => DatumColumnEncoder::U64(UInt64Builder::with_capacity(0)),
1779 SqlScalarType::Int16 => DatumColumnEncoder::I16(Int16Builder::with_capacity(0)),
1780 SqlScalarType::Int32 => DatumColumnEncoder::I32(Int32Builder::with_capacity(0)),
1781 SqlScalarType::Int64 => DatumColumnEncoder::I64(Int64Builder::with_capacity(0)),
1782 SqlScalarType::Float32 => DatumColumnEncoder::F32(Float32Builder::with_capacity(0)),
1783 SqlScalarType::Float64 => DatumColumnEncoder::F64(Float64Builder::with_capacity(0)),
1784 SqlScalarType::Numeric { .. } => DatumColumnEncoder::Numeric {
1785 approx_values: Float64Builder::with_capacity(0),
1786 binary_values: BinaryBuilder::with_capacity(0, 0),
1787 numeric_context: crate::adt::numeric::cx_datum().clone(),
1788 },
1789 SqlScalarType::String
1790 | SqlScalarType::PgLegacyName
1791 | SqlScalarType::Char { .. }
1792 | SqlScalarType::VarChar { .. } => {
1793 DatumColumnEncoder::String(StringBuilder::with_capacity(0, 0))
1794 }
1795 SqlScalarType::Bytes => DatumColumnEncoder::Bytes(BinaryBuilder::with_capacity(0, 0)),
1796 SqlScalarType::Date => DatumColumnEncoder::Date(Int32Builder::with_capacity(0)),
1797 SqlScalarType::Time => {
1798 DatumColumnEncoder::Time(FixedSizeBinaryBuilder::with_capacity(0, TIME_FIXED_BYTES))
1799 }
1800 SqlScalarType::Timestamp { .. } => DatumColumnEncoder::Timestamp(
1801 FixedSizeBinaryBuilder::with_capacity(0, TIMESTAMP_FIXED_BYTES),
1802 ),
1803 SqlScalarType::TimestampTz { .. } => DatumColumnEncoder::TimestampTz(
1804 FixedSizeBinaryBuilder::with_capacity(0, TIMESTAMP_FIXED_BYTES),
1805 ),
1806 SqlScalarType::MzTimestamp => {
1807 DatumColumnEncoder::MzTimestamp(UInt64Builder::with_capacity(0))
1808 }
1809 SqlScalarType::Interval => DatumColumnEncoder::Interval(
1810 FixedSizeBinaryBuilder::with_capacity(0, INTERVAL_FIXED_BYTES),
1811 ),
1812 SqlScalarType::Uuid => {
1813 DatumColumnEncoder::Uuid(FixedSizeBinaryBuilder::with_capacity(0, UUID_FIXED_BYTES))
1814 }
1815 SqlScalarType::AclItem => DatumColumnEncoder::AclItem(
1816 FixedSizeBinaryBuilder::with_capacity(0, ACL_ITEM_FIXED_BYTES),
1817 ),
1818 SqlScalarType::MzAclItem => {
1819 DatumColumnEncoder::MzAclItem(BinaryBuilder::with_capacity(0, 0))
1820 }
1821 SqlScalarType::Range { .. } => {
1822 DatumColumnEncoder::Range(BinaryBuilder::with_capacity(0, 0))
1823 }
1824 SqlScalarType::Jsonb => DatumColumnEncoder::Jsonb {
1825 offsets: vec![0],
1826 buf: Vec::new(),
1827 nulls: None,
1828 },
1829 s @ SqlScalarType::Array(_) | s @ SqlScalarType::Int2Vector => {
1830 let element_type = match s {
1831 SqlScalarType::Array(inner) => inner,
1832 SqlScalarType::Int2Vector => &SqlScalarType::Int16,
1833 _ => unreachable!("checked above"),
1834 };
1835 let inner = scalar_type_to_encoder(element_type)?;
1836 DatumColumnEncoder::Array {
1837 dims: ListBuilder::new(FixedSizeBinaryBuilder::with_capacity(
1838 0,
1839 ARRAY_DIMENSION_FIXED_BYTES,
1840 )),
1841 val_lengths: Vec::new(),
1842 vals: Box::new(inner),
1843 nulls: None,
1844 }
1845 }
1846 SqlScalarType::List { element_type, .. } => {
1847 let inner = scalar_type_to_encoder(&*element_type)?;
1848 DatumColumnEncoder::List {
1849 lengths: Vec::new(),
1850 values: Box::new(inner),
1851 nulls: None,
1852 }
1853 }
1854 SqlScalarType::Map { value_type, .. } => {
1855 let inner = scalar_type_to_encoder(&*value_type)?;
1856 DatumColumnEncoder::Map {
1857 lengths: Vec::new(),
1858 keys: StringBuilder::with_capacity(0, 0),
1859 vals: Box::new(inner),
1860 nulls: None,
1861 }
1862 }
1863 SqlScalarType::Record { fields, .. } if fields.is_empty() => {
1864 DatumColumnEncoder::RecordEmpty(BooleanBuilder::with_capacity(0))
1865 }
1866 SqlScalarType::Record { fields, .. } => {
1867 let encoders = fields
1868 .iter()
1869 .map(|(_name, ty)| {
1870 scalar_type_to_encoder(&ty.scalar_type).map(|e| DatumEncoder {
1871 nullable: ty.nullable,
1872 encoder: e,
1873 })
1874 })
1875 .collect::<Result<_, _>>()?;
1876
1877 DatumColumnEncoder::Record {
1878 fields: encoders,
1879 nulls: None,
1880 length: 0,
1881 }
1882 }
1883 };
1884 Ok(encoder)
1885}
1886
1887impl Codec for Row {
1888 type Storage = ProtoRow;
1889 type Schema = RelationDesc;
1890
1891 fn codec_name() -> String {
1892 "protobuf[Row]".into()
1893 }
1894
1895 fn encode<B>(&self, buf: &mut B)
1901 where
1902 B: BufMut,
1903 {
1904 self.into_proto()
1905 .encode(buf)
1906 .expect("no required fields means no initialization errors");
1907 }
1908
1909 fn decode(buf: &[u8], schema: &RelationDesc) -> Result<Row, String> {
1915 let mut row = Row::with_capacity(buf.len());
1923 <Self as Codec>::decode_from(&mut row, buf, &mut None, schema)?;
1924 Ok(row)
1925 }
1926
1927 fn decode_from<'a>(
1928 &mut self,
1929 buf: &'a [u8],
1930 storage: &mut Option<ProtoRow>,
1931 schema: &RelationDesc,
1932 ) -> Result<(), String> {
1933 let mut proto = storage.take().unwrap_or_default();
1934 proto.clear();
1935 proto.merge(buf).map_err(|err| err.to_string())?;
1936 let ret = self.decode_from_proto(&proto, schema);
1937 storage.replace(proto);
1938 ret
1939 }
1940
1941 fn validate(row: &Self, desc: &Self::Schema) -> Result<(), String> {
1942 for x in Itertools::zip_longest(desc.iter_types(), row.iter()) {
1943 match x {
1944 EitherOrBoth::Both(typ, datum) if datum.is_instance_of_sql(typ) => continue,
1945 _ => return Err(format!("row {:?} did not match desc {:?}", row, desc)),
1946 };
1947 }
1948 Ok(())
1949 }
1950
1951 fn encode_schema(schema: &Self::Schema) -> Bytes {
1952 schema.into_proto().encode_to_vec().into()
1953 }
1954
1955 fn decode_schema(buf: &Bytes) -> Self::Schema {
1956 let proto = ProtoRelationDesc::decode(buf.as_ref()).expect("valid schema");
1957 proto.into_rust().expect("valid schema")
1958 }
1959}
1960
1961impl<'a> From<Datum<'a>> for ProtoDatum {
1962 fn from(x: Datum<'a>) -> Self {
1963 let datum_type = match x {
1964 Datum::False => DatumType::Other(ProtoDatumOther::False.into()),
1965 Datum::True => DatumType::Other(ProtoDatumOther::True.into()),
1966 Datum::Int16(x) => DatumType::Int16(x.into()),
1967 Datum::Int32(x) => DatumType::Int32(x),
1968 Datum::UInt8(x) => DatumType::Uint8(x.into()),
1969 Datum::UInt16(x) => DatumType::Uint16(x.into()),
1970 Datum::UInt32(x) => DatumType::Uint32(x),
1971 Datum::UInt64(x) => DatumType::Uint64(x),
1972 Datum::Int64(x) => DatumType::Int64(x),
1973 Datum::Float32(x) => DatumType::Float32(x.into_inner()),
1974 Datum::Float64(x) => DatumType::Float64(x.into_inner()),
1975 Datum::Date(x) => DatumType::Date(x.into_proto()),
1976 Datum::Time(x) => DatumType::Time(ProtoNaiveTime {
1977 secs: x.num_seconds_from_midnight(),
1978 frac: x.nanosecond(),
1979 }),
1980 Datum::Timestamp(x) => DatumType::Timestamp(x.into_proto()),
1981 Datum::TimestampTz(x) => DatumType::TimestampTz(x.into_proto()),
1982 Datum::Interval(x) => DatumType::Interval(x.into_proto()),
1983 Datum::Bytes(x) => DatumType::Bytes(Bytes::copy_from_slice(x)),
1984 Datum::String(x) => DatumType::String(x.to_owned()),
1985 Datum::Array(x) => DatumType::Array(ProtoArray {
1986 elements: Some(ProtoRow {
1987 datums: x.elements().iter().map(|x| x.into()).collect(),
1988 }),
1989 dims: x
1990 .dims()
1991 .into_iter()
1992 .map(|x| ProtoArrayDimension {
1993 lower_bound: i64::cast_from(x.lower_bound),
1994 length: u64::cast_from(x.length),
1995 })
1996 .collect(),
1997 }),
1998 Datum::List(x) => DatumType::List(ProtoRow {
1999 datums: x.iter().map(|x| x.into()).collect(),
2000 }),
2001 Datum::Map(x) => DatumType::Dict(ProtoDict {
2002 elements: x
2003 .iter()
2004 .map(|(k, v)| ProtoDictElement {
2005 key: k.to_owned(),
2006 val: Some(v.into()),
2007 })
2008 .collect(),
2009 }),
2010 Datum::Numeric(x) => {
2011 let mut x = x.0.clone();
2013 if let Some((bcd, scale)) = x.to_packed_bcd() {
2014 DatumType::Numeric(ProtoNumeric { bcd, scale })
2015 } else if x.is_nan() {
2016 DatumType::Other(ProtoDatumOther::NumericNaN.into())
2017 } else if x.is_infinite() {
2018 if x.is_negative() {
2019 DatumType::Other(ProtoDatumOther::NumericNegInf.into())
2020 } else {
2021 DatumType::Other(ProtoDatumOther::NumericPosInf.into())
2022 }
2023 } else if x.is_special() {
2024 panic!("internal error: unhandled special numeric value: {}", x);
2025 } else {
2026 panic!(
2027 "internal error: to_packed_bcd returned None for non-special value: {}",
2028 x
2029 )
2030 }
2031 }
2032 Datum::JsonNull => DatumType::Other(ProtoDatumOther::JsonNull.into()),
2033 Datum::Uuid(x) => DatumType::Uuid(x.as_bytes().to_vec()),
2034 Datum::MzTimestamp(x) => DatumType::MzTimestamp(x.into()),
2035 Datum::Dummy => DatumType::Other(ProtoDatumOther::Dummy.into()),
2036 Datum::Null => DatumType::Other(ProtoDatumOther::Null.into()),
2037 Datum::Range(super::Range { inner }) => DatumType::Range(Box::new(ProtoRange {
2038 inner: inner.map(|RangeInner { lower, upper }| {
2039 Box::new(ProtoRangeInner {
2040 lower_inclusive: lower.inclusive,
2041 lower: lower.bound.map(|bound| Box::new(bound.datum().into())),
2042 upper_inclusive: upper.inclusive,
2043 upper: upper.bound.map(|bound| Box::new(bound.datum().into())),
2044 })
2045 }),
2046 })),
2047 Datum::MzAclItem(x) => DatumType::MzAclItem(x.into_proto()),
2048 Datum::AclItem(x) => DatumType::AclItem(x.into_proto()),
2049 };
2050 ProtoDatum {
2051 datum_type: Some(datum_type),
2052 }
2053 }
2054}
2055
2056impl RowPacker<'_> {
2057 pub(crate) fn try_push_proto(&mut self, x: &ProtoDatum) -> Result<(), String> {
2058 match &x.datum_type {
2059 Some(DatumType::Other(o)) => match ProtoDatumOther::try_from(*o) {
2060 Ok(ProtoDatumOther::Unknown) => return Err("unknown datum type".into()),
2061 Ok(ProtoDatumOther::Null) => self.push(Datum::Null),
2062 Ok(ProtoDatumOther::False) => self.push(Datum::False),
2063 Ok(ProtoDatumOther::True) => self.push(Datum::True),
2064 Ok(ProtoDatumOther::JsonNull) => self.push(Datum::JsonNull),
2065 Ok(ProtoDatumOther::Dummy) => {
2066 #[cfg(feature = "tracing")]
2070 tracing::error!("protobuf decoding found Dummy datum");
2071 self.push(Datum::Dummy);
2072 }
2073 Ok(ProtoDatumOther::NumericPosInf) => self.push(Datum::from(Numeric::infinity())),
2074 Ok(ProtoDatumOther::NumericNegInf) => self.push(Datum::from(-Numeric::infinity())),
2075 Ok(ProtoDatumOther::NumericNaN) => self.push(Datum::from(Numeric::nan())),
2076 Err(_) => return Err(format!("unknown datum type: {}", o)),
2077 },
2078 Some(DatumType::Int16(x)) => {
2079 let x = i16::try_from(*x)
2080 .map_err(|_| format!("int16 field stored with out of range value: {}", *x))?;
2081 self.push(Datum::Int16(x))
2082 }
2083 Some(DatumType::Int32(x)) => self.push(Datum::Int32(*x)),
2084 Some(DatumType::Int64(x)) => self.push(Datum::Int64(*x)),
2085 Some(DatumType::Uint8(x)) => {
2086 let x = u8::try_from(*x)
2087 .map_err(|_| format!("uint8 field stored with out of range value: {}", *x))?;
2088 self.push(Datum::UInt8(x))
2089 }
2090 Some(DatumType::Uint16(x)) => {
2091 let x = u16::try_from(*x)
2092 .map_err(|_| format!("uint16 field stored with out of range value: {}", *x))?;
2093 self.push(Datum::UInt16(x))
2094 }
2095 Some(DatumType::Uint32(x)) => self.push(Datum::UInt32(*x)),
2096 Some(DatumType::Uint64(x)) => self.push(Datum::UInt64(*x)),
2097 Some(DatumType::Float32(x)) => self.push(Datum::Float32((*x).into())),
2098 Some(DatumType::Float64(x)) => self.push(Datum::Float64((*x).into())),
2099 Some(DatumType::Bytes(x)) => self.push(Datum::Bytes(x)),
2100 Some(DatumType::String(x)) => self.push(Datum::String(x)),
2101 Some(DatumType::Uuid(x)) => {
2102 let u = Uuid::from_slice(x).map_err(|err| err.to_string())?;
2107 self.push(Datum::Uuid(u));
2108 }
2109 Some(DatumType::Date(x)) => self.push(Datum::Date(x.clone().into_rust()?)),
2110 Some(DatumType::Time(x)) => self.push(Datum::Time(x.clone().into_rust()?)),
2111 Some(DatumType::Timestamp(x)) => self.push(Datum::Timestamp(x.clone().into_rust()?)),
2112 Some(DatumType::TimestampTz(x)) => {
2113 self.push(Datum::TimestampTz(x.clone().into_rust()?))
2114 }
2115 Some(DatumType::Interval(x)) => self.push(Datum::Interval(
2116 x.clone()
2117 .into_rust()
2118 .map_err(|e: TryFromProtoError| e.to_string())?,
2119 )),
2120 Some(DatumType::List(x)) => self.push_list_with(|row| -> Result<(), String> {
2121 for d in x.datums.iter() {
2122 row.try_push_proto(d)?;
2123 }
2124 Ok(())
2125 })?,
2126 Some(DatumType::Array(x)) => {
2127 let dims = x
2128 .dims
2129 .iter()
2130 .map(|x| ArrayDimension {
2131 lower_bound: isize::cast_from(x.lower_bound),
2132 length: usize::cast_from(x.length),
2133 })
2134 .collect::<Vec<_>>();
2135 match x.elements.as_ref() {
2136 None => self.try_push_array(&dims, [].iter()),
2137 Some(elements) => {
2138 let elements_row = Row::try_from(elements)?;
2141 self.try_push_array(&dims, elements_row.iter())
2142 }
2143 }
2144 .map_err(|err| err.to_string())?
2145 }
2146 Some(DatumType::Dict(x)) => self.push_dict_with(|row| -> Result<(), String> {
2147 let mut prev_key: Option<&str> = None;
2148 for e in x.elements.iter() {
2149 if let Some(prev) = prev_key
2153 && e.key.as_str() <= prev
2154 {
2155 return Err(format!(
2156 "dict keys must be unique and in ascending order, \
2157 but {:?} came after {:?}",
2158 e.key, prev,
2159 ));
2160 }
2161 prev_key = Some(e.key.as_str());
2162 row.push(Datum::from(e.key.as_str()));
2163 let val = e
2164 .val
2165 .as_ref()
2166 .ok_or_else(|| format!("missing val for key: {}", e.key))?;
2167 row.try_push_proto(val)?;
2168 }
2169 Ok(())
2170 })?,
2171 Some(DatumType::Numeric(x)) => {
2172 const MAX_BCD_LEN: usize =
2184 mz_ore::cast::u8_to_usize(NUMERIC_DATUM_MAX_PRECISION) / 2 + 1;
2185 if x.bcd.is_empty() || x.bcd.len() > MAX_BCD_LEN {
2186 return Err(format!(
2187 "ProtoNumeric.bcd length {} not in 1..={}",
2188 x.bcd.len(),
2189 MAX_BCD_LEN,
2190 ));
2191 }
2192 let n = Decimal::from_packed_bcd(&x.bcd, x.scale).map_err(|err| err.to_string())?;
2193 self.push(Datum::from(n))
2194 }
2195 Some(DatumType::MzTimestamp(x)) => self.push(Datum::MzTimestamp((*x).into())),
2196 Some(DatumType::Range(inner)) => {
2197 let ProtoRange { inner } = &**inner;
2198 match inner {
2199 None => self.push_range(Range { inner: None }).unwrap(),
2200 Some(inner) => {
2201 let ProtoRangeInner {
2202 lower_inclusive,
2203 lower,
2204 upper_inclusive,
2205 upper,
2206 } = &**inner;
2207
2208 let is_null_proto = |d: &ProtoDatum| {
2213 matches!(
2214 d.datum_type,
2215 Some(DatumType::Other(o))
2216 if ProtoDatumOther::try_from(o)
2217 == Ok(ProtoDatumOther::Null)
2218 )
2219 };
2220 if lower.as_deref().is_some_and(is_null_proto)
2221 || upper.as_deref().is_some_and(is_null_proto)
2222 {
2223 return Err("range bound cannot be Null".into());
2224 }
2225
2226 self.push_range_with(
2227 RangeLowerBound {
2228 inclusive: *lower_inclusive,
2229 bound: lower
2230 .as_ref()
2231 .map(|d| |row: &mut RowPacker| row.try_push_proto(&*d)),
2232 },
2233 RangeUpperBound {
2234 inclusive: *upper_inclusive,
2235 bound: upper
2236 .as_ref()
2237 .map(|d| |row: &mut RowPacker| row.try_push_proto(&*d)),
2238 },
2239 )
2240 .map_err(|err| err.to_string())?;
2241 }
2242 }
2243 }
2244 Some(DatumType::MzAclItem(x)) => self.push(Datum::MzAclItem(x.clone().into_rust()?)),
2245 Some(DatumType::AclItem(x)) => self.push(Datum::AclItem(x.clone().into_rust()?)),
2246 None => return Err("unknown datum type".into()),
2247 };
2248 Ok(())
2249 }
2250}
2251
2252impl TryFrom<&ProtoRow> for Row {
2254 type Error = String;
2255
2256 fn try_from(x: &ProtoRow) -> Result<Self, Self::Error> {
2257 let mut row = Row::default();
2260 let mut packer = row.packer();
2261 for d in x.datums.iter() {
2262 packer.try_push_proto(d)?;
2263 }
2264 Ok(row)
2265 }
2266}
2267
2268impl RustType<ProtoRow> for Row {
2269 fn into_proto(&self) -> ProtoRow {
2270 let datums = self.iter().map(|x| x.into()).collect();
2271 ProtoRow { datums }
2272 }
2273
2274 fn from_proto(proto: ProtoRow) -> Result<Self, TryFromProtoError> {
2275 let mut row = Row::default();
2278 let mut packer = row.packer();
2279 for d in proto.datums.iter() {
2280 packer
2281 .try_push_proto(d)
2282 .map_err(TryFromProtoError::RowConversionError)?;
2283 }
2284 Ok(row)
2285 }
2286}
2287
2288#[cfg(test)]
2289mod tests {
2290 use std::collections::BTreeSet;
2291
2292 use arrow::array::{ArrayData, make_array};
2293 use arrow::compute::SortOptions;
2294 use arrow::datatypes::ArrowNativeType;
2295 use arrow::row::SortField;
2296 use chrono::{DateTime, NaiveDate, NaiveTime, Utc};
2297 use mz_ore::assert_err;
2298 use mz_ore::collections::CollectionExt;
2299 use mz_persist::indexed::columnar::arrow::realloc_array;
2300 use mz_persist::metrics::ColumnarMetrics;
2301 use mz_persist_types::Codec;
2302 use mz_persist_types::arrow::{ArrayBound, ArrayOrd};
2303 use mz_persist_types::columnar::{codec_to_schema, schema_to_codec};
2304 use mz_proto::{ProtoType, RustType};
2305 use proptest::prelude::*;
2306 use proptest::strategy::Strategy;
2307 use uuid::Uuid;
2308
2309 use super::*;
2310 use crate::adt::array::ArrayDimension;
2311 use crate::adt::interval::Interval;
2312 use crate::adt::numeric::Numeric;
2313 use crate::adt::timestamp::CheckedTimestamp;
2314 use crate::relation::arb_relation_desc;
2315 use crate::{ColumnName, RowArena, SqlColumnType, arb_datum_for_column, arb_row_for_relation};
2316 use crate::{Datum, RelationDesc, Row, SqlScalarType};
2317
2318 #[mz_ore::test]
2319 fn proto_row_invalid_range_is_error() {
2320 use prost::Message;
2325 let bytes: &[u8] = &[
2326 0x0a, 0x03, 0xaa, 0x01, 0x00, 0x0a, 0x03, 0xaa, 0x01, 0x00, 0x0a, 0x03, 0xa2, 0x01,
2327 0x00, 0x0a, 0x03, 0xaa, 0x01, 0x00, 0x0a, 0x20, 0xfa, 0x01, 0x1d, 0x1d, 0x9f, 0x00,
2328 0x00, 0x00, 0xaa, 0x01, 0x00, 0x0a, 0x13, 0xf8, 0x01, 0x08, 0xaa, 0x0a, 0x03, 0xba,
2329 0x01, 0x00, 0x22, 0x03, 0xba, 0x01, 0x00, 0x12, 0x03, 0xaa, 0x01, 0x00, 0x0a, 0x03,
2330 0xaa, 0x01, 0x00,
2331 ];
2332 let proto = ProtoRow::decode(bytes).expect("crash input decodes as a proto");
2333 let result: Result<Row, _> = proto.into_rust();
2334 assert_err!(result);
2335 }
2336
2337 #[mz_ore::test]
2338 fn proto_row_unordered_dict_keys_is_error() {
2339 use prost::Message;
2344 let bytes: &[u8] = &[
2345 0x0a, 0x32, 0x18, 0x4e, 0x18, 0x18, 0x68, 0x4e, 0xe8, 0x68, 0x57, 0xba, 0x01, 0x0a,
2346 0x0a, 0x08, 0x60, 0xff, 0xff, 0x10, 0x12, 0x02, 0x10, 0x10, 0x99, 0x68, 0x0a, 0x18,
2347 0x18, 0x4e, 0x18, 0x18, 0x68, 0x4e, 0xe8, 0x5b, 0x18, 0x68, 0x57, 0xba, 0x01, 0x0a,
2348 0x0a, 0x08, 0x60, 0xff, 0xff, 0x10, 0x12, 0x02, 0x18, 0x10,
2349 ];
2350 let proto = ProtoRow::decode(bytes).expect("crash input decodes as a proto");
2351 let result: Result<Row, _> = proto.into_rust();
2352 assert_err!(result);
2353 }
2354
2355 #[mz_ore::test]
2356 #[cfg_attr(miri, ignore)] fn proto_row_oversized_numeric_bcd_is_error() {
2358 for len in [21, 64, 4096] {
2364 let mut bcd = vec![0x99u8; len];
2365 *bcd.last_mut().unwrap() = 0x9c;
2367 let proto = ProtoRow {
2368 datums: vec![ProtoDatum {
2369 datum_type: Some(DatumType::Numeric(ProtoNumeric { bcd, scale: 0 })),
2370 }],
2371 };
2372 assert_err!(Row::try_from(&proto));
2373 }
2374
2375 let mut cx = crate::adt::numeric::cx_datum();
2378 let digits = "9".repeat(usize::from(NUMERIC_DATUM_MAX_PRECISION));
2379 let n: Numeric = cx.parse(digits.as_str()).unwrap();
2380 let row = Row::pack_slice(&[Datum::Numeric(OrderedDecimal(n))]);
2381 assert_eq!(Row::try_from(&row.into_proto()).unwrap(), row);
2382 }
2383
2384 #[track_caller]
2385 fn roundtrip_datum<'a>(
2386 ty: SqlColumnType,
2387 datum: impl Iterator<Item = Datum<'a>>,
2388 metrics: &ColumnarMetrics,
2389 ) {
2390 let desc = RelationDesc::builder().with_column("a", ty).finish();
2391 let rows = datum.map(|d| Row::pack_slice(&[d])).collect();
2392 roundtrip_rows(&desc, rows, metrics)
2393 }
2394
2395 #[track_caller]
2396 fn roundtrip_rows(desc: &RelationDesc, rows: Vec<Row>, metrics: &ColumnarMetrics) {
2397 let mut encoder = <RelationDesc as Schema<Row>>::encoder(desc).unwrap();
2398 for row in &rows {
2399 encoder.append(row);
2400 }
2401 let col = encoder.finish();
2402
2403 let col = realloc_array(&col, metrics);
2405 {
2407 let proto = col.to_data().into_proto();
2408 let bytes = proto.encode_to_vec();
2409 let proto = mz_persist_types::arrow::ProtoArrayData::decode(&bytes[..]).unwrap();
2410 let array_data: ArrayData = proto.into_rust().unwrap();
2411
2412 let col_rnd = StructArray::from(array_data.clone());
2413 assert_eq!(col, col_rnd);
2414
2415 let col_dyn = arrow::array::make_array(array_data);
2416 let col_dyn = col_dyn.as_any().downcast_ref::<StructArray>().unwrap();
2417 assert_eq!(&col, col_dyn);
2418 }
2419
2420 let decoder = <RelationDesc as Schema<Row>>::decoder(desc, col.clone()).unwrap();
2421 let stats = decoder.stats();
2422
2423 let arena = RowArena::new();
2425 let (stats, stat_nulls): (Vec<_>, Vec<_>) = desc
2426 .iter()
2427 .map(|(name, ty)| {
2428 let col_stats = stats.cols.get(name.as_str()).unwrap();
2429 let lower_upper =
2430 crate::stats::col_values(&ty.scalar_type, &col_stats.values, &arena);
2431 let null_count = col_stats.nulls.map_or(0, |n| n.count);
2432
2433 (lower_upper, null_count)
2434 })
2435 .unzip();
2436 let mut actual_nulls = vec![0usize; stats.len()];
2438
2439 let mut rnd_row = Row::default();
2440 for (idx, og_row) in rows.iter().enumerate() {
2441 decoder.decode(idx, &mut rnd_row);
2442 assert_eq!(og_row, &rnd_row);
2443
2444 for (c_idx, (rnd_datum, ty)) in rnd_row.iter().zip_eq(desc.typ().columns()).enumerate()
2446 {
2447 let lower_upper = stats[c_idx];
2448
2449 if rnd_datum.is_null() {
2451 actual_nulls[c_idx] += 1;
2452 } else if let Some((lower, upper)) = lower_upper {
2453 assert!(rnd_datum >= lower, "{rnd_datum:?} is not >= {lower:?}");
2454 assert!(rnd_datum <= upper, "{rnd_datum:?} is not <= {upper:?}");
2455 } else {
2456 match &ty.scalar_type {
2457 SqlScalarType::Jsonb => (),
2459 SqlScalarType::AclItem
2461 | SqlScalarType::MzAclItem
2462 | SqlScalarType::Range { .. }
2463 | SqlScalarType::Array(_)
2464 | SqlScalarType::Map { .. }
2465 | SqlScalarType::List { .. }
2466 | SqlScalarType::Record { .. }
2467 | SqlScalarType::Int2Vector => (),
2468 other => panic!("should have collected stats for {other:?}"),
2469 }
2470 }
2471 }
2472 }
2473
2474 for (col_idx, (stats_count, actual_count)) in
2476 stat_nulls.iter().zip_eq(actual_nulls.iter()).enumerate()
2477 {
2478 assert_eq!(
2479 stats_count, actual_count,
2480 "column {col_idx} has incorrect number of nulls!"
2481 );
2482 }
2483
2484 let codec = schema_to_codec::<Row>(desc, &col).unwrap();
2486 let col2 = codec_to_schema::<Row>(desc, &codec).unwrap();
2487 assert_eq!(col2.as_ref(), &col);
2488
2489 let converter = arrow::row::RowConverter::new(vec![SortField::new_with_options(
2491 col.data_type().clone(),
2492 SortOptions {
2493 descending: false,
2494 nulls_first: false,
2495 },
2496 )])
2497 .expect("sortable");
2498 let rows = converter
2499 .convert_columns(&[Arc::new(col.clone())])
2500 .expect("convertible");
2501 let mut row_vec = rows.iter().collect::<Vec<_>>();
2502 row_vec.sort();
2503 let row_col = converter
2504 .convert_rows(row_vec)
2505 .expect("convertible")
2506 .into_element();
2507 assert_eq!(row_col.len(), col.len());
2508
2509 let ord = ArrayOrd::new(&col);
2510 let mut indices = (0..u64::usize_as(col.len())).collect::<Vec<_>>();
2511 indices.sort_by_key(|i| ord.at(i.as_usize()));
2512 let indices = UInt64Array::from(indices);
2513 let ord_col = ::arrow::compute::take(&col, &indices, None).expect("takeable");
2514 assert_eq!(row_col.as_ref(), ord_col.as_ref());
2515
2516 let ordered_prefix_len = desc
2518 .iter()
2519 .take_while(|(_, c)| preserves_order(&c.scalar_type))
2520 .count();
2521 let decoder = <RelationDesc as Schema<Row>>::decoder_any(desc, ord_col.as_ref()).unwrap();
2522 let rows = (0..ord_col.len()).map(|i| {
2523 let mut row = Row::default();
2524 decoder.decode(i, &mut row);
2525 row
2526 });
2527 for (a, b) in rows.tuple_windows() {
2528 let a_prefix = a.iter().take(ordered_prefix_len);
2529 let b_prefix = b.iter().take(ordered_prefix_len);
2530 assert!(
2531 a_prefix.cmp(b_prefix).is_le(),
2532 "ordering should be consistent on preserves_order columns: {:#?}\n{:?}\n{:?}",
2533 desc.iter().take(ordered_prefix_len).collect_vec(),
2534 a.iter().take(ordered_prefix_len).collect_vec(),
2535 b.iter().take(ordered_prefix_len).collect_vec()
2536 );
2537 }
2538
2539 assert_eq!(
2541 ord.goodbytes(),
2542 (0..col.len()).map(|i| ord.at(i).goodbytes()).sum::<usize>(),
2543 "total size should match the sum of the sizes at each index"
2544 );
2545
2546 if !ord_col.is_empty() {
2548 let min_idx = indices.values()[0].as_usize();
2549 let lower_bound = ArrayBound::new(ord_col, min_idx);
2550 let max_encoded_len = 1000;
2551 if let Some(proto) = lower_bound.to_proto_lower(max_encoded_len) {
2552 assert!(
2553 proto.encoded_len() <= max_encoded_len,
2554 "should respect the max len"
2555 );
2556 let array_data = proto.into_rust().expect("valid array");
2557 let new_lower_bound = ArrayBound::new(make_array(array_data), 0);
2558 assert!(
2559 new_lower_bound.get() <= lower_bound.get(),
2560 "proto-roundtripped bound should be <= the original"
2561 );
2562 }
2563 }
2564 }
2565
2566 #[mz_ore::test]
2567 #[cfg_attr(miri, ignore)] fn proptest_datums() {
2569 let strat = any::<SqlColumnType>().prop_flat_map(|ty| {
2570 proptest::collection::vec(arb_datum_for_column(ty.clone()), 0..16)
2571 .prop_map(move |d| (ty.clone(), d))
2572 });
2573 let metrics = ColumnarMetrics::disconnected();
2574
2575 proptest!(|((ty, datums) in strat)| {
2576 roundtrip_datum(ty.clone(), datums.iter().map(Datum::from), &metrics);
2577 })
2578 }
2579
2580 #[mz_ore::test]
2581 #[cfg_attr(miri, ignore)] fn proptest_non_empty_relation_descs() {
2583 let strat = arb_relation_desc(1..8).prop_flat_map(|desc| {
2584 proptest::collection::vec(arb_row_for_relation(&desc), 0..12)
2585 .prop_map(move |rows| (desc.clone(), rows))
2586 });
2587 let metrics = ColumnarMetrics::disconnected();
2588
2589 proptest!(|((desc, rows) in strat)| {
2590 roundtrip_rows(&desc, rows, &metrics)
2591 })
2592 }
2593
2594 #[mz_ore::test]
2595 fn empty_relation_desc_returns_error() {
2596 let empty_desc = RelationDesc::empty();
2597 let result = <RelationDesc as Schema<Row>>::encoder(&empty_desc);
2598 assert_err!(result);
2599 }
2600
2601 #[mz_ore::test]
2602 fn smoketest_collections() {
2603 let mut row = Row::default();
2604 let mut packer = row.packer();
2605 let metrics = ColumnarMetrics::disconnected();
2606
2607 packer
2608 .try_push_array(
2609 &[ArrayDimension {
2610 lower_bound: 0,
2611 length: 3,
2612 }],
2613 [Datum::UInt32(4), Datum::UInt32(5), Datum::UInt32(6)],
2614 )
2615 .unwrap();
2616
2617 let array = row.unpack_first();
2618 roundtrip_datum(
2619 SqlScalarType::Array(Box::new(SqlScalarType::UInt32)).nullable(true),
2620 [array].into_iter(),
2621 &metrics,
2622 );
2623 }
2624
2625 #[mz_ore::test]
2626 fn smoketest_row() {
2627 let desc = RelationDesc::builder()
2628 .with_column("a", SqlScalarType::Int64.nullable(true))
2629 .with_column("b", SqlScalarType::String.nullable(true))
2630 .with_column("c", SqlScalarType::Bool.nullable(true))
2631 .with_column(
2632 "d",
2633 SqlScalarType::List {
2634 element_type: Box::new(SqlScalarType::UInt32),
2635 custom_id: None,
2636 }
2637 .nullable(true),
2638 )
2639 .with_column(
2640 "e",
2641 SqlScalarType::Map {
2642 value_type: Box::new(SqlScalarType::Int16),
2643 custom_id: None,
2644 }
2645 .nullable(true),
2646 )
2647 .finish();
2648 let mut encoder = <RelationDesc as Schema<Row>>::encoder(&desc).unwrap();
2649
2650 let mut og_row = Row::default();
2651 {
2652 let mut packer = og_row.packer();
2653 packer.push(Datum::Int64(100));
2654 packer.push(Datum::String("hello world"));
2655 packer.push(Datum::True);
2656 packer.push_list([Datum::UInt32(1), Datum::UInt32(2), Datum::UInt32(3)]);
2657 packer.push_dict([("bar", Datum::Int16(9)), ("foo", Datum::Int16(3))]);
2658 }
2659 let mut og_row_2 = Row::default();
2660 {
2661 let mut packer = og_row_2.packer();
2662 packer.push(Datum::Null);
2663 packer.push(Datum::Null);
2664 packer.push(Datum::Null);
2665 packer.push(Datum::Null);
2666 packer.push(Datum::Null);
2667 }
2668
2669 encoder.append(&og_row);
2670 encoder.append(&og_row_2);
2671 let col = encoder.finish();
2672
2673 let decoder = <RelationDesc as Schema<Row>>::decoder(&desc, col).unwrap();
2674
2675 let mut rnd_row = Row::default();
2676 decoder.decode(0, &mut rnd_row);
2677 assert_eq!(og_row, rnd_row);
2678
2679 let mut rnd_row = Row::default();
2680 decoder.decode(1, &mut rnd_row);
2681 assert_eq!(og_row_2, rnd_row);
2682 }
2683
2684 #[mz_ore::test]
2685 fn test_nested_list() {
2686 let desc = RelationDesc::builder()
2687 .with_column(
2688 "a",
2689 SqlScalarType::List {
2690 element_type: Box::new(SqlScalarType::List {
2691 element_type: Box::new(SqlScalarType::Int64),
2692 custom_id: None,
2693 }),
2694 custom_id: None,
2695 }
2696 .nullable(false),
2697 )
2698 .finish();
2699 let mut encoder = <RelationDesc as Schema<Row>>::encoder(&desc).unwrap();
2700
2701 let mut og_row = Row::default();
2702 {
2703 let mut packer = og_row.packer();
2704 packer.push_list_with(|inner| {
2705 inner.push_list([Datum::Int64(1), Datum::Int64(2)]);
2706 inner.push_list([Datum::Int64(5)]);
2707 inner.push_list([Datum::Int64(9), Datum::Int64(99), Datum::Int64(999)]);
2708 });
2709 }
2710
2711 encoder.append(&og_row);
2712 let col = encoder.finish();
2713
2714 let decoder = <RelationDesc as Schema<Row>>::decoder(&desc, col).unwrap();
2715 let mut rnd_row = Row::default();
2716 decoder.decode(0, &mut rnd_row);
2717
2718 assert_eq!(og_row, rnd_row);
2719 }
2720
2721 #[mz_ore::test]
2722 fn test_record() {
2723 let desc = RelationDesc::builder()
2724 .with_column(
2725 "a",
2726 SqlScalarType::Record {
2727 fields: [
2728 (
2729 ColumnName::from("foo"),
2730 SqlScalarType::Int64.nullable(false),
2731 ),
2732 (
2733 ColumnName::from("bar"),
2734 SqlScalarType::String.nullable(true),
2735 ),
2736 (
2737 ColumnName::from("baz"),
2738 SqlScalarType::List {
2739 element_type: Box::new(SqlScalarType::UInt32),
2740 custom_id: None,
2741 }
2742 .nullable(false),
2743 ),
2744 ]
2745 .into(),
2746 custom_id: None,
2747 }
2748 .nullable(true),
2749 )
2750 .finish();
2751 let mut encoder = <RelationDesc as Schema<Row>>::encoder(&desc).unwrap();
2752
2753 let mut og_row = Row::default();
2754 {
2755 let mut packer = og_row.packer();
2756 packer.push_list_with(|inner| {
2757 inner.push(Datum::Int64(42));
2758 inner.push(Datum::Null);
2759 inner.push_list([Datum::UInt32(1), Datum::UInt32(2), Datum::UInt32(3)]);
2760 });
2761 }
2762 let null_row = Row::pack_slice(&[Datum::Null]);
2763
2764 encoder.append(&og_row);
2765 encoder.append(&null_row);
2766 let col = encoder.finish();
2767
2768 let decoder = <RelationDesc as Schema<Row>>::decoder(&desc, col).unwrap();
2769 let mut rnd_row = Row::default();
2770
2771 decoder.decode(0, &mut rnd_row);
2772 assert_eq!(og_row, rnd_row);
2773
2774 rnd_row.packer();
2775 decoder.decode(1, &mut rnd_row);
2776 assert_eq!(null_row, rnd_row);
2777 }
2778
2779 #[mz_ore::test]
2780 #[cfg_attr(miri, ignore)] fn roundtrip() {
2782 let mut row = Row::default();
2783 let mut packer = row.packer();
2784 packer.extend([
2785 Datum::False,
2786 Datum::True,
2787 Datum::Int16(1),
2788 Datum::Int32(2),
2789 Datum::Int64(3),
2790 Datum::Float32(4f32.into()),
2791 Datum::Float64(5f64.into()),
2792 Datum::Date(
2793 NaiveDate::from_ymd_opt(6, 7, 8)
2794 .unwrap()
2795 .try_into()
2796 .unwrap(),
2797 ),
2798 Datum::Time(NaiveTime::from_hms_opt(9, 10, 11).unwrap()),
2799 Datum::Timestamp(
2800 CheckedTimestamp::from_timestamplike(
2801 NaiveDate::from_ymd_opt(12, 13 % 12, 14)
2802 .unwrap()
2803 .and_time(NaiveTime::from_hms_opt(15, 16, 17).unwrap()),
2804 )
2805 .unwrap(),
2806 ),
2807 Datum::TimestampTz(
2808 CheckedTimestamp::from_timestamplike(DateTime::from_naive_utc_and_offset(
2809 NaiveDate::from_ymd_opt(18, 19 % 12, 20)
2810 .unwrap()
2811 .and_time(NaiveTime::from_hms_opt(21, 22, 23).unwrap()),
2812 Utc,
2813 ))
2814 .unwrap(),
2815 ),
2816 Datum::Interval(Interval {
2817 months: 24,
2818 days: 42,
2819 micros: 25,
2820 }),
2821 Datum::Bytes(&[26, 27]),
2822 Datum::String("28"),
2823 Datum::from(Numeric::from(29)),
2824 Datum::from(Numeric::infinity()),
2825 Datum::from(-Numeric::infinity()),
2826 Datum::from(Numeric::nan()),
2827 Datum::JsonNull,
2828 Datum::Uuid(Uuid::from_u128(30)),
2829 Datum::Dummy,
2830 Datum::Null,
2831 ]);
2832 packer
2833 .try_push_array(
2834 &[ArrayDimension {
2835 lower_bound: 2,
2836 length: 2,
2837 }],
2838 vec![Datum::Int32(31), Datum::Int32(32)],
2839 )
2840 .expect("valid array");
2841 packer.push_list_with(|packer| {
2842 packer.push(Datum::String("33"));
2843 packer.push_list_with(|packer| {
2844 packer.push(Datum::String("34"));
2845 packer.push(Datum::String("35"));
2846 });
2847 packer.push(Datum::String("36"));
2848 packer.push(Datum::String("37"));
2849 });
2850 packer.push_dict_with(|row| {
2851 let mut i = 38;
2854 for _ in 0..20 {
2855 row.push(Datum::String(&i.to_string()));
2856 row.push(Datum::Int32(i + 1));
2857 i += 2;
2858 }
2859 });
2860
2861 let mut desc = RelationDesc::builder();
2862 for (idx, _) in row.iter().enumerate() {
2863 desc = desc.with_column(idx.to_string(), SqlScalarType::Int32.nullable(true));
2866 }
2867 let desc = desc.finish();
2868
2869 let encoded = row.encode_to_vec();
2870 assert_eq!(Row::decode(&encoded, &desc), Ok(row));
2871 }
2872
2873 #[mz_ore::test]
2874 fn smoketest_projection() {
2875 let desc = RelationDesc::builder()
2876 .with_column("a", SqlScalarType::Int64.nullable(true))
2877 .with_column("b", SqlScalarType::String.nullable(true))
2878 .with_column("c", SqlScalarType::Bool.nullable(true))
2879 .finish();
2880 let mut encoder = <RelationDesc as Schema<Row>>::encoder(&desc).unwrap();
2881
2882 let mut og_row = Row::default();
2883 {
2884 let mut packer = og_row.packer();
2885 packer.push(Datum::Int64(100));
2886 packer.push(Datum::String("hello world"));
2887 packer.push(Datum::True);
2888 }
2889 let mut og_row_2 = Row::default();
2890 {
2891 let mut packer = og_row_2.packer();
2892 packer.push(Datum::Null);
2893 packer.push(Datum::Null);
2894 packer.push(Datum::Null);
2895 }
2896
2897 encoder.append(&og_row);
2898 encoder.append(&og_row_2);
2899 let col = encoder.finish();
2900
2901 let projected_desc = desc.apply_demand(&BTreeSet::from([0, 2]));
2902
2903 let decoder = <RelationDesc as Schema<Row>>::decoder(&projected_desc, col).unwrap();
2904
2905 let mut rnd_row = Row::default();
2906 decoder.decode(0, &mut rnd_row);
2907 let expected_row = Row::pack_slice(&[Datum::Int64(100), Datum::True]);
2908 assert_eq!(expected_row, rnd_row);
2909
2910 let mut rnd_row = Row::default();
2911 decoder.decode(1, &mut rnd_row);
2912 let expected_row = Row::pack_slice(&[Datum::Null, Datum::Null]);
2913 assert_eq!(expected_row, rnd_row);
2914 }
2915}