Add Map compute and compression support

fuzz/src/array/filter.rs

@@ -11,6 +11,7 @@ use vortex_array::arrays::StructArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::arrays::bool::BoolArrayExt;
 use vortex_array::arrays::struct_::StructArrayExt;
+use vortex_array::builders::builder_with_capacity;
 use vortex_array::dtype::DType;
 use vortex_array::match_each_decimal_value_type;
 use vortex_array::match_each_native_ptype;
@@ -121,11 +122,17 @@ pub fn filter_canonical_array(
             )
             .map(|a| a.into_array())
         }
-        d @ (DType::Null
-        | DType::Map(..)
-        | DType::Union(..)
-        | DType::Variant(_)
-        | DType::Extension(_)) => {
+        DType::Map(..) => {
+            let mut builder =
+                builder_with_capacity(array.dtype(), filter.iter().filter(|b| **b).count());
+            for (idx, keep) in filter.iter().enumerate() {
+                if *keep {
+                    builder.append_scalar(&array.execute_scalar(idx, ctx)?)?;
+                }
+            }
+            Ok(builder.finish())
+        }
+        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/mask.rs

@@ -21,6 +21,7 @@ use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
 use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
 use vortex_array::arrays::listview::ListViewArraySlotsExt;
 use vortex_array::arrays::struct_::StructArrayExt;
+use vortex_array::builders::builder_with_capacity;
 use vortex_array::dtype::Nullability;
 use vortex_array::match_each_decimal_value_type;
 use vortex_array::validity::Validity;
@@ -139,6 +140,18 @@ pub fn mask_canonical_array(
             .vortex_expect("StructArray creation should succeed in fuzz test")
             .into_array()
         }
+        Canonical::Map(array) => {
+            let result_dtype = array.dtype().as_nullable();
+            let mut builder = builder_with_capacity(&result_dtype, array.len());
+            for idx in 0..array.len() {
+                if mask.value(idx) {
+                    builder.append_scalar(&array.execute_scalar(idx, ctx)?.cast(&result_dtype)?)?;
+                } else {
+                    builder.append_null();
+                }
+            }
+            builder.finish()
+        }
         Canonical::Extension(array) => {
             // Recursively mask the storage array
             let storage_canonical = array.storage_array().clone().execute::<Canonical>(ctx)?;
@@ -153,7 +166,6 @@ 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/mod.rs

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

fuzz/src/array/slice.rs

@@ -16,6 +16,7 @@ use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
 use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
 use vortex_array::arrays::listview::ListViewArraySlotsExt;
 use vortex_array::arrays::struct_::StructArrayExt;
+use vortex_array::builders::builder_with_capacity;
 use vortex_array::dtype::DType;
 use vortex_array::match_each_decimal_value_type;
 use vortex_array::match_each_native_ptype;
@@ -125,11 +126,14 @@ pub fn slice_canonical_array(
             )
             .map(|a| a.into_array())
         }
-        d @ (DType::Null
-        | DType::Map(..)
-        | DType::Union(..)
-        | DType::Variant(_)
-        | DType::Extension(_)) => {
+        DType::Map(..) => {
+            let mut builder = builder_with_capacity(array.dtype(), stop - start);
+            for idx in start..stop {
+                builder.append_scalar(&array.execute_scalar(idx, ctx)?)?;
+            }
+            Ok(builder.finish())
+        }
+        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/take.rs

@@ -148,11 +148,20 @@ pub fn take_canonical_array(
             )
             .map(|a| a.into_array())
         }
-        d @ (DType::Null
-        | DType::Map(..)
-        | DType::Union(..)
-        | DType::Variant(_)
-        | DType::Extension(_)) => {
+        DType::Map(..) => {
+            let result_dtype = array.dtype().union_nullability(nullable);
+            let mut builder = builder_with_capacity(&result_dtype, indices.len());
+            for idx in indices {
+                if let Some(idx) = idx {
+                    builder
+                        .append_scalar(&array.execute_scalar(*idx, ctx)?.cast(&result_dtype)?)?;
+                } else {
+                    builder.append_null();
+                }
+            }
+            Ok(builder.finish())
+        }
+        d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

vortex-array/src/aggregate_fn/fns/all_non_distinct/map.rs

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use super::list::check_list_identical;
+use crate::ExecutionCtx;
+use crate::arrays::ListView;
+use crate::arrays::MapArray;
+use crate::arrays::map::MapArrayExt;
+use crate::arrays::map::MapArraySlotsExt;
+
+pub(super) fn check_map_identical(
+    lhs: &MapArray,
+    rhs: &MapArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<bool> {
+    if lhs.map_dtype() != rhs.map_dtype() {
+        return Ok(false);
+    }
+
+    check_list_identical(
+        &lhs.entries().as_::<ListView>().into_owned(),
+        &rhs.entries().as_::<ListView>().into_owned(),
+        ctx,
+    )
+}

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

@@ -7,6 +7,7 @@ mod extension;
 mod filter;
 mod fixed_size_list;
 mod list;
+mod map;
 mod primitive;
 mod struct_;
 #[cfg(test)]
@@ -27,6 +28,7 @@ use self::extension::check_extension_identical;
 use self::filter::shared_validity_mask;
 use self::fixed_size_list::check_fixed_size_list_identical;
 use self::list::check_list_identical;
+use self::map::check_map_identical;
 use self::primitive::check_primitive_identical;
 use self::struct_::check_struct_identical;
 use self::varbin::check_varbinview_identical;
@@ -262,6 +264,7 @@ fn check_canonical_identical(
         }
         (Canonical::Struct(lhs), Canonical::Struct(rhs)) => check_struct_identical(lhs, rhs, ctx),
         (Canonical::List(lhs), Canonical::List(rhs)) => check_list_identical(lhs, rhs, ctx),
+        (Canonical::Map(lhs), Canonical::Map(rhs)) => check_map_identical(lhs, rhs, ctx),
         (Canonical::FixedSizeList(lhs), Canonical::FixedSizeList(rhs)) => {
             check_fixed_size_list_identical(lhs, rhs, ctx)
         }

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

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use super::list::check_listview_constant;
+use crate::ExecutionCtx;
+use crate::arrays::ListView;
+use crate::arrays::MapArray;
+use crate::arrays::map::MapArraySlotsExt;
+
+pub(super) fn check_map_constant(map: &MapArray, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
+    check_listview_constant(&map.entries().as_::<ListView>().into_owned(), ctx)
+}

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

@@ -6,6 +6,7 @@ mod decimal;
 mod extension;
 mod fixed_size_list;
 mod list;
+mod map;
 pub mod primitive;
 mod struct_;
 mod varbin;
@@ -20,6 +21,7 @@ use self::decimal::check_decimal_constant;
 use self::extension::check_extension_constant;
 use self::fixed_size_list::check_fixed_size_list_constant;
 use self::list::check_listview_constant;
+use self::map::check_map_constant;
 use self::primitive::check_primitive_constant;
 use self::struct_::check_struct_constant;
 use self::varbin::check_varbinview_constant;
@@ -402,9 +404,7 @@ 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::Map(m) => check_map_constant(m, ctx)?,
                     Canonical::FixedSizeList(f) => check_fixed_size_list_constant(f, ctx)?,
                     Canonical::Null(_) => true,
                     Canonical::Union(_) => {
@@ -456,14 +456,54 @@ mod tests {
     use crate::arrays::ListArray;
     use crate::arrays::PrimitiveArray;
     use crate::arrays::StructArray;
+    use crate::builders::MapBuilder;
     use crate::dtype::DType;
     use crate::dtype::DecimalDType;
     use crate::dtype::FieldNames;
+    use crate::dtype::MapDType;
     use crate::dtype::Nullability;
     use crate::dtype::PType;
     use crate::expr::stats::Stat;
+    use crate::scalar::Scalar;
     use crate::validity::Validity;
 
+    type MapEntryFixture<'a> = (i32, Option<&'a str>);
+    type MapRowFixture<'a> = Option<Vec<MapEntryFixture<'a>>>;
+
+    fn map_array_from_rows(rows: &[MapRowFixture<'_>]) -> VortexResult<crate::ArrayRef> {
+        let map_dtype = MapDType::try_new(
+            DType::Primitive(PType::I32, Nullability::NonNullable),
+            DType::Utf8(Nullability::Nullable),
+            false,
+        )?;
+        let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable);
+        let mut builder =
+            MapBuilder::<u64, u64>::with_capacity(map_dtype, Nullability::Nullable, rows.len());
+
+        for row in rows {
+            let scalar = match row {
+                Some(entries) => {
+                    let entries = entries
+                        .iter()
+                        .map(|(key, value)| {
+                            let key = Scalar::primitive(*key, Nullability::NonNullable);
+                            let value = value.map_or_else(
+                                || Scalar::null(DType::Utf8(Nullability::Nullable)),
+                                |value| Scalar::utf8(value, Nullability::Nullable),
+                            );
+                            (key, value)
+                        })
+                        .collect::<Vec<_>>();
+                    Scalar::try_map(dtype.clone(), entries)?
+                }
+                None => Scalar::null(dtype.clone()),
+            };
+            builder.append_value(scalar.as_map())?;
+        }
+
+        Ok(builder.finish_into_map().into_array())
+    }
+
     // Tests migrated from compute/is_constant.rs
     #[test]
     fn is_constant_min_max_no_nan() -> VortexResult<()> {
@@ -687,4 +727,26 @@ mod tests {
         assert_eq!(is_constant(&list_array.into_array(), &mut ctx)?, expected);
         Ok(())
     }
+
+    #[test]
+    fn test_map_is_constant() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+
+        let identical = map_array_from_rows(&[
+            Some(vec![(1, Some("one")), (2, None)]),
+            Some(vec![(1, Some("one")), (2, None)]),
+        ])?;
+        assert!(is_constant(&identical, &mut ctx)?);
+
+        let different = map_array_from_rows(&[
+            Some(vec![(1, Some("one")), (2, None)]),
+            Some(vec![(1, Some("one")), (3, None)]),
+        ])?;
+        assert!(!is_constant(&different, &mut ctx)?);
+
+        let all_null = map_array_from_rows(&[None, None])?;
+        assert!(is_constant(&all_null, &mut ctx)?);
+
+        Ok(())
+    }
 }

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

@@ -45,7 +45,8 @@ 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::ListView;
+use crate::arrays::map::MapArraySlotsExt;
 use crate::arrays::varbinview::BinaryView;
 use crate::dtype::DType;
 use crate::dtype::DecimalType;
@@ -200,9 +201,10 @@ 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::Map(array) => list_view_uncompressed_size_in_bytes(
+            &array.entries().as_::<ListView>().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),

vortex-array/src/arrays/arbitrary.rs

@@ -28,8 +28,10 @@ use crate::builders::ArrayBuilder;
 use crate::builders::DecimalBuilder;
 use crate::builders::FixedSizeListBuilder;
 use crate::builders::ListViewBuilder;
+use crate::builders::MapBuilder;
 use crate::dtype::DType;
 use crate::dtype::IntegerPType;
+use crate::dtype::MapDType;
 use crate::dtype::NativePType;
 use crate::dtype::Nullability;
 use crate::dtype::OffsetBuilderPType;
@@ -158,7 +160,9 @@ 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::Map(map_dtype, nullability) => {
+            random_map(u, map_dtype.clone(), *nullability, chunk_len)
+        }
         DType::Struct(sdt, n) => {
             let first_array = sdt
                 .fields()
@@ -199,6 +203,41 @@ fn random_array_chunk(
     }
 }
 
+fn random_map(
+    u: &mut Unstructured,
+    map_dtype: MapDType,
+    nullability: Nullability,
+    chunk_len: Option<usize>,
+) -> Result<ArrayRef> {
+    let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
+    let key_dtype = map_dtype.key_dtype();
+    let value_dtype = map_dtype.value_dtype();
+    let dtype = DType::Map(map_dtype.clone(), nullability);
+    let mut builder = MapBuilder::<u64, u64>::with_capacity(map_dtype, nullability, array_length);
+
+    for _ in 0..array_length {
+        if nullability == Nullability::Nullable && u.arbitrary::<bool>()? {
+            builder.append_null();
+        } else {
+            let entry_count = u.int_in_range(0..=20)?;
+            let entries = (0..entry_count)
+                .map(|_| {
+                    let key = random_scalar(u, &key_dtype)?;
+                    let value = random_scalar(u, &value_dtype)?;
+                    Ok((key, value))
+                })
+                .collect::<Result<Vec<_>>>()?;
+            let scalar = Scalar::try_map(dtype.clone(), entries)
+                .vortex_expect("generated map scalar should be valid");
+            builder
+                .append_scalar(&scalar)
+                .vortex_expect("generated map scalar should append");
+        }
+    }
+
+    Ok(builder.finish_into_map().into_array())
+}
+
 /// Creates a random fixed-size list array.
 ///
 /// If the `chunk_len` is specified, the length of the array will be equal to the chunk length.