better slots fix

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

Author: AdamGS

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

@@ -5,8 +5,10 @@ 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,
@@ -18,8 +20,8 @@ pub(super) fn check_map_identical(
     }
 
     check_list_identical(
-        &lhs.entries().into_owned(),
-        &rhs.entries().into_owned(),
+        &lhs.entries().as_::<ListView>().into_owned(),
+        &rhs.entries().as_::<ListView>().into_owned(),
         ctx,
     )
 }

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

@@ -5,9 +5,10 @@ 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::MapArrayExt;
+use crate::arrays::map::MapArraySlotsExt;
 
 pub(super) fn check_map_constant(map: &MapArray, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
-    check_listview_constant(&map.entries().into_owned(), ctx)
+    check_listview_constant(&map.entries().as_::<ListView>().into_owned(), ctx)
 }

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/map/array.rs

@@ -6,20 +6,20 @@ use std::fmt::Formatter;
 use std::hash::Hasher;
 use std::sync::Arc;
 
-use smallvec::smallvec;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 
 use crate::ArrayEq;
 use crate::ArrayHash;
 use crate::ArrayRef;
+use crate::ArraySlots;
 use crate::EqMode;
 use crate::IntoArray;
 use crate::array::Array;
 use crate::array::ArrayParts;
-use crate::array::ArrayView;
 use crate::array::TypedArrayRef;
+use crate::array_slots;
 use crate::arrays::ListView;
 use crate::arrays::ListViewArray;
 use crate::arrays::listview::ListViewArrayExt;
@@ -28,10 +28,12 @@ use crate::dtype::DType;
 use crate::dtype::MapDType;
 use crate::validity::Validity;
 
-/// The one child slot holding a [`ListViewArray`] of map entries.
-pub(super) const ENTRIES_SLOT: usize = 0;
-pub(super) const NUM_SLOTS: usize = 1;
-pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["entries"];
+#[array_slots(Map)]
+pub struct MapSlots {
+    /// The list-view storage of non-null `{key, value}` entry structs.
+    #[slot(0)]
+    pub entries: ArrayRef,
+}
 
 /// Encoding-specific metadata for [`crate::arrays::MapArray`].
 ///
@@ -56,6 +58,12 @@ impl ArrayHash for MapData {
     fn array_hash<H: Hasher>(&self, _state: &mut H, _accuracy: EqMode) {}
 }
 
+impl MapData {
+    pub(crate) fn make_slots(entries: ArrayRef) -> ArraySlots {
+        MapSlots { entries }.into_slots()
+    }
+}
+
 /// The logical and physical inputs used to construct a [`crate::arrays::MapArray`].
 pub struct MapDataParts {
     /// The key/value type and sortedness assertion for the map.
@@ -65,28 +73,20 @@ pub struct MapDataParts {
 }
 
 /// Accessors for the canonical map representation.
-pub trait MapArrayExt: TypedArrayRef<Map> {
-    /// Returns the list-view storage of map entry structs.
-    fn entries(&self) -> ArrayView<'_, ListView> {
-        self.as_ref().slots()[ENTRIES_SLOT]
-            .as_ref()
-            .vortex_expect("MapArray entries slot")
-            .as_::<ListView>()
-    }
-
+pub trait MapArrayExt: MapArraySlotsExt {
     /// Returns the entry structs for one map row.
     fn entries_at(&self, index: usize) -> VortexResult<ArrayRef> {
-        self.entries().list_elements_at(index)
+        self.entries().as_::<ListView>().list_elements_at(index)
     }
 
     /// Returns the number of entries in one map row.
     fn entry_count_at(&self, index: usize) -> usize {
-        self.entries().size_at(index)
+        self.entries().as_::<ListView>().size_at(index)
     }
 
     /// Returns the outer map validity delegated from the entries list-view.
     fn map_validity(&self) -> Validity {
-        self.entries().listview_validity()
+        self.entries().as_::<ListView>().listview_validity()
     }
 
     /// Returns this map's key/value type information.
@@ -125,8 +125,8 @@ impl Array<Map> {
         let nullability = entries.nullability();
         let dtype = DType::Map(map_dtype, nullability);
         let len = entries.len();
-        let parts = ArrayParts::new(Map, dtype, len, MapData)
-            .with_slots(smallvec![Some(entries.into_array())]);
+        let slots = MapData::make_slots(entries.into_array());
+        let parts = ArrayParts::new(Map, dtype, len, MapData).with_slots(slots);
         Self::try_from_parts(parts)
     }
 
@@ -141,8 +141,8 @@ impl Array<Map> {
         let nullability = entries.nullability();
         let dtype = DType::Map(map_dtype, nullability);
         let len = entries.len();
-        let parts = ArrayParts::new(Map, dtype, len, MapData)
-            .with_slots(smallvec![Some(entries.into_array())]);
+        let slots = MapData::make_slots(entries.into_array());
+        let parts = ArrayParts::new(Map, dtype, len, MapData).with_slots(slots);
         unsafe { Self::from_parts_unchecked(parts) }
     }
 
@@ -153,7 +153,7 @@ impl Array<Map> {
             .as_map_opt()
             .vortex_expect("MapArray requires a map dtype")
             .clone();
-        let entries = self.entries().into_owned();
+        let entries = self.entries().clone().downcast::<ListView>();
         MapDataParts { map_dtype, entries }
     }
 }

vortex-array/src/arrays/map/compute/cast.rs

@@ -11,6 +11,7 @@ use crate::array::ArrayView;
 use crate::arrays::ListView;
 use crate::arrays::map::Map;
 use crate::arrays::map::MapArrayExt;
+use crate::arrays::map::MapArraySlotsExt;
 use crate::arrays::map::compute::rebuild_map_from_array;
 use crate::dtype::DType;
 use crate::dtype::MapDType;
@@ -47,7 +48,10 @@ impl CastReduce for Map {
             return Ok(None);
         };
 
-        let Some(entries) = <ListView as CastReduce>::cast(array.entries(), &target_entries_dtype)?
+        let Some(entries) = <ListView as CastReduce>::cast(
+            array.entries().as_::<ListView>(),
+            &target_entries_dtype,
+        )?
         else {
             return Ok(None);
         };
@@ -67,8 +71,11 @@ impl CastKernel for Map {
             return Ok(None);
         };
 
-        let Some(entries) =
-            <ListView as CastKernel>::cast(array.entries(), &target_entries_dtype, ctx)?
+        let Some(entries) = <ListView as CastKernel>::cast(
+            array.entries().as_::<ListView>(),
+            &target_entries_dtype,
+            ctx,
+        )?
         else {
             return Ok(None);
         };

vortex-array/src/arrays/map/compute/filter.rs

@@ -5,19 +5,22 @@ use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::ArrayRef;
+use crate::IntoArray;
 use crate::array::ArrayView;
+use crate::arrays::ListView;
 use crate::arrays::ListViewArray;
+use crate::arrays::MapArray;
 use crate::arrays::filter::FilterKernel;
 use crate::arrays::filter::FilterReduce;
 use crate::arrays::listview::ListViewArraySlotsExt;
 use crate::arrays::map::Map;
 use crate::arrays::map::MapArrayExt;
-use crate::arrays::map::compute::rebuild_map;
+use crate::arrays::map::MapArraySlotsExt;
 use crate::executor::ExecutionCtx;
 
 impl FilterReduce for Map {
     fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult<Option<ArrayRef>> {
-        let entries = array.entries();
+        let entries = array.entries().as_::<ListView>();
 
         // SAFETY: filtering row metadata keeps offsets and sizes paired, preserves the original
         // elements, and filters validity to the same output length. The zero-copy-to-list flag is
@@ -31,7 +34,11 @@ impl FilterReduce for Map {
             )
         };
 
-        rebuild_map(array.map_dtype().clone(), filtered_entries).map(Some)
+        {
+            let map_dtype = array.map_dtype().clone();
+            MapArray::try_new(map_dtype, filtered_entries).map(IntoArray::into_array)
+        }
+        .map(Some)
     }
 }
 

vortex-array/src/arrays/map/compute/mask.rs

@@ -8,14 +8,17 @@ use crate::array::ArrayView;
 use crate::arrays::ListView;
 use crate::arrays::map::Map;
 use crate::arrays::map::MapArrayExt;
+use crate::arrays::map::MapArraySlotsExt;
 use crate::arrays::map::compute::rebuild_map_from_array;
 use crate::executor::ExecutionCtx;
 use crate::scalar_fn::fns::mask::MaskKernel;
 use crate::scalar_fn::fns::mask::MaskReduce;
 
 impl MaskReduce for Map {
     fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-        let Some(entries) = <ListView as MaskReduce>::mask(array.entries(), mask)? else {
+        let Some(entries) =
+            <ListView as MaskReduce>::mask(array.entries().as_::<ListView>(), mask)?
+        else {
             return Ok(None);
         };
 

vortex-array/src/arrays/map/compute/mod.rs

@@ -16,23 +16,16 @@ use vortex_error::vortex_err;
 use crate::ArrayRef;
 use crate::IntoArray;
 use crate::arrays::ListView;
-use crate::arrays::ListViewArray;
 use crate::arrays::MapArray;
 use crate::dtype::MapDType;
 
-fn rebuild_map(map_dtype: MapDType, entries: ListViewArray) -> VortexResult<ArrayRef> {
-    MapArray::try_new(map_dtype, entries).map(IntoArray::into_array)
-}
-
 fn rebuild_map_from_array(map_dtype: MapDType, entries: ArrayRef) -> VortexResult<ArrayRef> {
-    let entries = entries
-        .as_opt::<ListView>()
-        .ok_or_else(|| {
-            vortex_err!(
-                "Map entries operation expected vortex.listview/ListView, got {}",
-                entries.encoding_id()
-            )
-        })?
-        .into_owned();
-    rebuild_map(map_dtype, entries)
+    let map_entries = entries.try_downcast::<ListView>().map_err(|arr| {
+        vortex_err!(
+            "Map entries operation expected vortex.listview/ListView, got {}",
+            arr.encoding_id()
+        )
+    })?;
+
+    MapArray::try_new(map_dtype, map_entries).map(IntoArray::into_array)
 }

vortex-array/src/arrays/map/compute/slice.rs

@@ -10,14 +10,17 @@ use crate::array::ArrayView;
 use crate::arrays::ListView;
 use crate::arrays::map::Map;
 use crate::arrays::map::MapArrayExt;
+use crate::arrays::map::MapArraySlotsExt;
 use crate::arrays::map::compute::rebuild_map_from_array;
 use crate::arrays::slice::SliceKernel;
 use crate::arrays::slice::SliceReduce;
 use crate::executor::ExecutionCtx;
 
 impl SliceReduce for Map {
     fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
-        let Some(sliced_entries) = <ListView as SliceReduce>::slice(array.entries(), range)? else {
+        let Some(sliced_entries) =
+            <ListView as SliceReduce>::slice(array.entries().as_::<ListView>(), range)?
+        else {
             return Ok(None);
         };
 

vortex-array/src/arrays/map/compute/take.rs

@@ -10,12 +10,15 @@ use crate::arrays::dict::TakeExecute;
 use crate::arrays::dict::TakeReduce;
 use crate::arrays::map::Map;
 use crate::arrays::map::MapArrayExt;
+use crate::arrays::map::MapArraySlotsExt;
 use crate::arrays::map::compute::rebuild_map_from_array;
 use crate::executor::ExecutionCtx;
 
 impl TakeReduce for Map {
     fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-        let Some(entries) = <ListView as TakeReduce>::take(array.entries(), indices)? else {
+        let Some(entries) =
+            <ListView as TakeReduce>::take(array.entries().as_::<ListView>(), indices)?
+        else {
             return Ok(None);
         };
 
@@ -29,7 +32,9 @@ impl TakeExecute for Map {
         indices: &ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
-        let Some(entries) = <ListView as TakeExecute>::take(array.entries(), indices, ctx)? else {
+        let Some(entries) =
+            <ListView as TakeExecute>::take(array.entries().as_::<ListView>(), indices, ctx)?
+        else {
             return Ok(None);
         };