Add canonical Map arrays and builder

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

fuzz/src/array/fill_null.rs

@@ -47,6 +47,7 @@ pub fn fill_null_canonical_array(
         }
         Canonical::Struct(_)
         | Canonical::List(_)
+        | Canonical::Map(_)
         | Canonical::FixedSizeList(_)
         | Canonical::Extension(_) => canonical.into_array().fill_null(fill_value.clone())?,
         Canonical::Union(_) => {

fuzz/src/array/mask.rs

@@ -153,6 +153,7 @@ pub fn mask_canonical_array(
         Canonical::Union(_) => {
             todo!("TODO(connor)[Union]: support Union arrays in the mask fuzzer")
         }
+        Canonical::Map(_) => unreachable!("Map arrays are not fuzzed"),
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

fuzz/src/array/scalar_at.rs

@@ -107,6 +107,7 @@ pub fn scalar_at_canonical_array(
         Canonical::Union(_) => {
             todo!("TODO(connor)[Union]: support Union arrays in the scalar_at fuzzer")
         }
+        Canonical::Map(_) => unreachable!("Map arrays are not fuzzed"),
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

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

@@ -402,6 +402,9 @@ impl AggregateFnVTable for IsConstant {
                     Canonical::Struct(s) => check_struct_constant(s, ctx)?,
                     Canonical::Extension(e) => check_extension_constant(e, ctx)?,
                     Canonical::List(l) => check_listview_constant(l, ctx)?,
+                    Canonical::Map(_) => {
+                        vortex_bail!("Map arrays don't support IsConstant")
+                    }
                     Canonical::FixedSizeList(f) => check_fixed_size_list_constant(f, ctx)?,
                     Canonical::Null(_) => true,
                     Canonical::Union(_) => {

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

@@ -422,6 +422,7 @@ impl AggregateFnVTable for MinMax {
                 }
                 Canonical::Struct(_)
                 | Canonical::List(_)
+                | Canonical::Map(_)
                 | Canonical::FixedSizeList(_)
                 | Canonical::Variant(_) => {
                     vortex_bail!("Unsupported canonical type for min_max: {}", batch.dtype())

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

@@ -45,6 +45,7 @@ use crate::aggregate_fn::EmptyOptions;
 use crate::array::ArrayView;
 use crate::arrays::Constant;
 use crate::arrays::ConstantArray;
+use crate::arrays::map::MapArrayExt;
 use crate::arrays::varbinview::BinaryView;
 use crate::dtype::DType;
 use crate::dtype::DecimalType;
@@ -199,6 +200,9 @@ pub(crate) fn canonical_uncompressed_size_in_bytes(
         Canonical::Decimal(array) => decimal_uncompressed_size_in_bytes(array, ctx),
         Canonical::VarBinView(array) => varbinview_uncompressed_size_in_bytes(array, ctx),
         Canonical::List(array) => list_view_uncompressed_size_in_bytes(array, ctx),
+        Canonical::Map(array) => {
+            list_view_uncompressed_size_in_bytes(&array.entries().into_owned(), ctx)
+        }
         Canonical::FixedSizeList(array) => fixed_size_list_uncompressed_size_in_bytes(array, ctx),
         Canonical::Struct(array) => struct_uncompressed_size_in_bytes(array, ctx),
         Canonical::Union(array) => union_uncompressed_size_in_bytes(array, ctx),
@@ -232,13 +236,14 @@ pub(crate) fn constant_uncompressed_size_in_bytes(
             array.len(),
             array.scalar().as_binary().value().map(|value| value.len()),
         )?,
-        DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) | DType::Extension(_) => {
+        DType::List(..)
+        | DType::Map(..)
+        | DType::FixedSizeList(..)
+        | DType::Struct(..)
+        | DType::Extension(_) => {
             let canonical = array.array().clone().execute::<Canonical>(ctx)?;
             return canonical_uncompressed_size_in_bytes(&canonical, ctx);
         }
-        DType::Map(..) => {
-            vortex_bail!("UncompressedSizeInBytes is not supported for map arrays yet")
-        }
         DType::Union(..) => {
             todo!(
                 "TODO(connor)[Union]: support constant Union size accounting after constant Union \
@@ -297,7 +302,10 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool {
         DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => {
             supports_uncompressed_size_in_bytes(element_dtype)
         }
-        DType::Map(..) => false,
+        DType::Map(map_dtype, _) => {
+            supports_uncompressed_size_in_bytes(&map_dtype.key_dtype())
+                && supports_uncompressed_size_in_bytes(&map_dtype.value_dtype())
+        }
         DType::Struct(fields, _) => fields
             .fields()
             .all(|field| supports_uncompressed_size_in_bytes(&field)),

vortex-array/src/arrays/constant/vtable/canonical.rs

@@ -20,6 +20,7 @@ use crate::arrays::DecimalArray;
 use crate::arrays::ExtensionArray;
 use crate::arrays::FixedSizeListArray;
 use crate::arrays::ListViewArray;
+use crate::arrays::MapArray;
 use crate::arrays::NullArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::StructArray;
@@ -126,7 +127,16 @@ pub(crate) fn constant_canonicalize(
             ))
         }
         DType::List(..) => Canonical::List(constant_canonical_list_array(scalar, array.len())),
-        DType::Map(..) => vortex_error::vortex_bail!("canonical map arrays are not yet supported"),
+        DType::Map(map_dtype, nullability) => {
+            let entries_scalar = Scalar::try_new(
+                DType::List(Arc::new(map_dtype.entries_dtype()), *nullability),
+                scalar.value().cloned(),
+            )?;
+            Canonical::Map(MapArray::try_new(
+                map_dtype.clone(),
+                constant_canonical_list_array(&entries_scalar, array.len()),
+            )?)
+        }
         DType::FixedSizeList(element_dtype, list_size, _) => {
             let value = scalar.as_list();
 

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

@@ -49,6 +49,7 @@ pub(crate) fn take_canonical(
         Canonical::Decimal(a) => Canonical::Decimal(take_decimal(&a, codes, ctx)),
         Canonical::VarBinView(a) => Canonical::VarBinView(take_varbinview(&a, codes, ctx)),
         Canonical::List(a) => Canonical::List(take_listview(&a, codes, ctx)),
+        Canonical::Map(_) => vortex_error::vortex_bail!("Map arrays don't support take"),
         Canonical::FixedSizeList(a) => {
             Canonical::FixedSizeList(take_fixed_size_list(&a, codes, ctx))
         }

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_bail;
 use vortex_mask::Mask;
 use vortex_mask::MaskValues;
 
@@ -87,14 +88,18 @@ pub(super) fn execute_filter_fast_paths(
 }
 
 /// Filter a canonical array by a mask, returning a new canonical array.
-pub(super) fn execute_filter(canonical: Canonical, mask: &Arc<MaskValues>) -> Canonical {
-    match canonical {
+pub(super) fn execute_filter(
+    canonical: Canonical,
+    mask: &Arc<MaskValues>,
+) -> VortexResult<Canonical> {
+    Ok(match canonical {
         Canonical::Null(_) => Canonical::Null(NullArray::new(mask.true_count())),
         Canonical::Bool(a) => Canonical::Bool(bool::filter_bool(&a, mask)),
         Canonical::Primitive(a) => Canonical::Primitive(primitive::filter_primitive(&a, mask)),
         Canonical::Decimal(a) => Canonical::Decimal(decimal::filter_decimal(&a, mask)),
         Canonical::VarBinView(a) => Canonical::VarBinView(varbinview::filter_varbinview(&a, mask)),
         Canonical::List(a) => Canonical::List(listview::filter_listview(&a, mask)),
+        Canonical::Map(_) => vortex_bail!("Map arrays don't support filter"),
         Canonical::FixedSizeList(a) => {
             Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask))
         }
@@ -123,5 +128,5 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc<MaskValues>) -> Ca
                     .vortex_expect("filtered VariantArray children are row-aligned"),
             )
         }
-    }
+    })
 }

vortex-array/src/arrays/filter/vtable.rs

@@ -167,7 +167,7 @@ impl VTable for Filter {
         // TODO(joe): fix the ownership of AnyCanonical
         let child = Canonical::from(array.child().as_::<AnyCanonical>());
         Ok(ExecutionResult::done(
-            execute_filter(child, &mask_values).into_array(),
+            execute_filter(child, &mask_values)?.into_array(),
         ))
     }