diff --git a/fuzz/src/array/filter.rs b/fuzz/src/array/filter.rs index 746479aa46c..147b9c166ac 100644 --- a/fuzz/src/array/filter.rs +++ b/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") } } diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index 528d4576dfa..de8547c2306 100644 --- a/fuzz/src/array/mask.rs +++ b/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::(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"), }) } diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index 9bb6e12f66b..e513c5daf81 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -518,7 +518,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet { 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"), diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index b503cc86d98..86228a4e538 100644 --- a/fuzz/src/array/slice.rs +++ b/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") } } diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index 8e59bc085db..dacd13fbbb3 100644 --- a/fuzz/src/array/take.rs +++ b/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") } } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/map.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/map.rs new file mode 100644 index 00000000000..f1150b68213 --- /dev/null +++ b/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 { + if lhs.map_dtype() != rhs.map_dtype() { + return Ok(false); + } + + check_list_identical( + &lhs.entries().as_::().into_owned(), + &rhs.entries().as_::().into_owned(), + ctx, + ) +} diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index f030cc810e2..a8392f3fa74 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/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) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/map.rs b/vortex-array/src/aggregate_fn/fns/is_constant/map.rs new file mode 100644 index 00000000000..63d82b463d9 --- /dev/null +++ b/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 { + check_listview_constant(&map.entries().as_::().into_owned(), ctx) +} 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 ac7f7cb9ce3..ee0acd4462b 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/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>>; + + fn map_array_from_rows(rows: &[MapRowFixture<'_>]) -> VortexResult { + 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::::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::>(); + 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(()) + } } 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 49f68d4fa7f..ee9426c7919 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,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_::().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), diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs index 97b75f255f0..e77d196ba94 100644 --- a/vortex-array/src/arrays/arbitrary.rs +++ b/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, +) -> Result { + 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::::with_capacity(map_dtype, nullability, array_length); + + for _ in 0..array_length { + if nullability == Nullability::Nullable && u.arbitrary::()? { + 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::>>()?; + 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. diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index 6462ed38d16..a4afb7edf20 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -5,7 +5,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use crate::ArrayView; use crate::Canonical; @@ -21,6 +20,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; @@ -50,7 +51,7 @@ pub(crate) fn take_canonical( CanonicalView::Decimal(a) => Canonical::Decimal(take_decimal(a, codes, ctx)), CanonicalView::VarBinView(a) => Canonical::VarBinView(take_varbinview(a, codes, ctx)), CanonicalView::List(a) => Canonical::List(take_listview(a, codes, ctx)), - CanonicalView::Map(_) => vortex_bail!("Map arrays don't support take"), + CanonicalView::Map(a) => Canonical::Map(take_map(a, codes, ctx)), CanonicalView::FixedSizeList(a) => { Canonical::FixedSizeList(take_fixed_size_list(a, codes, ctx)) } @@ -143,6 +144,18 @@ fn take_listview( .into_owned() } +fn take_map( + array: ArrayView<'_, Map>, + codes: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> MapArray { + ::take(array, codes.array(), ctx) + .vortex_expect("take map execute") + .vortex_expect("Map TakeExecute should not return None") + .as_::() + .into_owned() +} + fn take_fixed_size_list( array: ArrayView<'_, FixedSizeList>, codes: ArrayView<'_, Primitive>, diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 3be654d213d..ca643d3bfab 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -9,7 +9,6 @@ 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; @@ -21,10 +20,13 @@ use crate::array::ArrayView; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::Filter; +use crate::arrays::Map; +use crate::arrays::MapArray; use crate::arrays::NullArray; use crate::arrays::VariantArray; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterArraySlotsExt; +use crate::arrays::filter::FilterReduce; use crate::arrays::variant::VariantArraySlotsExt; use crate::scalar::Scalar; use crate::validity::Validity; @@ -96,7 +98,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::Map(a) => Canonical::Map(filter_map(&a, mask)), Canonical::FixedSizeList(a) => { Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask)) } @@ -127,3 +129,11 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca } } } + +fn filter_map(array: &MapArray, mask: &Arc) -> MapArray { + let filter_mask = Mask::Values(Arc::clone(mask)); + let filtered = ::filter(array.as_view(), &filter_mask) + .vortex_expect("MapArray somehow could not be filtered") + .vortex_expect("Map filter reduce always produces an array"); + filtered.as_::().into_owned() +} diff --git a/vortex-array/src/arrays/map/array.rs b/vortex-array/src/arrays/map/array.rs index 21c69a4a0f9..1cdc362b606 100644 --- a/vortex-array/src/arrays/map/array.rs +++ b/vortex-array/src/arrays/map/array.rs @@ -6,7 +6,6 @@ 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; @@ -14,12 +13,13 @@ 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(&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 { - /// 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_::() - } - +pub trait MapArrayExt: MapArraySlotsExt { /// Returns the entry structs for one map row. fn entries_at(&self, index: usize) -> VortexResult { - self.entries().list_elements_at(index) + self.entries().as_::().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_::().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_validity() } /// Returns this map's key/value type information. @@ -125,8 +125,8 @@ impl Array { 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 { 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 { .as_map_opt() .vortex_expect("MapArray requires a map dtype") .clone(); - let entries = self.entries().into_owned(); + let entries = self.entries().clone().downcast::(); MapDataParts { map_dtype, entries } } } diff --git a/vortex-array/src/arrays/map/compute/cast.rs b/vortex-array/src/arrays/map/compute/cast.rs new file mode 100644 index 00000000000..bfbab6e17e9 --- /dev/null +++ b/vortex-array/src/arrays/map/compute/cast.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::ArrayRef; +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; +use crate::executor::ExecutionCtx; +use crate::scalar_fn::fns::cast::CastKernel; +use crate::scalar_fn::fns::cast::CastReduce; + +fn prepare_map_cast_target( + array: ArrayView<'_, Map>, + dtype: &DType, +) -> VortexResult> { + let Some(target_map_dtype) = dtype.as_map_opt() else { + return Ok(None); + }; + + if target_map_dtype.keys_sorted() && !array.keys_sorted() { + vortex_bail!( + "Cannot cast {} to {dtype}: source does not assert sorted map keys", + array.dtype() + ); + } + + let target_entries_dtype = DType::List( + Arc::new(target_map_dtype.entries_dtype()), + dtype.nullability(), + ); + Ok(Some((target_map_dtype.clone(), target_entries_dtype))) +} + +impl CastReduce for Map { + fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { + let Some((target_map_dtype, target_entries_dtype)) = prepare_map_cast_target(array, dtype)? + else { + return Ok(None); + }; + + let Some(entries) = ::cast( + array.entries().as_::(), + &target_entries_dtype, + )? + else { + return Ok(None); + }; + + rebuild_map_from_array(target_map_dtype, entries).map(Some) + } +} + +impl CastKernel for Map { + fn cast( + array: ArrayView<'_, Self>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some((target_map_dtype, target_entries_dtype)) = prepare_map_cast_target(array, dtype)? + else { + return Ok(None); + }; + + let Some(entries) = ::cast( + array.entries().as_::(), + &target_entries_dtype, + ctx, + )? + else { + return Ok(None); + }; + + rebuild_map_from_array(target_map_dtype, entries).map(Some) + } +} diff --git a/vortex-array/src/arrays/map/compute/filter.rs b/vortex-array/src/arrays/map/compute/filter.rs new file mode 100644 index 00000000000..53a862bb35a --- /dev/null +++ b/vortex-array/src/arrays/map/compute/filter.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::MapArraySlotsExt; +use crate::executor::ExecutionCtx; + +impl FilterReduce for Map { + fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { + let entries = array.entries().as_::(); + + // 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 + // not carried over: dropping a non-empty row leaves a gap in the referenced elements. + let filtered_entries = unsafe { + ListViewArray::new_unchecked( + entries.elements().clone(), + entries.offsets().filter(mask.clone())?, + entries.sizes().filter(mask.clone())?, + entries.validity()?.filter(mask)?, + ) + }; + + { + let map_dtype = array.map_dtype().clone(); + MapArray::try_new(map_dtype, filtered_entries).map(IntoArray::into_array) + } + .map(Some) + } +} + +impl FilterKernel for Map { + fn filter( + array: ArrayView<'_, Self>, + mask: &Mask, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + ::filter(array, mask) + } +} diff --git a/vortex-array/src/arrays/map/compute/mask.rs b/vortex-array/src/arrays/map/compute/mask.rs new file mode 100644 index 00000000000..41703a2a3ea --- /dev/null +++ b/vortex-array/src/arrays/map/compute/mask.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +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> { + let Some(entries) = + ::mask(array.entries().as_::(), mask)? + else { + return Ok(None); + }; + + rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some) + } +} + +impl MaskKernel for Map { + fn mask( + array: ArrayView<'_, Self>, + mask: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + ::mask(array, mask) + } +} diff --git a/vortex-array/src/arrays/map/compute/mod.rs b/vortex-array/src/arrays/map/compute/mod.rs new file mode 100644 index 00000000000..ea0be66f0e4 --- /dev/null +++ b/vortex-array/src/arrays/map/compute/mod.rs @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Map array compute-kernel namespace. + +mod cast; +mod filter; +mod mask; +pub(crate) mod rules; +mod slice; +mod take; + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::ListView; +use crate::arrays::MapArray; +use crate::dtype::MapDType; + +fn rebuild_map_from_array(map_dtype: MapDType, entries: ArrayRef) -> VortexResult { + let map_entries = entries.try_downcast::().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) +} diff --git a/vortex-array/src/arrays/map/compute/rules.rs b/vortex-array/src/arrays/map/compute/rules.rs new file mode 100644 index 00000000000..27622ff1c5f --- /dev/null +++ b/vortex-array/src/arrays/map/compute/rules.rs @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::arrays::Map; +use crate::arrays::dict::TakeReduceAdaptor; +use crate::arrays::filter::FilterReduceAdaptor; +use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleSet; +use crate::scalar_fn::fns::cast::CastReduceAdaptor; +use crate::scalar_fn::fns::mask::MaskReduceAdaptor; + +pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&FilterReduceAdaptor(Map)), + ParentRuleSet::lift(&CastReduceAdaptor(Map)), + ParentRuleSet::lift(&MaskReduceAdaptor(Map)), + ParentRuleSet::lift(&SliceReduceAdaptor(Map)), + ParentRuleSet::lift(&TakeReduceAdaptor(Map)), +]); diff --git a/vortex-array/src/arrays/map/compute/slice.rs b/vortex-array/src/arrays/map/compute/slice.rs new file mode 100644 index 00000000000..463dee74f60 --- /dev/null +++ b/vortex-array/src/arrays/map/compute/slice.rs @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +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) -> VortexResult> { + let Some(sliced_entries) = + ::slice(array.entries().as_::(), range)? + else { + return Ok(None); + }; + + rebuild_map_from_array(array.map_dtype().clone(), sliced_entries).map(Some) + } +} + +impl SliceKernel for Map { + fn slice( + array: ArrayView<'_, Self>, + range: Range, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + ::slice(array, range) + } +} diff --git a/vortex-array/src/arrays/map/compute/take.rs b/vortex-array/src/arrays/map/compute/take.rs new file mode 100644 index 00000000000..58be62e9645 --- /dev/null +++ b/vortex-array/src/arrays/map/compute/take.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::ArrayView; +use crate::arrays::ListView; +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> { + let Some(entries) = + ::take(array.entries().as_::(), indices)? + else { + return Ok(None); + }; + + rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some) + } +} + +impl TakeExecute for Map { + fn take( + array: ArrayView<'_, Self>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(entries) = + ::take(array.entries().as_::(), indices, ctx)? + else { + return Ok(None); + }; + + rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some) + } +} diff --git a/vortex-array/src/arrays/map/mod.rs b/vortex-array/src/arrays/map/mod.rs index 27eb138a1ed..ef1a04166bc 100644 --- a/vortex-array/src/arrays/map/mod.rs +++ b/vortex-array/src/arrays/map/mod.rs @@ -5,12 +5,21 @@ mod array; pub use array::MapArrayExt; +pub use array::MapArraySlotsExt; pub use array::MapData; pub use array::MapDataParts; +pub use array::MapSlots; +pub use array::MapSlotsView; + +pub(crate) mod compute; mod vtable; pub use vtable::Map; pub use vtable::MapArray; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/map/tests.rs b/vortex-array/src/arrays/map/tests.rs index 833e60f0194..2b5570a8d16 100644 --- a/vortex-array/src/arrays/map/tests.rs +++ b/vortex-array/src/arrays/map/tests.rs @@ -1,9 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use smallvec::smallvec; use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_mask::Mask; use vortex_session::registry::ReadContext; use crate::Array; @@ -14,17 +15,24 @@ use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; +use crate::arrays::FilterArray; +use crate::arrays::ListView; 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::MapArraySlotsExt; use crate::arrays::map::MapData; use crate::arrays::map::MapDataParts; +use crate::arrays::map::MapSlots; +use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::MapBuilder; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::MapDType; use crate::dtype::Nullability; @@ -65,14 +73,41 @@ fn sample_scalar(dtype: DType) -> VortexResult { ) } +fn map_array_from_rows( + map_dtype: MapDType, + nullability: Nullability, + rows: impl IntoIterator>>, +) -> VortexResult { + let rows = rows.into_iter().collect::>(); + let dtype = DType::Map(map_dtype.clone(), nullability); + let mut builder = MapBuilder::::with_capacity(map_dtype, nullability, rows.len()); + + for row in rows { + let scalar = match row { + Some(entries) => Scalar::try_map(dtype.clone(), entries)?, + None => Scalar::null(dtype.clone()), + }; + builder.append_scalar(&scalar)?; + } + + Ok(builder.finish_into_map()) +} + 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()) + map_array_from_rows( + map_dtype, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + Some(vec![]), + None, + ], + ) } #[test] @@ -144,7 +179,12 @@ fn rejects_malformed_entry_storage() -> VortexResult<()> { entries.len(), MapData, ) - .with_slots(smallvec![Some(entries.into_array())]); + .with_slots( + MapSlots { + entries: entries.into_array(), + } + .into_slots(), + ); assert!(Array::::try_from_parts(parts).is_err()); assert!( @@ -241,6 +281,366 @@ fn scalar_access_preserves_variable_entry_counts_and_utf8_pairs() -> VortexResul Ok(()) } +#[test] +fn slice_preserves_map_rows() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let expected = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + Some(vec![]), + ], + )?; + let sliced = source.slice(0..2)?; + let mut ctx = array_session().create_execution_ctx(); + + assert!(sliced.is::()); + assert_arrays_eq!(sliced, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn filter_handles_all_none_and_mixed_maps() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let map_dtype = map_dtype()?; + let mut ctx = array_session().create_execution_ctx(); + + let all = source.filter(Mask::from_iter([true, true, true]))?; + assert!(all.is::()); + assert_arrays_eq!(all, sample_array()?, &mut ctx); + + let none = source.filter(Mask::from_iter([false, false, false]))?; + let expected_none = map_array_from_rows( + map_dtype.clone(), + Nullability::Nullable, + Vec::>>::new(), + )?; + assert!(none.is::()); + assert_arrays_eq!(none, expected_none, &mut ctx); + + let mixed = source.filter(Mask::from_iter([true, false, true]))?; + let expected_mixed = map_array_from_rows( + map_dtype, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + None, + ], + )?; + assert!(mixed.is::()); + assert_arrays_eq!(mixed, expected_mixed, &mut ctx); + + Ok(()) +} + +#[test] +fn filter_dropping_nonempty_middle_row_clears_zero_copy_flag() -> VortexResult<()> { + let source = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [ + Some(vec![ + (key(1), value(Some("one"))), + (key(2), value(Some("two"))), + ]), + Some(vec![(key(3), value(Some("three")))]), + Some(vec![(key(4), value(Some("four")))]), + ], + )? + .into_array(); + let expected = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [ + Some(vec![ + (key(1), value(Some("one"))), + (key(2), value(Some("two"))), + ]), + Some(vec![(key(4), value(Some("four")))]), + ], + )?; + let mask = Mask::from_iter([true, false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let reduced = source.filter(mask.clone())?; + assert!(reduced.is::()); + // The dropped middle row leaves a gap in the referenced entry elements, so the filtered + // entries must not claim zero-copy-to-list convertibility. + assert!( + !reduced + .as_::() + .entries() + .as_::() + .into_owned() + .is_zero_copy_to_list() + ); + assert_arrays_eq!(reduced, expected, &mut ctx); + + let executed = FilterArray::new(source, mask) + .into_array() + .execute::(&mut ctx)?; + assert!( + !executed + .entries() + .as_::() + .into_owned() + .is_zero_copy_to_list() + ); + assert_arrays_eq!(executed, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn take_supports_reordered_duplicate_and_nullable_indices() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let map_dtype = map_dtype()?; + let mut ctx = array_session().create_execution_ctx(); + + let taken = source.take(PrimitiveArray::from_iter([2u64, 0, 0]).into_array())?; + let expected_taken = map_array_from_rows( + map_dtype.clone(), + Nullability::Nullable, + [ + None, + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + ], + )?; + assert!(taken.is::()); + assert_arrays_eq!(taken, expected_taken, &mut ctx); + + let nonnullable_source = map_array_from_rows( + map_dtype.clone(), + Nullability::NonNullable, + [ + Some(vec![(key(1), value(Some("one")))]), + Some(vec![]), + Some(vec![(key(3), value(None))]), + ], + )? + .into_array(); + let nullable_taken = nonnullable_source + .take(PrimitiveArray::from_option_iter([Some(1u64), None, Some(0)]).into_array())?; + let expected_nullable_taken = map_array_from_rows( + map_dtype, + Nullability::Nullable, + [Some(vec![]), None, Some(vec![(key(1), value(Some("one")))])], + )?; + assert!(nullable_taken.is::()); + assert_eq!(nullable_taken.dtype().nullability(), Nullability::Nullable); + assert_arrays_eq!(nullable_taken, expected_nullable_taken, &mut ctx); + + Ok(()) +} + +#[test] +fn mask_combines_with_existing_map_validity() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let mask = BoolArray::from_iter([true, false, true]).into_array(); + let masked = source.mask(mask)?; + let expected = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + None, + None, + ], + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert!(masked.is::()); + assert_arrays_eq!(masked, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn cast_widens_map_key_value_and_outer_nullability() -> VortexResult<()> { + let source_map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + true, + )?; + let source = map_array_from_rows( + source_map_dtype, + Nullability::NonNullable, + [ + Some(vec![ + ( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::NonNullable), + ), + ( + Scalar::primitive(2i32, Nullability::NonNullable), + Scalar::utf8("two", Nullability::NonNullable), + ), + ]), + Some(vec![]), + ], + )? + .into_array(); + let target_dtype = DType::map( + DType::Primitive(PType::I64, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + let target_map_dtype = target_dtype + .as_map_opt() + .vortex_expect("target dtype is map") + .clone(); + let cast = source.cast(target_dtype)?; + let expected = map_array_from_rows( + target_map_dtype, + Nullability::Nullable, + [ + Some(vec![ + ( + Scalar::primitive(1i64, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + ), + ( + Scalar::primitive(2i64, Nullability::NonNullable), + Scalar::utf8("two", Nullability::Nullable), + ), + ]), + Some(vec![]), + ], + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert!(cast.is::()); + assert_arrays_eq!(cast, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn cast_can_drop_but_not_create_sortedness_assertion() -> VortexResult<()> { + let sorted_source = sample_array()?.into_array(); + let unsorted_target = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + )?; + let cast = sorted_source.cast(unsorted_target)?; + let expected = map_array_from_rows( + MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + )?, + Nullability::Nullable, + [ + Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ]), + Some(vec![]), + None, + ], + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(cast, expected, &mut ctx); + + let unsorted_map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + )?; + let unsorted_source = map_array_from_rows( + unsorted_map_dtype, + Nullability::Nullable, + [Some(vec![(key(1), value(Some("one")))])], + )? + .into_array(); + let sorted_target = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + assert!(unsorted_source.cast(sorted_target).is_err()); + + Ok(()) +} + +#[test] +fn null_map_cast_cannot_create_sortedness_assertion() -> VortexResult<()> { + let unsorted_map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + )?; + let unsorted_dtype = DType::Map(unsorted_map_dtype.clone(), Nullability::Nullable); + let sorted_dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + let scalar = Scalar::null(unsorted_dtype); + + assert!(scalar.cast(&sorted_dtype).is_err()); + let constant_cast = ConstantArray::new(scalar, 2) + .into_array() + .cast(sorted_dtype.clone())?; + let mut ctx = array_session().create_execution_ctx(); + assert!(constant_cast.execute::(&mut ctx).is_err()); + + let all_null = + map_array_from_rows(unsorted_map_dtype, Nullability::Nullable, [None, None])?.into_array(); + assert!(all_null.cast(sorted_dtype).is_err()); + + Ok(()) +} + +#[test] +fn filter_preserves_duplicate_map_keys() -> VortexResult<()> { + let source = sample_array()?.into_array(); + let filtered = source.filter(Mask::from_iter([true, false, false]))?; + let expected = map_array_from_rows( + map_dtype()?, + Nullability::Nullable, + [Some(vec![ + (key(2), value(Some("two"))), + (key(1), value(None)), + (key(2), value(Some("duplicate"))), + ])], + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert_arrays_eq!(filtered, expected, &mut ctx); + + Ok(()) +} + #[test] fn builder_appends_existing_map_arrays() -> VortexResult<()> { let source = sample_array()?; diff --git a/vortex-array/src/arrays/map/vtable/kernel.rs b/vortex-array/src/arrays/map/vtable/kernel.rs new file mode 100644 index 00000000000..860496834d6 --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/kernel.rs @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Dict; +use crate::arrays::Filter; +use crate::arrays::Map; +use crate::arrays::Slice; +use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::filter::FilterExecuteAdaptor; +use crate::arrays::slice::SliceExecuteAdaptor; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::mask::Mask; +use crate::scalar_fn::fns::mask::MaskExecuteAdaptor; + +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), Map, CastExecuteAdaptor(Map)); + kernels.register_execute_parent_kernel(Dict.id(), Map, TakeExecuteAdaptor(Map)); + kernels.register_execute_parent_kernel(Filter.id(), Map, FilterExecuteAdaptor(Map)); + kernels.register_execute_parent_kernel(Mask.id(), Map, MaskExecuteAdaptor(Map)); + kernels.register_execute_parent_kernel(Slice.id(), Map, SliceExecuteAdaptor(Map)); +} diff --git a/vortex-array/src/arrays/map/vtable/mod.rs b/vortex-array/src/arrays/map/vtable/mod.rs index 4c8b19e3d0b..f11e3e72a89 100644 --- a/vortex-array/src/arrays/map/vtable/mod.rs +++ b/vortex-array/src/arrays/map/vtable/mod.rs @@ -1,7 +1,6 @@ // 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; @@ -21,22 +20,27 @@ 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::MapSlots; +use crate::arrays::map::MapSlotsView; use crate::arrays::map::array::validate_entries; +use crate::arrays::map::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::match_each_map_builder; use crate::serde::ArrayChildren; +mod kernel; mod operations; mod validity; /// A [`Map`]-encoded Vortex array. pub type MapArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + /// The canonical encoding for [`DType::Map`]. /// /// A map array has one `ListView>` child. Its outer dtype retains map-specific @@ -63,18 +67,17 @@ impl VTable for Map { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "MapArray expected {NUM_SLOTS} slot, found {}", + slots.len() == MapSlots::COUNT, + "MapArray expected {} slot, found {}", + MapSlots::COUNT, 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) + let slots = MapSlotsView::from_slots(slots); + validate_entries(map_dtype, *nullability, len, slots.entries) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -125,26 +128,27 @@ impl VTable for Map { vortex_bail!("Expected map dtype, got {dtype}"); }; vortex_ensure!( - children.len() == NUM_SLOTS, - "MapArray expected {NUM_SLOTS} child, found {}", + children.len() == MapSlots::COUNT, + "MapArray expected {} child, found {}", + MapSlots::COUNT, 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)?; + let entries = children.get(MapSlots::ENTRIES, &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)])) + let slots = MapData::make_slots(entries); + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, MapData).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + MapSlots::NAMES[idx].to_string() } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { @@ -165,4 +169,12 @@ impl VTable for Map { ), } } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } } diff --git a/vortex-array/src/arrays/map/vtable/validity.rs b/vortex-array/src/arrays/map/vtable/validity.rs index 0cdc1f886cb..e81c7602077 100644 --- a/vortex-array/src/arrays/map/vtable/validity.rs +++ b/vortex-array/src/arrays/map/vtable/validity.rs @@ -5,10 +5,10 @@ use crate::ArrayRef; use crate::array::ArrayView; use crate::array::ValidityChild; use crate::arrays::Map; -use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; impl ValidityChild for Map { fn validity_child(array: ArrayView<'_, Map>) -> ArrayRef { - array.entries().array().clone() + array.entries().clone() } } diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index b9d25946710..eebcf1a2335 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use crate::Canonical; use crate::IntoArray; @@ -14,7 +13,9 @@ use crate::arrays::BoolArray; use crate::arrays::DecimalArray; use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::ListView; use crate::arrays::ListViewArray; +use crate::arrays::MapArray; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -26,6 +27,8 @@ use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; use crate::arrays::struct_::StructArrayExt; use crate::arrays::union::UnionArrayExt; use crate::arrays::union::UnionArraySlotsExt; @@ -51,7 +54,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::Map(a) => Canonical::Map(mask_validity_map(a, validity)?), Canonical::FixedSizeList(a) => { Canonical::FixedSizeList(mask_validity_fixed_size_list(a, validity)?) } @@ -129,6 +132,11 @@ fn mask_validity_listview(array: ListViewArray, validity: Validity) -> VortexRes }) } +fn mask_validity_map(array: MapArray, validity: Validity) -> VortexResult { + let entries = mask_validity_listview(array.entries().as_::().into_owned(), validity)?; + MapArray::try_new(array.map_dtype().clone(), entries) +} + fn mask_validity_fixed_size_list( array: FixedSizeListArray, validity: Validity, diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 8e819094b6f..84729fd551d 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -143,6 +143,7 @@ pub(crate) fn initialize(session: &VortexSession) { fixed_size_list::initialize(session); list::initialize(session); listview::initialize(session); + map::initialize(session); patched::initialize(session); primitive::initialize(session); struct_::initialize(session); diff --git a/vortex-array/src/builders/map.rs b/vortex-array/src/builders/map.rs index 25083488c7e..362c2211a30 100644 --- a/vortex-array/src/builders/map.rs +++ b/vortex-array/src/builders/map.rs @@ -13,9 +13,10 @@ use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::ListView; use crate::arrays::Map; use crate::arrays::MapArray; -use crate::arrays::map::MapArrayExt; +use crate::arrays::map::MapArraySlotsExt; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::ListViewBuilder; @@ -103,7 +104,7 @@ impl MapBuilder { array.dtype() ); self.entries_builder - .append_listview_array(array.entries(), ctx) + .append_listview_array(array.entries().as_::(), ctx) } } diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 1d8487c4eed..5330aa2f1ed 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -50,6 +50,7 @@ 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::map::MapArraySlotsExt; use crate::arrays::primitive::PrimitiveDataParts; use crate::arrays::struct_::StructDataParts; use crate::arrays::union::UnionDataParts; @@ -293,6 +294,7 @@ impl Canonical { array.map_dtype().clone(), array .entries() + .as_::() .into_owned() .rebuild(ListViewRebuildMode::TrimElements, ctx)?, ))), @@ -707,14 +709,10 @@ impl Executable for CanonicalValidity { } Canonical::Map(map) => { let map_dtype = map.map_dtype().clone(); - let entries = map.entries().into_owned(); + let entries = map.entries().clone(); Ok(CanonicalValidity(Canonical::Map(MapArray::new( map_dtype, - entries - .into_array() - .execute::(ctx)? - .0 - .into_listview(), + entries.execute::(ctx)?.0.into_listview(), )))) } Canonical::FixedSizeList(fsl) => { @@ -893,11 +891,10 @@ impl Executable for RecursiveCanonical { } Canonical::Map(map) => { let map_dtype = map.map_dtype().clone(); - let entries = map.entries().into_owned(); + let entries = map.entries().clone(); Ok(RecursiveCanonical(Canonical::Map(MapArray::new( map_dtype, entries - .into_array() .execute::(ctx)? .0 .into_listview(), diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 18cda1c425e..cc1d826bdcf 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -1092,6 +1092,26 @@ fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { } } +fn widened_primitive_dtype(dtype: &DType) -> Option { + let DType::Primitive(ptype, nullability) = dtype else { + return None; + }; + let widened = match ptype { + PType::U8 => PType::U16, + PType::U16 => PType::U32, + PType::U32 => PType::U64, + PType::U64 => return None, + PType::I8 => PType::I16, + PType::I16 => PType::I32, + PType::I32 => PType::I64, + PType::I64 => return None, + PType::F16 => PType::F32, + PType::F32 => PType::F64, + PType::F64 => return None, + }; + Some(DType::Primitive(widened, *nullability)) +} + /// Tests that cast operations preserve array properties when sliced. /// /// # Invariant @@ -1227,7 +1247,45 @@ fn test_cast_slice_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { opposite, )] } - DType::Map(..) => vec![], /* Map arrays are not materializable until their layout is chosen. */ + DType::Map(map_dtype, nullability) => { + let opposite = match nullability { + Nullability::NonNullable => Nullability::Nullable, + Nullability::Nullable => Nullability::NonNullable, + }; + let key_dtype = map_dtype.key_dtype(); + let value_dtype = map_dtype.value_dtype(); + let mut targets = vec![DType::Map(map_dtype.clone(), opposite)]; + + if let Some(widened_key) = widened_primitive_dtype(&key_dtype) + && let Ok(dtype) = DType::map( + widened_key, + value_dtype.clone(), + map_dtype.keys_sorted(), + *nullability, + ) + { + targets.push(dtype); + } + + if let Some(widened_value) = widened_primitive_dtype(&value_dtype) + && let Ok(dtype) = DType::map( + key_dtype.clone(), + widened_value, + map_dtype.keys_sorted(), + *nullability, + ) + { + targets.push(dtype); + } + + if map_dtype.keys_sorted() + && let Ok(dtype) = DType::map(key_dtype, value_dtype, false, *nullability) + { + targets.push(dtype); + } + + targets + } DType::Struct(fields, nullability) => { let opposite = match nullability { Nullability::NonNullable => Nullability::Nullable, diff --git a/vortex-array/src/dtype/arbitrary/mod.rs b/vortex-array/src/dtype/arbitrary/mod.rs index 9b95074ff6f..ddfe1712ec3 100644 --- a/vortex-array/src/dtype/arbitrary/mod.rs +++ b/vortex-array/src/dtype/arbitrary/mod.rs @@ -36,7 +36,7 @@ impl<'a> Arbitrary<'a> for FieldName { fn random_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result { const BASE_TYPE_COUNT: i32 = 5; - const CONTAINER_TYPE_COUNT: i32 = 3; + const CONTAINER_TYPE_COUNT: i32 = 4; let max_dtype_kind = if depth == 0 { BASE_TYPE_COUNT } else { @@ -59,12 +59,23 @@ fn random_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result { u.choose_index(3)?.try_into().vortex_expect("impossible"), u.arbitrary()?, ), + 9 => random_map_dtype(u, depth - 1)?, // Null, // Extension(ExtDType, Nullability), _ => unreachable!("Number out of range"), }) } +fn random_map_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result { + Ok(DType::map( + random_dtype(u, depth)?.as_nonnullable(), + random_dtype(u, depth)?, + false, + u.arbitrary()?, + ) + .vortex_expect("non-nullable generated map keys are always valid")) +} + impl<'a> Arbitrary<'a> for Nullability { fn arbitrary(u: &mut Unstructured<'a>) -> Result { Ok(if u.arbitrary()? { @@ -123,3 +134,23 @@ fn random_struct_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result>>()?; Ok(StructFields::new(names, dtypes)) } + +#[cfg(test)] +mod tests { + use arbitrary::Unstructured; + + use super::random_map_dtype; + use crate::dtype::DType; + + #[test] + fn random_map_dtype_never_asserts_sorted_keys() { + for byte in 0..=u8::MAX { + let bytes = [byte; 128]; + let mut u = Unstructured::new(&bytes); + let Ok(DType::Map(map_dtype, _)) = random_map_dtype(&mut u, 0) else { + continue; + }; + assert!(!map_dtype.keys_sorted()); + } + } +} diff --git a/vortex-array/src/dtype/coercion.rs b/vortex-array/src/dtype/coercion.rs index 9f1ccd9861e..5021c1ae013 100644 --- a/vortex-array/src/dtype/coercion.rs +++ b/vortex-array/src/dtype/coercion.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use crate::dtype::DType; use crate::dtype::PType; +use crate::dtype::StructFields; use crate::dtype::decimal::DecimalDType; impl PType { @@ -101,19 +102,32 @@ impl DType { } if let (DType::Map(lhs, _), DType::Map(rhs, _)) = (self, other) { - if lhs.key_dtype() != rhs.key_dtype() || lhs.value_dtype() != rhs.value_dtype() { - return None; - } - + let key_dtype = lhs.key_dtype().least_supertype(&rhs.key_dtype())?; + let value_dtype = lhs.value_dtype().least_supertype(&rhs.value_dtype())?; return DType::map( - lhs.key_dtype(), - lhs.value_dtype(), + key_dtype, + value_dtype, lhs.keys_sorted() && rhs.keys_sorted(), union_null, ) .ok(); } + if let (DType::Struct(lhs, _), DType::Struct(rhs, _)) = (self, other) { + if lhs.nfields() != rhs.nfields() || lhs.names() != rhs.names() { + return None; + } + let fields = lhs + .fields() + .zip(rhs.fields()) + .map(|(lhs, rhs)| lhs.least_supertype(&rhs)) + .collect::>>()?; + return Some(DType::Struct( + StructFields::new(lhs.names().clone(), fields), + union_null, + )); + } + // Identity (ignoring nullability): return self with union nullability if self.eq_ignore_nullability(other) { return Some(self.with_nullability(union_null)); @@ -204,8 +218,18 @@ impl DType { if let (DType::Map(target, _), DType::Map(source, _)) = (self, other) { return (self.is_nullable() || !other.is_nullable()) && (!target.keys_sorted() || source.keys_sorted()) - && target.key_dtype() == source.key_dtype() - && target.value_dtype() == source.value_dtype(); + && target.key_dtype().can_coerce_from(&source.key_dtype()) + && target.value_dtype().can_coerce_from(&source.value_dtype()); + } + + if let (DType::Struct(target, _), DType::Struct(source, _)) = (self, other) { + return (self.is_nullable() || !other.is_nullable()) + && target.nfields() == source.nfields() + && target.names() == source.names() + && target + .fields() + .zip(source.fields()) + .all(|(target, source)| target.can_coerce_from(&source)); } // Same type (ignoring nullability): check nullability compatibility @@ -332,7 +356,9 @@ mod tests { use std::sync::Arc; use crate::dtype::DType; + use crate::dtype::FieldNames; use crate::dtype::PType; + use crate::dtype::StructFields; use crate::dtype::UnionVariants; use crate::dtype::decimal::DecimalDType; use crate::dtype::nullability::Nullability::NonNullable; @@ -827,48 +853,195 @@ mod tests { } #[test] - fn map_least_supertype_unions_outer_nullability_and_intersects_sortedness() { + fn map_least_supertype_recursively_widens_key_and_value_dtypes() { let key = DType::Primitive(PType::I32, NonNullable); - let value = DType::Utf8(Nullable); - let sorted = DType::map(key.clone(), value.clone(), true, NonNullable).unwrap(); - let unsorted = DType::map(key.clone(), value.clone(), false, Nullable).unwrap(); + let wide_key = DType::Primitive(PType::I64, NonNullable); + let value = DType::Utf8(NonNullable); + let nullable_value = DType::Utf8(Nullable); + let sorted = DType::map(key, value, true, NonNullable).unwrap(); + let unsorted = + DType::map(wide_key.clone(), nullable_value.clone(), false, Nullable).unwrap(); assert_eq!( sorted.least_supertype(&unsorted), - Some(DType::map(key, value, false, Nullable).unwrap()) + Some(DType::map(wide_key, nullable_value, false, Nullable).unwrap()) ); } #[test] - fn map_least_supertype_requires_identical_key_and_value_dtypes() { - let i32_map = DType::map( + fn map_least_supertype_widens_child_nullability() { + let key = DType::Primitive(PType::I32, NonNullable); + let lhs = DType::map(key.clone(), DType::Utf8(NonNullable), true, NonNullable).unwrap(); + let rhs = DType::map(key.clone(), DType::Utf8(Nullable), true, NonNullable).unwrap(); + let expected = DType::map(key, DType::Utf8(Nullable), true, NonNullable).unwrap(); + + assert_eq!(lhs.least_supertype(&rhs), Some(expected)); + } + + #[test] + fn map_least_supertype_rejects_incompatible_child_dtypes() { + let primitive_key_map = DType::map( DType::Primitive(PType::I32, NonNullable), DType::Utf8(Nullable), false, NonNullable, ) .unwrap(); - let i64_map = DType::map( + let utf8_key_map = DType::map( + DType::Utf8(NonNullable), + DType::Utf8(Nullable), + false, + NonNullable, + ) + .unwrap(); + + assert_eq!(primitive_key_map.least_supertype(&utf8_key_map), None); + } + + #[test] + fn map_coercion_recursively_widens_children_and_preserves_sortedness_rules() { + let sorted_source = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(NonNullable), + true, + NonNullable, + ) + .unwrap(); + let unsorted_source = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(NonNullable), + false, + NonNullable, + ) + .unwrap(); + let unsorted_target = DType::map( DType::Primitive(PType::I64, NonNullable), DType::Utf8(Nullable), false, + Nullable, + ) + .unwrap(); + let sorted_target = DType::map( + DType::Primitive(PType::I64, NonNullable), + DType::Utf8(Nullable), + true, + Nullable, + ) + .unwrap(); + let incompatible_target = DType::map( + DType::Utf8(NonNullable), + DType::Utf8(Nullable), + false, + Nullable, + ) + .unwrap(); + + assert!(unsorted_target.can_coerce_from(&sorted_source)); + assert!(!sorted_target.can_coerce_from(&unsorted_source)); + assert!(!incompatible_target.can_coerce_from(&sorted_source)); + } + + #[test] + fn map_struct_least_supertype_is_symmetric_over_field_nullability() { + let key_dtype = DType::Primitive(PType::I32, NonNullable); + let nonnullable_struct = DType::Struct( + StructFields::new(FieldNames::from(["a"]), vec![DType::Utf8(NonNullable)]), + NonNullable, + ); + let nullable_field_struct = DType::Struct( + StructFields::new(FieldNames::from(["a"]), vec![DType::Utf8(Nullable)]), + NonNullable, + ); + let lhs = DType::map(key_dtype.clone(), nonnullable_struct, false, NonNullable).unwrap(); + let rhs = DType::map(key_dtype, nullable_field_struct, false, NonNullable).unwrap(); + let expected = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Struct( + StructFields::new(FieldNames::from(["a"]), vec![DType::Utf8(Nullable)]), + NonNullable, + ), + false, NonNullable, ) .unwrap(); - assert_eq!(i32_map.least_supertype(&i64_map), None); + assert_eq!(lhs.least_supertype(&rhs), Some(expected.clone())); + assert_eq!(rhs.least_supertype(&lhs), Some(expected)); } #[test] - fn map_coercion_does_not_create_a_sortedness_assertion() { - let key = DType::Primitive(PType::I32, NonNullable); - let value = DType::Utf8(Nullable); - let sorted = DType::map(key.clone(), value.clone(), true, Nullable).unwrap(); - let unsorted = DType::map(key.clone(), value, false, Nullable).unwrap(); - let different_value = DType::map(key, DType::Utf8(NonNullable), false, Nullable).unwrap(); - - assert!(!sorted.can_coerce_from(&unsorted)); - assert!(unsorted.can_coerce_from(&sorted)); - assert!(!unsorted.can_coerce_from(&different_value)); + fn map_struct_can_coerce_from_respects_field_nullability() { + let key_dtype = DType::Primitive(PType::I32, NonNullable); + let source = DType::map( + key_dtype.clone(), + DType::Struct( + StructFields::new(FieldNames::from(["a"]), vec![DType::Utf8(NonNullable)]), + NonNullable, + ), + false, + NonNullable, + ) + .unwrap(); + let nullable_field_target = DType::map( + key_dtype.clone(), + DType::Struct( + StructFields::new(FieldNames::from(["a"]), vec![DType::Utf8(Nullable)]), + NonNullable, + ), + false, + NonNullable, + ) + .unwrap(); + let nullable_field_source = DType::map( + key_dtype.clone(), + DType::Struct( + StructFields::new(FieldNames::from(["a"]), vec![DType::Utf8(Nullable)]), + NonNullable, + ), + false, + NonNullable, + ) + .unwrap(); + let nonnullable_field_target = DType::map( + key_dtype, + DType::Struct( + StructFields::new(FieldNames::from(["a"]), vec![DType::Utf8(NonNullable)]), + NonNullable, + ), + false, + NonNullable, + ) + .unwrap(); + + assert!(nullable_field_target.can_coerce_from(&source)); + assert!(!nonnullable_field_target.can_coerce_from(&nullable_field_source)); + } + + #[test] + fn map_struct_coercion_rejects_different_field_names() { + let lhs = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Struct( + StructFields::new(FieldNames::from(["a"]), vec![DType::Utf8(Nullable)]), + NonNullable, + ), + false, + NonNullable, + ) + .unwrap(); + let rhs = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Struct( + StructFields::new(FieldNames::from(["b"]), vec![DType::Utf8(Nullable)]), + NonNullable, + ), + false, + NonNullable, + ) + .unwrap(); + + assert_eq!(lhs.least_supertype(&rhs), None); + assert!(!lhs.can_coerce_from(&rhs)); + assert!(!rhs.can_coerce_from(&lhs)); } } diff --git a/vortex-array/src/scalar/cast.rs b/vortex-array/src/scalar/cast.rs index 69379ec2c6f..0ce8c331f1d 100644 --- a/vortex-array/src/scalar/cast.rs +++ b/vortex-array/src/scalar/cast.rs @@ -7,6 +7,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::dtype::DType; use crate::scalar::Scalar; @@ -31,6 +32,16 @@ impl Scalar { return Scalar::try_new(target_dtype.clone(), self.value().cloned()); } + if let (Some(source), Some(target)) = (self.dtype().as_map_opt(), target_dtype.as_map_opt()) + && target.keys_sorted() + && !source.keys_sorted() + { + return Err(vortex_err!( + "Cannot cast {} to {target_dtype}: source does not assert sorted map keys", + self.dtype() + )); + } + // Null can be cast into any nullable type as null. // Note that the `matches` clause is technically unnecessary here, just protective. if self.value().is_none() || matches!(self.dtype(), DType::Null) { diff --git a/vortex-array/src/scalar/typed_view/map.rs b/vortex-array/src/scalar/typed_view/map.rs index f2eadd5c2b9..9dadf074523 100644 --- a/vortex-array/src/scalar/typed_view/map.rs +++ b/vortex-array/src/scalar/typed_view/map.rs @@ -54,7 +54,7 @@ impl Eq for MapScalar<'_> {} impl Hash for MapScalar<'_> { fn hash(&self, state: &mut H) { - self.dtype.as_nonnullable().hash(state); + self.dtype.hash_ignore_nullability(state); self.entries.hash(state); } } @@ -141,7 +141,9 @@ impl<'a> MapScalar<'a> { /// # Errors /// /// Returns an error when `dtype` is not a map, its key/value dtypes cannot be cast, or the - /// target claims sorted keys when this scalar's dtype does not make that assertion. + /// target claims sorted keys when this scalar's dtype does not make that assertion. Also + /// returns an error for direct null-map casts; callers should use [`Scalar::cast`] so null + /// handling and sortedness checks stay centralized. pub(crate) fn cast(&self, dtype: &DType) -> VortexResult { let target = dtype .as_map_opt() @@ -155,7 +157,10 @@ impl<'a> MapScalar<'a> { } let Some(entries) = self.entries else { - return Ok(Scalar::null(dtype.clone())); + vortex_bail!( + "Cannot cast null map {} to {dtype}: Scalar::cast should handle nulls first", + self.dtype + ); }; let target_key = target.key_dtype(); @@ -194,7 +199,12 @@ impl<'a> MapScalar<'a> { #[cfg(test)] mod tests { + use std::collections::hash_map::DefaultHasher; + use std::hash::Hash; + use std::hash::Hasher; + use vortex_error::VortexResult; + use vortex_utils::aliases::hash_set::HashSet; use crate::dtype::DType; use crate::dtype::Nullability; @@ -331,4 +341,66 @@ mod tests { Ok(()) } + + #[test] + fn equal_maps_with_different_nested_nullability_hash_equal() -> VortexResult<()> { + let key_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let nullable_dtype = DType::map( + key_dtype.clone(), + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + )?; + let nonnullable_dtype = DType::map( + key_dtype, + DType::Utf8(Nullability::NonNullable), + false, + Nullability::NonNullable, + )?; + let nullable_scalar = Scalar::try_map( + nullable_dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + )], + )?; + let nonnullable_scalar = Scalar::try_map( + nonnullable_dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::NonNullable), + )], + )?; + let nullable_map = nullable_scalar.as_map(); + let nonnullable_map = nonnullable_scalar.as_map(); + + assert_eq!(nullable_map, nonnullable_map); + + let mut nullable_hash = DefaultHasher::new(); + nullable_map.hash(&mut nullable_hash); + let mut nonnullable_hash = DefaultHasher::new(); + nonnullable_map.hash(&mut nonnullable_hash); + assert_eq!(nullable_hash.finish(), nonnullable_hash.finish()); + + let mut set = HashSet::new(); + set.insert(nullable_map); + assert!(set.contains(&nonnullable_map)); + + Ok(()) + } + + #[test] + fn direct_null_map_cast_returns_error() -> VortexResult<()> { + let target_dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + Nullability::NonNullable, + )?; + let scalar = Scalar::null(dtype()?); + + assert!(scalar.as_map().cast(&target_dtype).is_err()); + + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 046e19a7c9f..8df6accece3 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -26,6 +26,7 @@ use crate::arrays::Decimal; use crate::arrays::Extension; use crate::arrays::FixedSizeList; use crate::arrays::ListView; +use crate::arrays::Map; use crate::arrays::Null; use crate::arrays::Primitive; use crate::arrays::VarBinView; @@ -197,7 +198,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::Map(a) => ::cast(a, dtype, ctx), CanonicalView::FixedSizeList(a) => ::cast(a, dtype, ctx), CanonicalView::Struct(a) => struct_cast(a, dtype, ctx), CanonicalView::Union(_) => { diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index e5dcf96ac7f..86d45d2c0d9 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -29,7 +29,6 @@ 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,7 +153,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::Map(map_array) => self.compress_map_array(map_array, compress_ctx, exec_ctx), Canonical::FixedSizeList(fsl_array) => { let compressed_elems = self.compress(fsl_array.elements(), exec_ctx)?; diff --git a/vortex-compressor/src/compressor/structural.rs b/vortex-compressor/src/compressor/structural.rs index 8b9cd999d15..e0d9462ede1 100644 --- a/vortex-compressor/src/compressor/structural.rs +++ b/vortex-compressor/src/compressor/structural.rs @@ -9,13 +9,16 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListView; use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::MapArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::list::ListArrayExt; use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_error::VortexResult; +use vortex_error::vortex_err; use super::ROOT_SCHEME_ID; use crate::CascadingCompressor; @@ -106,6 +109,23 @@ impl CascadingCompressor { .into_array()) } + /// Compresses a [`MapArray`] by recursively compressing its entries [`ListViewArray`]. + pub(super) fn compress_map_array( + &self, + map_array: MapArray, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let data_parts = map_array.into_data_parts(); + let entries = self.compress_list_view_array(data_parts.entries, compress_ctx, exec_ctx)?; + let entries = entries + .as_opt::() + .ok_or_else(|| vortex_err!("Compressed map entries became {}", entries.encoding_id()))? + .into_owned(); + + Ok(MapArray::try_new(data_parts.map_dtype, entries)?.into_array()) + } + /// Compress very child slot of the array, then re-build it from them. pub(super) fn compress_physical_slots( &self, diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index b23152d7ad3..ec14383ce36 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -12,8 +12,16 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; +use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builders::MapBuilder; +use vortex_array::dtype::DType; +use vortex_array::dtype::MapDType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -741,3 +749,95 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { assert!(matches!(score, EstimateScore::FiniteCompression(ratio) if ratio.is_finite())); Ok(()) } + +type MapEntryFixture<'a> = (i32, Option<&'a str>); +type MapRowFixture<'a> = Option>>; + +fn map_array_from_rows(rows: &[MapRowFixture<'_>], keys_sorted: bool) -> VortexResult { + let map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + keys_sorted, + )?; + let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); + let mut builder = + MapBuilder::::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::>(); + Scalar::try_map(dtype.clone(), entries)? + } + None => Scalar::null(dtype.clone()), + }; + builder.append_value(scalar.as_map())?; + } + + Ok(builder.finish_into_map().into_array()) +} + +#[test] +fn map_compression_preserves_mixed_rows() -> VortexResult<()> { + let array = map_array_from_rows( + &[ + Some(vec![(1, Some("one")), (2, None)]), + None, + Some(vec![]), + Some(vec![(1, Some("dup-old")), (1, Some("dup-new"))]), + ], + false, + )?; + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressed = compressor().compress(&array, &mut exec_ctx)?; + + assert!(compressed.is::()); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(&compressed, &array, &mut exec_ctx); + Ok(()) +} + +#[test] +fn all_null_map_compression_preserves_values() -> VortexResult<()> { + let array = map_array_from_rows(&[None, None, None], false)?; + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressed = compressor().compress(&array, &mut exec_ctx)?; + + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(&compressed, &array, &mut exec_ctx); + Ok(()) +} + +#[test] +fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { + let rows = (0..64) + .map(|idx| { + Some(vec![ + (idx % 4, Some("alpha")), + (idx % 4, Some("beta")), + (idx % 4, Some("alpha")), + ]) + }) + .collect::>(); + let array = map_array_from_rows(&rows, true)?; + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressed = compressor().compress(&array, &mut exec_ctx)?; + + assert!(compressed.is::()); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(&compressed, &array, &mut exec_ctx); + Ok(()) +}