From bc1872e05f0756dd10bf1ea243f4cad720c2d995 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Thu, 30 Jul 2026 11:13:06 +0100 Subject: [PATCH] Add canonical Map arrays and builder Signed-off-by: Adam Gutglick --- fuzz/src/array/fill_null.rs | 1 + fuzz/src/array/mask.rs | 1 + fuzz/src/array/scalar_at.rs | 1 + .../src/aggregate_fn/fns/is_constant/mod.rs | 3 + .../src/aggregate_fn/fns/min_max/mod.rs | 1 + .../fns/uncompressed_size_in_bytes/mod.rs | 18 +- .../src/arrays/constant/vtable/canonical.rs | 12 +- vortex-array/src/arrays/dict/execute.rs | 1 + vortex-array/src/arrays/filter/execute/mod.rs | 2 + vortex-array/src/arrays/map/array.rs | 190 +++++++++++++ vortex-array/src/arrays/map/mod.rs | 16 ++ vortex-array/src/arrays/map/tests.rs | 255 ++++++++++++++++++ vortex-array/src/arrays/map/vtable/mod.rs | 160 +++++++++++ .../src/arrays/map/vtable/operations.rs | 35 +++ .../src/arrays/map/vtable/validity.rs | 14 + vortex-array/src/arrays/masked/execute.rs | 2 + vortex-array/src/arrays/mod.rs | 6 +- vortex-array/src/builders/map.rs | 159 +++++++++++ vortex-array/src/builders/mod.rs | 27 +- vortex-array/src/canonical.rs | 99 ++++++- vortex-array/src/dtype/dtype_impl.rs | 6 +- vortex-array/src/scalar_fn/fns/cast/mod.rs | 1 + vortex-array/src/session/mod.rs | 2 + vortex-compressor/src/compressor/cascade.rs | 2 + vortex-duckdb/src/exporter/canonical.rs | 1 + vortex-flatbuffers/src/generated/array.rs | 2 +- vortex-flatbuffers/src/generated/dtype.rs | 2 +- vortex-flatbuffers/src/generated/footer.rs | 2 +- vortex-flatbuffers/src/generated/layout.rs | 2 +- vortex-flatbuffers/src/generated/message.rs | 2 +- 30 files changed, 1001 insertions(+), 24 deletions(-) create mode 100644 vortex-array/src/arrays/map/array.rs create mode 100644 vortex-array/src/arrays/map/mod.rs create mode 100644 vortex-array/src/arrays/map/tests.rs create mode 100644 vortex-array/src/arrays/map/vtable/mod.rs create mode 100644 vortex-array/src/arrays/map/vtable/operations.rs create mode 100644 vortex-array/src/arrays/map/vtable/validity.rs create mode 100644 vortex-array/src/builders/map.rs diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index 75adfa94847..fd8d8ec88c3 100644 --- a/fuzz/src/array/fill_null.rs +++ b/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(_) => { diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index 64013f336ed..528d4576dfa 100644 --- a/fuzz/src/array/mask.rs +++ b/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"), }) } diff --git a/fuzz/src/array/scalar_at.rs b/fuzz/src/array/scalar_at.rs index 887974c7d42..712dbccea7c 100644 --- a/fuzz/src/array/scalar_at.rs +++ b/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"), }) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 708c83debc4..ac7f7cb9ce3 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/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(_) => { diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index de950e59d7d..a5091cf4a9f 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/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()) diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 22482c60a45..fabe596f1cd 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/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::(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)), diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 9f51c545a6c..6929f30f815 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/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(); diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index 3349e01cd63..4a3ccda43c2 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/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)) } diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 2c712009d1d..3be654d213d 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/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; @@ -95,6 +96,7 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca 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_panic!("Map arrays don't support filter"), Canonical::FixedSizeList(a) => { Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask)) } diff --git a/vortex-array/src/arrays/map/array.rs b/vortex-array/src/arrays/map/array.rs new file mode 100644 index 00000000000..73a7e4db6bf --- /dev/null +++ b/vortex-array/src/arrays/map/array.rs @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +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::EqMode; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::TypedArrayRef; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::map::Map; +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"]; + +/// Encoding-specific metadata for [`MapArray`]. +/// +/// All map metadata is represented by the outer [`DType::Map`] and the entries child, so this +/// value is intentionally empty. +#[derive(Clone, Debug, Default)] +pub struct MapData; + +impl Display for MapData { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + Ok(()) + } +} + +impl ArrayEq for MapData { + fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { + true + } +} + +impl ArrayHash for MapData { + fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} +} + +/// The logical and physical inputs used to construct a [`MapArray`]. +pub struct MapDataParts { + /// The key/value type and sortedness assertion for the map. + pub map_dtype: MapDType, + /// The physical list-view storage of `{key, value}` entry structs. + pub entries: ListViewArray, +} + +/// Accessors for the canonical map representation. +pub trait MapArrayExt: TypedArrayRef { + /// 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_::() + } + + /// Returns the entry structs for one map row. + fn entries_at(&self, index: usize) -> VortexResult { + self.entries().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) + } + + /// Returns the outer map validity delegated from the entries list-view. + fn map_validity(&self) -> Validity { + self.entries().listview_validity() + } + + /// Returns this map's key/value type information. + fn map_dtype(&self) -> &MapDType { + self.as_ref() + .dtype() + .as_map_opt() + .vortex_expect("MapArray requires a map dtype") + } + + /// Returns whether producers assert sorted keys within each map value. + fn keys_sorted(&self) -> bool { + self.map_dtype().keys_sorted() + } +} +impl> MapArrayExt for T {} + +impl Array { + /// Creates a canonical map array from its map dtype and list-view entry storage. + /// + /// # Panics + /// + /// Panics if `entries` is not a list of the map dtype's non-nullable `{key, value}` entry + /// struct with matching outer nullability. + pub fn new(map_dtype: MapDType, entries: ListViewArray) -> Self { + Self::try_new(map_dtype, entries).vortex_expect("MapArray construction failed") + } + + /// Constructs a canonical map array from its map dtype and list-view entry storage. + /// + /// # Errors + /// + /// Returns an error when the entry child is not `ListView>`, has a + /// different outer nullability, or has a different length than the outer map array. + pub fn try_new(map_dtype: MapDType, entries: ListViewArray) -> VortexResult { + 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())]); + Self::try_from_parts(parts) + } + + /// Creates a canonical map array without validating its entry storage. + /// + /// # Safety + /// + /// The caller must ensure that `entries` has dtype + /// `List(Struct { key, value }, entries.nullability())`, where the struct exactly matches + /// `map_dtype.entries_dtype()`. + pub unsafe fn new_unchecked(map_dtype: MapDType, entries: ListViewArray) -> Self { + 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())]); + unsafe { Self::from_parts_unchecked(parts) } + } + + /// Decomposes this map array into its logical dtype and physical entries child. + pub fn into_data_parts(self) -> MapDataParts { + let map_dtype = self + .dtype() + .as_map_opt() + .vortex_expect("MapArray requires a map dtype") + .clone(); + let entries = self.entries().into_owned(); + MapDataParts { map_dtype, entries } + } +} + +fn expected_entries_dtype(map_dtype: &MapDType, nullability: crate::dtype::Nullability) -> DType { + DType::List(Arc::new(map_dtype.entries_dtype()), nullability) +} + +pub(super) fn validate_entries( + map_dtype: &MapDType, + nullability: crate::dtype::Nullability, + len: usize, + entries: &ArrayRef, +) -> VortexResult<()> { + vortex_ensure!( + entries.is::(), + "MapArray entries must use vortex.listview encoding, got {}", + entries.encoding_id() + ); + vortex_ensure!( + entries.len() == len, + "MapArray entries length {} does not match outer length {len}", + entries.len() + ); + + let expected_dtype = expected_entries_dtype(map_dtype, nullability); + vortex_ensure!( + entries.dtype() == &expected_dtype, + "MapArray entries dtype {} does not match expected {expected_dtype}", + entries.dtype() + ); + + Ok(()) +} diff --git a/vortex-array/src/arrays/map/mod.rs b/vortex-array/src/arrays/map/mod.rs new file mode 100644 index 00000000000..27eb138a1ed --- /dev/null +++ b/vortex-array/src/arrays/map/mod.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Canonical map arrays backed by [`ListView`](crate::arrays::ListView) entry storage. + +mod array; +pub use array::MapArrayExt; +pub use array::MapData; +pub use array::MapDataParts; + +mod vtable; +pub use vtable::Map; +pub use vtable::MapArray; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/map/tests.rs b/vortex-array/src/arrays/map/tests.rs new file mode 100644 index 00000000000..a7bc41b5292 --- /dev/null +++ b/vortex-array/src/arrays/map/tests.rs @@ -0,0 +1,255 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::smallvec; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexResult; +use vortex_session::registry::ReadContext; + +use crate::Array; +use crate::ArrayContext; +use crate::ArrayParts; +use crate::ArrayVTable; +use crate::Canonical; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::ListViewArray; +use crate::arrays::Map; +use crate::arrays::MapArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapData; +use crate::arrays::map::MapDataParts; +use crate::builders::ArrayBuilder; +use crate::builders::MapBuilder; +use crate::dtype::DType; +use crate::dtype::MapDType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; +use crate::session::ArraySessionExt; +use crate::validity::Validity; + +fn map_dtype() -> VortexResult { + MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + ) +} + +fn key(value: i32) -> Scalar { + Scalar::primitive(value, Nullability::NonNullable) +} + +fn value(value: Option<&str>) -> Scalar { + match value { + Some(value) => Scalar::utf8(value, Nullability::Nullable), + None => Scalar::null(DType::Utf8(Nullability::Nullable)), + } +} + +fn sample_scalar(dtype: DType) -> VortexResult { + Scalar::try_map( + dtype, + [ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ], + ) +} + +fn sample_array() -> VortexResult { + let map_dtype = map_dtype()?; + let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); + let mut builder = MapBuilder::::with_capacity(map_dtype, Nullability::Nullable, 3); + builder.append_scalar(&sample_scalar(dtype.clone())?)?; + builder.append_scalar(&Scalar::try_map(dtype.clone(), [])?)?; + builder.append_scalar(&Scalar::null(dtype))?; + Ok(builder.finish_into_map()) +} + +#[test] +fn constructs_map_with_listview_entries() -> VortexResult<()> { + let MapDataParts { map_dtype, entries } = sample_array()?.into_data_parts(); + let array = MapArray::try_new(map_dtype.clone(), entries)?; + + assert_eq!( + array.dtype(), + &DType::Map(map_dtype.clone(), Nullability::Nullable) + ); + assert!(array.keys_sorted()); + assert_eq!(array.entry_count_at(0), 3); + assert_eq!(array.entry_count_at(1), 0); + assert_eq!(array.entries_at(0)?.dtype(), &map_dtype.entries_dtype()); + + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array + .map_validity() + .execute_mask(array.len(), &mut ctx)? + .true_count(), + 2 + ); + + Ok(()) +} + +#[test] +fn accepts_duplicate_and_unsorted_keys() -> VortexResult<()> { + let array = sample_array()?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + sample_scalar(array.dtype().clone())? + ); + + Ok(()) +} + +#[test] +fn rejects_malformed_entry_storage() -> VortexResult<()> { + let map_dtype = map_dtype()?; + let offsets = PrimitiveArray::from_iter([0u64]).into_array(); + let sizes = PrimitiveArray::from_iter([1u64]).into_array(); + let non_struct_entries = ListViewArray::try_new( + PrimitiveArray::from_iter([1i32]).into_array(), + offsets, + sizes, + Validity::NonNullable, + )?; + assert!(MapArray::try_new(map_dtype.clone(), non_struct_entries).is_err()); + + let nullable_entry_struct = + ConstantArray::new(Scalar::null(map_dtype.entries_dtype().as_nullable()), 1).into_array(); + let null_entries = ListViewArray::try_new( + nullable_entry_struct, + PrimitiveArray::from_iter([0u64]).into_array(), + PrimitiveArray::from_iter([1u64]).into_array(), + Validity::NonNullable, + )?; + assert!(MapArray::try_new(map_dtype.clone(), null_entries).is_err()); + + let MapDataParts { entries, .. } = sample_array()?.into_data_parts(); + let parts = ArrayParts::new( + Map, + DType::Map(map_dtype, Nullability::NonNullable), + entries.len(), + MapData, + ) + .with_slots(smallvec![Some(entries.into_array())]); + assert!(Array::::try_from_parts(parts).is_err()); + + assert!( + MapDType::try_new( + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::Nullable), + false, + ) + .is_err() + ); + + Ok(()) +} + +#[test] +fn scalar_access_preserves_null_and_empty_maps() -> VortexResult<()> { + let array = sample_array()?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + sample_scalar(array.dtype().clone())? + ); + assert!(array.execute_scalar(1, &mut ctx)?.as_map().is_empty()); + assert!(array.execute_scalar(2, &mut ctx)?.is_null()); + + Ok(()) +} + +#[test] +fn builder_appends_existing_map_arrays() -> VortexResult<()> { + let source = sample_array()?; + let mut builder = MapBuilder::::with_capacity( + source.map_dtype().clone(), + source.dtype().nullability(), + 0, + ); + let mut ctx = array_session().create_execution_ctx(); + builder.append_map_array(source.as_view(), &mut ctx)?; + builder.append_map_array(source.as_view(), &mut ctx)?; + let array = builder.finish_into_map(); + + assert_eq!(array.len(), 6); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + source.execute_scalar(0, &mut ctx)? + ); + assert!(array.execute_scalar(2, &mut ctx)?.is_null()); + assert_eq!( + array.execute_scalar(3, &mut ctx)?, + source.execute_scalar(0, &mut ctx)? + ); + + Ok(()) +} + +#[test] +fn canonicalizes_empty_constant_and_chunked_maps() -> VortexResult<()> { + let map_dtype = map_dtype()?; + let dtype = DType::Map(map_dtype, Nullability::Nullable); + let empty = Canonical::empty(&dtype); + assert!(empty.as_map().is_empty()); + + let mut ctx = array_session().create_execution_ctx(); + let constant = ConstantArray::new(sample_scalar(dtype.clone())?, 2) + .into_array() + .execute::(&mut ctx)?; + assert_eq!(constant.as_map().len(), 2); + + let first = sample_array()?.into_array(); + let second = sample_array()?.into_array(); + let chunked = ChunkedArray::try_new(vec![first, second], dtype)?.into_array(); + let canonical = chunked.execute::(&mut ctx)?; + assert_eq!(canonical.as_map().len(), 6); + + Ok(()) +} + +#[test] +fn serde_roundtrip_uses_registered_map_vtable() -> VortexResult<()> { + let session = array_session(); + assert!(session.arrays().registry().contains_key(&Map.id())); + + let array = sample_array()?.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + let array_ctx = ArrayContext::empty(); + let serialized = array.serialize(&array_ctx, &session, &SerializeOptions::default())?; + let mut concat = ByteBufferMut::empty(); + for buffer in serialized { + concat.extend_from_slice(buffer.as_ref()); + } + + let serialized = SerializedArray::try_from(concat.freeze())?; + let decoded = + serialized.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)?; + assert!(decoded.is::()); + + let mut ctx = session.create_execution_ctx(); + for index in 0..len { + assert_eq!( + decoded.execute_scalar(index, &mut ctx)?, + array.execute_scalar(index, &mut ctx)? + ); + } + + Ok(()) +} diff --git a/vortex-array/src/arrays/map/vtable/mod.rs b/vortex-array/src/arrays/map/vtable/mod.rs new file mode 100644 index 00000000000..7e4e2f67d3d --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/mod.rs @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayParts; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::array::ValidityVTableFromChild; +use crate::array::with_empty_buffers; +use crate::arrays::ListView; +use crate::arrays::map::MapData; +use crate::arrays::map::array::ENTRIES_SLOT; +use crate::arrays::map::array::NUM_SLOTS; +use crate::arrays::map::array::SLOT_NAMES; +use crate::arrays::map::array::validate_entries; +use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::dtype::DType; +use crate::serde::ArrayChildren; + +mod operations; +mod validity; + +/// A [`Map`]-encoded Vortex array. +pub type MapArray = Array; + +/// The canonical encoding for [`DType::Map`]. +/// +/// A map array has one `ListView>` child. Its outer dtype retains map-specific +/// metadata such as the `keys_sorted` assertion. +#[derive(Clone, Debug, Default)] +pub struct Map; + +impl VTable for Map { + type TypedArrayData = MapData; + + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.map"); + *ID + } + + fn validate( + &self, + _data: &MapData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + slots.len() == NUM_SLOTS, + "MapArray expected {NUM_SLOTS} slot, found {}", + slots.len() + ); + + let DType::Map(map_dtype, nullability) = dtype else { + vortex_bail!("Expected map dtype, got {dtype}"); + }; + let entries = slots[ENTRIES_SLOT] + .as_ref() + .ok_or_else(|| vortex_error::vortex_err!("MapArray missing entries slot"))?; + validate_entries(map_dtype, *nullability, len, entries) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("MapArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + if !metadata.is_empty() { + vortex_bail!( + "MapArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + vortex_ensure!(buffers.is_empty(), "MapArray expects no buffers"); + + let DType::Map(map_dtype, nullability) = dtype else { + vortex_bail!("Expected map dtype, got {dtype}"); + }; + vortex_ensure!( + children.len() == NUM_SLOTS, + "MapArray expected {NUM_SLOTS} child, found {}", + children.len() + ); + + let expected_entries_dtype = + DType::List(std::sync::Arc::new(map_dtype.entries_dtype()), *nullability); + let entries = children.get(ENTRIES_SLOT, &expected_entries_dtype, len)?; + vortex_ensure!( + entries.is::(), + "MapArray entries must use vortex.listview encoding, got {}", + entries.encoding_id() + ); + + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, MapData) + .with_slots(smallvec![Some(entries)])) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + SLOT_NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + builder.append_map_array(array, ctx) + } +} diff --git a/vortex-array/src/arrays/map/vtable/operations.rs b/vortex-array/src/arrays/map/vtable/operations.rs new file mode 100644 index 00000000000..d6e8fe87f12 --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/operations.rs @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ExecutionCtx; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::arrays::Map; +use crate::arrays::StructArray; +use crate::arrays::map::MapArrayExt; +use crate::arrays::struct_::StructArrayExt; +use crate::scalar::Scalar; + +impl OperationsVTable for Map { + fn scalar_at( + array: ArrayView<'_, Map>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let entries = array.entries_at(index)?.execute::(ctx)?; + let keys = entries.unmasked_field(0); + let values = entries.unmasked_field(1); + let pairs = (0..entries.len()) + .map(|entry_index| { + Ok(( + keys.execute_scalar(entry_index, ctx)?, + values.execute_scalar(entry_index, ctx)?, + )) + }) + .collect::>>()?; + + Scalar::try_map(array.dtype().clone(), pairs) + } +} diff --git a/vortex-array/src/arrays/map/vtable/validity.rs b/vortex-array/src/arrays/map/vtable/validity.rs new file mode 100644 index 00000000000..0cdc1f886cb --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/validity.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::ArrayRef; +use crate::array::ArrayView; +use crate::array::ValidityChild; +use crate::arrays::Map; +use crate::arrays::map::MapArrayExt; + +impl ValidityChild for Map { + fn validity_child(array: ArrayView<'_, Map>) -> ArrayRef { + array.entries().array().clone() + } +} diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index a8d17f15e2d..b9d25946710 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::Canonical; use crate::IntoArray; @@ -50,6 +51,7 @@ pub fn mask_validity_canonical( Canonical::Decimal(a) => Canonical::Decimal(mask_validity_decimal(a, validity)?), Canonical::VarBinView(a) => Canonical::VarBinView(mask_validity_varbinview(a, validity)?), Canonical::List(a) => Canonical::List(mask_validity_listview(a, validity)?), + Canonical::Map(_) => vortex_bail!("Map arrays don't support masking"), Canonical::FixedSizeList(a) => { Canonical::FixedSizeList(mask_validity_fixed_size_list(a, validity)?) } diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 7c31a1b48b1..8e819094b6f 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -5,7 +5,7 @@ //! //! Canonical arrays are the default uncompressed representation for a logical dtype: //! [`NullArray`], [`BoolArray`], [`PrimitiveArray`], [`DecimalArray`], [`VarBinViewArray`], -//! [`ListViewArray`], [`FixedSizeListArray`], [`StructArray`], [`UnionArray`], +//! [`ListViewArray`], [`MapArray`], [`FixedSizeListArray`], [`StructArray`], [`UnionArray`], //! [`ExtensionArray`], and [`VariantArray`]. //! //! Utility and lazy arrays represent common transformations without immediately materializing @@ -76,6 +76,10 @@ pub mod listview; pub use listview::ListView; pub use listview::ListViewArray; +pub mod map; +pub use map::Map; +pub use map::MapArray; + pub mod masked; pub use masked::Masked; pub use masked::MaskedArray; diff --git a/vortex-array/src/builders/map.rs b/vortex-array/src/builders/map.rs new file mode 100644 index 00000000000..a95f9367f12 --- /dev/null +++ b/vortex-array/src/builders/map.rs @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; +use std::sync::Arc; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Map; +use crate::arrays::MapArray; +use crate::arrays::map::MapArrayExt; +use crate::builders::ArrayBuilder; +use crate::builders::DEFAULT_BUILDER_CAPACITY; +use crate::builders::ListViewBuilder; +use crate::dtype::DType; +use crate::dtype::IntegerPType; +use crate::dtype::MapDType; +use crate::dtype::Nullability; +use crate::scalar::MapScalar; +use crate::scalar::Scalar; + +/// A builder for canonical [`MapArray`] values. +/// +/// The builder owns a [`ListViewBuilder`] whose elements are non-nullable `{key, value}` structs. +/// It preserves the map dtype's `keys_sorted` assertion while delegating offsets, sizes, and outer +/// validity to that list-view builder. +pub struct MapBuilder { + dtype: DType, + map_dtype: MapDType, + entries_builder: ListViewBuilder, +} + +impl MapBuilder { + /// Creates a map builder with the default capacity. + pub fn new(map_dtype: MapDType, nullability: Nullability) -> Self { + Self::with_capacity(map_dtype, nullability, DEFAULT_BUILDER_CAPACITY) + } + + /// Creates a map builder with space for `capacity` map rows. + pub fn with_capacity(map_dtype: MapDType, nullability: Nullability, capacity: usize) -> Self { + let entries_builder = ListViewBuilder::with_capacity( + Arc::new(map_dtype.entries_dtype()), + nullability, + capacity.saturating_mul(2), + capacity, + ); + let dtype = DType::Map(map_dtype.clone(), nullability); + Self { + dtype, + map_dtype, + entries_builder, + } + } + + /// Appends one map scalar. + pub fn append_value(&mut self, value: MapScalar<'_>) -> VortexResult<()> { + vortex_ensure!( + value.dtype() == &self.dtype, + "MapBuilder expected map scalar with dtype {}, got {}", + self.dtype, + value.dtype() + ); + + if value.is_null() { + self.entries_builder.append_null(); + return Ok(()); + } + + let entry_dtype = self.map_dtype.entries_dtype(); + let entries = value + .entries() + .map(|(key, value)| Scalar::struct_(entry_dtype.clone(), vec![key, value])) + .collect(); + let entries = Scalar::list(Arc::new(entry_dtype), entries, self.dtype.nullability()); + self.entries_builder.append_value(entries.as_list()) + } + + /// Finishes the builder directly into a [`MapArray`]. + pub fn finish_into_map(&mut self) -> MapArray { + MapArray::new( + self.map_dtype.clone(), + self.entries_builder.finish_into_listview(), + ) + } +} + +impl ArrayBuilder for MapBuilder { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn len(&self) -> usize { + self.entries_builder.len() + } + + fn append_zeros(&mut self, n: usize) { + self.entries_builder.append_zeros(n); + } + + unsafe fn append_nulls_unchecked(&mut self, n: usize) { + unsafe { self.entries_builder.append_nulls_unchecked(n) }; + } + + fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> { + vortex_ensure!( + scalar.dtype() == self.dtype(), + "MapBuilder expected scalar with dtype {}, got {}", + self.dtype(), + scalar.dtype() + ); + self.append_value(scalar.as_map()) + } + + fn reserve_exact(&mut self, additional: usize) { + self.entries_builder.reserve_exact(additional); + } + + unsafe fn set_validity_unchecked(&mut self, validity: Mask) { + unsafe { self.entries_builder.set_validity_unchecked(validity) }; + } + + fn finish(&mut self) -> ArrayRef { + self.finish_into_map().into_array() + } + + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { + Canonical::Map(self.finish_into_map()) + } + + fn append_map_array( + &mut self, + array: ArrayView<'_, Map>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.dtype(), + "MapBuilder expected map array with dtype {}, got {}", + self.dtype(), + array.dtype() + ); + self.entries_builder + .append_listview_array(array.entries(), ctx) + } +} diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 9e4af7677c4..bf29482bc9c 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -42,6 +42,7 @@ use crate::ExecutionCtx; use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListView; +use crate::arrays::Map; use crate::canonical::Canonical; use crate::dtype::DType; use crate::match_each_decimal_value_type; @@ -59,6 +60,7 @@ mod extension; mod fixed_size_list; mod list; mod listview; +mod map; mod null; mod primitive; mod struct_; @@ -70,6 +72,7 @@ pub use extension::*; pub use fixed_size_list::*; pub use list::*; pub use listview::*; +pub use map::*; pub use null::*; pub use primitive::*; pub use struct_::*; @@ -230,6 +233,22 @@ pub trait ArrayBuilder: Send { self.dtype() ) } + + /// Appends the values of a [`Map`]-encoded `array` to this builder. + /// + /// Only map-typed builders support this; canonical map arrays dispatch through this hook so + /// the generic offset and size types of their nested list-view builders stay erased. + fn append_map_array( + &mut self, + array: ArrayView<'_, Map>, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_bail!( + "cannot append a Map array of dtype {} to a {} builder", + array.dtype(), + self.dtype() + ) + } } /// Construct a new canonical builder for the given [`DType`]. @@ -290,9 +309,11 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box { - vortex_error::vortex_panic!(InvalidArgument: "map builders are not yet supported") - } + DType::Map(map_dtype, nullability) => Box::new(MapBuilder::::with_capacity( + map_dtype.clone(), + *nullability, + capacity, + )), DType::FixedSizeList(elem_dtype, list_size, null) => { Box::new(FixedSizeListBuilder::with_capacity( Arc::clone(elem_dtype), diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 36a59ae1b7c..1d8487c4eed 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -29,6 +29,8 @@ use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::ListView; use crate::arrays::ListViewArray; +use crate::arrays::Map; +use crate::arrays::MapArray; use crate::arrays::Null; use crate::arrays::NullArray; use crate::arrays::Primitive; @@ -47,6 +49,7 @@ use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewDataParts; use crate::arrays::listview::ListViewRebuildMode; +use crate::arrays::map::MapArrayExt; use crate::arrays::primitive::PrimitiveDataParts; use crate::arrays::struct_::StructDataParts; use crate::arrays::union::UnionDataParts; @@ -83,9 +86,10 @@ use crate::validity::Validity; /// /// # Arrow interoperability /// -/// Vortex canonical encodings have equivalent Arrow encodings that can be built zero-copy, except -/// [`UnionArray`], whose independent top-level validity cannot be represented directly by an Arrow -/// union. The corresponding Arrow array types can also be built directly. +/// Most Vortex canonical encodings have an equivalent Arrow encoding that can be built zero-copy, +/// and the corresponding Arrow array types can also be built directly. Map array Arrow transport is +/// not implemented yet, and [`UnionArray`]'s independent top-level validity cannot be represented +/// directly by an Arrow union. /// /// The full list of canonical types and their equivalent Arrow array types are: /// @@ -95,6 +99,7 @@ use crate::validity::Validity; /// * `DecimalArray`: `arrow_array::Decimal128Array` and `arrow_array::Decimal256Array` /// * `VarBinViewArray`: `arrow_array::GenericByteViewArray` /// * `ListViewArray`: `arrow_array::ListViewArray` +/// * `MapArray`: Vortex `ListView>` storage /// * `FixedSizeListArray`: `arrow_array::FixedSizeListArray` /// * `StructArray`: `arrow_array::StructArray` /// @@ -130,6 +135,7 @@ pub enum Canonical { Decimal(DecimalArray), VarBinView(VarBinViewArray), List(ListViewArray), + Map(MapArray), FixedSizeList(FixedSizeListArray), Struct(StructArray), Union(UnionArray), @@ -149,6 +155,7 @@ macro_rules! match_each_canonical { Canonical::Decimal($ident) => $eval, Canonical::VarBinView($ident) => $eval, Canonical::List($ident) => $eval, + Canonical::Map($ident) => $eval, Canonical::FixedSizeList($ident) => $eval, Canonical::Struct($ident) => $eval, Canonical::Union($ident) => $eval, @@ -215,9 +222,14 @@ impl Canonical { // An empty list view is trivially copyable to a list. .with_zero_copy_to_list(true) }), - DType::Map(..) => { - vortex_panic!(InvalidArgument: "canonical map arrays are not yet supported") - } + DType::Map(map_dtype, nullability) => Canonical::Map(MapArray::new( + map_dtype.clone(), + Canonical::empty(&DType::List( + Arc::new(map_dtype.entries_dtype()), + *nullability, + )) + .into_listview(), + )), DType::FixedSizeList(elem_dtype, list_size, null) => Canonical::FixedSizeList(unsafe { FixedSizeListArray::new_unchecked( Canonical::empty(elem_dtype).into_array(), @@ -277,6 +289,13 @@ impl Canonical { Canonical::List(array) => Ok(Canonical::List( array.rebuild(ListViewRebuildMode::TrimElements, ctx)?, )), + Canonical::Map(array) => Ok(Canonical::Map(MapArray::new( + array.map_dtype().clone(), + array + .entries() + .into_owned() + .rebuild(ListViewRebuildMode::TrimElements, ctx)?, + ))), _ => Ok(self.clone()), } } @@ -380,6 +399,22 @@ impl Canonical { } } + pub fn as_map(&self) -> &MapArray { + if let Canonical::Map(a) = self { + a + } else { + vortex_panic!("Cannot get MapArray from {:?}", &self) + } + } + + pub fn into_map(self) -> MapArray { + if let Canonical::Map(a) = self { + a + } else { + vortex_panic!("Cannot unwrap MapArray from {:?}", &self) + } + } + pub fn as_fixed_size_list(&self) -> &FixedSizeListArray { if let Canonical::FixedSizeList(a) = self { a @@ -486,6 +521,10 @@ pub trait ToCanonical { #[deprecated(note = "use `array.execute::(ctx)` instead")] fn to_listview(&self) -> ListViewArray; + /// Canonicalize into a [`MapArray`] if the target is [`Map`](DType::Map) typed. + #[deprecated(note = "use `array.execute::(ctx)` instead")] + fn to_map(&self) -> MapArray; + /// Canonicalize into a [`FixedSizeListArray`] if the target is [`List`](DType::FixedSizeList) /// typed. #[deprecated(note = "use `array.execute::(ctx)` instead")] @@ -541,6 +580,12 @@ impl ToCanonical for ArrayRef { result.into_listview() } + fn to_map(&self) -> MapArray { + #[expect(deprecated)] + let result = self.to_canonical().vortex_expect("to_canonical failed"); + result.into_map() + } + fn to_fixed_size_list(&self) -> FixedSizeListArray { #[expect(deprecated)] let result = self.to_canonical().vortex_expect("to_canonical failed"); @@ -660,6 +705,18 @@ impl Executable for CanonicalValidity { .with_zero_copy_to_list(zctl) }))) } + Canonical::Map(map) => { + let map_dtype = map.map_dtype().clone(); + let entries = map.entries().into_owned(); + Ok(CanonicalValidity(Canonical::Map(MapArray::new( + map_dtype, + entries + .into_array() + .execute::(ctx)? + .0 + .into_listview(), + )))) + } Canonical::FixedSizeList(fsl) => { let list_size = fsl.list_size(); let len = fsl.len(); @@ -834,6 +891,18 @@ impl Executable for RecursiveCanonical { .with_zero_copy_to_list(zctl) }))) } + Canonical::Map(map) => { + let map_dtype = map.map_dtype().clone(); + let entries = map.entries().into_owned(); + Ok(RecursiveCanonical(Canonical::Map(MapArray::new( + map_dtype, + entries + .into_array() + .execute::(ctx)? + .0 + .into_listview(), + )))) + } Canonical::FixedSizeList(fsl) => { let list_size = fsl.list_size(); let len = fsl.len(); @@ -1041,6 +1110,18 @@ impl Executable for ListViewArray { } } +/// Execute the array to canonical form and unwrap as a [`MapArray`]. +/// +/// This will panic if the array's dtype is not map. +impl Executable for MapArray { + fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + match array.try_downcast::() { + Ok(map) => Ok(map), + Err(array) => Ok(Canonical::execute(array, ctx)?.into_map()), + } + } +} + /// Execute the array to canonical form and unwrap as a [`FixedSizeListArray`]. /// /// This will panic if the array's dtype is not fixed size list. @@ -1104,6 +1185,7 @@ pub enum CanonicalView<'a> { Decimal(ArrayView<'a, Decimal>), VarBinView(ArrayView<'a, VarBinView>), List(ArrayView<'a, ListView>), + Map(ArrayView<'a, Map>), FixedSizeList(ArrayView<'a, FixedSizeList>), Struct(ArrayView<'a, Struct>), Union(ArrayView<'a, Union>), @@ -1120,6 +1202,7 @@ impl From> for Canonical { CanonicalView::Decimal(a) => Canonical::Decimal(a.into_owned()), CanonicalView::VarBinView(a) => Canonical::VarBinView(a.into_owned()), CanonicalView::List(a) => Canonical::List(a.into_owned()), + CanonicalView::Map(a) => Canonical::Map(a.into_owned()), CanonicalView::FixedSizeList(a) => Canonical::FixedSizeList(a.into_owned()), CanonicalView::Struct(a) => Canonical::Struct(a.into_owned()), CanonicalView::Union(a) => Canonical::Union(a.into_owned()), @@ -1139,6 +1222,7 @@ impl CanonicalView<'_> { CanonicalView::Decimal(a) => a.array().clone(), CanonicalView::VarBinView(a) => a.array().clone(), CanonicalView::List(a) => a.array().clone(), + CanonicalView::Map(a) => a.array().clone(), CanonicalView::FixedSizeList(a) => a.array().clone(), CanonicalView::Struct(a) => a.array().clone(), CanonicalView::Union(a) => a.array().clone(), @@ -1162,6 +1246,7 @@ impl Matcher for AnyCanonical { || array.is::() || array.is::() || array.is::() + || array.is::() || array.is::() || array.is::() || array.is::() @@ -1184,6 +1269,8 @@ impl Matcher for AnyCanonical { Some(CanonicalView::Union(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::List(a)) + } else if let Some(a) = array.as_opt::() { + Some(CanonicalView::Map(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::FixedSizeList(a)) } else if let Some(a) = array.as_opt::() { diff --git a/vortex-array/src/dtype/dtype_impl.rs b/vortex-array/src/dtype/dtype_impl.rs index 6e3f2c263b4..3f3c70fcb96 100644 --- a/vortex-array/src/dtype/dtype_impl.rs +++ b/vortex-array/src/dtype/dtype_impl.rs @@ -643,16 +643,16 @@ mod tests { } #[test] - fn test_map_dtype_hash_ignores_value_nullability() -> VortexResult<()> { + fn test_map_dtype_hash_ignores_entry_nullability() -> VortexResult<()> { let lhs = DType::map( DType::Primitive(PType::I32, NonNullable), - DType::Utf8(Nullable), + DType::Utf8(NonNullable), true, NonNullable, )?; let rhs = DType::map( DType::Primitive(PType::I32, NonNullable), - DType::Utf8(NonNullable), + DType::Utf8(Nullable), true, Nullable, )?; diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 311b199cf76..046e19a7c9f 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -197,6 +197,7 @@ fn cast_canonical( CanonicalView::Decimal(a) => ::cast(a, dtype, ctx), CanonicalView::VarBinView(a) => ::cast(a, dtype, ctx), CanonicalView::List(a) => ::cast(a, dtype, ctx), + CanonicalView::Map(_) => vortex_bail!("Map arrays don't support casting"), CanonicalView::FixedSizeList(a) => ::cast(a, dtype, ctx), CanonicalView::Struct(a) => struct_cast(a, dtype, ctx), CanonicalView::Union(_) => { diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 988522e4f4a..2f3fbb9e4e7 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -24,6 +24,7 @@ use crate::arrays::Extension; use crate::arrays::FixedSizeList; use crate::arrays::List; use crate::arrays::ListView; +use crate::arrays::Map; use crate::arrays::Masked; use crate::arrays::Null; use crate::arrays::PiecewiseSequence; @@ -74,6 +75,7 @@ impl Default for ArraySession { this.register(Decimal); this.register(VarBinView); this.register(ListView); + this.register(Map); this.register(FixedSizeList); this.register(Struct); this.register(Union); diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index eb81e81241e..2df4a7e3ac5 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -30,6 +30,7 @@ use vortex_array::arrays::union::UnionArraySlotsExt; use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use super::CascadingCompressor; use super::constant; @@ -154,6 +155,7 @@ impl CascadingCompressor { self.compress_list_view_array(list_view_array, compress_ctx, exec_ctx) } } + Canonical::Map(_) => vortex_bail!("Map arrays are not yet supported by the compressor"), Canonical::FixedSizeList(fsl_array) => { let compressed_elems = self.compress(fsl_array.elements(), exec_ctx)?; diff --git a/vortex-duckdb/src/exporter/canonical.rs b/vortex-duckdb/src/exporter/canonical.rs index c6990811d6e..e7374298c26 100644 --- a/vortex-duckdb/src/exporter/canonical.rs +++ b/vortex-duckdb/src/exporter/canonical.rs @@ -30,6 +30,7 @@ pub(crate) fn new_exporter( Canonical::Decimal(array) => decimal::new_exporter(array, ctx), Canonical::VarBinView(array) => varbinview::new_exporter(array, ctx), Canonical::List(array) => list_view::new_exporter(array, cache, ctx), + Canonical::Map(_) => vortex_bail!("Map arrays can't be exported to DuckDB"), Canonical::FixedSizeList(array) => fixed_size_list::new_exporter(array, cache, ctx), Canonical::Struct(array) => struct_::new_exporter(array, cache, ctx), Canonical::Union(_) => { diff --git a/vortex-flatbuffers/src/generated/array.rs b/vortex-flatbuffers/src/generated/array.rs index 6d903a56aa5..d56c9f80c2d 100644 --- a/vortex-flatbuffers/src/generated/array.rs +++ b/vortex-flatbuffers/src/generated/array.rs @@ -1,7 +1,7 @@ // automatically generated by the FlatBuffers compiler, do not modify // @generated -extern crate alloc; +extern crate alloc; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] pub const ENUM_MIN_COMPRESSION: u8 = 0; diff --git a/vortex-flatbuffers/src/generated/dtype.rs b/vortex-flatbuffers/src/generated/dtype.rs index 44470d5ed33..1bf446ea99b 100644 --- a/vortex-flatbuffers/src/generated/dtype.rs +++ b/vortex-flatbuffers/src/generated/dtype.rs @@ -1,7 +1,7 @@ // automatically generated by the FlatBuffers compiler, do not modify // @generated -extern crate alloc; +extern crate alloc; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] pub const ENUM_MIN_PTYPE: u8 = 0; diff --git a/vortex-flatbuffers/src/generated/footer.rs b/vortex-flatbuffers/src/generated/footer.rs index 62ad85542e1..defa166f391 100644 --- a/vortex-flatbuffers/src/generated/footer.rs +++ b/vortex-flatbuffers/src/generated/footer.rs @@ -1,9 +1,9 @@ // automatically generated by the FlatBuffers compiler, do not modify // @generated -extern crate alloc; use crate::array::*; use crate::layout::*; +extern crate alloc; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] pub const ENUM_MIN_COMPRESSION_SCHEME: u8 = 0; diff --git a/vortex-flatbuffers/src/generated/layout.rs b/vortex-flatbuffers/src/generated/layout.rs index 0c7c7557b74..2c13c196f63 100644 --- a/vortex-flatbuffers/src/generated/layout.rs +++ b/vortex-flatbuffers/src/generated/layout.rs @@ -1,7 +1,7 @@ // automatically generated by the FlatBuffers compiler, do not modify // @generated -extern crate alloc; +extern crate alloc; pub enum LayoutOffset {} #[derive(Copy, Clone, PartialEq)] diff --git a/vortex-flatbuffers/src/generated/message.rs b/vortex-flatbuffers/src/generated/message.rs index b3e4be76ef7..2b49f7df47c 100644 --- a/vortex-flatbuffers/src/generated/message.rs +++ b/vortex-flatbuffers/src/generated/message.rs @@ -1,9 +1,9 @@ // automatically generated by the FlatBuffers compiler, do not modify // @generated -extern crate alloc; use crate::array::*; use crate::dtype::*; +extern crate alloc; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] pub const ENUM_MIN_MESSAGE_VERSION: u8 = 0;