diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index aa93dae0c45..2cfed7fd25b 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -153,6 +153,7 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort DType::FixedSizeList(.., nullability) => { execute_sparse_fixed_size_list(&patches, &fill_value, len, *nullability, ctx)? } + DType::Map(..) => vortex_bail!("Sparse canonicalization does not support Map arrays yet"), DType::Struct(struct_fields, ..) => execute_sparse_struct( struct_fields, fill_value.as_struct(), diff --git a/fuzz/src/array/compare.rs b/fuzz/src/array/compare.rs index f594a9b4102..03ce50a7170 100644 --- a/fuzz/src/array/compare.rs +++ b/fuzz/src/array/compare.rs @@ -186,7 +186,11 @@ pub fn compare_canonical_array( })) .into_array() } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/fuzz/src/array/filter.rs b/fuzz/src/array/filter.rs index aa54b0d9488..746479aa46c 100644 --- a/fuzz/src/array/filter.rs +++ b/fuzz/src/array/filter.rs @@ -121,7 +121,11 @@ pub fn filter_canonical_array( ) .map(|a| a.into_array()) } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index b6402cb1e1d..9bb6e12f66b 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -518,6 +518,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet { acc.intersection(&actions).copied().collect() }) } + DType::Map(..) => HashSet::new(), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), // Currently, no support at all DType::Variant(_) => unreachable!("Variant dtype shouldn't be fuzzed"), diff --git a/fuzz/src/array/search_sorted.rs b/fuzz/src/array/search_sorted.rs index 762182b35fb..6eb0d63e93a 100644 --- a/fuzz/src/array/search_sorted.rs +++ b/fuzz/src/array/search_sorted.rs @@ -149,7 +149,11 @@ pub fn search_sorted_canonical_array( .collect::>>()?; scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side) } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index 906200f1007..b503cc86d98 100644 --- a/fuzz/src/array/slice.rs +++ b/fuzz/src/array/slice.rs @@ -125,7 +125,11 @@ pub fn slice_canonical_array( ) .map(|a| a.into_array()) } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/fuzz/src/array/sort.rs b/fuzz/src/array/sort.rs index 3f00fccd06e..b0af8638898 100644 --- a/fuzz/src/array/sort.rs +++ b/fuzz/src/array/sort.rs @@ -102,7 +102,11 @@ pub fn sort_canonical_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexR }); take_canonical_array_non_nullable_indices(array, &sort_indices, ctx) } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index c67ae6184d8..8e59bc085db 100644 --- a/fuzz/src/array/take.rs +++ b/fuzz/src/array/take.rs @@ -148,7 +148,11 @@ pub fn take_canonical_array( ) .map(|a| a.into_array()) } - d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { + d @ (DType::Null + | DType::Map(..) + | DType::Union(..) + | DType::Variant(_) + | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index ada9e6eb2e5..99a7b99128a 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -246,6 +246,7 @@ impl AggregateFnVTable for IsSorted { DType::Null | DType::List(..) | DType::FixedSizeList(..) + | DType::Map(..) | DType::Struct(..) | DType::Union(..) | DType::Variant(..) @@ -263,6 +264,7 @@ impl AggregateFnVTable for IsSorted { DType::Null | DType::List(..) | DType::FixedSizeList(..) + | DType::Map(..) | DType::Struct(..) | DType::Union(..) | DType::Variant(..) 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 48635f2545b..1517d9047a9 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 @@ -236,6 +236,9 @@ pub(crate) fn constant_uncompressed_size_in_bytes( 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 \ @@ -294,6 +297,7 @@ 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::Struct(fields, _) => fields .fields() .all(|field| supports_uncompressed_size_in_bytes(&field)), diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs index a21ee369b16..97b75f255f0 100644 --- a/vortex-array/src/arrays/arbitrary.rs +++ b/vortex-array/src/arrays/arbitrary.rs @@ -158,6 +158,7 @@ fn random_array_chunk( DType::FixedSizeList(elem_dtype, list_size, null) => { random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len) } + DType::Map(..) => Err(IncorrectFormat), DType::Struct(sdt, n) => { let first_array = sdt .fields() diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 1477ba3b91b..9f51c545a6c 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -126,6 +126,7 @@ 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::FixedSizeList(element_dtype, list_size, _) => { let value = scalar.as_list(); diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 85dafb69cb8..f76316d9357 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -347,6 +347,9 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box { + vortex_error::vortex_panic!(InvalidArgument: "map builders are not yet supported") + } DType::FixedSizeList(elem_dtype, list_size, null) => { Box::new(FixedSizeListBuilder::with_capacity( Arc::clone(elem_dtype), diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 138e6941def..a9db688f239 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -631,6 +631,9 @@ fn create_test_scalars_for_dtype(dtype: &DType, count: usize) -> Vec { .collect(); Scalar::fixed_size_list(Arc::clone(element_dtype), elements, *n) } + DType::Map(..) => { + panic!("map builders are not supported until MapArray exists") + } DType::Struct(fields, n) => { // Create struct scalars with field values. let field_values: Vec = fields diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index bee2d52a1f8..36a59ae1b7c 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -215,6 +215,9 @@ 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::FixedSizeList(elem_dtype, list_size, null) => Canonical::FixedSizeList(unsafe { FixedSizeListArray::new_unchecked( Canonical::empty(elem_dtype).into_array(), diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 73f98cce4a7..18cda1c425e 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -1227,6 +1227,7 @@ 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::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..4a56cba98ec 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,6 +59,13 @@ fn random_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result { u.choose_index(3)?.try_into().vortex_expect("impossible"), u.arbitrary()?, ), + 9 => DType::map( + random_dtype(u, depth - 1)?.as_nonnullable(), + random_dtype(u, depth - 1)?, + u.arbitrary()?, + u.arbitrary()?, + ) + .vortex_expect("non-nullable generated map keys are always valid"), // Null, // Extension(ExtDType, Nullability), _ => unreachable!("Number out of range"), diff --git a/vortex-array/src/dtype/coercion.rs b/vortex-array/src/dtype/coercion.rs index 4d30a80c216..9f1ccd9861e 100644 --- a/vortex-array/src/dtype/coercion.rs +++ b/vortex-array/src/dtype/coercion.rs @@ -100,6 +100,20 @@ impl DType { return Some(DType::List(Arc::new(elem), union_null)); } + 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; + } + + return DType::map( + lhs.key_dtype(), + lhs.value_dtype(), + lhs.keys_sorted() && rhs.keys_sorted(), + union_null, + ) + .ok(); + } + // Identity (ignoring nullability): return self with union nullability if self.eq_ignore_nullability(other) { return Some(self.with_nullability(union_null)); @@ -187,6 +201,13 @@ impl DType { && target_elem.can_coerce_from(source_elem); } + 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(); + } + // Same type (ignoring nullability): check nullability compatibility if self.eq_ignore_nullability(other) { return self.is_nullable() || !other.is_nullable(); @@ -804,4 +825,50 @@ mod tests { DType::Decimal(DecimalDType::new(15, 5), NonNullable) ); } + + #[test] + fn map_least_supertype_unions_outer_nullability_and_intersects_sortedness() { + 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(); + + assert_eq!( + sorted.least_supertype(&unsorted), + Some(DType::map(key, value, false, Nullable).unwrap()) + ); + } + + #[test] + fn map_least_supertype_requires_identical_key_and_value_dtypes() { + let i32_map = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + false, + NonNullable, + ) + .unwrap(); + let i64_map = DType::map( + DType::Primitive(PType::I64, NonNullable), + DType::Utf8(Nullable), + false, + NonNullable, + ) + .unwrap(); + + assert_eq!(i32_map.least_supertype(&i64_map), None); + } + + #[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)); + } } diff --git a/vortex-array/src/dtype/dtype_impl.rs b/vortex-array/src/dtype/dtype_impl.rs index 2cfa0b01383..b2e53dcefae 100644 --- a/vortex-array/src/dtype/dtype_impl.rs +++ b/vortex-array/src/dtype/dtype_impl.rs @@ -10,11 +10,13 @@ use std::sync::Arc; use DType::*; use itertools::Itertools; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use super::DType; use crate::dtype::FieldDType; use crate::dtype::FieldName; +use crate::dtype::MapDType; use crate::dtype::PType; use crate::dtype::StructFields; use crate::dtype::UnionVariants; @@ -65,6 +67,7 @@ impl DType { | Binary(null) | List(_, null) | FixedSizeList(_, _, null) + | Map(_, null) | Struct(_, null) | Union(_, null) | Variant(null) => matches!(null, Nullability::Nullable), @@ -95,6 +98,7 @@ impl DType { Binary(_) => Binary(nullability), List(edt, _) => List(Arc::clone(edt), nullability), FixedSizeList(edt, size, _) => FixedSizeList(Arc::clone(edt), *size, nullability), + Map(map, _) => Map(map.clone(), nullability), Struct(sf, _) => Struct(sf.clone(), nullability), Union(vs, _) => Union(vs.clone(), nullability), Variant(_) => Variant(nullability), @@ -121,6 +125,7 @@ impl DType { (FixedSizeList(lhs_dtype, lhs_size, _), FixedSizeList(rhs_dtype, rhs_size, _)) => { lhs_size == rhs_size && lhs_dtype.eq_ignore_nullability(rhs_dtype) } + (Map(lhs, _), Map(rhs, _)) => lhs.eq_ignore_nullability(rhs), (Struct(lhs_dtype, _), Struct(rhs_dtype, _)) => { lhs_dtype.eq_ignore_nullability(rhs_dtype) } @@ -146,6 +151,7 @@ impl DType { element.hash_ignore_nullability(state); size.hash(state); } + Map(map, _) => map.hash_ignore_nullability(state), Struct(fields, _) => fields.hash_ignore_nullability(state), Union(variants, _) => variants.hash_ignore_nullability(state), Extension(ext) => ext.hash_ignore_nullability(state), @@ -262,6 +268,11 @@ impl DType { matches!(self, FixedSizeList(..)) } + /// Check if `self` is a [`DType::Map`]. + pub fn is_map(&self) -> bool { + matches!(self, Map(..)) + } + /// Check if `self` is a [`DType::Struct`] pub fn is_struct(&self) -> bool { matches!(self, Struct(_, _)) @@ -286,7 +297,7 @@ impl DType { /// recursive type. pub fn is_nested(&self) -> bool { match self { - List(..) | FixedSizeList(..) | Struct(..) | Union(..) | Variant(..) => true, + List(..) | FixedSizeList(..) | Map(..) | Struct(..) | Union(..) | Variant(..) => true, Extension(ext) => ext.storage_dtype().is_nested(), _ => false, } @@ -305,7 +316,7 @@ impl DType { Decimal(decimal, _) => { Some(DecimalType::smallest_decimal_value_type(decimal).byte_width()) } - Utf8(_) | Binary(_) | List(..) => None, + Utf8(_) | Binary(_) | List(..) | Map(..) => None, FixedSizeList(elem_dtype, list_size, _) => { elem_dtype.element_size().map(|s| s * *list_size as usize) } @@ -384,6 +395,24 @@ impl DType { } } + /// Get the [`MapDType`] if `self` is a [`DType::Map`], otherwise `None`. + pub fn as_map_opt(&self) -> Option<&MapDType> { + if let Map(map, _) = self { + Some(map) + } else { + None + } + } + + /// Owned version of [Self::as_map_opt]. + pub fn into_map_opt(self) -> Option { + if let Map(map, _) = self { + Some(map) + } else { + None + } + } + /// Get the inner element dtype if `self` is **either** a [`DType::List`] or a /// [`DType::FixedSizeList`], otherwise returns `None` pub fn as_any_size_list_element_opt(&self) -> Option<&Arc> { @@ -488,6 +517,23 @@ impl DType { List(Arc::new(dtype.into()), nullability) } + /// Convenience method for creating a [`DType::Map`]. + /// + /// # Errors + /// + /// Returns an error when the key dtype is nullable. + pub fn map( + key: impl Into, + value: impl Into, + keys_sorted: bool, + nullability: Nullability, + ) -> VortexResult { + Ok(Map( + MapDType::try_new(key.into(), value.into(), keys_sorted)?, + nullability, + )) + } + /// Convenience method for creating a [`DType::Struct`]. pub fn struct_, impl Into)>>( iter: I, @@ -508,6 +554,7 @@ impl Display for DType { Binary(null) => write!(f, "binary{null}"), List(edt, null) => write!(f, "list({edt}){null}"), FixedSizeList(edt, size, null) => write!(f, "fixed_size_list({edt})[{size}]{null}"), + Map(map, null) => write!(f, "{map}{null}"), Struct(sf, null) => write!( f, "{{{}}}{null}", @@ -595,6 +642,27 @@ mod tests { Ok(()) } + #[test] + fn test_map_dtype_hash_ignores_value_nullability() -> VortexResult<()> { + let lhs = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + true, + NonNullable, + )?; + let rhs = DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(NonNullable), + true, + Nullable, + )?; + + assert!(lhs.eq_ignore_nullability(&rhs)); + assert_eq!(hash_ignore_nullability(&lhs), hash_ignore_nullability(&rhs)); + + Ok(()) + } + #[test] fn element_size_null() { assert_eq!(DType::Null.element_size(), Some(0)); diff --git a/vortex-array/src/dtype/map.rs b/vortex-array/src/dtype/map.rs new file mode 100644 index 00000000000..e138cea6545 --- /dev/null +++ b/vortex-array/src/dtype/map.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::dtype::DType; +use crate::dtype::FieldDType; +use crate::dtype::Nullability; + +/// Logical type information for a map's entries. +/// +/// A map has ordered key/value entries. Keys must be non-nullable, while values may be nullable. +/// `keys_sorted` is a producer assertion matching Arrow's map type; it is not validated against +/// data at type construction time. +#[allow( + clippy::derived_hash_with_manual_eq, + reason = "manual PartialEq adds Arc::ptr_eq fast path only" +)] +#[derive(Clone, Eq, Hash)] +pub struct MapDType(Arc); + +impl PartialEq for MapDType { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) || self.0 == other.0 + } +} + +impl fmt::Debug for MapDType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MapDType") + .field("key", &self.0.key) + .field("value", &self.0.value) + .field("keys_sorted", &self.0.keys_sorted) + .finish() + } +} + +impl fmt::Display for MapDType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "map({}, {}, keys_sorted={})", + self.key_dtype(), + self.value_dtype(), + self.keys_sorted() + ) + } +} + +struct MapDTypeInner { + key: FieldDType, + value: FieldDType, + keys_sorted: bool, +} + +impl PartialEq for MapDTypeInner { + fn eq(&self, other: &Self) -> bool { + self.key == other.key && self.value == other.value && self.keys_sorted == other.keys_sorted + } +} + +impl Eq for MapDTypeInner {} + +impl Hash for MapDTypeInner { + fn hash(&self, state: &mut H) { + self.key.hash(state); + self.value.hash(state); + self.keys_sorted.hash(state); + } +} + +impl MapDType { + pub(crate) fn eq_ignore_nullability(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + || (self.key_dtype().eq_ignore_nullability(&other.key_dtype()) + && self + .value_dtype() + .eq_ignore_nullability(&other.value_dtype()) + && self.keys_sorted() == other.keys_sorted()) + } + + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + self.key_dtype().hash_ignore_nullability(state); + self.value_dtype().hash_ignore_nullability(state); + self.keys_sorted().hash(state); + } + + /// Creates a map dtype from its key and value dtypes. + /// + /// # Errors + /// + /// Returns an error when `key` is nullable. Arrow map keys cannot be null. + pub fn try_new(key: DType, value: DType, keys_sorted: bool) -> VortexResult { + Self::try_from_fields(key.into(), value.into(), keys_sorted) + } + + pub(crate) fn try_from_fields( + key: FieldDType, + value: FieldDType, + keys_sorted: bool, + ) -> VortexResult { + vortex_ensure!( + !key.value()?.is_nullable(), + "map key dtype must be non-nullable" + ); + + Ok(Self(Arc::new(MapDTypeInner { + key, + value, + keys_sorted, + }))) + } + + /// Returns the dtype of the map keys. + pub fn key_dtype(&self) -> DType { + self.0 + .key + .value() + .vortex_expect("map key dtype must be valid") + } + + /// Returns the dtype of the map values. + pub fn value_dtype(&self) -> DType { + self.0 + .value + .value() + .vortex_expect("map value dtype must be valid") + } + + /// Returns whether producers assert that keys are sorted within every map value. + pub fn keys_sorted(&self) -> bool { + self.0.keys_sorted + } + + /// Returns the non-nullable `{key, value}` struct dtype used for map entries. + pub fn entries_dtype(&self) -> DType { + DType::struct_( + [("key", self.key_dtype()), ("value", self.value_dtype())], + Nullability::NonNullable, + ) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::dtype::DType; + use crate::dtype::MapDType; + use crate::dtype::Nullability; + use crate::dtype::PType; + + #[test] + fn rejects_nullable_keys() { + let result = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Utf8(Nullability::Nullable), + false, + ); + + assert!(result.is_err()); + } + + #[test] + fn accepts_nullable_values() -> VortexResult<()> { + let dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + )?; + + assert_eq!( + dtype.entries_dtype(), + DType::struct_( + [ + ( + "key", + DType::Primitive(PType::I32, Nullability::NonNullable) + ), + ("value", DType::Utf8(Nullability::Nullable)), + ], + Nullability::NonNullable, + ) + ); + assert!(dtype.keys_sorted()); + assert_eq!(dtype.to_string(), "map(i32, utf8?, keys_sorted=true)"); + + Ok(()) + } + + #[test] + fn outer_nullability_is_independent_of_map_details() -> VortexResult<()> { + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + )?; + let nullable = dtype.as_nullable(); + + assert!(dtype.is_map()); + assert!(!dtype.is_nullable()); + assert!(nullable.is_nullable()); + assert!(dtype.eq_ignore_nullability(&nullable)); + assert_eq!(dtype.as_map_opt(), nullable.as_map_opt()); + assert!(nullable.as_map_opt().unwrap().keys_sorted()); + + Ok(()) + } +} diff --git a/vortex-array/src/dtype/mod.rs b/vortex-array/src/dtype/mod.rs index 8b579c70d39..df621293ab8 100644 --- a/vortex-array/src/dtype/mod.rs +++ b/vortex-array/src/dtype/mod.rs @@ -24,6 +24,7 @@ mod f16; mod field; mod field_mask; mod field_names; +mod map; mod native_dtype; mod nullability; mod ptype; @@ -102,6 +103,12 @@ pub enum DType { /// well as a `u32` size that determines the fixed length of each `FixedSizeList` scalar. FixedSizeList(Arc, u32, Nullability), + /// A logical map type. + /// + /// Map keys are non-nullable, values may be nullable, and [`MapDType`] stores Arrow's + /// `keys_sorted` assertion. + Map(MapDType, Nullability), + /// A logical struct type. /// /// A `Struct` type is composed of an ordered list of fields, each with a corresponding name and @@ -147,6 +154,7 @@ impl PartialEq for DType { (Self::FixedSizeList(da, sa, na), Self::FixedSizeList(db, sb, nb)) => { sa == sb && na == nb && (Arc::ptr_eq(da, db) || da == db) } + (Self::Map(ma, na), Self::Map(mb, nb)) => na == nb && ma == mb, // StructFields handles its own Arc::ptr_eq in its PartialEq impl. (Self::Struct(a, na), Self::Struct(b, nb)) => na == nb && a == b, // UnionVariants handles its own Arc::ptr_eq in its PartialEq impl. @@ -163,6 +171,7 @@ impl PartialEq for DType { | (Self::Binary(_), _) | (Self::List(..), _) | (Self::FixedSizeList(..), _) + | (Self::Map(..), _) | (Self::Struct(..), _) | (Self::Union(..), _) | (Self::Variant(_), _) @@ -179,6 +188,7 @@ pub use field::*; pub use field_mask::*; pub use field_names::*; pub use half; +pub use map::*; pub use nullability::*; pub use ptype::*; pub use struct_::*; diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index 6344973c2b3..0f7a16c3d65 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -21,6 +21,7 @@ use vortex_session::VortexSession; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldDType; +use crate::dtype::MapDType; use crate::dtype::PType; use crate::dtype::StructFields; use crate::dtype::UnionVariants; @@ -128,6 +129,32 @@ impl UnionVariants { } } +impl MapDType { + /// Creates a map dtype from a flatbuffer-defined object and its underlying buffer. + fn from_fb( + fb_map: fbd::Map<'_>, + buffer: FlatBuffer, + session: VortexSession, + ) -> VortexResult { + let key = fb_map + .key_type() + .ok_or_else(|| vortex_err!("failed to parse map key type from flatbuffer"))?; + let value = fb_map + .value_type() + .ok_or_else(|| vortex_err!("failed to parse map value type from flatbuffer"))?; + + MapDType::try_from_fields( + FieldDType::from(ViewedDType::from_fb_loc( + key._tab.loc(), + buffer.clone(), + session.clone(), + )), + FieldDType::from(ViewedDType::from_fb_loc(value._tab.loc(), buffer, session)), + fb_map.keys_sorted(), + ) + } +} + impl DType { /// Create a [`DType`] from a flatbuffer buffer. pub fn from_flatbuffer(buffer: FlatBuffer, session: &VortexSession) -> VortexResult { @@ -220,6 +247,13 @@ impl TryFrom for DType { fb_fixed_size_list.nullable().into(), )) } + fb::Type::Map => { + let fb_map = fb + .type__as_map() + .ok_or_else(|| vortex_err!("failed to parse map from flatbuffer"))?; + let map = MapDType::from_fb(fb_map, vfdt.buffer().clone(), vfdt.session.clone())?; + Ok(Self::Map(map, fb_map.nullable().into())) + } fb::Type::Struct_ => { let fb_struct = fb .type__as_struct_() @@ -357,6 +391,20 @@ impl WriteFlatBuffer for DType { ) .as_union_value() } + Self::Map(map, n) => { + let key_type = Some(map.key_dtype().write_flatbuffer(fbb)?); + let value_type = Some(map.value_dtype().write_flatbuffer(fbb)?); + fb::Map::create( + fbb, + &fb::MapArgs { + key_type, + value_type, + keys_sorted: map.keys_sorted(), + nullable: (*n).into(), + }, + ) + .as_union_value() + } Self::Struct(st, n) => { let names = st .names() @@ -448,6 +496,7 @@ impl WriteFlatBuffer for DType { Self::Binary(_) => fb::Type::Binary, Self::List(..) => fb::Type::List, Self::FixedSizeList(..) => fb::Type::FixedSizeList, + Self::Map(..) => fb::Type::Map, Self::Struct(..) => fb::Type::Struct_, Self::Union(..) => fb::Type::Union, Self::Variant(_) => fb::Type::Variant, @@ -569,6 +618,21 @@ mod test { ), Nullability::NonNullable, )); + let inner_map = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + ) + .unwrap(); + let map = DType::map( + inner_map, + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + ) + .unwrap(); + roundtrip_dtype(DType::struct_([("map", map)], Nullability::NonNullable)); roundtrip_dtype(DType::Variant(Nullability::Nullable)); } diff --git a/vortex-array/src/dtype/serde/mod.rs b/vortex-array/src/dtype/serde/mod.rs index c6837f4ab2d..f26912b5e53 100644 --- a/vortex-array/src/dtype/serde/mod.rs +++ b/vortex-array/src/dtype/serde/mod.rs @@ -203,4 +203,31 @@ mod test { assert_eq!(deserialized, dtype); assert_eq!(deserialized.nullability(), Nullability::Nullable); } + + #[test] + fn test_serde_nested_map_dtype_json_roundtrip() { + let inner_map = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + ) + .unwrap(); + let map = DType::map( + inner_map, + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + ) + .unwrap(); + let dtype = DType::struct_([("map", map)], Nullability::Nullable); + + let json = serde_json::to_string(&dtype).unwrap(); + let mut deserializer = serde_json::Deserializer::from_str(&json); + let deserialized: DType = DTypeSerde::::new(&SESSION) + .deserialize(&mut deserializer) + .unwrap(); + + assert_eq!(deserialized, dtype); + } } diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 485907e2e95..50367a90276 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -10,6 +10,7 @@ use vortex_session::VortexSession; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::MapDType; use crate::dtype::PType; use crate::dtype::StructFields; use crate::dtype::UnionVariants; @@ -77,6 +78,26 @@ impl DType { nullable, )) } + DtypeType::Map(map) => Ok(Self::Map( + MapDType::try_new( + DType::from_proto( + map.key_type + .as_ref() + .ok_or_else(|| vortex_err!(Serde: "Invalid map key type"))? + .as_ref(), + session, + )?, + DType::from_proto( + map.value_type + .as_ref() + .ok_or_else(|| vortex_err!(Serde: "Invalid map value type"))? + .as_ref(), + session, + )?, + map.keys_sorted, + )?, + map.nullable.into(), + )), DtypeType::Struct(s) => Ok(Self::Struct( StructFields::new( s.names.iter().map(|s| s.as_str()).collect(), @@ -165,6 +186,16 @@ impl TryFrom<&DType> for pb::DType { nullable: (*null).into(), })) } + DType::Map(map, null) => { + let key_dtype = map.key_dtype(); + let value_dtype = map.value_dtype(); + DtypeType::Map(Box::new(pb::Map { + key_type: Some(Box::new(Self::try_from(&key_dtype)?)), + value_type: Some(Box::new(Self::try_from(&value_dtype)?)), + keys_sorted: map.keys_sorted(), + nullable: (*null).into(), + })) + } DType::Struct(s, null) => DtypeType::Struct(pb::Struct { names: s.names().iter().map(|s| s.as_ref().to_string()).collect(), dtypes: s @@ -393,6 +424,27 @@ mod tests { } } + #[test] + fn test_nested_map_round_trip() { + let inner_map = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + ) + .unwrap(); + let map = DType::map( + inner_map, + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + ) + .unwrap(); + let dtype = DType::struct_([("map", map)], Nullability::NonNullable); + + assert_eq!(round_trip_dtype(&dtype), dtype); + } + #[test] fn test_extension_round_trip() { let ext_dtype = diff --git a/vortex-array/src/dtype/serde/serde.rs b/vortex-array/src/dtype/serde/serde.rs index 8fd52fda054..e941fb475bb 100644 --- a/vortex-array/src/dtype/serde/serde.rs +++ b/vortex-array/src/dtype/serde/serde.rs @@ -110,6 +110,14 @@ impl Serialize for DType { state.serialize_field(n)?; state.end() } + DType::Map(map, n) => { + let mut state = serializer.serialize_tuple_variant("DType", 12, "Map", 4)?; + state.serialize_field(&map.key_dtype())?; + state.serialize_field(&map.value_dtype())?; + state.serialize_field(&map.keys_sorted())?; + state.serialize_field(n)?; + state.end() + } DType::Struct(fields, n) => { let mut state = serializer.serialize_tuple_variant("DType", 8, "Struct", 2)?; state.serialize_field(&fields)?; @@ -181,6 +189,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { "Union", "Variant", "Extension", + "Map", ]; struct DTypeVisitor<'a> { @@ -234,6 +243,9 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { "FixedSizeList" => access.newtype_variant_seed(FixedSizeListFieldsSeed { session: self.session, }), + "Map" => access.newtype_variant_seed(MapFieldsSeed { + session: self.session, + }), "Struct" => access.newtype_variant_seed(StructFieldsSeed { session: self.session, }), @@ -264,6 +276,59 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { } } +struct MapFieldsSeed<'a> { + session: &'a VortexSession, +} + +impl<'de> DeserializeSeed<'de> for MapFieldsSeed<'_> { + type Value = DType; + + fn deserialize(self, deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct MapVisitor<'a> { + session: &'a VortexSession, + } + + impl<'de> Visitor<'de> for MapVisitor<'_> { + type Value = DType; + + fn expecting(&self, f: &mut Formatter) -> fmt::Result { + f.write_str("Map tuple (key_dtype, value_dtype, keys_sorted, nullability)") + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: SeqAccess<'de>, + { + let key = seq + .next_element_seed(DTypeSerde::::new(self.session))? + .ok_or_else(|| de::Error::invalid_length(0, &self))?; + let value = seq + .next_element_seed(DTypeSerde::::new(self.session))? + .ok_or_else(|| de::Error::invalid_length(1, &self))?; + let keys_sorted = seq + .next_element()? + .ok_or_else(|| de::Error::invalid_length(2, &self))?; + let nullability = seq + .next_element()? + .ok_or_else(|| de::Error::invalid_length(3, &self))?; + + DType::map(key, value, keys_sorted, nullability) + .map_err(|error| de::Error::custom(error.to_string())) + } + } + + deserializer.deserialize_tuple( + 4, + MapVisitor { + session: self.session, + }, + ) + } +} + // ============================================================================ // Helper seeds for nested DType variants (with session) // ============================================================================ diff --git a/vortex-array/src/scalar/arbitrary.rs b/vortex-array/src/scalar/arbitrary.rs index 49528f05a68..9781d3b9531 100644 --- a/vortex-array/src/scalar/arbitrary.rs +++ b/vortex-array/src/scalar/arbitrary.rs @@ -86,6 +86,23 @@ pub fn random_scalar(u: &mut Unstructured, dtype: &DType) -> Result { )), ) .vortex_expect("unable to construct random `Scalar`_"), + DType::Map(map, _) => Scalar::try_new( + dtype.clone(), + Some(ScalarValue::Tuple( + iter::from_fn(|| { + u.arbitrary().unwrap_or(false).then(|| { + let key = random_scalar(u, &map.key_dtype())?; + let value = random_scalar(u, &map.value_dtype())?; + Ok(Some(ScalarValue::Tuple(vec![ + key.into_value(), + value.into_value(), + ]))) + }) + }) + .collect::>>()?, + )), + ) + .vortex_expect("unable to construct random `Scalar`_"), DType::Struct(sdt, _) => Scalar::try_new( dtype.clone(), Some(ScalarValue::Tuple( diff --git a/vortex-array/src/scalar/cast.rs b/vortex-array/src/scalar/cast.rs index a4bb04e05e0..69379ec2c6f 100644 --- a/vortex-array/src/scalar/cast.rs +++ b/vortex-array/src/scalar/cast.rs @@ -57,6 +57,7 @@ impl Scalar { DType::Utf8(_) => self.as_utf8().cast(target_dtype), DType::Binary(_) => self.as_binary().cast(target_dtype), DType::List(..) | DType::FixedSizeList(..) => self.as_list().cast(target_dtype), + DType::Map(..) => self.as_map().cast(target_dtype), DType::Struct(..) => self.as_struct().cast(target_dtype), DType::Union(..) => vortex_bail!( "union scalar cast from {} to {target_dtype} is not supported (yet)", diff --git a/vortex-array/src/scalar/constructor.rs b/vortex-array/src/scalar/constructor.rs index 234e95441a7..3ea25be816e 100644 --- a/vortex-array/src/scalar/constructor.rs +++ b/vortex-array/src/scalar/constructor.rs @@ -9,6 +9,7 @@ use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; @@ -139,6 +140,58 @@ impl Scalar { Self::create_list(element_dtype, children, nullability, ListKind::FixedSize) } + /// Creates a map scalar with the given dtype and ordered key/value entries. + /// + /// # Panics + /// + /// Panics when `dtype` is not a map or an entry has an incompatible key or value dtype. + pub fn map(dtype: DType, entries: impl IntoIterator) -> Self { + Self::try_map(dtype, entries).vortex_expect("unable to construct a map `Scalar`") + } + + /// Attempts to create a map scalar with the given dtype and ordered key/value entries. + /// + /// # Errors + /// + /// Returns an error when `dtype` is not a map or an entry has an incompatible key or value + /// dtype. + pub fn try_map( + dtype: DType, + entries: impl IntoIterator, + ) -> VortexResult { + let map = dtype + .as_map_opt() + .ok_or_else(|| vortex_error::vortex_err!("Expected map dtype, found {dtype}"))?; + let key_dtype = map.key_dtype(); + let value_dtype = map.value_dtype(); + + let entries = entries + .into_iter() + .enumerate() + .map(|(index, (key, value))| { + if key.dtype() != &key_dtype { + vortex_bail!( + "map entry {index} expected key dtype {key_dtype}, got {}", + key.dtype() + ); + } + if value.dtype() != &value_dtype { + vortex_bail!( + "map entry {index} expected value dtype {value_dtype}, got {}", + value.dtype() + ); + } + + Ok(Some(ScalarValue::Tuple(vec![ + key.into_value(), + value.into_value(), + ]))) + }) + .collect::>>()?; + + Self::try_new(dtype, Some(ScalarValue::Tuple(entries))) + } + /// Creates a list [`Scalar`] from an element dtype, children, nullability, and list kind. fn create_list( element_dtype: impl Into>, diff --git a/vortex-array/src/scalar/convert/from_scalar.rs b/vortex-array/src/scalar/convert/from_scalar.rs index 32753dea1dc..77b35ad3f28 100644 --- a/vortex-array/src/scalar/convert/from_scalar.rs +++ b/vortex-array/src/scalar/convert/from_scalar.rs @@ -14,6 +14,7 @@ use crate::scalar::BoolScalar; use crate::scalar::DecimalScalar; use crate::scalar::ExtScalar; use crate::scalar::ListScalar; +use crate::scalar::MapScalar; use crate::scalar::PrimitiveScalar; use crate::scalar::Scalar; use crate::scalar::StructScalar; @@ -95,6 +96,16 @@ impl<'a> TryFrom<&'a Scalar> for ListScalar<'a> { } } +impl<'a> TryFrom<&'a Scalar> for MapScalar<'a> { + type Error = VortexError; + + fn try_from(value: &'a Scalar) -> VortexResult { + value + .as_map_opt() + .ok_or_else(|| vortex_err!("Expected map scalar, found {}", value.dtype())) + } +} + impl<'a> TryFrom<&'a Scalar> for ExtScalar<'a> { type Error = VortexError; diff --git a/vortex-array/src/scalar/display.rs b/vortex-array/src/scalar/display.rs index 2947afe5a51..6cd64c7008d 100644 --- a/vortex-array/src/scalar/display.rs +++ b/vortex-array/src/scalar/display.rs @@ -19,6 +19,7 @@ impl Display for Scalar { DType::Utf8(_) => write!(f, "{}", self.as_utf8()), DType::Binary(_) => write!(f, "{}", self.as_binary()), DType::List(..) | DType::FixedSizeList(..) => write!(f, "{}", self.as_list()), + DType::Map(..) => write!(f, "{}", self.as_map()), DType::Struct(..) => write!(f, "{}", self.as_struct()), DType::Union(..) => write!(f, "{}", self.as_union()), DType::Variant(_) => write!(f, "{}", self.as_variant()), diff --git a/vortex-array/src/scalar/downcast.rs b/vortex-array/src/scalar/downcast.rs index 2cdb67fade2..91d17e14109 100644 --- a/vortex-array/src/scalar/downcast.rs +++ b/vortex-array/src/scalar/downcast.rs @@ -14,6 +14,7 @@ use crate::scalar::DecimalScalar; use crate::scalar::DecimalValue; use crate::scalar::ExtScalar; use crate::scalar::ListScalar; +use crate::scalar::MapScalar; use crate::scalar::PValue; use crate::scalar::PrimitiveScalar; use crate::scalar::Scalar; @@ -137,6 +138,21 @@ impl Scalar { ListScalar::try_new(self.dtype(), self.value()).ok() } + /// Returns a view of the scalar as a map scalar. + /// + /// # Panics + /// + /// Panics if the scalar does not have a [`Map`](crate::dtype::DType::Map) type. + pub fn as_map(&self) -> MapScalar<'_> { + self.as_map_opt() + .vortex_expect("Failed to convert scalar to map") + } + + /// Returns a view of the scalar as a map scalar if it has a map type. + pub fn as_map_opt(&self) -> Option> { + MapScalar::try_new(self.dtype(), self.value()).ok() + } + /// Returns a view of the scalar as an extension scalar. /// /// # Panics diff --git a/vortex-array/src/scalar/proto.rs b/vortex-array/src/scalar/proto.rs index d3f7894f86e..7e1b5ac85db 100644 --- a/vortex-array/src/scalar/proto.rs +++ b/vortex-array/src/scalar/proto.rs @@ -471,13 +471,17 @@ fn list_from_proto( .map(|(value, field_dtype)| ScalarValue::from_proto(value, &field_dtype, session)) .collect::>>()? } - _ => { - vortex_bail!( - Serde: "expected List, FixedSizeList, or Struct dtype for ListValue, got {dtype}" - ) + DType::Map(map, _) => { + let entry_dtype = map.entries_dtype(); + v.values + .iter() + .map(|entry| ScalarValue::from_proto(entry, &entry_dtype, session)) + .collect::>>()? } + _ => vortex_bail!( + Serde: "expected a tuple-backed dtype for ListValue, got {dtype}" + ), }; - Ok(ScalarValue::Tuple(values)) } @@ -606,6 +610,34 @@ mod tests { )); } + #[test] + fn test_map() { + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + ) + .unwrap(); + round_trip( + Scalar::try_map( + dtype.clone(), + [ + ( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + ), + ( + Scalar::primitive(2i32, Nullability::NonNullable), + Scalar::null(DType::Utf8(Nullability::Nullable)), + ), + ], + ) + .unwrap(), + ); + round_trip(Scalar::null(dtype)); + } + #[test] fn test_f16() { round_trip(Scalar::primitive( diff --git a/vortex-array/src/scalar/scalar_impl.rs b/vortex-array/src/scalar/scalar_impl.rs index a4351932152..4d9f5a3a867 100644 --- a/vortex-array/src/scalar/scalar_impl.rs +++ b/vortex-array/src/scalar/scalar_impl.rs @@ -135,6 +135,7 @@ impl Scalar { /// - `Utf8`: `""` /// - `Binary`: An empty buffer /// - `List`: An empty list + /// - `Map`: An empty map /// - `FixedSizeList`: A list (with correct size) of zero values, which is determined by the /// element [`DType`] /// - `Struct`: A struct where each field has a zero value, which is determined by the field @@ -212,6 +213,7 @@ impl Scalar { DType::Utf8(_) => value.as_utf8().is_empty(), DType::Binary(_) => value.as_binary().is_empty(), DType::List(..) => value.as_list().is_empty(), + DType::Map(..) => self.as_map().is_empty(), // A fixed-size list is zero only if it has the expected number of elements and every // element is itself a non-null zero value.1 DType::FixedSizeList(_, list_size, _) => { @@ -298,6 +300,11 @@ impl Scalar { .elements() .map(|fields| fields.into_iter().map(|f| f.approx_nbytes()).sum::()) .unwrap_or_default(), + DType::Map(..) => self + .as_map() + .entries() + .map(|(key, value)| key.approx_nbytes() + value.approx_nbytes()) + .sum(), DType::Struct(..) => self .as_struct() .fields_iter() @@ -444,6 +451,7 @@ fn partial_cmp_tuple_values( partial_cmp_list_values(element_dtype, lhs, rhs) } DType::Struct(fields, _) => partial_cmp_struct_values(fields, lhs, rhs), + DType::Map(..) => None, DType::Extension(ext_dtype) => { partial_cmp_tuple_values(ext_dtype.storage_dtype(), lhs, rhs) } diff --git a/vortex-array/src/scalar/scalar_value.rs b/vortex-array/src/scalar/scalar_value.rs index c898832d96c..eb74c1a0a0e 100644 --- a/vortex-array/src/scalar/scalar_value.rs +++ b/vortex-array/src/scalar/scalar_value.rs @@ -36,7 +36,7 @@ pub enum ScalarValue { Binary(ByteBuffer), /// A tuple of potentially null scalar values. /// - /// Used as the underlying representation for list, fixed-size list, and struct scalars. + /// Used as the underlying representation for list, fixed-size list, map, and struct scalars. Tuple(Vec>), /// A present union value carrying its selected type ID and raw child value. Union(UnionValue), @@ -61,6 +61,7 @@ impl ScalarValue { DType::Utf8(_) => Self::Utf8(BufferString::empty()), DType::Binary(_) => Self::Binary(ByteBuffer::empty()), DType::List(..) => Self::Tuple(vec![]), + DType::Map(..) => Self::Tuple(vec![]), DType::FixedSizeList(edt, size, _) => { let elements = (0..*size) .map(|_| Self::try_zero_value(edt).map(Some)) @@ -112,6 +113,7 @@ impl ScalarValue { DType::Utf8(_) => Self::Utf8(BufferString::empty()), DType::Binary(_) => Self::Binary(ByteBuffer::empty()), DType::List(..) => Self::Tuple(vec![]), + DType::Map(..) => Self::Tuple(vec![]), DType::FixedSizeList(edt, size, _) => { let elements = (0..*size) .map(|_| Self::try_default_value(edt)) diff --git a/vortex-array/src/scalar/typed_view/map.rs b/vortex-array/src/scalar/typed_view/map.rs new file mode 100644 index 00000000000..f2eadd5c2b9 --- /dev/null +++ b/vortex-array/src/scalar/typed_view/map.rs @@ -0,0 +1,334 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`MapScalar`] typed view implementation. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::dtype::DType; +use crate::dtype::MapDType; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; + +/// A scalar value representing an ordered sequence of map key/value entries. +/// +/// Map keys are non-null and each entry has exactly one key and one value. Duplicate keys are +/// preserved because map logical types do not enforce key uniqueness. +#[derive(Debug, Clone, Copy)] +pub struct MapScalar<'a> { + dtype: &'a DType, + entries: Option<&'a [Option]>, +} + +impl Display for MapScalar<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + if self.is_null() { + return write!(f, "null"); + } + + write!( + f, + "{{{}}}", + self.entries() + .map(|(key, value)| format!("{key}: {value}")) + .format(", ") + ) + } +} + +impl PartialEq for MapScalar<'_> { + fn eq(&self, other: &Self) -> bool { + self.dtype.eq_ignore_nullability(other.dtype) && self.entries == other.entries + } +} + +impl Eq for MapScalar<'_> {} + +impl Hash for MapScalar<'_> { + fn hash(&self, state: &mut H) { + self.dtype.as_nonnullable().hash(state); + self.entries.hash(state); + } +} + +impl<'a> MapScalar<'a> { + /// Creates a map scalar view from a dtype and optional scalar value. + /// + /// # Errors + /// + /// Returns an error when `dtype` is not [`DType::Map`]. + pub fn try_new(dtype: &'a DType, value: Option<&'a ScalarValue>) -> VortexResult { + if !dtype.is_map() { + vortex_bail!("Expected map scalar, found {dtype}") + } + + Ok(Self { + dtype, + entries: value.map(ScalarValue::as_list), + }) + } + + /// Returns the map dtype. + #[inline] + pub fn dtype(&self) -> &'a DType { + self.dtype + } + + /// Returns the map type details. + #[inline] + pub fn map_dtype(&self) -> &'a MapDType { + self.dtype + .as_map_opt() + .vortex_expect("MapScalar always has a map dtype") + } + + /// Returns the number of entries, or zero for a null map. + #[inline] + pub fn len(&self) -> usize { + self.entries.map_or(0, <[Option]>::len) + } + + /// Returns whether the map has no entries or is null. + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns whether the entire map scalar is null. + #[inline] + pub fn is_null(&self) -> bool { + self.entries.is_none() + } + + /// Returns the entry at `index`, or `None` when the map is null or the index is out of bounds. + pub fn entry(&self, index: usize) -> Option<(Scalar, Scalar)> { + let values = self.entries?.get(index)?.as_ref()?.as_list(); + Some(self.entry_scalars(values)) + } + + /// Iterates over `(key, value)` entries. A null map yields no entries. + pub fn entries(&self) -> impl Iterator + '_ { + self.entries.into_iter().flatten().map(|entry| { + self.entry_scalars( + entry + .as_ref() + .vortex_expect("map entry is non-null") + .as_list(), + ) + }) + } + + /// Iterates over map keys. A null map yields no keys. + pub fn keys(&self) -> impl Iterator + '_ { + self.entries().map(|(key, _)| key) + } + + /// Iterates over map values. A null map yields no values. + pub fn values(&self) -> impl Iterator + '_ { + self.entries().map(|(_, value)| value) + } + + /// Casts this map scalar to another map dtype. + /// + /// # 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. + pub(crate) fn cast(&self, dtype: &DType) -> VortexResult { + let target = dtype + .as_map_opt() + .ok_or_else(|| vortex_err!("Cannot cast map to {dtype}: target must be a map"))?; + + if target.keys_sorted() && !self.map_dtype().keys_sorted() { + vortex_bail!( + "Cannot cast {} to {dtype}: source does not assert sorted map keys", + self.dtype + ); + } + + let Some(entries) = self.entries else { + return Ok(Scalar::null(dtype.clone())); + }; + + let target_key = target.key_dtype(); + let target_value = target.value_dtype(); + let entries = entries + .iter() + .map(|entry| { + let (key, value) = self.entry_scalars( + entry + .as_ref() + .vortex_expect("map entry is non-null") + .as_list(), + ); + Ok(Some(ScalarValue::Tuple(vec![ + key.cast(&target_key)?.into_value(), + value.cast(&target_value)?.into_value(), + ]))) + }) + .collect::>>()?; + + Scalar::try_new(dtype.clone(), Some(ScalarValue::Tuple(entries))) + } + + fn entry_scalars(&self, values: &[Option]) -> (Scalar, Scalar) { + let key = values.first().vortex_expect("map entry has a key").clone(); + let value = values.get(1).vortex_expect("map entry has a value").clone(); + + // SAFETY: MapScalar only views a Scalar that has passed Scalar::validate, which enforces + // the entry shape and the key/value dtypes. + ( + unsafe { Scalar::new_unchecked(self.map_dtype().key_dtype(), key) }, + unsafe { Scalar::new_unchecked(self.map_dtype().value_dtype(), value) }, + ) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::scalar::Scalar; + use crate::scalar::ScalarValue; + + fn dtype() -> VortexResult { + DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + ) + } + + #[test] + fn entries_and_display() -> VortexResult<()> { + let scalar = Scalar::try_map( + dtype()?, + [ + ( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + ), + ( + Scalar::primitive(2i32, Nullability::NonNullable), + Scalar::null(DType::Utf8(Nullability::Nullable)), + ), + ], + )?; + + let map = scalar.as_map(); + assert_eq!(map.len(), 2); + assert_eq!(map.keys().count(), 2); + assert_eq!(map.values().count(), 2); + assert_eq!( + map.entry(0).unwrap().0, + Scalar::primitive(1i32, Nullability::NonNullable) + ); + assert_eq!(format!("{map}"), "{1i32: \"one\", 2i32: null}"); + + Ok(()) + } + + #[test] + fn null_map_is_distinct_from_empty_map() -> VortexResult<()> { + let dtype = dtype()?; + let empty = Scalar::try_map(dtype.clone(), [])?; + let null = Scalar::null(dtype); + + assert!(empty.as_map().is_empty()); + assert!(!empty.as_map().is_null()); + assert!(null.as_map().is_null()); + assert_ne!(empty, null); + + Ok(()) + } + + #[test] + fn rejects_malformed_entries() -> VortexResult<()> { + let dtype = dtype()?; + let malformed = Scalar::try_new( + dtype, + Some(ScalarValue::Tuple(vec![Some(ScalarValue::Tuple(vec![ + Some(ScalarValue::Primitive(1i32.into())), + ]))])), + ); + + assert!(malformed.is_err()); + Ok(()) + } + + #[test] + fn rejects_null_keys() -> VortexResult<()> { + let malformed = Scalar::try_new( + dtype()?, + Some(ScalarValue::Tuple(vec![Some(ScalarValue::Tuple(vec![ + None, + Some(ScalarValue::Utf8("value".into())), + ]))])), + ); + + assert!(malformed.is_err()); + Ok(()) + } + + #[test] + fn cast_can_drop_a_sortedness_assertion() -> VortexResult<()> { + let source_dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::NonNullable, + )?; + let target_dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + Nullability::Nullable, + )?; + let scalar = Scalar::try_map( + source_dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + )], + )?; + + let cast = scalar.cast(&target_dtype)?; + assert_eq!(cast.dtype(), &target_dtype); + assert_eq!(cast.as_map().entry(0), scalar.as_map().entry(0)); + + Ok(()) + } + + #[test] + fn cast_cannot_create_a_sortedness_assertion() -> VortexResult<()> { + let target_dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + let scalar = Scalar::try_map( + dtype()?, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + )], + )?; + + assert!(scalar.cast(&target_dtype).is_err()); + + Ok(()) + } +} diff --git a/vortex-array/src/scalar/typed_view/mod.rs b/vortex-array/src/scalar/typed_view/mod.rs index 0715027a4ae..4ebb9f533cf 100644 --- a/vortex-array/src/scalar/typed_view/mod.rs +++ b/vortex-array/src/scalar/typed_view/mod.rs @@ -20,6 +20,7 @@ mod bool; mod decimal; mod extension; mod list; +mod map; mod primitive; mod struct_; mod union; @@ -31,6 +32,7 @@ pub use bool::*; pub use decimal::*; pub use extension::*; pub use list::*; +pub use map::*; pub use primitive::*; pub use struct_::*; pub use union::*; diff --git a/vortex-array/src/scalar/validate.rs b/vortex-array/src/scalar/validate.rs index f81147376b8..d6dba2c34b4 100644 --- a/vortex-array/src/scalar/validate.rs +++ b/vortex-array/src/scalar/validate.rs @@ -102,6 +102,37 @@ impl Scalar { })?; } } + DType::Map(map, _) => { + let ScalarValue::Tuple(entries) = value else { + vortex_bail!("map dtype expected Tuple value, got {value}"); + }; + let key_dtype = map.key_dtype(); + let value_dtype = map.value_dtype(); + + for (index, entry) in entries.iter().enumerate() { + let entry = entry.as_ref().ok_or_else(|| { + vortex_error::vortex_err!("map entry at index {index} cannot be null") + })?; + let ScalarValue::Tuple(values) = entry else { + vortex_bail!( + "map entry at index {index} expected Tuple value, got {entry}" + ); + }; + vortex_ensure_eq!( + values.len(), + 2, + "map entry at index {index} expected 2 values, got {}", + values.len(), + ); + + Self::validate(&key_dtype, values[0].as_ref()).map_err(|error| { + vortex_error::vortex_err!("map key at entry {index}: {error}") + })?; + Self::validate(&value_dtype, values[1].as_ref()).map_err(|error| { + vortex_error::vortex_err!("map value at entry {index}: {error}") + })?; + } + } DType::Struct(fields, _) => { let ScalarValue::Tuple(values) = value else { vortex_bail!("struct dtype expected Tuple value, got {value}"); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index 70e98639b54..0452f4a3156 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -217,7 +217,7 @@ fn compare_arrays( DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { nested::compare_nested(lhs, rhs, op, nullability, ctx) } - DType::Union(..) | DType::Variant(_) | DType::Extension(_) => { + DType::Map(..) | DType::Union(..) | DType::Variant(_) | DType::Extension(_) => { vortex_bail!("compare is not supported for dtype {}", lhs.dtype()) } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs index 4ccd1fbe363..e174ce17fa7 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs @@ -214,7 +214,7 @@ fn build_values_comparator( let rhs = rhs.clone().execute::(ctx)?; build_comparator(lhs.storage_array(), rhs.storage_array(), ctx)? } - DType::Union(..) | DType::Variant(_) => { + DType::Map(..) | DType::Union(..) | DType::Variant(_) => { vortex_bail!("compare is not supported for dtype {}", lhs.dtype()) } }) diff --git a/vortex-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index e227b401633..40d46025bd7 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -546,6 +546,9 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef { DataType::ListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable), + DataType::Map(..) => { + vortex_bail!("Arrow MapArray conversion is not yet supported") + } DataType::Null => Self::from_arrow(as_null_array(array), nullable), DataType::Timestamp(u, _) => match u { ArrowTimeUnit::Second => { diff --git a/vortex-arrow/src/dtype.rs b/vortex-arrow/src/dtype.rs index a87c931f3df..9c8be4840ae 100644 --- a/vortex-arrow/src/dtype.rs +++ b/vortex-arrow/src/dtype.rs @@ -38,6 +38,8 @@ use vortex_array::extension::datetime::Timestamp; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; @@ -224,6 +226,35 @@ impl TryFromArrowType<(&DataType, Nullability)> for DType { nullability, ), DataType::Struct(f) => DType::Struct(StructFields::try_from_arrow(f)?, nullability), + DataType::Map(entries, keys_sorted) => { + vortex_ensure!( + !entries.is_nullable(), + "Arrow map entries field must be non-nullable" + ); + let DataType::Struct(fields) = entries.data_type() else { + vortex_bail!( + "Arrow map entries field must have Struct type, got {:?}", + entries.data_type() + ); + }; + vortex_ensure_eq!( + fields.len(), + 2, + InvalidArgument: "Arrow map entries struct must contain exactly two fields" + ); + let key = &fields[0]; + let value = &fields[1]; + vortex_ensure!( + !key.is_nullable(), + "Arrow map key field must be non-nullable" + ); + DType::map( + Self::try_from_arrow(key.as_ref())?, + Self::try_from_arrow(value.as_ref())?, + *keys_sorted, + nullability, + )? + } DataType::Dictionary(_, value_type) => { Self::try_from_arrow((value_type.as_ref(), nullability))? } @@ -368,6 +399,17 @@ pub(crate) fn to_data_type_naive(dtype: &DType) -> VortexResult { )), *size as i32, ), + DType::Map(map_dtype, _) => { + let key = Field::new("key", to_data_type_naive(&map_dtype.key_dtype())?, false); + let value_dtype = map_dtype.value_dtype(); + let value = Field::new( + "value", + to_data_type_naive(&value_dtype)?, + value_dtype.is_nullable(), + ); + let entries = Field::new_struct("entries", Fields::from(vec![key, value]), false); + DataType::Map(FieldRef::new(entries), map_dtype.keys_sorted()) + } DType::Struct(struct_dtype, _) => { let mut fields = Vec::with_capacity(struct_dtype.names().len()); for (field_name, field_dt) in struct_dtype.names().iter().zip(struct_dtype.fields()) { @@ -638,6 +680,92 @@ mod test { assert_eq!(original_dtype, roundtripped_dtype); } + #[test] + fn map_dtype_roundtrip_uses_conventional_export_names_and_positional_import() -> VortexResult<()> + { + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + + let arrow = dtype.to_arrow_dtype()?; + let DataType::Map(entries, keys_sorted) = &arrow else { + panic!("expected Map, got {arrow:?}"); + }; + assert!(*keys_sorted); + assert_eq!(entries.name(), "entries"); + assert!(!entries.is_nullable()); + let DataType::Struct(fields) = entries.data_type() else { + panic!("expected map entries to be a struct"); + }; + assert_eq!(fields[0].name(), "key"); + assert!(!fields[0].is_nullable()); + assert_eq!(fields[1].name(), "value"); + assert!(fields[1].is_nullable()); + assert_eq!( + DType::try_from_arrow((&arrow, Nullability::Nullable))?, + dtype + ); + + let positional = DataType::Map( + Arc::new(Field::new_struct( + "anything", + Fields::from(vec![ + Field::new("first", DataType::Int32, false), + Field::new("second", DataType::Utf8, true), + ]), + false, + )), + false, + ); + assert_eq!( + DType::try_from_arrow((&positional, Nullability::NonNullable))?, + DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + Nullability::NonNullable, + )? + ); + + Ok(()) + } + + #[test] + fn map_dtype_import_rejects_invalid_arrow_shape() { + let invalid_entries = [ + Field::new_struct( + "entries", + Fields::from(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, true), + ]), + true, + ), + Field::new("entries", DataType::Int32, false), + Field::new_struct( + "entries", + Fields::from(vec![Field::new("key", DataType::Int32, false)]), + false, + ), + Field::new_struct( + "entries", + Fields::from(vec![ + Field::new("key", DataType::Int32, true), + Field::new("value", DataType::Utf8, true), + ]), + false, + ), + ]; + + for entries in invalid_entries { + let data_type = DataType::Map(Arc::new(entries), false); + assert!(DType::try_from_arrow((&data_type, Nullability::NonNullable)).is_err()); + } + } + // Regression test for https://github.com/vortex-data/vortex/issues/8346: unsupported Arrow // types must return an error instead of panicking with `unimplemented!`. #[rstest] diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index 854a713f746..cae67f3143b 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -191,8 +191,8 @@ pub(crate) fn execute_arrow_naive( dt @ (DataType::Date32 | DataType::Date64) => to_arrow_date(array, dt, ctx), dt @ (DataType::Time32(_) | DataType::Time64(_)) => to_arrow_time(array, dt, ctx), dt @ DataType::Timestamp(..) => to_arrow_timestamp(array, dt, ctx), + DataType::Map(..) => vortex_bail!("Arrow MapArray conversion is not yet supported"), DataType::FixedSizeBinary(_) - | DataType::Map(..) | DataType::Duration(_) | DataType::Interval(_) | DataType::Union(..) => { diff --git a/vortex-arrow/src/scalar.rs b/vortex-arrow/src/scalar.rs index aeac56f43d3..1b8f2c01ed1 100644 --- a/vortex-arrow/src/scalar.rs +++ b/vortex-arrow/src/scalar.rs @@ -7,6 +7,10 @@ use std::sync::Arc; use arrow_array::Scalar as ArrowScalar; use arrow_array::*; +use arrow_buffer::NullBuffer; +use arrow_buffer::OffsetBuffer; +use arrow_schema::Field; +use arrow_schema::Fields; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::extension::datetime::AnyTemporal; @@ -17,6 +21,7 @@ use vortex_array::scalar::BoolScalar; use vortex_array::scalar::DecimalScalar; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::ExtScalar; +use vortex_array::scalar::MapScalar; use vortex_array::scalar::PrimitiveScalar; use vortex_array::scalar::Scalar; use vortex_array::scalar::Utf8Scalar; @@ -24,6 +29,8 @@ use vortex_error::VortexError; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use crate::dtype::to_data_type_naive; + /// Arrow represents scalars as single-element arrays. This constant is the length of those arrays. const SCALAR_ARRAY_LEN: usize = 1; @@ -70,6 +77,7 @@ impl ToArrowDatum for Scalar { DType::Binary(_) => binary_to_arrow(value.as_binary()), DType::List(..) => unimplemented!("list scalar conversion"), DType::FixedSizeList(..) => unimplemented!("fixed-size list scalar conversion"), + DType::Map(..) => map_to_arrow(value.as_map()), DType::Struct(..) => unimplemented!("struct scalar conversion"), DType::Union(..) => unimplemented!("union scalar conversion"), DType::Variant(_) => unimplemented!("Variant scalar conversion"), @@ -126,6 +134,69 @@ fn binary_to_arrow(scalar: BinaryScalar<'_>) -> Result, VortexErr value_to_arrow_scalar!(scalar.value(), BinaryViewArray) } +/// Convert a [`MapScalar`] to an Arrow [`Datum`]. +fn map_to_arrow(scalar: MapScalar<'_>) -> Result, VortexError> { + let map_dtype = scalar.map_dtype(); + let key_dtype = map_dtype.key_dtype(); + let value_dtype = map_dtype.value_dtype(); + let key_field = Field::new("key", to_data_type_naive(&key_dtype)?, false); + let value_field = Field::new( + "value", + to_data_type_naive(&value_dtype)?, + value_dtype.is_nullable(), + ); + let fields = Fields::from(vec![key_field, value_field]); + + let entries = scalar.entries().collect::>(); + let keys = entries + .iter() + .map(|(key, _)| key.to_arrow_datum()) + .collect::, _>>()?; + let values = entries + .iter() + .map(|(_, value)| value.to_arrow_datum()) + .collect::, _>>()?; + + let key_array = concat_scalar_arrays(&keys, &key_dtype)?; + let value_array = concat_scalar_arrays(&values, &value_dtype)?; + let entries = StructArray::new(fields.clone(), vec![key_array, value_array], None); + + let entries_len = entries.len(); + let entries_len = i32::try_from(entries_len).map_err(|_| { + vortex_err!( + "Cannot convert map scalar with {entries_len} entries to Arrow: MapArray offsets are i32" + ) + })?; + let offsets = OffsetBuffer::new(vec![0_i32, entries_len].into()); + let entries_field = Arc::new(Field::new_struct("entries", fields, false)); + let nulls = scalar + .is_null() + .then(|| NullBuffer::new_null(SCALAR_ARRAY_LEN)); + let map = MapArray::try_new( + entries_field, + offsets, + entries, + nulls, + map_dtype.keys_sorted(), + )?; + Ok(Arc::new(ArrowScalar::new(map))) +} + +fn concat_scalar_arrays( + scalars: &[Arc], + dtype: &DType, +) -> Result { + if scalars.is_empty() { + return Ok(new_empty_array(&to_data_type_naive(dtype)?)); + } + + let arrays = scalars + .iter() + .map(|scalar| scalar.get().0) + .collect::>(); + Ok(arrow_select::concat::concat(&arrays)?) +} + /// Convert an [`ExtScalar`] to an Arrow [`Datum`]. /// /// Currently only temporal extension types (timestamps, dates, and times) are supported. @@ -199,6 +270,10 @@ fn extension_to_arrow(scalar: ExtScalar<'_>) -> Result, VortexErr mod tests { use std::sync::Arc; + use arrow_array::Array; + use arrow_array::Int32Array; + use arrow_array::MapArray; + use arrow_array::StringViewArray; use rstest::rstest; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; @@ -419,6 +494,57 @@ mod tests { assert!(result.is_ok()); } + #[test] + fn test_map_scalar_to_arrow() -> VortexResult<()> { + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + )?; + let scalar = Scalar::try_map( + dtype, + [ + ( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::utf8("one", Nullability::Nullable), + ), + ( + Scalar::primitive(2i32, Nullability::NonNullable), + Scalar::null(DType::Utf8(Nullability::Nullable)), + ), + ], + )?; + + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + let map = array + .as_any() + .downcast_ref::() + .expect("map scalar should convert to MapArray"); + assert_eq!(map.len(), 1); + assert_eq!(map.value_offsets(), &[0, 2]); + assert!(map.is_valid(0)); + assert_eq!( + map.keys() + .as_any() + .downcast_ref::() + .expect("map key array should be Int32") + .values(), + &[1, 2] + ); + let values = map + .values() + .as_any() + .downcast_ref::() + .expect("map value array should be StringView"); + assert_eq!(values.value(0), "one"); + assert!(values.is_null(1)); + + Ok(()) + } + #[test] #[should_panic(expected = "struct scalar conversion")] fn test_struct_scalar_to_arrow_todo() { diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index e95feead1f7..87a08a548ea 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -247,6 +247,16 @@ impl ArrowSession { nullability.is_nullable(), )) } + DType::Map(map_dtype, nullability) => { + let key = self.to_arrow_field("key", &map_dtype.key_dtype())?; + let value = self.to_arrow_field("value", &map_dtype.value_dtype())?; + let entries = Field::new_struct("entries", Fields::from(vec![key, value]), false); + Ok(Field::new( + name, + DataType::Map(Arc::new(entries), map_dtype.keys_sorted()), + nullability.is_nullable(), + )) + } DType::Struct(fields, nullability) => { let arrow_fields = Fields::from_iter( fields @@ -346,6 +356,32 @@ impl ArrowSession { *size as u32, nullability, ), + DataType::Map(entries, keys_sorted) => { + vortex_ensure!( + !entries.is_nullable(), + "Arrow map entries field must be non-nullable" + ); + let DataType::Struct(fields) = entries.data_type() else { + vortex_bail!( + "Arrow map entries field must have Struct type, got {:?}", + entries.data_type() + ); + }; + vortex_ensure!( + fields.len() == 2, + "Arrow map entries struct must contain exactly two fields" + ); + vortex_ensure!( + !fields[0].is_nullable(), + "Arrow map key field must be non-nullable" + ); + DType::map( + self.from_arrow_field(fields[0].as_ref())?, + self.from_arrow_field(fields[1].as_ref())?, + *keys_sorted, + nullability, + )? + } DataType::Struct(fields) => { let entries = fields .iter() @@ -584,6 +620,9 @@ impl ArrowSession { let validity = nulls(list.nulls(), field.is_nullable())?; Ok(ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array()) } + DataType::Map(..) => { + vortex_bail!("Arrow MapArray conversion is not yet supported") + } _ => ArrayRef::from_arrow(array.as_ref(), field.is_nullable()), } } @@ -709,6 +748,40 @@ mod tests { Ok(()) } + #[test] + fn schema_roundtrip_preserves_map_uuid_fields() -> VortexResult<()> { + let session = ArrowSession::default(); + let map = DType::map( + uuid_dtype(false), + uuid_dtype(true), + true, + Nullability::Nullable, + )?; + let dtype = DType::Struct( + StructFields::from_iter([(FieldName::from("ids"), map)]), + Nullability::NonNullable, + ); + + let schema = session.to_arrow_schema(&dtype)?; + let field = schema.field(0); + let DataType::Map(entries, keys_sorted) = field.data_type() else { + panic!("expected Map, got {:?}", field.data_type()); + }; + assert!(*keys_sorted); + assert_eq!(entries.name(), "entries"); + assert!(!entries.is_nullable()); + let DataType::Struct(fields) = entries.data_type() else { + panic!("expected map entries struct, got {:?}", entries.data_type()); + }; + assert!(has_valid_extension_type::(&fields[0])); + assert!(has_valid_extension_type::(&fields[1])); + assert!(!fields[0].is_nullable()); + assert!(fields[1].is_nullable()); + + assert_eq!(session.from_arrow_schema(&schema)?, dtype); + Ok(()) + } + #[test] fn to_arrow_schema_struct_of_struct_uuid() -> VortexResult<()> { let session = ArrowSession::default(); diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index ec262b479d9..d1606374872 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -125,6 +125,9 @@ impl TryToDataFusion for Scalar { dtype @ DType::FixedSizeList(..) => vortex_bail!( "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" ), + dtype @ DType::Map(..) => vortex_bail!( + "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" + ), DType::Struct(..) => struct_to_df(self)?, dtype @ DType::Union(..) => vortex_bail!( "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 8e08d531844..36f3b5f4cd6 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::Schema; @@ -180,6 +182,52 @@ fn calculate_physical_field_type( } } + // Map field names and child metadata come from the reference schema, while child + // types are recursively reconciled so nested extension metadata is preserved. + DataType::Map(logical_entries, keys_sorted) => { + let DType::Map(map_dtype, _) = dtype else { + return Err(exec_datafusion_err!( + "Failed to convert dtype to arrow: Vortex DType is {dtype} which is not compatible with {logical_type}" + )); + }; + let DataType::Struct(logical_fields) = logical_entries.data_type() else { + return Err(exec_datafusion_err!( + "Failed to convert dtype to arrow: Arrow Map entries must be a Struct, got {:?}", + logical_entries.data_type() + )); + }; + if logical_fields.len() != 2 { + return Err(exec_datafusion_err!( + "Failed to convert dtype to arrow: Arrow Map entries must contain exactly two fields" + )); + } + + let key = Field::new( + logical_fields[0].name(), + calculate_physical_field_type( + &map_dtype.key_dtype(), + logical_fields[0].data_type(), + arrow_session, + )?, + false, + ) + .with_metadata(logical_fields[0].metadata().clone()); + let value = Field::new( + logical_fields[1].name(), + calculate_physical_field_type( + &map_dtype.value_dtype(), + logical_fields[1].data_type(), + arrow_session, + )?, + logical_fields[1].is_nullable(), + ) + .with_metadata(logical_fields[1].metadata().clone()); + let entries = Field::new_struct(logical_entries.name(), vec![key, value], false) + .with_metadata(logical_entries.metadata().clone()); + + DataType::Map(Arc::new(entries), *keys_sorted) + } + // For list view types, recursively check the element type DataType::ListView(logical_elem) | DataType::LargeListView(logical_elem) => { if let DType::List(elem_dtype, _) = dtype { @@ -440,6 +488,65 @@ mod tests { } } + #[test] + fn test_map_schema_conversion_preserves_reference_fields() { + let key = Field::new("custom_key", DataType::Int32, false) + .with_metadata([("key_metadata".to_owned(), "key_value".to_owned())].into()); + let value = Field::new("custom_value", DataType::Utf8, true) + .with_metadata([("value_metadata".to_owned(), "value_value".to_owned())].into()); + let entries = Field::new_struct("custom_entries", vec![key, value], false) + .with_metadata([("entries_metadata".to_owned(), "entries_value".to_owned())].into()); + let logical_schema = Schema::new(vec![Field::new( + "map_col", + DataType::Map(Arc::new(entries), true), + true, + )]); + let dtype = DType::Struct( + StructFields::from_iter([( + "map_col", + DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + true, + Nullability::Nullable, + ) + .unwrap(), + )]), + Nullability::NonNullable, + ); + + let physical_schema = + calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()).unwrap(); + let field = physical_schema.field(0); + assert!(field.is_nullable()); + let DataType::Map(entries, keys_sorted) = field.data_type() else { + panic!("expected Map type, got {:?}", field.data_type()); + }; + assert!(*keys_sorted); + assert_eq!(entries.name(), "custom_entries"); + assert_eq!( + entries.metadata().get("entries_metadata"), + Some(&"entries_value".to_owned()) + ); + let DataType::Struct(fields) = entries.data_type() else { + panic!("expected map entries struct, got {:?}", entries.data_type()); + }; + assert_eq!(fields[0].name(), "custom_key"); + assert_eq!(fields[0].data_type(), &DataType::Int32); + assert!(!fields[0].is_nullable()); + assert_eq!( + fields[0].metadata().get("key_metadata"), + Some(&"key_value".to_owned()) + ); + assert_eq!(fields[1].name(), "custom_value"); + assert_eq!(fields[1].data_type(), &DataType::Utf8); + assert!(fields[1].is_nullable()); + assert_eq!( + fields[1].metadata().get("value_metadata"), + Some(&"value_value".to_owned()) + ); + } + #[test] fn test_non_struct_dtype_error() { // Test that non-struct DType produces an error diff --git a/vortex-duckdb/src/convert/dtype.rs b/vortex-duckdb/src/convert/dtype.rs index dd26dfbcbba..6b0568d4e5c 100644 --- a/vortex-duckdb/src/convert/dtype.rs +++ b/vortex-duckdb/src/convert/dtype.rs @@ -244,6 +244,7 @@ impl TryFrom<&DType> for LogicalType { DType::Struct(struct_type, _) => { return LogicalType::try_from(struct_type); } + DType::Map(..) => vortex_bail!("Vortex Map isn't supported"), // TODO(connor): Union DType::Union(..) => vortex_bail!("Vortex Union isn't supported"), DType::Variant(_) => vortex_bail!("Vortex Variant array aren't supported"), diff --git a/vortex-duckdb/src/convert/scalar.rs b/vortex-duckdb/src/convert/scalar.rs index 7ad12e89ca3..c9cc6100a1d 100644 --- a/vortex-duckdb/src/convert/scalar.rs +++ b/vortex-duckdb/src/convert/scalar.rs @@ -82,6 +82,7 @@ impl ToDuckDBScalar for Scalar { DType::FixedSizeList(..) => { vortex_bail!("Vortex FixedSizeList scalars aren't supported") } + DType::Map(..) => vortex_bail!("Vortex Map scalars aren't supported"), DType::Variant(_) => vortex_bail!("Vortex Variant scalars aren't supported"), DType::Struct(..) => vortex_bail!("Vortex Struct scalars aren't supported"), // TODO(connor): Union diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index 74a6f56fb68..2569340be3f 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -74,6 +74,7 @@ impl From<&DType> for vx_dtype_variant { DType::List(..) => vx_dtype_variant::DTYPE_LIST, DType::FixedSizeList(..) => vx_dtype_variant::DTYPE_FIXED_SIZE_LIST, DType::Struct(..) => vx_dtype_variant::DTYPE_STRUCT, + DType::Map(..) => vortex_panic!("Map is not supported in FFI yet"), DType::Union(..) => vortex_panic!("Union is not supported in FFI yet"), DType::Variant(_) => vortex_panic!("Variant is not supported in FFI yet"), DType::Extension(_) => vx_dtype_variant::DTYPE_EXTENSION, diff --git a/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs b/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs index fb3b19513de..7de10c7a70d 100644 --- a/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs +++ b/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs @@ -74,6 +74,13 @@ table Union { nullable: bool; } +table Map { + key_type: DType; + value_type: DType; + keys_sorted: bool; + nullable: bool; +} + union Type { Null = 1, Bool = 2, @@ -87,6 +94,7 @@ union Type { FixedSizeList = 10, // This is after `Extension` for backwards compatibility. Variant = 11, Union = 12, + Map = 13, } table DType { diff --git a/vortex-flatbuffers/src/generated/dtype.rs b/vortex-flatbuffers/src/generated/dtype.rs index 56536ad0b94..44470d5ed33 100644 --- a/vortex-flatbuffers/src/generated/dtype.rs +++ b/vortex-flatbuffers/src/generated/dtype.rs @@ -126,10 +126,10 @@ impl ::flatbuffers::SimpleToVerifyInSlice for PType {} #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] pub const ENUM_MIN_TYPE: u8 = 0; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] -pub const ENUM_MAX_TYPE: u8 = 12; +pub const ENUM_MAX_TYPE: u8 = 13; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_TYPE: [Type; 13] = [ +pub const ENUM_VALUES_TYPE: [Type; 14] = [ Type::NONE, Type::Null, Type::Bool, @@ -143,6 +143,7 @@ pub const ENUM_VALUES_TYPE: [Type; 13] = [ Type::FixedSizeList, Type::Variant, Type::Union, + Type::Map, ]; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] @@ -163,9 +164,10 @@ impl Type { pub const FixedSizeList: Self = Self(10); pub const Variant: Self = Self(11); pub const Union: Self = Self(12); + pub const Map: Self = Self(13); pub const ENUM_MIN: u8 = 0; - pub const ENUM_MAX: u8 = 12; + pub const ENUM_MAX: u8 = 13; pub const ENUM_VALUES: &'static [Self] = &[ Self::NONE, Self::Null, @@ -180,6 +182,7 @@ impl Type { Self::FixedSizeList, Self::Variant, Self::Union, + Self::Map, ]; /// Returns the variant's name or "" if unknown. pub fn variant_name(self) -> Option<&'static str> { @@ -197,6 +200,7 @@ impl Type { Self::FixedSizeList => Some("FixedSizeList"), Self::Variant => Some("Variant"), Self::Union => Some("Union"), + Self::Map => Some("Map"), _ => None, } } @@ -1608,6 +1612,153 @@ impl ::core::fmt::Debug for Union<'_> { ds.finish() } } +pub enum MapOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct Map<'a> { + pub _tab: ::flatbuffers::Table<'a>, +} + +impl<'a> ::flatbuffers::Follow<'a> for Map<'a> { + type Inner = Map<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } + } +} + +impl<'a> Map<'a> { + pub const VT_KEY_TYPE: ::flatbuffers::VOffsetT = 4; + pub const VT_VALUE_TYPE: ::flatbuffers::VOffsetT = 6; + pub const VT_KEYS_SORTED: ::flatbuffers::VOffsetT = 8; + pub const VT_NULLABLE: ::flatbuffers::VOffsetT = 10; + + #[inline] + pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { + Map { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args MapArgs<'args> + ) -> ::flatbuffers::WIPOffset> { + let mut builder = MapBuilder::new(_fbb); + if let Some(x) = args.value_type { builder.add_value_type(x); } + if let Some(x) = args.key_type { builder.add_key_type(x); } + builder.add_nullable(args.nullable); + builder.add_keys_sorted(args.keys_sorted); + builder.finish() + } + + + #[inline] + pub fn key_type(&self) -> Option> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset>(Map::VT_KEY_TYPE, None)} + } + #[inline] + pub fn value_type(&self) -> Option> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset>(Map::VT_VALUE_TYPE, None)} + } + #[inline] + pub fn keys_sorted(&self) -> bool { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Map::VT_KEYS_SORTED, Some(false)).unwrap()} + } + #[inline] + pub fn nullable(&self) -> bool { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Map::VT_NULLABLE, Some(false)).unwrap()} + } +} + +impl ::flatbuffers::Verifiable for Map<'_> { + #[inline] + fn run_verifier( + v: &mut ::flatbuffers::Verifier, pos: usize + ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { + v.visit_table(pos)? + .visit_field::<::flatbuffers::ForwardsUOffset>("key_type", Self::VT_KEY_TYPE, false)? + .visit_field::<::flatbuffers::ForwardsUOffset>("value_type", Self::VT_VALUE_TYPE, false)? + .visit_field::("keys_sorted", Self::VT_KEYS_SORTED, false)? + .visit_field::("nullable", Self::VT_NULLABLE, false)? + .finish(); + Ok(()) + } +} +pub struct MapArgs<'a> { + pub key_type: Option<::flatbuffers::WIPOffset>>, + pub value_type: Option<::flatbuffers::WIPOffset>>, + pub keys_sorted: bool, + pub nullable: bool, +} +impl<'a> Default for MapArgs<'a> { + #[inline] + fn default() -> Self { + MapArgs { + key_type: None, + value_type: None, + keys_sorted: false, + nullable: false, + } + } +} + +pub struct MapBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { + fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, + start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, +} +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> MapBuilder<'a, 'b, A> { + #[inline] + pub fn add_key_type(&mut self, key_type: ::flatbuffers::WIPOffset>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset>(Map::VT_KEY_TYPE, key_type); + } + #[inline] + pub fn add_value_type(&mut self, value_type: ::flatbuffers::WIPOffset>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset>(Map::VT_VALUE_TYPE, value_type); + } + #[inline] + pub fn add_keys_sorted(&mut self, keys_sorted: bool) { + self.fbb_.push_slot::(Map::VT_KEYS_SORTED, keys_sorted, false); + } + #[inline] + pub fn add_nullable(&mut self, nullable: bool) { + self.fbb_.push_slot::(Map::VT_NULLABLE, nullable, false); + } + #[inline] + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> MapBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + MapBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> ::flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + ::flatbuffers::WIPOffset::new(o.value()) + } +} + +impl ::core::fmt::Debug for Map<'_> { + fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { + let mut ds = f.debug_struct("Map"); + ds.field("key_type", &self.key_type()); + ds.field("value_type", &self.value_type()); + ds.field("keys_sorted", &self.keys_sorted()); + ds.field("nullable", &self.nullable()); + ds.finish() + } +} pub enum DTypeOffset {} #[derive(Copy, Clone, PartialEq)] @@ -1837,6 +1988,21 @@ impl<'a> DType<'a> { } } + #[inline] + #[allow(non_snake_case)] + pub fn type__as_map(&self) -> Option> { + if self.type_type() == Type::Map { + self.type_().map(|t| { + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + unsafe { Map::init_from_table(t) } + }) + } else { + None + } + } + } impl ::flatbuffers::Verifiable for DType<'_> { @@ -1859,6 +2025,7 @@ impl ::flatbuffers::Verifiable for DType<'_> { Type::FixedSizeList => v.verify_union_variant::<::flatbuffers::ForwardsUOffset>("Type::FixedSizeList", pos), Type::Variant => v.verify_union_variant::<::flatbuffers::ForwardsUOffset>("Type::Variant", pos), Type::Union => v.verify_union_variant::<::flatbuffers::ForwardsUOffset>("Type::Union", pos), + Type::Map => v.verify_union_variant::<::flatbuffers::ForwardsUOffset>("Type::Map", pos), _ => Ok(()), } })? @@ -1997,6 +2164,13 @@ impl ::core::fmt::Debug for DType<'_> { ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") } }, + Type::Map => { + if let Some(x) = self.type__as_map() { + ds.field("type_", &x) + } else { + ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") + } + }, _ => { let x: Option<()> = None; ds.field("type_", &x) diff --git a/vortex-flatbuffers/src/generated/message.rs b/vortex-flatbuffers/src/generated/message.rs index 9cc48fe6be6..b3e4be76ef7 100644 --- a/vortex-flatbuffers/src/generated/message.rs +++ b/vortex-flatbuffers/src/generated/message.rs @@ -2,8 +2,8 @@ // @generated extern crate alloc; -use crate::dtype::*; use crate::array::*; +use crate::dtype::*; #[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; diff --git a/vortex-proto/proto/dtype.proto b/vortex-proto/proto/dtype.proto index 11feef63a8c..3993962824f 100644 --- a/vortex-proto/proto/dtype.proto +++ b/vortex-proto/proto/dtype.proto @@ -81,6 +81,13 @@ message Union { bool nullable = 4; } +message Map { + DType key_type = 1; + DType value_type = 2; + bool keys_sorted = 3; + bool nullable = 4; +} + message DType { oneof dtype_type { Null null = 1; @@ -95,6 +102,7 @@ message DType { FixedSizeList fixed_size_list = 10; // This is after `Extension` for backwards compatibility. Variant variant = 11; Union union = 12; + Map map = 13; } } diff --git a/vortex-proto/src/generated/vortex.dtype.rs b/vortex-proto/src/generated/vortex.dtype.rs index 16188687ee5..37ffec0a426 100644 --- a/vortex-proto/src/generated/vortex.dtype.rs +++ b/vortex-proto/src/generated/vortex.dtype.rs @@ -84,8 +84,22 @@ pub struct Union { pub nullable: bool, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct Map { + #[prost(message, optional, boxed, tag = "1")] + pub key_type: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub value_type: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(bool, tag = "3")] + pub keys_sorted: bool, + #[prost(bool, tag = "4")] + pub nullable: bool, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct DType { - #[prost(oneof = "d_type::DtypeType", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12")] + #[prost( + oneof = "d_type::DtypeType", + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13" + )] pub dtype_type: ::core::option::Option, } /// Nested message and enum types in `DType`. @@ -117,6 +131,8 @@ pub mod d_type { Variant(super::Variant), #[prost(message, tag = "12")] Union(super::Union), + #[prost(message, tag = "13")] + Map(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/vortex-python/src/dtype/mod.rs b/vortex-python/src/dtype/mod.rs index ce18ceb6d83..8f391b89698 100644 --- a/vortex-python/src/dtype/mod.rs +++ b/vortex-python/src/dtype/mod.rs @@ -134,6 +134,9 @@ impl PyDType { DType::Binary(..) => Self::with_subclass(py, dtype, PyBinaryDType), DType::List(..) => Self::with_subclass(py, dtype, PyListDType), DType::FixedSizeList(..) => Self::with_subclass(py, dtype, PyFixedSizeListDType), + DType::Map(..) => Err(PyValueError::new_err( + "Map dtypes are not supported in Python", + )), DType::Struct(..) => Self::with_subclass(py, dtype, PyStructDType), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => Err(PyValueError::new_err( diff --git a/vortex-python/src/python_repr.rs b/vortex-python/src/python_repr.rs index 56819184dca..1e804b046cc 100644 --- a/vortex-python/src/python_repr.rs +++ b/vortex-python/src/python_repr.rs @@ -80,6 +80,7 @@ impl Display for DTypePythonRepr<'_> { size, n.python_repr() ), + DType::Map(..) => write!(f, "{dtype}"), DType::Struct(st, n) => write!( f, "struct({{{}}}, nullable={})", diff --git a/vortex-python/src/scalar/into_py.rs b/vortex-python/src/scalar/into_py.rs index a3b81ab6b2e..44931f16145 100644 --- a/vortex-python/src/scalar/into_py.rs +++ b/vortex-python/src/scalar/into_py.rs @@ -83,6 +83,9 @@ impl<'py> IntoPyObject<'py> for PyVortex<&'_ Scalar> { DType::List(..) | DType::FixedSizeList(..) => { PyVortex(self.0.as_list()).into_pyobject(py) } + DType::Map(..) => Err(PyValueError::new_err( + "Map scalars are not supported in Python", + )), DType::Struct(..) => PyVortex(self.0.as_struct()).into_pyobject(py), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => Err(PyValueError::new_err( diff --git a/vortex-python/src/scalar/mod.rs b/vortex-python/src/scalar/mod.rs index eea66dccc52..f032fa85c0e 100644 --- a/vortex-python/src/scalar/mod.rs +++ b/vortex-python/src/scalar/mod.rs @@ -116,6 +116,9 @@ impl PyScalar { // of "fixed-size" only applies to full arrays, not scalars. Self::with_subclass(py, scalar, PyListScalar) } + DType::Map(..) => Err(PyValueError::new_err( + "Map scalars are not supported in Python", + )), DType::Struct(..) => Self::with_subclass(py, scalar, PyStructScalar), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => Err(PyValueError::new_err( diff --git a/vortex-row/src/codec.rs b/vortex-row/src/codec.rs index 5a0f62f974e..58fbc3ccbbb 100644 --- a/vortex-row/src/codec.rs +++ b/vortex-row/src/codec.rs @@ -224,9 +224,9 @@ pub(crate) fn row_width_for_dtype(dtype: &DType) -> VortexResult { } Ok(RowWidth::Fixed(total)) } - DType::List(..) => { + DType::List(..) | DType::Map(..) => { vortex_bail!( - "row encoding does not support variable-size List arrays (no well-defined ordering)" + "row encoding does not support variable-size List or Map arrays (no well-defined ordering)" ) } DType::Variant(_) => {