1use std::sync::Arc;
5
6use crate::Literal;
7use arrow::array::{
8 Array, ArrayRef, BooleanArray, Float64Array, Int8Array, Int64Array, ListArray, StringArray,
9 StructArray,
10};
11use arrow::buffer::{NullBuffer, OffsetBuffer};
12use arrow::datatypes::{DataType, Field, Fields};
13
14pub const TAG: &str = "__het_tag";
16pub const KEY: &str = "__het_key";
18pub const INT: &str = "__het_int";
20pub const FLOAT: &str = "__het_float";
22pub const STR: &str = "__het_str";
24pub const BOOL: &str = "__het_bool";
26pub const LIST: &str = "__het_list";
28pub const MAP: &str = "__het_map";
30pub const MAP_KEY: &str = "__het_mkey";
32pub const MAP_VALUE: &str = "__het_mval";
34pub const DYNAMIC_PREFIX: &str = "__het_value_";
36
37#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
39pub enum ValueError {
40 #[error("GF_VALUE_SCHEMA: heterogeneous schema is not a supported layout")]
42 Schema,
43 #[error("GF_VALUE_TAG: invalid heterogeneous tag {0}")]
45 Tag(i8),
46 #[error("GF_VALUE_NULL: heterogeneous tag or selected payload is null")]
48 NullPayload,
49 #[error("GF_VALUE_PAYLOAD: heterogeneous row has conflicting payloads")]
51 ConflictingPayload,
52 #[error("GF_VALUE_BOUNDS: heterogeneous row or nesting exceeds its layout")]
54 Bounds,
55 #[error("GF_VALUE_KIND: value cannot be encoded in this heterogeneous layout")]
57 Kind,
58 #[error("GF_VALUE_ARROW: {0}")]
60 Arrow(String),
61}
62
63impl ValueError {
64 #[must_use]
66 pub const fn code(&self) -> &'static str {
67 match self {
68 Self::Schema => "GF_VALUE_SCHEMA",
69 Self::Tag(_) => "GF_VALUE_TAG",
70 Self::NullPayload => "GF_VALUE_NULL",
71 Self::ConflictingPayload => "GF_VALUE_PAYLOAD",
72 Self::Bounds => "GF_VALUE_BOUNDS",
73 Self::Kind => "GF_VALUE_KIND",
74 Self::Arrow(_) => "GF_VALUE_ARROW",
75 }
76 }
77}
78
79#[derive(Debug, Clone, PartialEq, Eq)]
81pub enum Layout {
82 ScalarV1,
84 ConstantV1 {
86 depth: usize,
88 },
89 DynamicV1 {
91 payload_types: Vec<DataType>,
93 },
94}
95
96#[must_use]
97pub fn payload_field(index: impl std::fmt::Display) -> String {
99 format!("{DYNAMIC_PREFIX}{index}")
100}
101
102#[must_use]
103pub fn scalar_fields() -> Fields {
105 vec![
106 Field::new(TAG, DataType::Int8, false),
107 Field::new(INT, DataType::Int64, true),
108 Field::new(FLOAT, DataType::Float64, true),
109 Field::new(STR, DataType::Utf8, true),
110 Field::new(BOOL, DataType::Boolean, true),
111 ]
112 .into()
113}
114
115#[must_use]
116pub fn map_entry_fields(depth: usize) -> Fields {
118 vec![
119 Field::new(MAP_KEY, DataType::Utf8, false),
120 Field::new(MAP_VALUE, DataType::Struct(constant_fields(depth)), true),
121 ]
122 .into()
123}
124
125#[must_use]
126pub fn constant_fields(depth: usize) -> Fields {
128 let mut fields = vec![Field::new(KEY, DataType::Float64, true)];
129 fields.extend(scalar_fields().iter().map(|field| field.as_ref().clone()));
130 if depth > 0 {
131 fields.push(Field::new(
132 LIST,
133 DataType::new_list(DataType::Struct(constant_fields(depth - 1)), true),
134 true,
135 ));
136 fields.push(Field::new(
137 MAP,
138 DataType::new_list(DataType::Struct(map_entry_fields(depth - 1)), true),
139 true,
140 ));
141 }
142 fields.into()
143}
144
145#[must_use]
147pub fn dynamic_fields(types: &[DataType]) -> Fields {
148 let mut fields = vec![Field::new(TAG, DataType::Int8, false)];
149 fields.extend(
150 types
151 .iter()
152 .enumerate()
153 .map(|(i, ty)| Field::new(payload_field(i), ty.clone(), true)),
154 );
155 fields.into()
156}
157
158pub fn recognize(data_type: &DataType) -> Result<Option<Layout>, ValueError> {
163 Ok(recognize_layout(data_type)?.map(|layout| match layout {
164 RowLayout::Scalar => Layout::ScalarV1,
165 RowLayout::Constant { depth } => Layout::ConstantV1 { depth },
166 RowLayout::Dynamic { .. } => {
167 let DataType::Struct(fields) = data_type else {
168 unreachable!("recognized struct")
169 };
170 Layout::DynamicV1 {
171 payload_types: fields
172 .iter()
173 .skip(1)
174 .map(|field| field.data_type().clone())
175 .collect(),
176 }
177 }
178 }))
179}
180
181#[derive(Clone, Copy)]
182enum RowLayout {
183 Scalar,
184 Constant { depth: usize },
185 Dynamic { width: usize },
186}
187
188fn field_header(field: &Field, name: &str, nullable: bool) -> bool {
189 field.name() == name && field.is_nullable() == nullable && field.metadata().is_empty()
190}
191
192fn field_matches(field: &Field, name: &str, data_type: &DataType, nullable: bool) -> bool {
193 field_header(field, name, nullable) && field.data_type() == data_type
194}
195
196fn names_match(fields: &Fields, names: &[&str]) -> bool {
197 fields.len() == names.len()
198 && fields
199 .iter()
200 .zip(names)
201 .all(|(field, name)| field.name() == name)
202}
203
204fn dynamic_name_matches(name: &str, index: usize) -> bool {
205 let Some(suffix) = name.strip_prefix(DYNAMIC_PREFIX) else {
206 return false;
207 };
208 !suffix.is_empty()
210 && (suffix == "0" || !suffix.starts_with('0'))
211 && suffix.bytes().all(|byte| byte.is_ascii_digit())
212 && suffix.parse::<usize>() == Ok(index)
213}
214
215fn scalar_payloads_match(fields: &[Arc<Field>]) -> bool {
216 fields.len() == 5
217 && [
218 (TAG, DataType::Int8, false),
219 (INT, DataType::Int64, true),
220 (FLOAT, DataType::Float64, true),
221 (STR, DataType::Utf8, true),
222 (BOOL, DataType::Boolean, true),
223 ]
224 .iter()
225 .zip(fields)
226 .all(|((name, data_type, nullable), field)| {
227 field_matches(field, name, data_type, *nullable)
228 })
229}
230
231fn constant_depth(fields: &Fields) -> Result<usize, ValueError> {
232 if !matches!(fields.len(), 6 | 8)
233 || !field_matches(&fields[0], KEY, &DataType::Float64, true)
234 || !scalar_payloads_match(&fields[1..6])
235 {
236 return Err(ValueError::Schema);
237 }
238 if fields.len() == 6 {
239 return Ok(0);
240 }
241 if !field_header(&fields[6], LIST, true) || !field_header(&fields[7], MAP, true) {
242 return Err(ValueError::Schema);
243 }
244 let (DataType::List(list_item), DataType::List(map_item)) =
245 (fields[6].data_type(), fields[7].data_type())
246 else {
247 return Err(ValueError::Schema);
248 };
249 if !field_header(list_item, "item", true) || !field_header(map_item, "item", true) {
250 return Err(ValueError::Schema);
251 }
252 let (DataType::Struct(child), DataType::Struct(entries)) =
253 (list_item.data_type(), map_item.data_type())
254 else {
255 return Err(ValueError::Schema);
256 };
257 if entries.len() != 2
258 || !field_matches(&entries[0], MAP_KEY, &DataType::Utf8, false)
259 || !field_header(&entries[1], MAP_VALUE, true)
260 || entries[1].data_type() != list_item.data_type()
261 {
262 return Err(ValueError::Schema);
263 }
264 constant_depth(child)?
265 .checked_add(1)
266 .ok_or(ValueError::Bounds)
267}
268
269#[cfg(test)]
270thread_local! {
271 static SCHEMA_VISITS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
272}
273
274fn recognize_layout(data_type: &DataType) -> Result<Option<RowLayout>, ValueError> {
275 #[cfg(test)]
276 SCHEMA_VISITS.with(|visits| visits.set(visits.get() + 1));
277 let DataType::Struct(fields) = data_type else {
278 return Ok(None);
279 };
280 let Some(tag) = fields.iter().find(|field| field.name() == TAG) else {
281 return Ok(None);
282 };
283 let dynamic_names = fields.first().is_some_and(|field| field.name() == TAG)
284 && fields.len() > 1
285 && fields
286 .iter()
287 .skip(1)
288 .enumerate()
289 .all(|(index, field)| dynamic_name_matches(field.name(), index));
290 let complete_names = names_match(fields, &[TAG, INT, FLOAT, STR, BOOL])
291 || names_match(fields, &[KEY, TAG, INT, FLOAT, STR, BOOL])
292 || names_match(fields, &[KEY, TAG, INT, FLOAT, STR, BOOL, LIST, MAP])
293 || dynamic_names;
294 if fields.len() == 1 || (tag.data_type() != &DataType::Int8 && !complete_names) {
295 return Ok(None);
296 }
297 if scalar_payloads_match(fields) {
298 return Ok(Some(RowLayout::Scalar));
299 }
300 if fields.first().is_some_and(|field| field.name() == KEY) {
301 return constant_depth(fields).map(|depth| Some(RowLayout::Constant { depth }));
302 }
303 if dynamic_names
304 && i8::try_from(fields.len() - 1).is_ok()
305 && field_matches(&fields[0], TAG, &DataType::Int8, false)
306 && fields
307 .iter()
308 .skip(1)
309 .all(|field| field.is_nullable() && field.metadata().is_empty())
310 {
311 return Ok(Some(RowLayout::Dynamic {
312 width: fields.len() - 1,
313 }));
314 }
315 Err(ValueError::Schema)
316}
317
318pub enum Decoded<'a> {
320 Null,
322 Payload(&'a ArrayRef),
324 Map(&'a ArrayRef),
326}
327
328fn payload_is_null(array: &dyn Array, row: usize) -> bool {
329 if array.data_type() == &DataType::Null || array.is_null(row) {
330 return true;
331 }
332 matches!(
335 array.data_type(),
336 DataType::Dictionary(_, _) | DataType::RunEndEncoded(_, _)
337 ) && array.slice(row, 1).logical_null_count() == 1
338}
339
340pub fn decode_row(array: &StructArray, row: usize) -> Result<Decoded<'_>, ValueError> {
345 let layout = recognize_layout(array.data_type())?.ok_or(ValueError::Schema)?;
346 Ok(match select_row(array, row, layout)? {
347 None => Decoded::Null,
348 Some((index, true)) => Decoded::Map(array.column(index)),
349 Some((index, false)) => Decoded::Payload(array.column(index)),
350 })
351}
352
353fn select_row(
354 array: &StructArray,
355 row: usize,
356 layout: RowLayout,
357) -> Result<Option<(usize, bool)>, ValueError> {
358 if row >= array.len() {
359 return Err(ValueError::Bounds);
360 }
361 if array.is_null(row) {
362 return Ok(None);
363 }
364 let tag_index = usize::from(matches!(layout, RowLayout::Constant { .. }));
365 let tags = array
366 .column(tag_index)
367 .as_any()
368 .downcast_ref::<Int8Array>()
369 .ok_or(ValueError::Schema)?;
370 if tags.is_null(row) {
371 return Err(ValueError::NullPayload);
372 }
373 let tag = tags.value(row);
374 let (selected, map, first_payload) = match layout {
375 RowLayout::Scalar => match tag {
376 0..=3 => (
377 Some(usize::try_from(tag).map_err(|_| ValueError::Tag(tag))? + 1),
378 false,
379 1,
380 ),
381 4 => (None, false, 1),
382 _ => return Err(ValueError::Tag(tag)),
383 },
384 RowLayout::Constant { depth } => match tag {
385 0..=3 => (
386 Some(usize::try_from(tag).map_err(|_| ValueError::Tag(tag))? + 2),
387 false,
388 2,
389 ),
390 4..=5 if depth > 0 => (
391 Some(usize::try_from(tag).map_err(|_| ValueError::Tag(tag))? + 2),
392 tag == 5,
393 2,
394 ),
395 _ => return Err(ValueError::Tag(tag)),
396 },
397 RowLayout::Dynamic { width } => {
398 let index = usize::try_from(tag).map_err(|_| ValueError::Tag(tag))?;
399 if index >= width {
400 return Err(ValueError::Tag(tag));
401 }
402 (Some(index + 1), false, 1)
403 }
404 };
405 for index in first_payload..array.num_columns() {
406 if Some(index) != selected && !payload_is_null(array.column(index).as_ref(), row) {
407 return Err(ValueError::ConflictingPayload);
408 }
409 }
410 let Some(index) = selected else {
411 return Ok(None);
412 };
413 let payload = array.column(index);
414 if payload_is_null(payload.as_ref(), row) {
415 return Err(ValueError::NullPayload);
416 }
417 Ok(Some((index, map)))
418}
419
420pub fn contains_heterogeneous(data_type: &DataType) -> Result<bool, ValueError> {
425 let own = recognize_layout(data_type)?.is_some();
426 let mut nested = false;
427 match data_type {
428 DataType::Struct(fields) => {
429 for field in fields {
430 nested |= contains_heterogeneous(field.data_type())?;
431 }
432 }
433 DataType::List(field)
434 | DataType::LargeList(field)
435 | DataType::FixedSizeList(field, _)
436 | DataType::Map(field, _) => nested = contains_heterogeneous(field.data_type())?,
437 DataType::Dictionary(_, value) => nested = contains_heterogeneous(value)?,
438 _ => {}
439 }
440 Ok(own || nested)
441}
442
443pub fn validate_array(array: &dyn Array) -> Result<(), ValueError> {
448 let validator = ArrayValidator::new(array)?;
449 validator.rows(0..array.len())
450}
451
452enum ArrayValidator<'a> {
456 Skip,
457 Struct {
458 array: &'a StructArray,
459 layout: Option<RowLayout>,
460 children: Vec<ArrayValidator<'a>>,
461 },
462 List {
463 array: &'a dyn Array,
464 offsets: ChildOffsets<'a>,
465 child: Box<ArrayValidator<'a>>,
466 },
467 Dictionary {
468 array: &'a dyn Array,
469 value_type: &'a DataType,
470 },
471}
472
473enum ChildOffsets<'a> {
474 Small(&'a [i32]),
475 Large(&'a [i64]),
476 Fixed(usize),
477}
478
479impl ChildOffsets<'_> {
480 fn range(&self, row: usize) -> Result<std::ops::Range<usize>, ValueError> {
481 match self {
482 Self::Small(offsets) => Ok(usize::try_from(offsets[row])
483 .map_err(|_| ValueError::Bounds)?
484 ..usize::try_from(offsets[row + 1]).map_err(|_| ValueError::Bounds)?),
485 Self::Large(offsets) => Ok(usize::try_from(offsets[row])
486 .map_err(|_| ValueError::Bounds)?
487 ..usize::try_from(offsets[row + 1]).map_err(|_| ValueError::Bounds)?),
488 Self::Fixed(width) => {
489 let start = row.checked_mul(*width).ok_or(ValueError::Bounds)?;
490 Ok(start..start.checked_add(*width).ok_or(ValueError::Bounds)?)
491 }
492 }
493 }
494}
495
496impl<'a> ArrayValidator<'a> {
497 fn new(array: &'a dyn Array) -> Result<Self, ValueError> {
498 use arrow::array::{FixedSizeListArray, LargeListArray, MapArray};
499 if !contains_heterogeneous(array.data_type())? {
500 return Ok(Self::Skip);
501 }
502 if let DataType::Dictionary(_, value_type) = array.data_type() {
503 return Ok(Self::Dictionary { array, value_type });
504 }
505 if let Some(values) = array.as_any().downcast_ref::<StructArray>() {
506 return Ok(Self::Struct {
507 array: values,
508 layout: recognize_layout(values.data_type())?,
509 children: values
510 .columns()
511 .iter()
512 .map(|value| Self::new(value.as_ref()))
513 .collect::<Result<_, _>>()?,
514 });
515 }
516 let (offsets, child): (ChildOffsets<'a>, &'a dyn Array) =
517 if let Some(values) = array.as_any().downcast_ref::<ListArray>() {
518 (
519 ChildOffsets::Small(values.value_offsets()),
520 values.values().as_ref(),
521 )
522 } else if let Some(values) = array.as_any().downcast_ref::<LargeListArray>() {
523 (
524 ChildOffsets::Large(values.value_offsets()),
525 values.values().as_ref(),
526 )
527 } else if let Some(values) = array.as_any().downcast_ref::<FixedSizeListArray>() {
528 (
529 ChildOffsets::Fixed(
530 usize::try_from(values.value_length()).map_err(|_| ValueError::Bounds)?,
531 ),
532 values.values().as_ref(),
533 )
534 } else if let Some(values) = array.as_any().downcast_ref::<MapArray>() {
535 (
536 ChildOffsets::Small(values.value_offsets()),
537 values.entries(),
538 )
539 } else {
540 return Err(ValueError::Schema);
541 };
542 Ok(Self::List {
543 array,
544 offsets,
545 child: Box::new(Self::new(child)?),
546 })
547 }
548
549 fn rows(&self, rows: std::ops::Range<usize>) -> Result<(), ValueError> {
550 if matches!(self, Self::Skip) {
551 return Ok(());
552 }
553 for row in rows {
554 self.row(row)?;
555 }
556 Ok(())
557 }
558
559 fn row(&self, row: usize) -> Result<(), ValueError> {
560 match self {
561 Self::Skip => {}
562 Self::Struct {
563 array,
564 layout,
565 children,
566 } => {
567 if array.is_null(row) {
568 return Ok(());
569 }
570 if let Some(layout) = layout {
571 if let Some((index, _)) = select_row(array, row, *layout)? {
572 children[index].row(row)?;
573 }
574 } else {
575 for child in children {
576 child.row(row)?;
577 }
578 }
579 }
580 Self::List {
581 array,
582 offsets,
583 child,
584 } => {
585 if !array.is_null(row) {
586 child.rows(offsets.range(row)?)?;
587 }
588 }
589 Self::Dictionary { array, value_type } => {
590 if !payload_is_null(*array, row) {
591 let decoded = arrow::compute::cast(array.slice(row, 1).as_ref(), value_type)
592 .map_err(|error| ValueError::Arrow(error.to_string()))?;
593 validate_array(decoded.as_ref())?;
594 }
595 }
596 }
597 Ok(())
598 }
599}
600
601#[allow(
606 clippy::too_many_lines,
607 reason = "one exhaustive wire-kind table keeps payload columns and validity aligned"
608)]
609pub fn encode_constant(values: &[Literal], depth: usize) -> Result<StructArray, ValueError> {
610 let mut keys = Vec::new();
611 let mut tags = Vec::new();
612 let mut ints = Vec::new();
613 let mut floats = Vec::new();
614 let mut strings = Vec::new();
615 let mut bools = Vec::new();
616 let mut valid = Vec::new();
617 let mut children = Vec::new();
618 let mut offsets = vec![0_i32];
619 let mut child_valid = Vec::new();
620 let mut map_keys = Vec::new();
621 let mut map_values = Vec::new();
622 let mut map_offsets = vec![0_i32];
623 let mut map_valid = Vec::new();
624 for value in values {
625 let (mut key, mut tag, mut int, mut float, mut string, mut boolean) =
626 (None, 0, None, None, None, None);
627 match value {
628 Literal::Null => {}
629 Literal::Int(v) => {
630 #[allow(clippy::cast_precision_loss)]
631 {
632 key = Some(*v as f64);
633 }
634 int = Some(*v);
635 }
636 Literal::Float(v) => {
637 key = Some(*v);
638 float = Some(*v);
639 tag = 1;
640 }
641 Literal::Str(v) => {
642 string = Some(v.clone());
643 tag = 2;
644 }
645 Literal::Bool(v) => {
646 boolean = Some(*v);
647 tag = 3;
648 }
649 Literal::List(v) if depth > 0 => {
650 children.extend(v.iter().cloned());
651 tag = 4;
652 }
653 Literal::Map(v) if depth > 0 => {
654 for (k, v) in v {
655 map_keys.push(k.clone());
656 map_values.push(v.clone());
657 }
658 tag = 5;
659 }
660 _ => return Err(ValueError::Kind),
661 }
662 keys.push(key);
663 tags.push(tag);
664 ints.push(int);
665 floats.push(float);
666 strings.push(string);
667 bools.push(boolean);
668 valid.push(!matches!(value, Literal::Null));
669 child_valid.push(tag == 4);
670 map_valid.push(tag == 5);
671 offsets.push(i32::try_from(children.len()).map_err(|_| ValueError::Bounds)?);
672 map_offsets.push(i32::try_from(map_keys.len()).map_err(|_| ValueError::Bounds)?);
673 }
674 let mut columns: Vec<ArrayRef> = vec![
675 Arc::new(Float64Array::from(keys)),
676 Arc::new(Int8Array::from(tags)),
677 Arc::new(Int64Array::from(ints)),
678 Arc::new(Float64Array::from(floats)),
679 Arc::new(StringArray::from(strings)),
680 Arc::new(BooleanArray::from(bools)),
681 ];
682 if depth > 0 {
683 columns.push(Arc::new(ListArray::new(
684 Arc::new(Field::new(
685 "item",
686 DataType::Struct(constant_fields(depth - 1)),
687 true,
688 )),
689 OffsetBuffer::new(offsets.into()),
690 Arc::new(encode_constant(&children, depth - 1)?),
691 Some(NullBuffer::from(child_valid)),
692 )));
693 let entries = StructArray::new(
694 map_entry_fields(depth - 1),
695 vec![
696 Arc::new(StringArray::from(map_keys)),
697 Arc::new(encode_constant(&map_values, depth - 1)?),
698 ],
699 None,
700 );
701 columns.push(Arc::new(ListArray::new(
702 Arc::new(Field::new(
703 "item",
704 DataType::Struct(map_entry_fields(depth - 1)),
705 true,
706 )),
707 OffsetBuffer::new(map_offsets.into()),
708 Arc::new(entries),
709 Some(NullBuffer::from(map_valid)),
710 )));
711 }
712 StructArray::try_new(
713 constant_fields(depth),
714 columns,
715 Some(NullBuffer::from(valid)),
716 )
717 .map_err(|e| ValueError::Arrow(e.to_string()))
718}
719
720#[derive(Clone, Copy)]
722pub enum Scalar<'a> {
723 Null,
725 Int(i64),
727 Float(f64),
729 Str(&'a str),
731 Bool(bool),
733}
734
735#[must_use]
736pub fn encode_scalar<'a>(values: impl IntoIterator<Item = Option<Scalar<'a>>>) -> StructArray {
738 let mut tags = Vec::new();
739 let mut ints = Vec::new();
740 let mut floats = Vec::new();
741 let mut strings = Vec::new();
742 let mut bools = Vec::new();
743 let mut valid = Vec::new();
744 for value in values {
745 let (tag, int, float, string, boolean) = match value {
746 Some(Scalar::Int(v)) => (0, Some(v), None, None, None),
747 Some(Scalar::Float(v)) => (1, None, Some(v), None, None),
748 Some(Scalar::Str(v)) => (2, None, None, Some(v), None),
749 Some(Scalar::Bool(v)) => (3, None, None, None, Some(v)),
750 Some(Scalar::Null) => (4, None, None, None, None),
751 None => (0, None, None, None, None),
752 };
753 tags.push(tag);
754 ints.push(int);
755 floats.push(float);
756 strings.push(string);
757 bools.push(boolean);
758 valid.push(value.is_some());
759 }
760 StructArray::new(
761 scalar_fields(),
762 vec![
763 Arc::new(Int8Array::from(tags)),
764 Arc::new(Int64Array::from(ints)),
765 Arc::new(Float64Array::from(floats)),
766 Arc::new(StringArray::from(strings)),
767 Arc::new(BooleanArray::from(bools)),
768 ],
769 Some(NullBuffer::from(valid)),
770 )
771}
772
773pub fn decode_scalar(array: &StructArray, row: usize) -> Result<Literal, ValueError> {
778 if recognize(array.data_type())? != Some(Layout::ScalarV1) {
779 return Err(ValueError::Schema);
780 }
781 match decode_row(array, row)? {
782 Decoded::Null => Ok(Literal::Null),
783 Decoded::Payload(value) => match value.data_type() {
784 DataType::Int64 => Ok(Literal::Int(
785 value
786 .as_any()
787 .downcast_ref::<Int64Array>()
788 .ok_or(ValueError::Schema)?
789 .value(row),
790 )),
791 DataType::Float64 => Ok(Literal::Float(
792 value
793 .as_any()
794 .downcast_ref::<Float64Array>()
795 .ok_or(ValueError::Schema)?
796 .value(row),
797 )),
798 DataType::Utf8 => Ok(Literal::Str(
799 value
800 .as_any()
801 .downcast_ref::<StringArray>()
802 .ok_or(ValueError::Schema)?
803 .value(row)
804 .to_owned(),
805 )),
806 DataType::Boolean => Ok(Literal::Bool(
807 value
808 .as_any()
809 .downcast_ref::<BooleanArray>()
810 .ok_or(ValueError::Schema)?
811 .value(row),
812 )),
813 _ => Err(ValueError::Schema),
814 },
815 Decoded::Map(_) => Err(ValueError::Schema),
816 }
817}
818
819pub fn encode_dynamic_rows(
824 tags: Int8Array,
825 payloads: Vec<ArrayRef>,
826 validity: Option<NullBuffer>,
827) -> Result<StructArray, ValueError> {
828 if payloads.len() > i8::MAX as usize {
829 return Err(ValueError::Bounds);
830 }
831 let types = payloads
832 .iter()
833 .map(|value| value.data_type().clone())
834 .collect::<Vec<_>>();
835 let mut columns = vec![Arc::new(tags) as ArrayRef];
836 columns.extend(payloads);
837 let result = StructArray::try_new(dynamic_fields(&types), columns, validity)
838 .map_err(|error| ValueError::Arrow(error.to_string()))?;
839 validate_array(&result)?;
840 Ok(result)
841}
842
843pub fn encode_dynamic(values: &[ArrayRef], rows: usize) -> Result<ListArray, ValueError> {
848 use arrow::array::Int32Array;
849 use arrow::compute::take;
850 let width = values.len();
851 let width_i8 = i8::try_from(width).map_err(|_| ValueError::Bounds)?;
852 if values.iter().any(|value| value.len() != rows) {
853 return Err(ValueError::Bounds);
854 }
855 rows.checked_mul(width)
856 .filter(|length| i32::try_from(*length).is_ok())
857 .ok_or(ValueError::Bounds)?;
858 let types = values
859 .iter()
860 .map(|value| value.data_type().clone())
861 .collect::<Vec<_>>();
862 let fields = dynamic_fields(&types);
863 let tags = Int8Array::from_iter_values((0..rows).flat_map(|_| 0..width_i8));
864 let mut columns: Vec<ArrayRef> = vec![Arc::new(tags)];
865 for (index, value) in values.iter().enumerate() {
866 let indices = (0..rows)
867 .flat_map(|row| {
868 (0..width).map(move |element| {
869 (element == index)
870 .then(|| i32::try_from(row).ok())
871 .flatten()
872 })
873 })
874 .collect::<Int32Array>();
875 columns.push(
876 take(value.as_ref(), &indices, None).map_err(|e| ValueError::Arrow(e.to_string()))?,
877 );
878 }
879 let validity = (0..rows)
880 .flat_map(|row| {
881 values
882 .iter()
883 .map(move |value| !payload_is_null(value.as_ref(), row))
884 })
885 .collect::<NullBuffer>();
886 let tags = columns.remove(0);
887 let tags = tags
888 .as_any()
889 .downcast_ref::<Int8Array>()
890 .ok_or(ValueError::Schema)?
891 .clone();
892 let elements = encode_dynamic_rows(tags, columns, Some(validity))?;
893 Ok(ListArray::new(
894 Arc::new(Field::new("item", DataType::Struct(fields), true)),
895 OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows)),
896 Arc::new(elements),
897 None,
898 ))
899}
900
901#[cfg(test)]
902mod tests {
903 use super::*;
904
905 #[test]
906 fn scalar_v1_golden_fields_tags_and_null_states() {
907 let expected: Fields = vec![
908 Field::new("__het_tag", DataType::Int8, false),
909 Field::new("__het_int", DataType::Int64, true),
910 Field::new("__het_float", DataType::Float64, true),
911 Field::new("__het_str", DataType::Utf8, true),
912 Field::new("__het_bool", DataType::Boolean, true),
913 ]
914 .into();
915 assert_eq!(scalar_fields(), expected);
916 let values = encode_scalar([
917 Some(Scalar::Int(i64::MAX)),
918 Some(Scalar::Float(-1.5)),
919 Some(Scalar::Str("text")),
920 Some(Scalar::Bool(true)),
921 Some(Scalar::Null),
922 None,
923 ]);
924 let tags = values
925 .column(0)
926 .as_any()
927 .downcast_ref::<Int8Array>()
928 .unwrap();
929 assert_eq!(tags.values().as_ref(), &[0, 1, 2, 3, 4, 0]);
930 assert!(!values.is_null(4));
931 assert!(values.is_null(5));
932 let expected = [
933 Literal::Int(i64::MAX),
934 Literal::Float(-1.5),
935 Literal::Str("text".into()),
936 Literal::Bool(true),
937 Literal::Null,
938 Literal::Null,
939 ];
940 for (row, value) in expected.iter().enumerate() {
941 assert_eq!(&decode_scalar(&values, row).unwrap(), value);
942 }
943 validate_array(&values).unwrap();
944 }
945
946 #[test]
947 fn constant_v1_golden_nested_tags_and_dynamic_index_are_distinct() {
948 let values = encode_constant(
949 &[
950 Literal::Int(1),
951 Literal::Float(2.5),
952 Literal::Str("x".into()),
953 Literal::Bool(false),
954 Literal::List(vec![Literal::Int(9)]),
955 Literal::Map(vec![("key".into(), Literal::Null)]),
956 Literal::Null,
957 ],
958 1,
959 )
960 .unwrap();
961 assert_eq!(
962 values
963 .fields()
964 .iter()
965 .map(|f| f.name().as_str())
966 .collect::<Vec<_>>(),
967 [
968 "__het_key",
969 "__het_tag",
970 "__het_int",
971 "__het_float",
972 "__het_str",
973 "__het_bool",
974 "__het_list",
975 "__het_map"
976 ]
977 );
978 let tags = values
979 .column(1)
980 .as_any()
981 .downcast_ref::<Int8Array>()
982 .unwrap();
983 assert_eq!(tags.values().as_ref(), &[0, 1, 2, 3, 4, 5, 0]);
984 assert!(
985 matches!(decode_row(&values, 4).unwrap(), Decoded::Payload(value) if matches!(value.data_type(), DataType::List(_)))
986 );
987 assert!(matches!(decode_row(&values, 5).unwrap(), Decoded::Map(_)));
988 validate_array(&values).unwrap();
989 let columns: Vec<ArrayRef> = (0..6)
990 .map(|i| Arc::new(Int64Array::from(vec![i])) as ArrayRef)
991 .collect();
992 let dynamic = encode_dynamic(&columns, 1).unwrap();
993 let items = dynamic.value(0);
994 let items = items.as_any().downcast_ref::<StructArray>().unwrap();
995 assert_eq!(items.fields()[5].name(), "__het_value_4");
996 assert!(
997 matches!(decode_row(items, 4).unwrap(), Decoded::Payload(value) if value.data_type() == &DataType::Int64)
998 );
999 validate_array(&dynamic).unwrap();
1000 }
1001
1002 #[test]
1003 fn nested_batch_schema_work_is_independent_of_row_count() {
1004 let value = Literal::List(vec![Literal::Map(vec![("key".into(), Literal::Int(7))])]);
1005 let mut expected_visits = None;
1006 for rows in [1, 16, 256, 4096] {
1007 let array = encode_constant(&vec![value.clone(); rows], 2).unwrap();
1008 SCHEMA_VISITS.with(|visits| visits.set(0));
1009 validate_array(&array).unwrap();
1010 let visits = SCHEMA_VISITS.with(std::cell::Cell::get);
1011 assert!(visits > 0);
1012 assert_eq!(
1013 *expected_visits.get_or_insert(visits),
1014 visits,
1015 "schema admission must scale with the type, not {rows} visible rows"
1016 );
1017 let mut tags = vec![4_i8; rows];
1018 tags[rows - 1] = 99;
1019 let mut columns = array.columns().to_vec();
1020 columns[1] = Arc::new(Int8Array::from(tags));
1021 let malformed = StructArray::new(array.fields().clone(), columns, None);
1022 assert_eq!(
1023 validate_array(&malformed),
1024 Err(ValueError::Tag(99)),
1025 "every visible row remains checked after the bounded schema pass"
1026 );
1027 }
1028 }
1029
1030 #[test]
1031 fn dynamic_null_payload_is_logically_null_without_a_physical_bitmap() {
1032 let values: Vec<ArrayRef> = vec![
1033 Arc::new(StringArray::from(vec!["text"])),
1034 Arc::new(arrow::array::NullArray::new(1)),
1035 Arc::new(Int64Array::from(vec![7])),
1036 ];
1037 let lists = encode_dynamic(&values, 1).unwrap();
1038 validate_array(&lists).unwrap();
1039 let items = lists.value(0);
1040 let items = items.as_any().downcast_ref::<StructArray>().unwrap();
1041 assert!(matches!(decode_row(items, 0).unwrap(), Decoded::Payload(_)));
1042 assert!(matches!(decode_row(items, 1).unwrap(), Decoded::Null));
1043 assert!(matches!(decode_row(items, 2).unwrap(), Decoded::Payload(_)));
1044 }
1045
1046 #[test]
1047 fn malformed_tags_payloads_and_schema_have_distinct_errors() {
1048 let original = encode_scalar([Some(Scalar::Int(3))]);
1049 let mut columns = original.columns().to_vec();
1050 columns[0] = Arc::new(Int8Array::from(vec![99]));
1051 let invalid = StructArray::new(scalar_fields(), columns.clone(), None);
1052 assert!(matches!(decode_row(&invalid, 0), Err(ValueError::Tag(99))));
1053 columns[0] = Arc::new(Int8Array::from(vec![0]));
1054 columns[1] = Arc::new(Int64Array::from(vec![None]));
1055 assert!(matches!(
1056 decode_row(&StructArray::new(scalar_fields(), columns.clone(), None), 0),
1057 Err(ValueError::NullPayload)
1058 ));
1059 columns[1] = Arc::new(Int64Array::from(vec![3]));
1060 columns[2] = Arc::new(Float64Array::from(vec![1.0]));
1061 assert!(matches!(
1062 decode_row(&StructArray::new(scalar_fields(), columns, None), 0),
1063 Err(ValueError::ConflictingPayload)
1064 ));
1065 let mut fields = scalar_fields()
1066 .iter()
1067 .map(|f| f.as_ref().clone())
1068 .collect::<Vec<_>>();
1069 fields[1] = Field::new("__het_int", DataType::UInt64, true);
1070 assert_eq!(
1071 recognize(&DataType::Struct(fields.into())),
1072 Err(ValueError::Schema)
1073 );
1074 let mut fields = scalar_fields()
1075 .iter()
1076 .map(|field| field.as_ref().clone())
1077 .collect::<Vec<_>>();
1078 fields[0] = Field::new(TAG, DataType::Int64, false);
1079 assert_eq!(
1080 recognize(&DataType::Struct(fields.into())),
1081 Err(ValueError::Schema)
1082 );
1083 }
1084
1085 #[test]
1086 fn ordinary_property_names_do_not_claim_a_heterogeneous_layout() {
1087 for name in ["__het_int", "__het_key", "__het_value_0", "__het_tag"] {
1088 let ordinary = StructArray::new(
1089 vec![Field::new(name, DataType::Int64, false)].into(),
1090 vec![Arc::new(Int64Array::from(vec![7]))],
1091 None,
1092 );
1093 assert_eq!(recognize(ordinary.data_type()), Ok(None));
1094 validate_array(&ordinary).unwrap();
1095 }
1096 }
1097
1098 #[test]
1099 fn nested_validation_respects_null_parent_masks() {
1100 let original = encode_scalar([Some(Scalar::Int(3))]);
1101 let mut columns = original.columns().to_vec();
1102 columns[0] = Arc::new(Int8Array::from(vec![99]));
1103 let invalid = Arc::new(StructArray::new(scalar_fields(), columns, None)) as ArrayRef;
1104 let field = Arc::new(Field::new("item", invalid.data_type().clone(), true));
1105 let masked = ListArray::new(
1106 field.clone(),
1107 OffsetBuffer::new(vec![0, 1].into()),
1108 invalid.clone(),
1109 Some(NullBuffer::from(vec![false])),
1110 );
1111 validate_array(&masked).unwrap();
1112 let ordinary_fields =
1113 vec![Field::new("property", invalid.data_type().clone(), true)].into();
1114 let masked_struct = StructArray::new(
1115 ordinary_fields,
1116 vec![invalid.clone()],
1117 Some(NullBuffer::from(vec![false])),
1118 );
1119 validate_array(&masked_struct).unwrap();
1120 let masked_large = arrow::array::LargeListArray::new(
1121 field.clone(),
1122 OffsetBuffer::new(vec![0_i64, 1].into()),
1123 invalid.clone(),
1124 Some(NullBuffer::from(vec![false])),
1125 );
1126 validate_array(&masked_large).unwrap();
1127 let masked_fixed = arrow::array::FixedSizeListArray::new(
1128 field.clone(),
1129 1,
1130 invalid.clone(),
1131 Some(NullBuffer::from(vec![false])),
1132 );
1133 validate_array(&masked_fixed).unwrap();
1134 let unused_values = ListArray::new(
1135 field.clone(),
1136 OffsetBuffer::new(vec![0, 0].into()),
1137 invalid.clone(),
1138 None,
1139 );
1140 validate_array(&unused_values).unwrap();
1141 let visible = ListArray::new(field, OffsetBuffer::new(vec![0, 1].into()), invalid, None);
1142 assert_eq!(validate_array(&visible), Err(ValueError::Tag(99)));
1143 }
1144}