Fix scalar eq, cmp, hash, etc. (#2091)

Previously we were being a little bit lazy and delegating to
ScalarValue. We've been bitten by this in the past!

It is never safe to delegate to ScalarValue since it doesn't have enough
information to know what "equals" means, i.e. it's missing the
surrounding scalar's DType.

This PR removes PartialEq, PartialOrd, and Hash from ScalarValue, in
favor of forcing us to wrap the value into a Scalar and then compare
each dtype-specific scalar type.

Author: gatesn

pyvortex/src/scalar.rs

@@ -9,6 +9,7 @@ use pyo3::types::PyDict;
 use vortex::buffer::{BufferString, ByteBuffer};
 use vortex::dtype::half::f16;
 use vortex::dtype::{DType, PType};
+use vortex::error::VortexExpect;
 use vortex::scalar::{ListScalar, Scalar, StructScalar};
 
 pub fn scalar_into_py(py: Python, x: Scalar, copy_into_python: bool) -> PyResult<PyObject> {
@@ -190,6 +191,8 @@ impl PyVortexList {
 fn to_python_list(py: Python, scalar: ListScalar<'_>, recursive: bool) -> PyResult<PyObject> {
     Ok(scalar
         .elements()
+        .vortex_expect("non-null")
+        .into_iter()
         .map(|x| scalar_into_py(py, x, recursive))
         .collect::<PyResult<Vec<_>>>()?
         .into_py(py))

vortex-array/src/array/constant/canonical.rs

@@ -105,10 +105,12 @@ fn canonical_byte_view(
 
 #[cfg(test)]
 mod tests {
-    use vortex_dtype::{DType, Nullability};
+    use vortex_dtype::half::f16;
+    use vortex_dtype::{DType, Nullability, PType};
     use vortex_scalar::Scalar;
 
     use crate::array::ConstantArray;
+    use crate::canonical::IntoArrayVariant;
     use crate::compute::scalar_at;
     use crate::stats::{ArrayStatistics as _, StatsSet};
     use crate::{ArrayLen, IntoArrayData as _, IntoCanonical};
@@ -151,4 +153,17 @@ mod tests {
         assert_eq!(canonical_stats, StatsSet::constant(&scalar, 4));
         assert_eq!(canonical_stats, stats);
     }
+
+    #[test]
+    fn test_canonicalize_scalar_values() {
+        let f16_scalar = Scalar::primitive(f16::from_f32(5.722046e-6), Nullability::NonNullable);
+        let scalar = Scalar::new(
+            DType::Primitive(PType::F16, Nullability::NonNullable),
+            Scalar::primitive(96u8, Nullability::NonNullable).into_value(),
+        );
+        let const_array = ConstantArray::new(scalar.clone(), 1).into_array();
+        let canonical_const = const_array.into_primitive().unwrap();
+        assert_eq!(scalar_at(&canonical_const, 0).unwrap(), scalar);
+        assert_eq!(scalar_at(&canonical_const, 0).unwrap(), f16_scalar);
+    }
 }

vortex-array/src/builders/list.rs

@@ -49,20 +49,23 @@ where
     }
 
     pub fn append_value(&mut self, value: ListScalar) -> VortexResult<()> {
-        if value.is_null() {
-            if self.nullability == Nullability::NonNullable {
-                vortex_bail!("Cannot append null value to non-nullable list");
+        match value.elements() {
+            None => {
+                if self.nullability == Nullability::NonNullable {
+                    vortex_bail!("Cannot append null value to non-nullable list");
+                }
+                self.append_null();
+                Ok(())
             }
-            self.append_null();
-            Ok(())
-        } else {
-            for scalar in value.elements() {
-                // TODO(joe): This is slow, we should be able to append multiple values at once,
-                // or the list scalar should hold an ArrayData
-                self.value_builder.append_scalar(&scalar)?;
+            Some(elements) => {
+                for scalar in elements {
+                    // TODO(joe): This is slow, we should be able to append multiple values at once,
+                    // or the list scalar should hold an ArrayData
+                    self.value_builder.append_scalar(&scalar)?;
+                }
+                self.validity.append_value(true);
+                self.append_index(self.value_builder.len().as_())
             }
-            self.validity.append_value(true);
-            self.append_index(self.value_builder.len().as_())
         }
     }
 

vortex-scalar/src/binary.rs

@@ -7,11 +7,27 @@ use vortex_error::{vortex_bail, vortex_err, VortexError, VortexExpect as _, Vort
 use crate::value::{InnerScalarValue, ScalarValue};
 use crate::Scalar;
 
+#[derive(Debug, Hash)]
 pub struct BinaryScalar<'a> {
     dtype: &'a DType,
     value: Option<ByteBuffer>,
 }
 
+impl PartialEq for BinaryScalar<'_> {
+    fn eq(&self, other: &Self) -> bool {
+        self.dtype == other.dtype && self.value == other.value
+    }
+}
+
+impl Eq for BinaryScalar<'_> {}
+
+/// Ord is not implemented since it's undefined for different nullability
+impl PartialOrd for BinaryScalar<'_> {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        self.value.partial_cmp(&other.value)
+    }
+}
+
 impl<'a> BinaryScalar<'a> {
     #[inline]
     pub fn dtype(&self) -> &'a DType {

vortex-scalar/src/bool.rs

@@ -1,15 +1,35 @@
+use std::cmp::Ordering;
+
 use vortex_dtype::Nullability::NonNullable;
 use vortex_dtype::{DType, Nullability};
 use vortex_error::{vortex_bail, vortex_err, VortexError, VortexExpect as _, VortexResult};
 
 use crate::value::ScalarValue;
 use crate::{InnerScalarValue, Scalar};
 
+#[derive(Debug, Hash)]
 pub struct BoolScalar<'a> {
     dtype: &'a DType,
     value: Option<bool>,
 }
 
+impl PartialEq for BoolScalar<'_> {
+    fn eq(&self, other: &Self) -> bool {
+        self.value == other.value
+    }
+}
+
+impl Eq for BoolScalar<'_> {}
+
+impl PartialOrd for BoolScalar<'_> {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        if self.dtype != other.dtype {
+            return None;
+        }
+        self.value.partial_cmp(&other.value)
+    }
+}
+
 impl<'a> BoolScalar<'a> {
     #[inline]
     pub fn dtype(&self) -> &'a DType {

vortex-scalar/src/display.rs

@@ -61,11 +61,11 @@ impl Display for Scalar {
             }
             DType::List(..) => {
                 let v = ListScalar::try_from(self).map_err(|_| std::fmt::Error)?;
-
-                if v.is_null() {
-                    write!(f, "null")
-                } else {
-                    write!(f, "[{}]", v.elements().format(","))
+                match v.elements() {
+                    None => write!(f, "null"),
+                    Some(elems) => {
+                        write!(f, "[{}]", elems.iter().format(","))
+                    }
                 }
             }
             // Specialized handling for date/time/timestamp builtin extension types.

vortex-scalar/src/extension.rs

@@ -1,3 +1,4 @@
+use std::hash::Hash;
 use std::sync::Arc;
 
 use vortex_dtype::{DType, ExtDType};
@@ -11,6 +12,30 @@ pub struct ExtScalar<'a> {
     value: &'a ScalarValue,
 }
 
+impl PartialEq for ExtScalar<'_> {
+    fn eq(&self, other: &Self) -> bool {
+        self.ext_dtype == other.ext_dtype && self.storage() == other.storage()
+    }
+}
+
+impl Eq for ExtScalar<'_> {}
+
+impl PartialOrd for ExtScalar<'_> {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        if self.ext_dtype != other.ext_dtype {
+            return None;
+        }
+        self.storage().partial_cmp(&other.storage())
+    }
+}
+
+impl Hash for ExtScalar<'_> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.ext_dtype.hash(state);
+        self.storage().hash(state);
+    }
+}
+
 impl<'a> ExtScalar<'a> {
     pub fn try_new(dtype: &'a DType, value: &'a ScalarValue) -> VortexResult<Self> {
         let DType::Extension(ext_dtype) = dtype else {

vortex-scalar/src/lib.rs

@@ -1,6 +1,5 @@
 use std::cmp::Ordering;
 use std::hash::Hash;
-use std::mem::discriminant;
 use std::sync::Arc;
 
 pub use scalar_type::ScalarType;
@@ -42,8 +41,9 @@ use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 /// A [`ScalarValue`] is opaque, and should be accessed via one of the type-specific scalar wrappers
 /// for example [`BoolScalar`], [`PrimitiveScalar`], etc.
 ///
-/// Note: [`PartialEq`] and [`PartialOrd`] are implemented only for an exact match of the scalar's
-/// dtype, including nullability.
+/// Note that [`PartialOrd`] is implemented only for an exact match of the scalar's dtype,
+/// including nullability. When the DType does match, ordering is nulls first (lowest), then the
+/// natural ordering of the scalar value.
 #[derive(Debug, Clone)]
 #[cfg_attr(feature = "serde", derive(::serde::Serialize, ::serde::Deserialize))]
 pub struct Scalar {
@@ -203,28 +203,56 @@ impl Scalar {
 
 impl PartialEq for Scalar {
     fn eq(&self, other: &Self) -> bool {
-        self.dtype == other.dtype && self.value.0 == other.value.0
+        if self.dtype != other.dtype {
+            return false;
+        }
+
+        match self.dtype() {
+            DType::Null => true,
+            DType::Bool(_) => self.as_bool() == other.as_bool(),
+            DType::Primitive(..) => self.as_primitive() == other.as_primitive(),
+            DType::Utf8(_) => self.as_utf8() == other.as_utf8(),
+            DType::Binary(_) => self.as_binary() == other.as_binary(),
+            DType::Struct(..) => self.as_struct() == other.as_struct(),
+            DType::List(..) => self.as_list() == other.as_list(),
+            DType::Extension(_) => self.as_extension() == other.as_extension(),
+        }
     }
 }
 
 impl Eq for Scalar {}
 
 impl PartialOrd for Scalar {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        // We check for DType equality, ignoring nullability, and allowing us to compare all
-        // primitive types to all other primitive types.
-        if discriminant(self.dtype()) == discriminant(other.dtype()) {
-            self.value.0.partial_cmp(&other.value.0)
-        } else {
-            None
+        if self.dtype() != other.dtype() {
+            return None;
+        }
+
+        match self.dtype() {
+            DType::Null => Some(Ordering::Equal),
+            DType::Bool(_) => self.as_bool().partial_cmp(&other.as_bool()),
+            DType::Primitive(..) => self.as_primitive().partial_cmp(&other.as_primitive()),
+            DType::Utf8(_) => self.as_utf8().partial_cmp(&other.as_utf8()),
+            DType::Binary(_) => self.as_binary().partial_cmp(&other.as_binary()),
+            DType::Struct(..) => self.as_struct().partial_cmp(&other.as_struct()),
+            DType::List(..) => self.as_list().partial_cmp(&other.as_list()),
+            DType::Extension(_) => self.as_extension().partial_cmp(&other.as_extension()),
         }
     }
 }
 
 impl Hash for Scalar {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        discriminant(self.dtype()).hash(state);
-        self.value.0.hash(state);
+        match self.dtype() {
+            DType::Null => self.dtype().hash(state), // Hash the dtype instead of the value
+            DType::Bool(_) => self.as_bool().hash(state),
+            DType::Primitive(..) => self.as_primitive().hash(state),
+            DType::Utf8(_) => self.as_utf8().hash(state),
+            DType::Binary(_) => self.as_binary().hash(state),
+            DType::Struct(..) => self.as_struct().hash(state),
+            DType::List(..) => self.as_list().hash(state),
+            DType::Extension(_) => self.as_extension().hash(state),
+        }
     }
 }
 

vortex-scalar/src/list.rs

@@ -1,9 +1,12 @@
+use std::hash::Hash;
 use std::ops::Deref;
 use std::sync::Arc;
 
 use itertools::Itertools as _;
 use vortex_dtype::{DType, Nullability};
-use vortex_error::{vortex_bail, vortex_panic, VortexError, VortexExpect as _, VortexResult};
+use vortex_error::{
+    vortex_bail, vortex_err, vortex_panic, VortexError, VortexExpect as _, VortexResult,
+};
 
 use crate::value::{InnerScalarValue, ScalarValue};
 use crate::Scalar;
@@ -14,6 +17,34 @@ pub struct ListScalar<'a> {
     elements: Option<Arc<[ScalarValue]>>,
 }
 
+impl PartialEq for ListScalar<'_> {
+    fn eq(&self, other: &Self) -> bool {
+        if self.dtype != other.dtype {
+            return false;
+        }
+        self.elements() == other.elements()
+    }
+}
+
+impl Eq for ListScalar<'_> {}
+
+/// Ord is not implemented since it's undefined for different DTypes
+impl PartialOrd for ListScalar<'_> {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        if self.dtype() != other.dtype() {
+            return None;
+        }
+        self.elements().partial_cmp(&other.elements())
+    }
+}
+
+impl Hash for ListScalar<'_> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.dtype.hash(state);
+        self.elements().hash(state);
+    }
+}
+
 impl<'a> ListScalar<'a> {
     #[inline]
     pub fn dtype(&self) -> &'a DType {
@@ -55,16 +86,13 @@ impl<'a> ListScalar<'a> {
             })
     }
 
-    pub fn elements(&self) -> impl Iterator<Item = Scalar> + '_ {
-        self.elements
-            .as_ref()
-            .map(AsRef::as_ref)
-            .unwrap_or_else(|| &[] as &[ScalarValue])
-            .iter()
-            .map(|e| Scalar {
-                dtype: self.element_dtype(),
-                value: e.clone(),
-            })
+    pub fn elements(&self) -> Option<Vec<Scalar>> {
+        self.elements.as_ref().map(|elems| {
+            elems
+                .iter()
+                .map(|e| Scalar::new(self.element_dtype(), e.clone()))
+                .collect_vec()
+        })
     }
 
     pub(crate) fn cast(&self, dtype: &DType) -> VortexResult<Scalar> {
@@ -143,7 +171,10 @@ impl<'a, T: for<'b> TryFrom<&'b Scalar, Error = VortexError>> TryFrom<&'a Scalar
     fn try_from(value: &'a Scalar) -> Result<Self, Self::Error> {
         let value = ListScalar::try_from(value)?;
         let mut elems = vec![];
-        for e in value.elements() {
+        for e in value
+            .elements()
+            .ok_or_else(|| vortex_err!("Expected non-null list"))?
+        {
             elems.push(T::try_from(&e)?);
         }
         Ok(elems)