Part 2 of making Variant a canonical array

fuzz/src/array/fill_null.rs

@@ -44,6 +44,7 @@ pub fn fill_null_canonical_array(
         | Canonical::List(_)
         | Canonical::FixedSizeList(_)
         | Canonical::Extension(_) => canonical.into_array().fill_null(fill_value.clone())?,
+        Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }
 

fuzz/src/array/mask.rs

@@ -136,6 +136,7 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult<A
                 .with_nullability(masked_storage.dtype().nullability());
             ExtensionArray::new(ext_dtype, masked_storage).into_array()
         }
+        Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }
 

fuzz/src/array/scalar_at.rs

@@ -95,5 +95,6 @@ pub fn scalar_at_canonical_array(canonical: Canonical, index: usize) -> VortexRe
                 scalar_at_canonical_array(array.storage_array().to_canonical()?, index)?;
             Scalar::extension_ref(array.ext_dtype().clone(), storage_scalar)
         }
+        Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

vortex-array/public-api.lock

@@ -5020,7 +5020,7 @@ impl vortex_array::arrays::variant::VariantArray
 
 pub fn vortex_array::arrays::variant::VariantArray::child(&self) -> &vortex_array::ArrayRef
 
-pub fn vortex_array::arrays::variant::VariantArray::new(child: vortex_array::ArrayRef, nullability: vortex_array::dtype::Nullability) -> Self
+pub fn vortex_array::arrays::variant::VariantArray::new(child: vortex_array::ArrayRef) -> Self
 
 impl vortex_array::arrays::variant::VariantArray
 
@@ -8272,7 +8272,7 @@ impl vortex_array::arrays::variant::VariantArray
 
 pub fn vortex_array::arrays::variant::VariantArray::child(&self) -> &vortex_array::ArrayRef
 
-pub fn vortex_array::arrays::variant::VariantArray::new(child: vortex_array::ArrayRef, nullability: vortex_array::dtype::Nullability) -> Self
+pub fn vortex_array::arrays::variant::VariantArray::new(child: vortex_array::ArrayRef) -> Self
 
 impl vortex_array::arrays::variant::VariantArray
 
@@ -22096,6 +22096,8 @@ pub vortex_array::Canonical::Struct(vortex_array::arrays::StructArray)
 
 pub vortex_array::Canonical::VarBinView(vortex_array::arrays::VarBinViewArray)
 
+pub vortex_array::Canonical::Variant(vortex_array::arrays::variant::VariantArray)
+
 impl vortex_array::Canonical
 
 pub fn vortex_array::Canonical::as_bool(&self) -> &vortex_array::arrays::BoolArray
@@ -22196,6 +22198,8 @@ pub vortex_array::CanonicalView::Struct(&'a vortex_array::arrays::StructArray)
 
 pub vortex_array::CanonicalView::VarBinView(&'a vortex_array::arrays::VarBinViewArray)
 
+pub vortex_array::CanonicalView::Variant(&'a vortex_array::arrays::variant::VariantArray)
+
 impl core::convert::AsRef<dyn vortex_array::DynArray> for vortex_array::CanonicalView<'_>
 
 pub fn vortex_array::CanonicalView<'_>::as_ref(&self) -> &dyn vortex_array::DynArray

vortex-array/src/aggregate_fn/fns/is_constant/mod.rs

@@ -12,6 +12,7 @@ mod varbin;
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 use vortex_mask::Mask;
 
 use self::bool::check_bool_constant;
@@ -273,14 +274,14 @@ impl AggregateFnVTable for IsConstant {
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
         match input_dtype {
-            DType::Null => None,
+            DType::Null | DType::Variant(..) => None,
             _ => Some(DType::Bool(Nullability::NonNullable)),
         }
     }
 
     fn partial_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
         match input_dtype {
-            DType::Null => None,
+            DType::Null | DType::Variant(..) => None,
             _ => Some(make_is_constant_partial_dtype(input_dtype)),
         }
     }
@@ -408,6 +409,9 @@ impl AggregateFnVTable for IsConstant {
                     Canonical::List(l) => check_listview_constant(l, ctx)?,
                     Canonical::FixedSizeList(f) => check_fixed_size_list_constant(f, ctx)?,
                     Canonical::Null(_) => true,
+                    Canonical::Variant(_) => {
+                        vortex_bail!("Variant arrays don't support IsConstant")
+                    }
                 };
 
                 if !batch_is_constant {

vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs

@@ -257,14 +257,22 @@ impl AggregateFnVTable for IsSorted {
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
         match input_dtype {
-            DType::Null | DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => None,
+            DType::Null
+            | DType::Struct(..)
+            | DType::List(..)
+            | DType::FixedSizeList(..)
+            | DType::Variant(..) => None,
             _ => Some(DType::Bool(Nullability::NonNullable)),
         }
     }
 
     fn partial_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
         match input_dtype {
-            DType::Null | DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => None,
+            DType::Null
+            | DType::Struct(..)
+            | DType::List(..)
+            | DType::FixedSizeList(..)
+            | DType::Variant(..) => None,
             _ => Some(make_is_sorted_partial_dtype(input_dtype)),
         }
     }

vortex-array/src/aggregate_fn/fns/min_max/mod.rs

@@ -271,7 +271,10 @@ impl AggregateFnVTable for MinMax {
                 Canonical::Decimal(d) => accumulate_decimal(partial, d),
                 Canonical::Extension(e) => accumulate_extension(partial, e, ctx),
                 Canonical::Null(_) => Ok(()),
-                Canonical::Struct(_) | Canonical::List(_) | Canonical::FixedSizeList(_) => {
+                Canonical::Struct(_)
+                | Canonical::List(_)
+                | Canonical::FixedSizeList(_)
+                | Canonical::Variant(_) => {
                     vortex_bail!("Unsupported canonical type for min_max: {}", batch.dtype())
                 }
             },

vortex-array/src/arrays/dict/execute.rs

@@ -5,6 +5,7 @@
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 
 use crate::Canonical;
 use crate::ExecutionCtx;
@@ -50,6 +51,9 @@ pub fn take_canonical(
         }
         Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)),
         Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)),
+        Canonical::Variant(_) => {
+            vortex_bail!("Variant arrays don't support Take")
+        }
     })
 }
 

vortex-array/src/arrays/filter/execute/mod.rs

@@ -9,6 +9,7 @@ use std::sync::Arc;
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 use vortex_mask::MaskValues;
 
@@ -94,5 +95,8 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc<MaskValues>) -> Ca
                 .vortex_expect("ExtensionArray storage type somehow could not be filtered");
             Canonical::Extension(ExtensionArray::new(a.ext_dtype().clone(), filtered_storage))
         }
+        Canonical::Variant(_) => {
+            vortex_panic!("Variant arrays don't support filtering")
+        }
     }
 }

vortex-array/src/arrays/masked/execute.rs

@@ -6,6 +6,7 @@
 use std::ops::BitAnd;
 
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 use vortex_mask::Mask;
 
 use crate::Canonical;
@@ -53,6 +54,9 @@ pub fn mask_validity_canonical(
         Canonical::Extension(a) => {
             Canonical::Extension(mask_validity_extension(a, validity_mask, ctx)?)
         }
+        Canonical::Variant(_) => {
+            vortex_bail!("Variant arrays don't masking validity")
+        }
     })
 }
 

vortex-array/src/arrays/variant/mod.rs

@@ -5,26 +5,24 @@ mod vtable;
 
 pub use self::vtable::Variant;
 use crate::ArrayRef;
-use crate::dtype::DType;
-use crate::dtype::Nullability;
 
 /// The canonical in-memory representation of variant (semi-structured) data.
 ///
 /// Wraps a single child array that contains the actual variant-encoded data
 /// (e.g. a `ParquetVariantArray` or any other variant encoding).
+///
+/// Nullability is delegated to the child array: `VariantArray`'s dtype is
+/// always the child's dtype. The child's validity determines which rows are
+/// null.
 #[derive(Clone, Debug)]
 pub struct VariantArray {
-    dtype: DType,
     child: ArrayRef,
 }
 
 impl VariantArray {
-    /// Creates a new VariantArray with the given nullability.
-    pub fn new(child: ArrayRef, nullability: Nullability) -> Self {
-        Self {
-            dtype: DType::Variant(nullability),
-            child,
-        }
+    /// Creates a new VariantArray. Nullability comes from the child's dtype.
+    pub fn new(child: ArrayRef) -> Self {
+        Self { child }
     }
 
     /// Returns a reference to the underlying child array.

vortex-array/src/arrays/variant/vtable/mod.rs

@@ -18,7 +18,6 @@ use crate::ArrayRef;
 use crate::EmptyMetadata;
 use crate::ExecutionCtx;
 use crate::ExecutionResult;
-use crate::IntoArray;
 use crate::Precision;
 use crate::arrays::VariantArray;
 use crate::buffer::BufferHandle;
@@ -60,7 +59,7 @@ impl VTable for Variant {
     }
 
     fn dtype(array: &Self::Array) -> &DType {
-        &array.dtype
+        array.child.dtype()
     }
 
     fn stats(array: &Self::Array) -> StatsSetRef<'_> {
@@ -136,13 +135,9 @@ impl VTable for Variant {
             "Expected 1 child, got {}",
             children.len()
         );
-        // The child can be any variant encoding, so we use DType::Variant.
-        let child = children.get(
-            0,
-            &DType::Variant(crate::dtype::Nullability::NonNullable),
-            len,
-        )?;
-        Ok(VariantArray::new(child, dtype.nullability()))
+        // The child carries the nullability for the whole VariantArray.
+        let child = children.get(0, dtype, len)?;
+        Ok(VariantArray::new(child))
     }
 
     fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
@@ -151,12 +146,14 @@ impl VTable for Variant {
             "VariantArray expects exactly 1 child, got {}",
             children.len()
         );
-        array.child = children.into_iter().next().vortex_expect("must exist");
+        array.child = children
+            .into_iter()
+            .next()
+            .vortex_expect("VariantArray must have 1 child");
         Ok(())
     }
 
     fn execute(array: Arc<Self::Array>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
-        // VariantArray is the canonical variant representation.
-        Ok(ExecutionResult::done(array.as_ref().clone().into_array()))
+        Ok(ExecutionResult::done_upcast::<Self>(array))
     }
 }

vortex-array/src/canonical.rs

@@ -35,6 +35,8 @@ use crate::arrays::Struct;
 use crate::arrays::StructArray;
 use crate::arrays::VarBinView;
 use crate::arrays::VarBinViewArray;
+use crate::arrays::Variant;
+use crate::arrays::VariantArray;
 use crate::arrays::bool::BoolArrayParts;
 use crate::arrays::decimal::DecimalArrayParts;
 use crate::arrays::listview::ListViewArrayParts;
@@ -122,6 +124,7 @@ pub enum Canonical {
     FixedSizeList(FixedSizeListArray),
     Struct(StructArray),
     Extension(ExtensionArray),
+    Variant(VariantArray),
 }
 
 /// Match on every canonical variant and evaluate a code block on all variants
@@ -136,6 +139,7 @@ macro_rules! match_each_canonical {
             Canonical::List($ident) => $eval,
             Canonical::FixedSizeList($ident) => $eval,
             Canonical::Struct($ident) => $eval,
+            Canonical::Variant($ident) => $eval,
             Canonical::Extension($ident) => $eval,
         }
     }};
@@ -641,6 +645,16 @@ impl Executable for CanonicalValidity {
                         .into_array(),
                 ),
             ))),
+            Canonical::Variant(variant) => {
+                Ok(CanonicalValidity(Canonical::Variant(VariantArray::new(
+                    variant
+                        .child()
+                        .clone()
+                        .execute::<CanonicalValidity>(ctx)?
+                        .0
+                        .into_array(),
+                ))))
+            }
         }
     }
 }
@@ -771,6 +785,16 @@ impl Executable for RecursiveCanonical {
                         .into_array(),
                 ),
             ))),
+            Canonical::Variant(variant) => {
+                Ok(RecursiveCanonical(Canonical::Variant(VariantArray::new(
+                    variant
+                        .child()
+                        .clone()
+                        .execute::<RecursiveCanonical>(ctx)?
+                        .0
+                        .into_array(),
+                ))))
+            }
         }
     }
 }
@@ -925,6 +949,7 @@ pub enum CanonicalView<'a> {
     FixedSizeList(&'a FixedSizeListArray),
     Struct(&'a StructArray),
     Extension(&'a ExtensionArray),
+    Variant(&'a VariantArray),
 }
 
 impl From<CanonicalView<'_>> for Canonical {
@@ -939,6 +964,7 @@ impl From<CanonicalView<'_>> for Canonical {
             CanonicalView::FixedSizeList(a) => Canonical::FixedSizeList(a.clone()),
             CanonicalView::Struct(a) => Canonical::Struct(a.clone()),
             CanonicalView::Extension(a) => Canonical::Extension(a.clone()),
+            CanonicalView::Variant(a) => Canonical::Variant(a.clone()),
         }
     }
 }
@@ -955,6 +981,7 @@ impl AsRef<dyn DynArray> for CanonicalView<'_> {
             CanonicalView::FixedSizeList(a) => a.as_ref(),
             CanonicalView::Struct(a) => a.as_ref(),
             CanonicalView::Extension(a) => a.as_ref(),
+            CanonicalView::Variant(a) => a.as_ref(),
         }
     }
 }
@@ -973,6 +1000,7 @@ impl Matcher for AnyCanonical {
             || array.is::<ListView>()
             || array.is::<FixedSizeList>()
             || array.is::<VarBinView>()
+            || array.is::<Variant>()
             || array.is::<Extension>()
     }
 
@@ -993,6 +1021,8 @@ impl Matcher for AnyCanonical {
             Some(CanonicalView::FixedSizeList(a))
         } else if let Some(a) = array.as_opt::<VarBinView>() {
             Some(CanonicalView::VarBinView(a))
+        } else if let Some(a) = array.as_opt::<Variant>() {
+            Some(CanonicalView::Variant(a))
         } else {
             array.as_opt::<Extension>().map(CanonicalView::Extension)
         }

vortex-array/src/scalar_fn/fns/cast/mod.rs

@@ -214,6 +214,9 @@ fn cast_canonical(
         CanonicalView::FixedSizeList(a) => <FixedSizeList as CastReduce>::cast(a, dtype),
         CanonicalView::Struct(a) => <Struct as CastKernel>::cast(a, dtype, ctx),
         CanonicalView::Extension(a) => <Extension as CastReduce>::cast(a, dtype),
+        CanonicalView::Variant(_) => {
+            vortex_bail!("Variant arrays don't support casting")
+        }
     }
 }