Implement compare for map types

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> {
+        DType::map(
+            DType::Primitive(PType::I32, Nullability::NonNullable),
+            DType::Utf8(Nullability::Nullable),
+            false,
+            nullability,
+        )
+    }
+
+    fn map_scalar(entries: Vec<(i32, Option<&str>)>) -> VortexResult<Scalar> {
+        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(

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())
         }
     }

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,13 +29,15 @@ 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;
 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::<ExtensionArray>(ctx)?;
             build_comparator(lhs.storage_array(), rhs.storage_array(), ctx)?
         }
-        DType::Map(..) | DType::Union(..) | DType::Variant(_) => {
+        DType::Map(..) => {
+            let lhs = lhs.clone().execute::<MapArray>(ctx)?;
+            let rhs = rhs.clone().execute::<MapArray>(ctx)?;
+            build_values_comparator(lhs.entries(), rhs.entries(), ctx)?
+        }
+        DType::Union(..) | DType::Variant(_) => {
             vortex_bail!("compare is not supported for dtype {}", lhs.dtype())
         }
     })

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> {
+    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> {
+    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<Item = Option<MapRow>>,
+) -> VortexResult<ArrayRef> {
+    let rows = rows.into_iter().collect::<Vec<_>>();
+    let dtype = map_dtype(nullability)?;
+    let map_dtype = dtype.as_map_opt().vortex_expect("map dtype").clone();
+    let mut builder = MapBuilder::<u64, u64>::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::<BoolArray>(&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(())
+}