[WIP] Map DType

encodings/sparse/src/canonical.rs

@@ -151,6 +151,7 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort
         DType::FixedSizeList(.., nullability) => {
             execute_sparse_fixed_size_list(&patches, &fill_value, len, *nullability, ctx)?
         }
+        DType::Map(..) => vortex_bail!("Sparse canonicalization does not support Map arrays yet"),
         DType::Struct(struct_fields, ..) => execute_sparse_struct(
             struct_fields,
             fill_value.as_struct(),

fuzz/src/array/compare.rs

@@ -186,7 +186,11 @@ pub fn compare_canonical_array(
             }))
             .into_array()
         }
-        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
+        d @ (DType::Null
+        | DType::Map(..)
+        | DType::Union(..)
+        | DType::Variant(_)
+        | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

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::Variant(_) => unreachable!("Variant arrays are not fuzzed"),

fuzz/src/array/filter.rs

@@ -121,7 +121,11 @@ pub fn filter_canonical_array(
             )
             .map(|a| a.into_array())
         }
-        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
+        d @ (DType::Null
+        | DType::Map(..)
+        | DType::Union(..)
+        | DType::Variant(_)
+        | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/mask.rs

@@ -149,6 +149,7 @@ pub fn mask_canonical_array(
                 .with_nullability(masked_storage.dtype().nullability());
             ExtensionArray::new(ext_dtype, masked_storage).into_array()
         }
+        Canonical::Map(_) => unreachable!("Map arrays are not fuzzed"),
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

fuzz/src/array/mod.rs

@@ -518,6 +518,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet<ActionType> {
                     acc.intersection(&actions).copied().collect()
                 })
         }
+        DType::Map(..) => HashSet::new(),
         DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
         // Currently, no support at all
         DType::Variant(_) => unreachable!("Variant dtype shouldn't be fuzzed"),

fuzz/src/array/scalar_at.rs

@@ -104,6 +104,7 @@ pub fn scalar_at_canonical_array(
             let storage_scalar = scalar_at_canonical_array(storage_canonical, index, ctx)?;
             Scalar::extension_ref(array.ext_dtype().clone(), storage_scalar)
         }
+        Canonical::Map(_) => unreachable!("Map arrays are not fuzzed"),
         Canonical::Variant(_) => unreachable!("Variant arrays are not fuzzed"),
     })
 }

fuzz/src/array/search_sorted.rs

@@ -149,7 +149,11 @@ pub fn search_sorted_canonical_array(
                 .collect::<VortexResult<Vec<_>>>()?;
             scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side)
         }
-        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
+        d @ (DType::Null
+        | DType::Map(..)
+        | DType::Union(..)
+        | DType::Variant(_)
+        | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/slice.rs

@@ -124,7 +124,11 @@ pub fn slice_canonical_array(
             )
             .map(|a| a.into_array())
         }
-        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
+        d @ (DType::Null
+        | DType::Map(..)
+        | DType::Union(..)
+        | DType::Variant(_)
+        | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/sort.rs

@@ -102,7 +102,11 @@ pub fn sort_canonical_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexR
             });
             take_canonical_array_non_nullable_indices(array, &sort_indices, ctx)
         }
-        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
+        d @ (DType::Null
+        | DType::Map(..)
+        | DType::Union(..)
+        | DType::Variant(_)
+        | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/take.rs

@@ -148,7 +148,11 @@ pub fn take_canonical_array(
             )
             .map(|a| a.into_array())
         }
-        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
+        d @ (DType::Null
+        | DType::Map(..)
+        | DType::Union(..)
+        | DType::Variant(_)
+        | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

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::Variant(_) => {

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

@@ -246,6 +246,7 @@ impl AggregateFnVTable for IsSorted {
             DType::Null
             | DType::List(..)
             | DType::FixedSizeList(..)
+            | DType::Map(..)
             | DType::Struct(..)
             | DType::Union(..)
             | DType::Variant(..)
@@ -263,6 +264,7 @@ impl AggregateFnVTable for IsSorted {
             DType::Null
             | DType::List(..)
             | DType::FixedSizeList(..)
+            | DType::Map(..)
             | DType::Struct(..)
             | DType::Union(..)
             | DType::Variant(..)

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

@@ -419,6 +419,7 @@ impl AggregateFnVTable for MinMax {
                 Canonical::Null(_) => Ok(()),
                 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

@@ -43,6 +43,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;
@@ -197,6 +198,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::Extension(array) => extension_uncompressed_size_in_bytes(array, ctx),
@@ -229,7 +233,11 @@ 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);
         }
@@ -286,6 +294,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(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/arbitrary.rs

@@ -157,6 +157,7 @@ fn random_array_chunk(
         DType::FixedSizeList(elem_dtype, list_size, null) => {
             random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len)
         }
+        DType::Map(..) => Err(IncorrectFormat),
         DType::Struct(sdt, n) => {
             let first_array = sdt
                 .fields()

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,6 +127,16 @@ pub(crate) fn constant_canonicalize(
             ))
         }
         DType::List(..) => Canonical::List(constant_canonical_list_array(scalar, array.len())),
+        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;
 
@@ -83,14 +84,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))
         }
@@ -118,5 +123,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

@@ -168,7 +168,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(),
         ))
     }