diff --git a/vortex-array/src/scalar/scalar_impl.rs b/vortex-array/src/scalar/scalar_impl.rs index 4d9f5a3a867..eaf83cec3bb 100644 --- a/vortex-array/src/scalar/scalar_impl.rs +++ b/vortex-array/src/scalar/scalar_impl.rs @@ -451,7 +451,8 @@ 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, + // A map compares as the list of its `{key, value}` entry structs. + DType::Map(map_dtype, _) => partial_cmp_list_values(&map_dtype.entries_dtype(), lhs, rhs), DType::Extension(ext_dtype) => { partial_cmp_tuple_values(ext_dtype.storage_dtype(), lhs, rhs) } @@ -497,9 +498,11 @@ fn partial_cmp_struct_values( #[cfg(test)] mod tests { + use std::cmp::Ordering; use std::sync::Arc; use rstest::rstest; + use vortex_error::VortexResult; use crate::dtype::DType; use crate::dtype::Nullability; @@ -518,6 +521,61 @@ mod tests { } } + fn map_dtype(nullability: Nullability) -> VortexResult { + DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + nullability, + ) + } + + fn map_scalar(entries: Vec<(i32, Option<&str>)>) -> VortexResult { + Scalar::try_map( + map_dtype(Nullability::Nullable)?, + entries.into_iter().map(|(key, value)| { + ( + i32_scalar(key), + match value { + Some(value) => Scalar::utf8(value, Nullability::Nullable), + None => Scalar::null(DType::Utf8(Nullability::Nullable)), + }, + ) + }), + ) + } + + /// Maps order entry-wise, then by entry count. + #[rstest] + #[case(vec![(1, Some("a"))], vec![(1, Some("a"))], Ordering::Equal)] + #[case(vec![(1, Some("a"))], vec![(1, Some("b"))], Ordering::Less)] + #[case(vec![(2, Some("a"))], vec![(1, Some("z"))], Ordering::Greater)] + #[case(vec![(1, Some("a"))], vec![(1, Some("a")), (2, Some("b"))], Ordering::Less)] + #[case(vec![], vec![], Ordering::Equal)] + #[case(vec![], vec![(1, Some("a"))], Ordering::Less)] + #[case(vec![(1, Some("a"))], vec![(1, None)], Ordering::Greater)] + fn map_ordering( + #[case] lhs: Vec<(i32, Option<&str>)>, + #[case] rhs: Vec<(i32, Option<&str>)>, + #[case] expected: Ordering, + ) -> VortexResult<()> { + assert_eq!( + map_scalar(lhs)?.partial_cmp(&map_scalar(rhs)?), + Some(expected) + ); + Ok(()) + } + + #[test] + fn null_map_orders_before_every_non_null_map() -> VortexResult<()> { + let null = Scalar::null(map_dtype(Nullability::Nullable)?); + + assert_eq!(null.partial_cmp(&map_scalar(vec![])?), Some(Ordering::Less)); + assert_eq!(null.partial_cmp(&null), Some(Ordering::Equal)); + + Ok(()) + } + fn ab_struct_dtype(nullability: Nullability) -> DType { DType::Struct( StructFields::new( 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 0452f4a3156..d25a652ee57 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -214,10 +214,10 @@ fn compare_arrays( DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), - DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { + DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Map(..) => { nested::compare_nested(lhs, rhs, op, nullability, ctx) } - DType::Map(..) | DType::Union(..) | DType::Variant(_) | DType::Extension(_) => { + 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 e174ce17fa7..7fd71262030 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs @@ -1,13 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Row-wise comparison of nested (struct, list, fixed-size list) arrays. +//! Row-wise comparison of nested (struct, list, fixed-size list, map) arrays. //! //! Nested comparisons canonicalize both operands recursively and build a tree of row //! comparators, one per nested level. The ordering semantics match [`Scalar`] comparison: //! structs compare field-by-field in declaration order, lists compare element-by-element and -//! then by length, and a null value (at any nesting depth) orders before every non-null value. -//! Only top-level nulls make the comparison result null. +//! then by length, maps compare as the list of their `{key, value}` entries, and a null value +//! (at any nesting depth) orders before every non-null value. Only top-level nulls make the +//! comparison result null. //! //! [`Scalar`]: crate::scalar::Scalar @@ -28,6 +29,7 @@ use crate::arrays::DecimalArray; use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; +use crate::arrays::MapArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; @@ -35,6 +37,7 @@ use crate::arrays::decimal::widened_buffer; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::map::MapArraySlotsExt; use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; use crate::dtype::NativePType; @@ -214,7 +217,12 @@ fn build_values_comparator( let rhs = rhs.clone().execute::(ctx)?; build_comparator(lhs.storage_array(), rhs.storage_array(), ctx)? } - DType::Map(..) | DType::Union(..) | DType::Variant(_) => { + DType::Map(..) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + build_values_comparator(lhs.entries(), rhs.entries(), ctx)? + } + DType::Union(..) | DType::Variant(_) => { vortex_bail!("compare is not supported for dtype {}", lhs.dtype()) } }) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index c2af83561a4..9831a963354 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -7,6 +7,7 @@ use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; @@ -24,6 +25,8 @@ use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; +use crate::builders::ArrayBuilder; +use crate::builders::MapBuilder; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -671,3 +674,189 @@ fn binary_compare() { let result = execute_compare_test(lhs, rhs, Operator::Eq); assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); } + +/// A `map(i32, utf8?)` dtype that makes no sortedness assertion. +fn map_dtype(nullability: Nullability) -> VortexResult { + DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + false, + nullability, + ) +} + +type MapRow = Vec<(i32, Option<&'static str>)>; + +fn map_scalar(nullability: Nullability, entries: MapRow) -> VortexResult { + Scalar::try_map( + map_dtype(nullability)?, + entries.into_iter().map(|(key, value)| { + ( + Scalar::primitive(key, Nullability::NonNullable), + match value { + Some(value) => Scalar::utf8(value, Nullability::Nullable), + None => Scalar::null(DType::Utf8(Nullability::Nullable)), + }, + ) + }), + ) +} + +fn map_array( + nullability: Nullability, + rows: impl IntoIterator>, +) -> VortexResult { + let rows = rows.into_iter().collect::>(); + let dtype = map_dtype(nullability)?; + let map_dtype = dtype.as_map_opt().vortex_expect("map dtype").clone(); + let mut builder = MapBuilder::::with_capacity(map_dtype, nullability, rows.len()); + for row in rows { + let scalar = match row { + Some(entries) => map_scalar(nullability, entries)?, + None => Scalar::null(dtype.clone()), + }; + builder.append_scalar(&scalar)?; + } + Ok(builder.finish_into_map().into_array()) +} + +/// Maps compare as the ordered sequence of their `{key, value}` entries: entry-wise first, then +/// by entry count. A null map value orders before every non-null one. +#[rstest] +#[case(Operator::Eq, [true, false, false, false, true])] +#[case(Operator::NotEq, [false, true, true, true, false])] +#[case(Operator::Lt, [false, true, false, false, false])] +#[case(Operator::Lte, [true, true, false, false, true])] +#[case(Operator::Gt, [false, false, true, true, false])] +#[case(Operator::Gte, [true, false, true, true, true])] +fn map_compare_all_operators( + #[case] op: Operator, + #[case] expected: [bool; 5], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = map_array( + Nullability::NonNullable, + [ + // Identical entries. + Some(vec![(1, Some("a")), (2, Some("b"))]), + // A strict prefix of the right-hand row, so it orders first. + Some(vec![(1, Some("a"))]), + // Keys break the tie before values do. + Some(vec![(2, Some("a"))]), + // A null map value orders before a non-null one. + Some(vec![(1, Some("a"))]), + // Two empty maps. + Some(vec![]), + ], + )?; + let rhs = map_array( + Nullability::NonNullable, + [ + Some(vec![(1, Some("a")), (2, Some("b"))]), + Some(vec![(1, Some("a")), (2, Some("b"))]), + Some(vec![(1, Some("z"))]), + Some(vec![(1, None)]), + Some(vec![]), + ], + )?; + + let result = execute_compare_test(lhs, rhs, op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +} + +/// Only top-level nulls make a map comparison null; an empty map is not a null map. +#[test] +fn map_compare_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = map_array( + Nullability::Nullable, + [None, None, Some(vec![]), Some(vec![(1, Some("a"))])], + )?; + let rhs = map_array( + Nullability::Nullable, + [None, Some(vec![]), Some(vec![]), Some(vec![(1, Some("a"))])], + )?; + + let result = execute_compare_test(lhs, rhs, Operator::Eq) + .execute::(&mut ctx) + .vortex_expect("bool array"); + let expected = BoolArray::from_iter([None, None, Some(true), Some(true)]); + assert_arrays_eq!(result, expected, &mut ctx); + + Ok(()) +} + +/// A map array compared against a constant map, and two constant maps compared through +/// [`scalar_cmp`], agree with the row-wise kernel. +#[test] +fn map_constant_compare() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = map_array( + Nullability::NonNullable, + [ + Some(vec![(1, Some("a"))]), + Some(vec![(1, Some("a")), (2, Some("b"))]), + Some(vec![(9, Some("a"))]), + ], + )?; + let needle = map_scalar(Nullability::NonNullable, vec![(1, Some("a"))])?; + let constant = ConstantArray::new(needle.clone(), array.len()).into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(array, constant, Operator::Lt); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, false]), + &mut ctx + ); + + // Constant-vs-constant folds through `scalar_cmp`. + let bigger = map_scalar( + Nullability::NonNullable, + vec![(1, Some("a")), (2, Some("b"))], + )?; + assert_eq!( + scalar_cmp(&needle, &bigger, CompareOperator::Lt)?, + Scalar::bool(true, Nullability::NonNullable) + ); + assert_eq!( + scalar_cmp(&needle, &needle, CompareOperator::Eq)?, + Scalar::bool(true, Nullability::NonNullable) + ); + + Ok(()) +} + +/// Maps nested inside another nested type compare through the same row comparator tree. +#[test] +fn struct_of_map_compare() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = StructArray::from_fields(&[( + "m", + map_array( + Nullability::NonNullable, + [Some(vec![(1, Some("a"))]), Some(vec![(1, Some("a"))])], + )?, + )])? + .into_array(); + let rhs = StructArray::from_fields(&[( + "m", + map_array( + Nullability::NonNullable, + [Some(vec![(1, Some("a"))]), Some(vec![(1, Some("b"))])], + )?, + )])? + .into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([true, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([false, true]), &mut ctx); + + Ok(()) +}