1use std::any::TypeId;
7use std::sync::Arc;
8
9use vortex_buffer::BitBuffer;
10use vortex_buffer::Buffer;
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_error::vortex_ensure;
14use vortex_error::vortex_panic;
15
16use crate::ArrayRef;
17use crate::ArraySlots;
18use crate::Executable;
19use crate::ExecutionCtx;
20use crate::IntoArray;
21use crate::array::ArrayView;
22use crate::array::VTable;
23use crate::array::child_to_validity;
24use crate::arrays::Bool;
25use crate::arrays::BoolArray;
26use crate::arrays::Decimal;
27use crate::arrays::DecimalArray;
28use crate::arrays::Extension;
29use crate::arrays::ExtensionArray;
30use crate::arrays::FixedSizeList;
31use crate::arrays::FixedSizeListArray;
32use crate::arrays::ListView;
33use crate::arrays::ListViewArray;
34use crate::arrays::Map;
35use crate::arrays::MapArray;
36use crate::arrays::Null;
37use crate::arrays::NullArray;
38use crate::arrays::Primitive;
39use crate::arrays::PrimitiveArray;
40use crate::arrays::Struct;
41use crate::arrays::StructArray;
42use crate::arrays::Union;
43use crate::arrays::UnionArray;
44use crate::arrays::VarBinView;
45use crate::arrays::VarBinViewArray;
46use crate::arrays::Variant;
47use crate::arrays::VariantArray;
48use crate::arrays::bool::BoolDataParts;
49use crate::arrays::decimal::DecimalDataParts;
50use crate::arrays::extension::ExtensionArrayExt;
51use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
52use crate::arrays::listview::ListViewDataParts;
53use crate::arrays::listview::ListViewRebuildMode;
54use crate::arrays::map::MapArrayExt;
55use crate::arrays::map::MapArraySlotsExt;
56use crate::arrays::primitive::PrimitiveDataParts;
57use crate::arrays::struct_::StructDataParts;
58use crate::arrays::union::UnionDataParts;
59use crate::arrays::varbinview::VarBinViewDataParts;
60use crate::arrays::variant::VariantArraySlotsExt;
61use crate::dtype::DType;
62use crate::dtype::NativePType;
63use crate::dtype::Nullability;
64use crate::dtype::PType;
65use crate::match_each_decimal_value_type;
66use crate::match_each_native_ptype;
67use crate::matcher::Matcher;
68use crate::validity::Validity;
69
70#[derive(Debug, Clone)]
134pub enum Canonical {
135 Null(NullArray),
136 Bool(BoolArray),
137 Primitive(PrimitiveArray),
138 Decimal(DecimalArray),
139 VarBinView(VarBinViewArray),
140 List(ListViewArray),
141 Map(MapArray),
142 FixedSizeList(FixedSizeListArray),
143 Struct(StructArray),
144 Union(UnionArray),
145 Extension(ExtensionArray),
147 Variant(VariantArray),
149}
150
151macro_rules! match_each_canonical {
153 ($self:expr, | $ident:ident | $eval:expr) => {{
154 match $self {
155 Canonical::Null($ident) => $eval,
156 Canonical::Bool($ident) => $eval,
157 Canonical::Primitive($ident) => $eval,
158 Canonical::Decimal($ident) => $eval,
159 Canonical::VarBinView($ident) => $eval,
160 Canonical::List($ident) => $eval,
161 Canonical::Map($ident) => $eval,
162 Canonical::FixedSizeList($ident) => $eval,
163 Canonical::Struct($ident) => $eval,
164 Canonical::Union($ident) => $eval,
165 Canonical::Variant($ident) => $eval,
166 Canonical::Extension($ident) => $eval,
167 }
168 }};
169}
170
171impl Canonical {
172 pub fn empty(dtype: &DType) -> Canonical {
174 match dtype {
175 DType::Null => Canonical::Null(NullArray::new(0)),
176 DType::Bool(n) => Canonical::Bool(unsafe {
177 BoolArray::new_unchecked(BitBuffer::empty(), Validity::from(n))
178 }),
179 DType::Primitive(ptype, n) => {
180 match_each_native_ptype!(ptype, |P| {
181 Canonical::Primitive(unsafe {
182 PrimitiveArray::new_unchecked(Buffer::<P>::empty(), Validity::from(n))
183 })
184 })
185 }
186 DType::Decimal(decimal_type, n) => {
187 match_each_decimal_value_type!(
188 DecimalType::smallest_decimal_value_type(decimal_type),
189 |D| {
190 Canonical::Decimal(unsafe {
191 DecimalArray::new_unchecked::<D>(
192 Buffer::empty(),
193 *decimal_type,
194 Validity::from(n),
195 )
196 })
197 }
198 )
199 }
200 DType::Utf8(n) => Canonical::VarBinView(unsafe {
201 VarBinViewArray::new_unchecked(
202 Buffer::empty(),
203 Arc::new([]),
204 dtype.clone(),
205 Validity::from(n),
206 )
207 }),
208 DType::Binary(n) => Canonical::VarBinView(unsafe {
209 VarBinViewArray::new_unchecked(
210 Buffer::empty(),
211 Arc::new([]),
212 dtype.clone(),
213 Validity::from(n),
214 )
215 }),
216 DType::List(dtype, n) => Canonical::List(unsafe {
217 ListViewArray::new_unchecked(
218 Canonical::empty(dtype).into_array(),
219 Canonical::empty(&DType::Primitive(PType::U8, Nullability::NonNullable))
220 .into_array(),
221 Canonical::empty(&DType::Primitive(PType::U8, Nullability::NonNullable))
222 .into_array(),
223 Validity::from(n),
224 )
225 .with_zero_copy_to_list(true)
227 }),
228 DType::Map(map_dtype, nullability) => Canonical::Map(MapArray::new(
229 map_dtype.clone(),
230 Canonical::empty(&DType::List(
231 Arc::new(map_dtype.entries_dtype()),
232 *nullability,
233 ))
234 .into_listview(),
235 )),
236 DType::FixedSizeList(elem_dtype, list_size, null) => Canonical::FixedSizeList(unsafe {
237 FixedSizeListArray::new_unchecked(
238 Canonical::empty(elem_dtype).into_array(),
239 *list_size,
240 Validity::from(null),
241 0,
242 )
243 }),
244 DType::Struct(struct_dtype, n) => Canonical::Struct(unsafe {
245 StructArray::new_unchecked(
246 struct_dtype
247 .fields()
248 .map(|f| Canonical::empty(&f).into_array())
249 .collect::<Vec<_>>(),
250 struct_dtype.clone(),
251 0,
252 Validity::from(n),
253 )
254 }),
255 DType::Union(variants, nullability) => {
256 Canonical::Union(UnionArray::empty(variants.clone(), *nullability))
257 }
258 DType::Variant(_) => {
259 vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant")
260 }
261 DType::Extension(ext_dtype) => Canonical::Extension(ExtensionArray::new(
262 ext_dtype.clone(),
263 Canonical::empty(ext_dtype.storage_dtype()).into_array(),
264 )),
265 }
266 }
267
268 pub fn len(&self) -> usize {
269 match_each_canonical!(self, |arr| arr.len())
270 }
271
272 pub fn dtype(&self) -> &DType {
273 match_each_canonical!(self, |arr| arr.dtype())
274 }
275
276 pub fn is_empty(&self) -> bool {
277 match_each_canonical!(self, |arr| arr.is_empty())
278 }
279}
280
281impl Canonical {
282 pub fn compact(&self, ctx: &mut ExecutionCtx) -> VortexResult<Canonical> {
290 match self {
291 Canonical::Bool(array) => Ok(Canonical::Bool(array.trim_bits()?)),
292 Canonical::VarBinView(array) => Ok(Canonical::VarBinView(array.compact_buffers(ctx)?)),
293 Canonical::List(array) => Ok(Canonical::List(
294 array.rebuild(ListViewRebuildMode::TrimElements, ctx)?,
295 )),
296 Canonical::Map(array) => Ok(Canonical::Map(MapArray::new(
297 array.map_dtype().clone(),
298 array
299 .entries()
300 .as_::<ListView>()
301 .into_owned()
302 .rebuild(ListViewRebuildMode::TrimElements, ctx)?,
303 ))),
304 _ => Ok(self.clone()),
305 }
306 }
307}
308
309impl Canonical {
311 pub fn as_null(&self) -> &NullArray {
312 if let Canonical::Null(a) = self {
313 a
314 } else {
315 vortex_panic!("Cannot get NullArray from {:?}", &self)
316 }
317 }
318
319 pub fn into_null(self) -> NullArray {
320 if let Canonical::Null(a) = self {
321 a
322 } else {
323 vortex_panic!("Cannot unwrap NullArray from {:?}", &self)
324 }
325 }
326
327 pub fn as_bool(&self) -> &BoolArray {
328 if let Canonical::Bool(a) = self {
329 a
330 } else {
331 vortex_panic!("Cannot get BoolArray from {:?}", &self)
332 }
333 }
334
335 pub fn into_bool(self) -> BoolArray {
336 if let Canonical::Bool(a) = self {
337 a
338 } else {
339 vortex_panic!("Cannot unwrap BoolArray from {:?}", &self)
340 }
341 }
342
343 pub fn as_primitive(&self) -> &PrimitiveArray {
344 if let Canonical::Primitive(a) = self {
345 a
346 } else {
347 vortex_panic!("Cannot get PrimitiveArray from {:?}", &self)
348 }
349 }
350
351 pub fn into_primitive(self) -> PrimitiveArray {
352 if let Canonical::Primitive(a) = self {
353 a
354 } else {
355 vortex_panic!("Cannot unwrap PrimitiveArray from {:?}", &self)
356 }
357 }
358
359 pub fn as_decimal(&self) -> &DecimalArray {
360 if let Canonical::Decimal(a) = self {
361 a
362 } else {
363 vortex_panic!("Cannot get DecimalArray from {:?}", &self)
364 }
365 }
366
367 pub fn into_decimal(self) -> DecimalArray {
368 if let Canonical::Decimal(a) = self {
369 a
370 } else {
371 vortex_panic!("Cannot unwrap DecimalArray from {:?}", &self)
372 }
373 }
374
375 pub fn as_varbinview(&self) -> &VarBinViewArray {
376 if let Canonical::VarBinView(a) = self {
377 a
378 } else {
379 vortex_panic!("Cannot get VarBinViewArray from {:?}", &self)
380 }
381 }
382
383 pub fn into_varbinview(self) -> VarBinViewArray {
384 if let Canonical::VarBinView(a) = self {
385 a
386 } else {
387 vortex_panic!("Cannot unwrap VarBinViewArray from {:?}", &self)
388 }
389 }
390
391 pub fn as_listview(&self) -> &ListViewArray {
392 if let Canonical::List(a) = self {
393 a
394 } else {
395 vortex_panic!("Cannot get ListArray from {:?}", &self)
396 }
397 }
398
399 pub fn into_listview(self) -> ListViewArray {
400 if let Canonical::List(a) = self {
401 a
402 } else {
403 vortex_panic!("Cannot unwrap ListArray from {:?}", &self)
404 }
405 }
406
407 pub fn as_map(&self) -> &MapArray {
408 if let Canonical::Map(a) = self {
409 a
410 } else {
411 vortex_panic!("Cannot get MapArray from {:?}", &self)
412 }
413 }
414
415 pub fn into_map(self) -> MapArray {
416 if let Canonical::Map(a) = self {
417 a
418 } else {
419 vortex_panic!("Cannot unwrap MapArray from {:?}", &self)
420 }
421 }
422
423 pub fn as_fixed_size_list(&self) -> &FixedSizeListArray {
424 if let Canonical::FixedSizeList(a) = self {
425 a
426 } else {
427 vortex_panic!("Cannot get FixedSizeListArray from {:?}", &self)
428 }
429 }
430
431 pub fn into_fixed_size_list(self) -> FixedSizeListArray {
432 if let Canonical::FixedSizeList(a) = self {
433 a
434 } else {
435 vortex_panic!("Cannot unwrap FixedSizeListArray from {:?}", &self)
436 }
437 }
438
439 pub fn as_struct(&self) -> &StructArray {
440 if let Canonical::Struct(a) = self {
441 a
442 } else {
443 vortex_panic!("Cannot get StructArray from {:?}", &self)
444 }
445 }
446
447 pub fn into_struct(self) -> StructArray {
448 if let Canonical::Struct(a) = self {
449 a
450 } else {
451 vortex_panic!("Cannot unwrap StructArray from {:?}", &self)
452 }
453 }
454
455 pub fn as_union(&self) -> &UnionArray {
457 if let Canonical::Union(a) = self {
458 a
459 } else {
460 vortex_panic!("Cannot get UnionArray from {:?}", &self)
461 }
462 }
463
464 pub fn into_union(self) -> UnionArray {
466 if let Canonical::Union(a) = self {
467 a
468 } else {
469 vortex_panic!("Cannot unwrap UnionArray from {:?}", &self)
470 }
471 }
472
473 pub fn as_extension(&self) -> &ExtensionArray {
474 if let Canonical::Extension(a) = self {
475 a
476 } else {
477 vortex_panic!("Cannot get ExtensionArray from {:?}", &self)
478 }
479 }
480
481 pub fn into_extension(self) -> ExtensionArray {
482 if let Canonical::Extension(a) = self {
483 a
484 } else {
485 vortex_panic!("Cannot unwrap ExtensionArray from {:?}", &self)
486 }
487 }
488}
489
490impl IntoArray for Canonical {
491 fn into_array(self) -> ArrayRef {
492 match_each_canonical!(self, |arr| arr.into_array())
493 }
494}
495
496#[deprecated(note = "use `array.execute::<T>(ctx)` instead")]
502pub trait ToCanonical {
503 #[deprecated(note = "use `array.execute::<NullArray>(ctx)` instead")]
505 fn to_null(&self) -> NullArray;
506
507 #[deprecated(note = "use `array.execute::<BoolArray>(ctx)` instead")]
509 fn to_bool(&self) -> BoolArray;
510
511 #[deprecated(note = "use `array.execute::<PrimitiveArray>(ctx)` instead")]
514 fn to_primitive(&self) -> PrimitiveArray;
515
516 #[deprecated(note = "use `array.execute::<DecimalArray>(ctx)` instead")]
519 fn to_decimal(&self) -> DecimalArray;
520
521 #[deprecated(note = "use `array.execute::<StructArray>(ctx)` instead")]
523 fn to_struct(&self) -> StructArray;
524
525 #[deprecated(note = "use `array.execute::<ListViewArray>(ctx)` instead")]
527 fn to_listview(&self) -> ListViewArray;
528
529 #[deprecated(note = "use `array.execute::<MapArray>(ctx)` instead")]
531 fn to_map(&self) -> MapArray;
532
533 #[deprecated(note = "use `array.execute::<FixedSizeListArray>(ctx)` instead")]
536 fn to_fixed_size_list(&self) -> FixedSizeListArray;
537
538 #[deprecated(note = "use `array.execute::<VarBinViewArray>(ctx)` instead")]
541 fn to_varbinview(&self) -> VarBinViewArray;
542
543 #[deprecated(note = "use `array.execute::<ExtensionArray>(ctx)` instead")]
546 fn to_extension(&self) -> ExtensionArray;
547}
548
549#[expect(deprecated)]
551impl ToCanonical for ArrayRef {
552 fn to_null(&self) -> NullArray {
553 #[expect(deprecated)]
554 let result = self.to_canonical().vortex_expect("to_canonical failed");
555 result.into_null()
556 }
557
558 fn to_bool(&self) -> BoolArray {
559 #[expect(deprecated)]
560 let result = self.to_canonical().vortex_expect("to_canonical failed");
561 result.into_bool()
562 }
563
564 fn to_primitive(&self) -> PrimitiveArray {
565 #[expect(deprecated)]
566 let result = self.to_canonical().vortex_expect("to_canonical failed");
567 result.into_primitive()
568 }
569
570 fn to_decimal(&self) -> DecimalArray {
571 #[expect(deprecated)]
572 let result = self.to_canonical().vortex_expect("to_canonical failed");
573 result.into_decimal()
574 }
575
576 fn to_struct(&self) -> StructArray {
577 #[expect(deprecated)]
578 let result = self.to_canonical().vortex_expect("to_canonical failed");
579 result.into_struct()
580 }
581
582 fn to_listview(&self) -> ListViewArray {
583 #[expect(deprecated)]
584 let result = self.to_canonical().vortex_expect("to_canonical failed");
585 result.into_listview()
586 }
587
588 fn to_map(&self) -> MapArray {
589 #[expect(deprecated)]
590 let result = self.to_canonical().vortex_expect("to_canonical failed");
591 result.into_map()
592 }
593
594 fn to_fixed_size_list(&self) -> FixedSizeListArray {
595 #[expect(deprecated)]
596 let result = self.to_canonical().vortex_expect("to_canonical failed");
597 result.into_fixed_size_list()
598 }
599
600 fn to_varbinview(&self) -> VarBinViewArray {
601 #[expect(deprecated)]
602 let result = self.to_canonical().vortex_expect("to_canonical failed");
603 result.into_varbinview()
604 }
605
606 fn to_extension(&self) -> ExtensionArray {
607 #[expect(deprecated)]
608 let result = self.to_canonical().vortex_expect("to_canonical failed");
609 result.into_extension()
610 }
611}
612
613impl From<Canonical> for ArrayRef {
614 fn from(value: Canonical) -> Self {
615 match_each_canonical!(value, |arr| arr.into_array())
616 }
617}
618
619impl Executable for Canonical {
625 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
626 let result = array.execute_until::<AnyCanonical>(ctx)?;
627 Ok(result
628 .as_opt::<AnyCanonical>()
629 .map(Canonical::from)
630 .vortex_expect("execute_until::<AnyCanonical> must return a canonical array"))
631 }
632}
633
634pub struct CanonicalValidity(pub Canonical);
639
640impl Executable for CanonicalValidity {
641 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
642 match array.execute::<Canonical>(ctx)? {
643 n @ Canonical::Null(_) => Ok(CanonicalValidity(n)),
644 Canonical::Bool(b) => {
645 let validity = child_to_validity(b.slots()[0].as_ref(), b.dtype().nullability());
646 let len = b.len();
647 let BoolDataParts { bits, meta } = b.into_data().into_parts(len);
648 Ok(CanonicalValidity(Canonical::Bool(
649 BoolArray::try_new_from_handle(
650 bits,
651 meta.offset(),
652 meta.len(),
653 validity.execute(ctx)?,
654 )?,
655 )))
656 }
657 Canonical::Primitive(p) => {
658 let PrimitiveDataParts {
659 ptype,
660 buffer,
661 validity,
662 } = p.into_data_parts();
663 Ok(CanonicalValidity(Canonical::Primitive(unsafe {
664 PrimitiveArray::new_unchecked_from_handle(buffer, ptype, validity.execute(ctx)?)
665 })))
666 }
667 Canonical::Decimal(d) => {
668 let DecimalDataParts {
669 decimal_dtype,
670 values,
671 values_type,
672 validity,
673 } = d.into_data_parts();
674 Ok(CanonicalValidity(Canonical::Decimal(unsafe {
675 DecimalArray::new_unchecked_handle(
676 values,
677 values_type,
678 decimal_dtype,
679 validity.execute(ctx)?,
680 )
681 })))
682 }
683 Canonical::VarBinView(vbv) => {
684 let VarBinViewDataParts {
685 dtype,
686 buffers,
687 views,
688 validity,
689 } = vbv.into_data_parts();
690 Ok(CanonicalValidity(Canonical::VarBinView(unsafe {
691 VarBinViewArray::new_handle_unchecked(
692 views,
693 buffers,
694 dtype,
695 validity.execute(ctx)?,
696 )
697 })))
698 }
699 Canonical::List(l) => {
700 let zctl = l.is_zero_copy_to_list();
701 let ListViewDataParts {
702 elements,
703 offsets,
704 sizes,
705 validity,
706 ..
707 } = l.into_data_parts();
708 Ok(CanonicalValidity(Canonical::List(unsafe {
709 ListViewArray::new_unchecked(elements, offsets, sizes, validity.execute(ctx)?)
710 .with_zero_copy_to_list(zctl)
711 })))
712 }
713 Canonical::Map(map) => {
714 let map_dtype = map.map_dtype().clone();
715 let entries = map.entries().clone();
716 Ok(CanonicalValidity(Canonical::Map(MapArray::new(
717 map_dtype,
718 entries.execute::<CanonicalValidity>(ctx)?.0.into_listview(),
719 ))))
720 }
721 Canonical::FixedSizeList(fsl) => {
722 let list_size = fsl.list_size();
723 let len = fsl.len();
724 let parts = fsl.into_data_parts();
725 let elements = parts.elements;
726 let validity = parts.validity;
727 Ok(CanonicalValidity(Canonical::FixedSizeList(
728 FixedSizeListArray::new(elements, list_size, validity.execute(ctx)?, len),
729 )))
730 }
731 Canonical::Struct(st) => {
732 let len = st.len();
733 let StructDataParts {
734 struct_fields,
735 fields,
736 validity,
737 } = st.into_data_parts();
738 Ok(CanonicalValidity(Canonical::Struct(unsafe {
739 StructArray::new_unchecked(fields, struct_fields, len, validity.execute(ctx)?)
740 })))
741 }
742 Canonical::Union(union) => {
743 let UnionDataParts {
744 variants,
745 type_ids,
746 children,
747 } = union.into_data_parts();
748 let type_ids = type_ids.execute::<CanonicalValidity>(ctx)?.0.into_array();
749
750 Ok(CanonicalValidity(Canonical::Union(unsafe {
751 UnionArray::new_unchecked(type_ids, variants, children.iter().cloned())
752 })))
753 }
754 Canonical::Extension(ext) => Ok(CanonicalValidity(Canonical::Extension(
755 ExtensionArray::new(
756 ext.ext_dtype().clone(),
757 ext.storage_array()
758 .clone()
759 .execute::<CanonicalValidity>(ctx)?
760 .0
761 .into_array(),
762 ),
763 ))),
764 Canonical::Variant(variant) => {
765 let core_storage = recursively_canonicalize_slots(variant.core_storage(), ctx)?;
766 let shredded = variant
767 .shredded()
768 .map(|shredded| {
769 if shredded.is::<Variant>() {
770 recursively_canonicalize_slots(shredded, ctx)
771 } else {
772 shredded
773 .clone()
774 .execute::<CanonicalValidity>(ctx)
775 .map(|canonical| canonical.0.into_array())
776 }
777 })
778 .transpose()?;
779 Ok(CanonicalValidity(Canonical::Variant(
780 VariantArray::try_new(core_storage, shredded)?,
781 )))
782 }
783 }
784 }
785}
786
787pub struct RecursiveCanonical(pub Canonical);
792
793fn recursively_canonicalize_slots(
796 array: &ArrayRef,
797 ctx: &mut ExecutionCtx,
798) -> VortexResult<ArrayRef> {
799 let slots = array
800 .slots()
801 .iter()
802 .map(|slot| {
803 slot.as_ref()
804 .map(|child| {
805 child
806 .clone()
807 .execute::<RecursiveCanonical>(ctx)
808 .map(|canonical| canonical.0.into_array())
809 })
810 .transpose()
811 })
812 .collect::<VortexResult<ArraySlots>>()?;
813 unsafe { array.clone().with_slots(slots) }
816}
817impl Executable for RecursiveCanonical {
818 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
819 match array.execute::<Canonical>(ctx)? {
820 n @ Canonical::Null(_) => Ok(RecursiveCanonical(n)),
821 Canonical::Bool(b) => {
822 let validity = child_to_validity(b.slots()[0].as_ref(), b.dtype().nullability());
823 let len = b.len();
824 let BoolDataParts { bits, meta } = b.into_data().into_parts(len);
825 Ok(RecursiveCanonical(Canonical::Bool(
826 BoolArray::try_new_from_handle(
827 bits,
828 meta.offset(),
829 meta.len(),
830 validity.execute(ctx)?,
831 )?,
832 )))
833 }
834 Canonical::Primitive(p) => {
835 let PrimitiveDataParts {
836 ptype,
837 buffer,
838 validity,
839 } = p.into_data_parts();
840 Ok(RecursiveCanonical(Canonical::Primitive(unsafe {
841 PrimitiveArray::new_unchecked_from_handle(buffer, ptype, validity.execute(ctx)?)
842 })))
843 }
844 Canonical::Decimal(d) => {
845 let DecimalDataParts {
846 decimal_dtype,
847 values,
848 values_type,
849 validity,
850 } = d.into_data_parts();
851 Ok(RecursiveCanonical(Canonical::Decimal(unsafe {
852 DecimalArray::new_unchecked_handle(
853 values,
854 values_type,
855 decimal_dtype,
856 validity.execute(ctx)?,
857 )
858 })))
859 }
860 Canonical::VarBinView(vbv) => {
861 let VarBinViewDataParts {
862 dtype,
863 buffers,
864 views,
865 validity,
866 } = vbv.into_data_parts();
867 Ok(RecursiveCanonical(Canonical::VarBinView(unsafe {
868 VarBinViewArray::new_handle_unchecked(
869 views,
870 buffers,
871 dtype,
872 validity.execute(ctx)?,
873 )
874 })))
875 }
876 Canonical::List(l) => {
877 let zctl = l.is_zero_copy_to_list();
878 let ListViewDataParts {
879 elements,
880 offsets,
881 sizes,
882 validity,
883 ..
884 } = l.into_data_parts();
885 Ok(RecursiveCanonical(Canonical::List(unsafe {
886 ListViewArray::new_unchecked(
887 elements.execute::<RecursiveCanonical>(ctx)?.0.into_array(),
888 offsets.execute::<RecursiveCanonical>(ctx)?.0.into_array(),
889 sizes.execute::<RecursiveCanonical>(ctx)?.0.into_array(),
890 validity.execute(ctx)?,
891 )
892 .with_zero_copy_to_list(zctl)
893 })))
894 }
895 Canonical::Map(map) => {
896 let map_dtype = map.map_dtype().clone();
897 let entries = map.entries().clone();
898 Ok(RecursiveCanonical(Canonical::Map(MapArray::new(
899 map_dtype,
900 entries
901 .execute::<RecursiveCanonical>(ctx)?
902 .0
903 .into_listview(),
904 ))))
905 }
906 Canonical::FixedSizeList(fsl) => {
907 let list_size = fsl.list_size();
908 let len = fsl.len();
909 let parts = fsl.into_data_parts();
910 let elements = parts.elements;
911 let validity = parts.validity;
912 Ok(RecursiveCanonical(Canonical::FixedSizeList(
913 FixedSizeListArray::new(
914 elements.execute::<RecursiveCanonical>(ctx)?.0.into_array(),
915 list_size,
916 validity.execute(ctx)?,
917 len,
918 ),
919 )))
920 }
921 Canonical::Struct(st) => {
922 let len = st.len();
923 let StructDataParts {
924 struct_fields,
925 fields,
926 validity,
927 } = st.into_data_parts();
928 let executed_fields = fields
929 .into_iter()
930 .map(|f| Ok(f.execute::<RecursiveCanonical>(ctx)?.0.into_array()))
931 .collect::<VortexResult<Vec<_>>>()?;
932
933 Ok(RecursiveCanonical(Canonical::Struct(unsafe {
934 StructArray::new_unchecked(
935 executed_fields,
936 struct_fields,
937 len,
938 validity.execute(ctx)?,
939 )
940 })))
941 }
942 Canonical::Union(union) => {
943 let UnionDataParts {
944 variants,
945 type_ids,
946 children,
947 } = union.into_data_parts();
948 let type_ids = type_ids.execute::<RecursiveCanonical>(ctx)?.0.into_array();
949 let children = children
950 .iter()
951 .cloned()
952 .map(|child| {
953 child
954 .execute::<RecursiveCanonical>(ctx)
955 .map(|canonical| canonical.0.into_array())
956 })
957 .collect::<VortexResult<Vec<_>>>()?;
958
959 Ok(RecursiveCanonical(Canonical::Union(unsafe {
960 UnionArray::new_unchecked(type_ids, variants, children)
961 })))
962 }
963 Canonical::Extension(ext) => Ok(RecursiveCanonical(Canonical::Extension(
964 ExtensionArray::new(
965 ext.ext_dtype().clone(),
966 ext.storage_array()
967 .clone()
968 .execute::<RecursiveCanonical>(ctx)?
969 .0
970 .into_array(),
971 ),
972 ))),
973 Canonical::Variant(variant) => {
974 let core_storage = recursively_canonicalize_slots(variant.core_storage(), ctx)?;
975 let shredded = variant
976 .shredded()
977 .map(|shredded| {
978 if shredded.is::<Variant>() {
979 recursively_canonicalize_slots(shredded, ctx)
980 } else {
981 shredded
982 .clone()
983 .execute::<RecursiveCanonical>(ctx)
984 .map(|canonical| canonical.0.into_array())
985 }
986 })
987 .transpose()?;
988 Ok(RecursiveCanonical(Canonical::Variant(
989 VariantArray::try_new(core_storage, shredded)?,
990 )))
991 }
992 }
993 }
994}
995
996impl<T: NativePType> Executable for Buffer<T> {
1002 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1003 let array = PrimitiveArray::execute(array, ctx)?;
1004 vortex_ensure!(
1005 matches!(
1006 array.validity()?,
1007 Validity::NonNullable | Validity::AllValid
1008 ),
1009 "Cannot execute to native buffer: array is not all-valid."
1010 );
1011 Ok(array.into_buffer())
1012 }
1013}
1014
1015impl Executable for PrimitiveArray {
1019 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1020 match array.try_downcast::<Primitive>() {
1021 Ok(primitive) => Ok(primitive),
1022 Err(array) => Ok(Canonical::execute(array, ctx)?.into_primitive()),
1023 }
1024 }
1025}
1026
1027impl Executable for BoolArray {
1031 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1032 match array.try_downcast::<Bool>() {
1033 Ok(bool_array) => Ok(bool_array),
1034 Err(array) => Ok(Canonical::execute(array, ctx)?.into_bool()),
1035 }
1036 }
1037}
1038
1039impl Executable for BitBuffer {
1043 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1044 let bool = BoolArray::execute(array, ctx)?;
1045 assert!(
1046 !bool.dtype().is_nullable(),
1047 "bit buffer execute only works with non-nullable bool arrays"
1048 );
1049 Ok(bool.into_bit_buffer())
1050 }
1051}
1052
1053impl Executable for NullArray {
1057 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1058 match array.try_downcast::<Null>() {
1059 Ok(null_array) => Ok(null_array),
1060 Err(array) => Ok(Canonical::execute(array, ctx)?.into_null()),
1061 }
1062 }
1063}
1064
1065impl Executable for VarBinViewArray {
1069 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1070 match array.try_downcast::<VarBinView>() {
1071 Ok(varbinview) => Ok(varbinview),
1072 Err(array) => Ok(Canonical::execute(array, ctx)?.into_varbinview()),
1073 }
1074 }
1075}
1076
1077impl Executable for ExtensionArray {
1081 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1082 match array.try_downcast::<Extension>() {
1083 Ok(ext_array) => Ok(ext_array),
1084 Err(array) => Ok(Canonical::execute(array, ctx)?.into_extension()),
1085 }
1086 }
1087}
1088
1089impl Executable for DecimalArray {
1093 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1094 match array.try_downcast::<Decimal>() {
1095 Ok(decimal) => Ok(decimal),
1096 Err(array) => Ok(Canonical::execute(array, ctx)?.into_decimal()),
1097 }
1098 }
1099}
1100
1101impl Executable for ListViewArray {
1105 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1106 match array.try_downcast::<ListView>() {
1107 Ok(list) => Ok(list),
1108 Err(array) => Ok(Canonical::execute(array, ctx)?.into_listview()),
1109 }
1110 }
1111}
1112
1113impl Executable for MapArray {
1117 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1118 match array.try_downcast::<Map>() {
1119 Ok(map) => Ok(map),
1120 Err(array) => Ok(Canonical::execute(array, ctx)?.into_map()),
1121 }
1122 }
1123}
1124
1125impl Executable for FixedSizeListArray {
1129 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1130 match array.try_downcast::<FixedSizeList>() {
1131 Ok(fsl) => Ok(fsl),
1132 Err(array) => Ok(Canonical::execute(array, ctx)?.into_fixed_size_list()),
1133 }
1134 }
1135}
1136
1137impl Executable for StructArray {
1141 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1142 match array.try_downcast::<Struct>() {
1143 Ok(struct_array) => Ok(struct_array),
1144 Err(array) => Ok(Canonical::execute(array, ctx)?.into_struct()),
1145 }
1146 }
1147}
1148
1149impl Executable for UnionArray {
1153 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1154 match array.try_downcast::<Union>() {
1155 Ok(union_array) => Ok(union_array),
1156 Err(array) => Ok(Canonical::execute(array, ctx)?.into_union()),
1157 }
1158 }
1159}
1160
1161impl Executable for VariantArray {
1165 fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
1166 match array.try_downcast::<Variant>() {
1167 Ok(variant_array) => Ok(variant_array),
1168 Err(array) => match Canonical::execute(array, ctx)? {
1169 Canonical::Variant(variant_array) => Ok(variant_array),
1170 canonical => vortex_panic!("Cannot unwrap VariantArray from {:?}", canonical),
1171 },
1172 }
1173 }
1174}
1175
1176#[derive(Debug, Clone, Copy)]
1181pub enum CanonicalView<'a> {
1182 Null(ArrayView<'a, Null>),
1183 Bool(ArrayView<'a, Bool>),
1184 Primitive(ArrayView<'a, Primitive>),
1185 Decimal(ArrayView<'a, Decimal>),
1186 VarBinView(ArrayView<'a, VarBinView>),
1187 List(ArrayView<'a, ListView>),
1188 Map(ArrayView<'a, Map>),
1189 FixedSizeList(ArrayView<'a, FixedSizeList>),
1190 Struct(ArrayView<'a, Struct>),
1191 Union(ArrayView<'a, Union>),
1192 Extension(ArrayView<'a, Extension>),
1193 Variant(ArrayView<'a, Variant>),
1194}
1195
1196impl From<CanonicalView<'_>> for Canonical {
1197 fn from(value: CanonicalView<'_>) -> Self {
1198 match value {
1199 CanonicalView::Null(a) => Canonical::Null(a.into_owned()),
1200 CanonicalView::Bool(a) => Canonical::Bool(a.into_owned()),
1201 CanonicalView::Primitive(a) => Canonical::Primitive(a.into_owned()),
1202 CanonicalView::Decimal(a) => Canonical::Decimal(a.into_owned()),
1203 CanonicalView::VarBinView(a) => Canonical::VarBinView(a.into_owned()),
1204 CanonicalView::List(a) => Canonical::List(a.into_owned()),
1205 CanonicalView::Map(a) => Canonical::Map(a.into_owned()),
1206 CanonicalView::FixedSizeList(a) => Canonical::FixedSizeList(a.into_owned()),
1207 CanonicalView::Struct(a) => Canonical::Struct(a.into_owned()),
1208 CanonicalView::Union(a) => Canonical::Union(a.into_owned()),
1209 CanonicalView::Extension(a) => Canonical::Extension(a.into_owned()),
1210 CanonicalView::Variant(a) => Canonical::Variant(a.into_owned()),
1211 }
1212 }
1213}
1214
1215impl CanonicalView<'_> {
1216 pub fn to_array_ref(&self) -> ArrayRef {
1218 match self {
1219 CanonicalView::Null(a) => a.array().clone(),
1220 CanonicalView::Bool(a) => a.array().clone(),
1221 CanonicalView::Primitive(a) => a.array().clone(),
1222 CanonicalView::Decimal(a) => a.array().clone(),
1223 CanonicalView::VarBinView(a) => a.array().clone(),
1224 CanonicalView::List(a) => a.array().clone(),
1225 CanonicalView::Map(a) => a.array().clone(),
1226 CanonicalView::FixedSizeList(a) => a.array().clone(),
1227 CanonicalView::Struct(a) => a.array().clone(),
1228 CanonicalView::Union(a) => a.array().clone(),
1229 CanonicalView::Extension(a) => a.array().clone(),
1230 CanonicalView::Variant(a) => a.array().clone(),
1231 }
1232 }
1233}
1234
1235macro_rules! canonical_kinds {
1237 ($($kind:ident => $vtable:ty),+ $(,)?) => {
1238 #[derive(Clone, Copy)]
1240 pub(crate) enum CanonicalKind {
1241 $($kind),+
1242 }
1243
1244 impl CanonicalKind {
1245 #[inline]
1247 pub(crate) fn of<V: VTable>() -> Option<Self> {
1248 $(if TypeId::of::<V>() == TypeId::of::<$vtable>() {
1249 return Some(Self::$kind);
1250 })+
1251 None
1252 }
1253
1254 #[inline]
1257 unsafe fn view(self, array: &ArrayRef) -> CanonicalView<'_> {
1258 match self {
1259 $(Self::$kind => CanonicalView::$kind(unsafe {
1261 array.as_typed_unchecked::<$vtable>()
1262 })),+
1263 }
1264 }
1265 }
1266 };
1267}
1268
1269canonical_kinds! {
1270 Null => Null,
1271 Bool => Bool,
1272 Primitive => Primitive,
1273 Decimal => Decimal,
1274 Struct => Struct,
1275 Union => Union,
1276 List => ListView,
1277 Map => Map,
1278 FixedSizeList => FixedSizeList,
1279 VarBinView => VarBinView,
1280 Variant => Variant,
1281 Extension => Extension,
1282}
1283
1284pub struct AnyCanonical;
1286impl Matcher for AnyCanonical {
1287 type Match<'a> = CanonicalView<'a>;
1288
1289 #[inline]
1290 fn matches(array: &ArrayRef) -> bool {
1291 array.dyn_array().canonical_kind().is_some()
1292 }
1293
1294 #[inline]
1295 fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
1296 let kind = array.dyn_array().canonical_kind()?;
1297 Some(unsafe { kind.view(array) })
1299 }
1300}
1301
1302#[cfg(test)]
1303mod tests {
1304 use std::sync::Arc;
1305 use std::sync::LazyLock;
1306
1307 use rstest::rstest;
1308 use vortex_error::VortexResult;
1309 use vortex_error::vortex_err;
1310 use vortex_session::VortexSession;
1311
1312 use crate::ArrayRef;
1313 use crate::Canonical;
1314 use crate::CanonicalValidity;
1315 use crate::IntoArray;
1316 use crate::VortexSessionExecute;
1317 use crate::array::new_foreign_array;
1318 use crate::arrays::Constant;
1319 use crate::arrays::ConstantArray;
1320 use crate::arrays::Primitive;
1321 use crate::arrays::Struct;
1322 use crate::arrays::VarBinArray;
1323 use crate::arrays::Variant;
1324 use crate::arrays::VariantArray;
1325 use crate::arrays::struct_::StructArrayExt;
1326 use crate::arrays::variant::VariantArraySlotsExt;
1327 use crate::canonical::AnyCanonical;
1328 use crate::canonical::StructArray;
1329 use crate::dtype::DType;
1330 use crate::dtype::DecimalDType;
1331 use crate::dtype::MapDType;
1332 use crate::dtype::Nullability;
1333 use crate::dtype::PType;
1334 use crate::dtype::StructFields;
1335 use crate::dtype::UnionVariants;
1336 use crate::extension::datetime::TimeUnit;
1337 use crate::extension::datetime::Timestamp;
1338 use crate::scalar::Scalar;
1339
1340 static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
1342
1343 fn variant_core_storage(len: usize) -> ArrayRef {
1344 ConstantArray::new(
1345 Scalar::variant(Scalar::primitive(1i32, Nullability::NonNullable)),
1346 len,
1347 )
1348 .into_array()
1349 }
1350
1351 fn canonical_arrays(nullability: Nullability) -> VortexResult<Vec<ArrayRef>> {
1352 let i32_dtype = DType::Primitive(PType::I32, nullability);
1353 let dtypes = [
1354 DType::Null,
1355 DType::Bool(nullability),
1356 i32_dtype.clone(),
1357 DType::Decimal(DecimalDType::try_new(10, 2)?, nullability),
1358 DType::Utf8(nullability),
1359 DType::Binary(nullability),
1360 DType::List(Arc::new(i32_dtype.clone()), nullability),
1361 DType::FixedSizeList(Arc::new(i32_dtype.clone()), 2, nullability),
1362 DType::Map(
1363 MapDType::try_new(
1364 DType::Utf8(Nullability::NonNullable),
1365 i32_dtype.clone(),
1366 false,
1367 )?,
1368 nullability,
1369 ),
1370 DType::Struct(
1371 StructFields::from_iter([("a", i32_dtype.clone())]),
1372 nullability,
1373 ),
1374 DType::Union(
1375 UnionVariants::new(["a"].into(), vec![i32_dtype])?,
1376 nullability,
1377 ),
1378 DType::Extension(Timestamp::new(TimeUnit::Microseconds, nullability).erased()),
1379 ];
1380 let mut arrays: Vec<_> = dtypes
1381 .iter()
1382 .map(|dtype| Canonical::empty(dtype).into_array())
1383 .collect();
1384 let value = Scalar::variant(Scalar::primitive(1i32, nullability))
1385 .cast(&DType::Variant(nullability))?;
1386 arrays.push(
1387 VariantArray::try_new(ConstantArray::new(value, 2).into_array(), None)?.into_array(),
1388 );
1389 Ok(arrays)
1390 }
1391
1392 #[rstest]
1393 fn any_canonical_matches_all_encodings(
1394 #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability,
1395 ) -> VortexResult<()> {
1396 for array in canonical_arrays(nullability)? {
1397 assert!(array.is::<AnyCanonical>(), "{array}");
1398 let canonical = Canonical::from(array.as_::<AnyCanonical>()).into_array();
1399 assert!(ArrayRef::ptr_eq(&array, &canonical));
1400 }
1401 Ok(())
1402 }
1403
1404 #[test]
1405 fn any_canonical_rejects_noncanonical_encodings() {
1406 let arrays = [
1407 ConstantArray::new(1i32, 4).into_array(),
1408 ConstantArray::new(true, 4).into_array(),
1409 ConstantArray::new("value", 4).into_array(),
1410 VarBinArray::from_iter([Some("value")], DType::Utf8(Nullability::Nullable))
1411 .into_array(),
1412 ];
1413 for array in arrays {
1414 assert!(!array.is::<AnyCanonical>(), "{array}");
1415 assert!(array.as_opt::<AnyCanonical>().is_none(), "{array}");
1416 }
1417 }
1418
1419 #[rstest]
1420 fn any_canonical_rejects_foreign_arrays_with_canonical_ids(
1421 #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability,
1422 ) -> VortexResult<()> {
1423 let mut ctx = SESSION.create_execution_ctx();
1424 for canonical in canonical_arrays(nullability)? {
1425 let foreign = new_foreign_array(
1426 canonical.encoding_id(),
1427 canonical.dtype().clone(),
1428 canonical.len(),
1429 vec![],
1430 vec![],
1431 Default::default(),
1432 )?;
1433 assert!(!foreign.is::<AnyCanonical>(), "{foreign}");
1434 assert!(foreign.as_opt::<AnyCanonical>().is_none(), "{foreign}");
1435 assert!(foreign.execute::<Canonical>(&mut ctx).is_err());
1436 }
1437 Ok(())
1438 }
1439
1440 #[test]
1441 fn canonical_validity_canonicalizes_variant_shredded_physical_slots() -> VortexResult<()> {
1442 let len = 2;
1443 let nested_shredded =
1444 StructArray::try_from_iter([("value", ConstantArray::new(10i32, len).into_array())])?;
1445 let inner_variant = VariantArray::try_new(
1446 variant_core_storage(len),
1447 Some(nested_shredded.into_array()),
1448 )?;
1449 let outer_variant =
1450 VariantArray::try_new(variant_core_storage(len), Some(inner_variant.into_array()))?;
1451
1452 let mut ctx = SESSION.create_execution_ctx();
1453 let Canonical::Variant(canonical) = outer_variant
1454 .into_array()
1455 .execute::<CanonicalValidity>(&mut ctx)?
1456 .0
1457 else {
1458 return Err(vortex_err!("expected canonical variant"));
1459 };
1460
1461 let nested_variant = canonical
1462 .shredded()
1463 .and_then(|shredded| shredded.as_opt::<Variant>())
1464 .ok_or_else(|| vortex_err!("expected nested variant shredded child"))?;
1465 let nested_struct = nested_variant
1466 .shredded()
1467 .and_then(|shredded| shredded.as_opt::<Struct>())
1468 .ok_or_else(|| vortex_err!("expected nested struct shredded child"))?;
1469 let value = nested_struct.unmasked_field_by_name("value")?;
1470
1471 assert!(value.is::<Primitive>());
1472 assert!(!value.is::<Constant>());
1473
1474 Ok(())
1475 }
1476}