Require matching dtypes for extension comparisons

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

Cargo.lock

@@ -18,6 +18,7 @@ use vortex_array::scalar_fn::fns::operators::CompareOperator;
 use vortex_array::scalar_fn::fns::operators::Operator;
 use vortex_array::validity::Validity;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
 use vortex_mask::AllOr;
 
 use crate::array::DateTimeParts;
@@ -31,6 +32,13 @@ impl CompareKernel for DateTimeParts {
         operator: CompareOperator,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
+        vortex_ensure!(
+            lhs.dtype().eq_ignore_nullability(rhs.dtype()),
+            "Comparison requires equal dtypes ignoring nullability, got {} and {}",
+            lhs.dtype(),
+            rhs.dtype()
+        );
+
         let Some(rhs_const) = rhs.as_constant() else {
             return Ok(None);
         };
@@ -284,7 +292,7 @@ fn compare_dtp(
 }
 
 #[cfg(test)]
-mod test {
+mod tests {
     use std::sync::LazyLock;
 
     use rstest::rstest;
@@ -303,6 +311,8 @@ mod test {
     use vortex_array::extension::datetime::TimestampOptions;
     use vortex_array::scalar::Scalar;
     use vortex_array::scalar_fn::fns::binary::CompareKernel;
+    use vortex_array::test_harness::ExpectedOutcome;
+    use vortex_array::test_harness::assert_execution_paths;
     use vortex_array::validity::Validity;
     use vortex_buffer::buffer;
     use vortex_session::VortexSession;
@@ -683,6 +693,49 @@ mod test {
         assert_matches_temporal(nullable, unit, constants)
     }
 
+    #[rstest]
+    #[case::storage(None)]
+    #[case::different_unit(Some(TimeUnit::Nanoseconds))]
+    fn test_comparison_rejects_mismatched_constant_dtype(
+        #[case] unit: Option<TimeUnit>,
+    ) -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let temporal = TemporalArray::new_timestamp(
+            buffer![1600i64, 1400].into_array(),
+            TimeUnit::Milliseconds,
+            None,
+        );
+        let array = DateTimeParts::try_from_temporal(temporal, &mut ctx)?;
+        let scalar = match unit {
+            Some(unit) => Scalar::extension::<Timestamp>(
+                TimestampOptions { unit, tz: None },
+                Scalar::from(1_500_000_000i64),
+            ),
+            None => Scalar::from(1500i64),
+        };
+        let rhs = ConstantArray::new(scalar, array.len()).into_array();
+
+        for operator in ALL_OPERATORS {
+            assert!(
+                <DateTimeParts as CompareKernel>::compare(
+                    array.as_view(),
+                    &rhs,
+                    operator,
+                    &mut ctx,
+                )
+                .is_err()
+            );
+            assert_execution_paths(
+                &[array.clone().into_array(), rhs.clone()],
+                |inputs| inputs[0].binary(inputs[1].clone(), operator.into()),
+                ExpectedOutcome::Error,
+                &mut ctx,
+            );
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn null_rows_stay_null_when_parts_differ() -> VortexResult<()> {
         let mut ctx = SESSION.create_execution_ctx();

encodings/datetime-parts/src/compute/compare.rs

@@ -12,7 +12,6 @@ use vortex_session::VortexSession;
 use crate::EqMode;
 use crate::IntoArray;
 use crate::VortexSessionExecute;
-use crate::arrays::BoolArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::datetime::TemporalData;
 use crate::assert_arrays_eq;
@@ -218,7 +217,7 @@ fn test_validity_preservation(#[case] validity: Validity) {
 }
 
 #[test]
-fn comparison_rejects_mismatched_timestamp_units() -> VortexResult<()> {
+fn test_comparison_rejects_mismatched_timestamp_units_at_binding() -> VortexResult<()> {
     let ts_array = PrimitiveArray::from_iter(vec![1704067200000i64, 1704153600000, 1704240000000])
         .into_array();
     let temporal = TemporalData::new_timestamp(ts_array, TimeUnit::Milliseconds, None);
@@ -236,13 +235,12 @@ fn comparison_rejects_mismatched_timestamp_units() -> VortexResult<()> {
 
     let error = temporal
         .into_array()
-        .apply(&filter_expr)?
-        .execute::<BoolArray>(&mut SESSION.create_execution_ctx())
-        .unwrap_err();
+        .apply(&filter_expr)
+        .expect_err("comparisons require matching timestamp units when the expression is bound");
     assert!(
         error
             .to_string()
-            .contains("Cannot compare extension dtypes")
+            .contains("Comparison requires equal dtypes ignoring nullability")
     );
 
     Ok(())

vortex-array/src/arrays/datetime/test.rs

@@ -247,7 +247,7 @@ mod tests {
     }
 
     #[test]
-    fn test_scalar_fn_no_pushdown_different_ext_types() {
+    fn test_comparison_rejects_different_extension_types_at_binding() {
         #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
         struct TestExt2;
         impl ExtVTable for TestExt2 {
@@ -289,21 +289,15 @@ mod tests {
         let storage = buffer![10i64, 20, 30].into_array();
         let ext_array = ExtensionArray::new(ext_dtype1, storage).into_array();
 
-        // Create constant with different extension type
         let const_scalar = Scalar::extension::<TestExt2>(EmptyMetadata, Scalar::from(25i64));
         let const_array = ConstantArray::new(const_scalar, 3).into_array();
 
-        let scalar_fn_array = Binary::try_new(ext_array, const_array, Operator::Lt)
-            .unwrap()
-            .into_array();
-
-        let optimized = scalar_fn_array.optimize().unwrap();
-
-        // The first child should still be an ExtensionArray (no pushdown happened)
-        let scalar_fn = optimized.as_opt::<ScalarFn>().unwrap();
+        let error = Binary::try_new(ext_array, const_array, Operator::Lt)
+            .expect_err("comparisons require matching extension dtypes at binding");
         assert!(
-            scalar_fn.children()[0].as_opt::<Extension>().is_some(),
-            "Expected first child to remain ExtensionArray when ext types differ"
+            error
+                .to_string()
+                .contains("Comparison requires equal dtypes ignoring nullability")
         );
     }
 
@@ -333,25 +327,19 @@ mod tests {
     }
 
     #[test]
-    fn test_scalar_fn_no_pushdown_non_extension_constant() {
+    fn test_comparison_rejects_storage_constant_at_binding() {
         let ext_dtype = test_ext_dtype();
         let storage = buffer![10i64, 20, 30].into_array();
         let ext_array = ExtensionArray::new(ext_dtype, storage).into_array();
 
-        // Create a non-extension constant (plain primitive)
         let const_array = ConstantArray::new(Scalar::from(25i64), 3).into_array();
 
-        let scalar_fn_array = Binary::try_new(ext_array, const_array, Operator::Lt)
-            .unwrap()
-            .into_array();
-
-        let optimized = scalar_fn_array.optimize().unwrap();
-
-        // No pushdown should happen because constant is not an extension scalar
-        let scalar_fn = optimized.as_opt::<ScalarFn>().unwrap();
+        let error = Binary::try_new(ext_array, const_array, Operator::Lt)
+            .expect_err("an extension cannot be compared with its storage dtype at binding");
         assert!(
-            scalar_fn.children()[0].as_opt::<Extension>().is_some(),
-            "Expected first child to remain ExtensionArray when constant is not extension"
+            error
+                .to_string()
+                .contains("Comparison requires equal dtypes ignoring nullability")
         );
     }
 }

vortex-array/src/arrays/extension/compute/rules.rs

@@ -15,6 +15,7 @@ use std::cmp::Ordering;
 
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
 
 use crate::ArrayRef;
@@ -101,6 +102,8 @@ where
             _ => return Ok(None),
         };
 
+        validate_comparison_dtypes(array.dtype(), other.dtype())?;
+
         let len = array.len();
         let nullable = array.dtype().is_nullable() || other.dtype().is_nullable();
 
@@ -133,6 +136,8 @@ pub(crate) fn execute_compare(
     op: CompareOperator,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
+    validate_comparison_dtypes(lhs.dtype(), rhs.dtype())?;
+
     let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
 
     if lhs.len() != rhs.len() {
@@ -174,32 +179,14 @@ fn compare_arrays(
 ) -> VortexResult<ArrayRef> {
     let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable());
 
-    // Extension arrays compare through their storage. When both sides are extensions they must
-    // agree on the full extension dtype (a timestamp in milliseconds must not compare its raw
-    // storage against one in seconds); a lone extension side may compare against its raw storage.
-    if lhs.dtype().is_extension() || rhs.dtype().is_extension() {
-        if lhs.dtype().is_extension()
-            && rhs.dtype().is_extension()
-            && !lhs.dtype().eq_ignore_nullability(rhs.dtype())
-        {
-            vortex_bail!(
-                "Cannot compare extension dtypes {} and {}",
-                lhs.dtype(),
-                rhs.dtype()
-            );
-        }
+    if lhs.dtype().is_extension() {
         let lhs = extension_storage(lhs, ctx)?;
         let rhs = extension_storage(rhs, ctx)?;
+
         return compare_arrays(&lhs, &rhs, op, ctx);
     }
 
-    if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) {
-        vortex_bail!(
-            "Cannot compare different DTypes {} and {}",
-            lhs.dtype(),
-            rhs.dtype()
-        );
-    }
+    validate_comparison_dtypes(lhs.dtype(), rhs.dtype())?;
 
     match lhs.dtype() {
         // Every value of the null dtype is null, so every comparison result is null.
@@ -221,12 +208,10 @@ fn compare_arrays(
     }
 }
 
-/// Replace an extension array (or extension constant) with its storage. Non-extension arrays are
-/// returned unchanged.
+/// Returns the storage of an extension array or constant.
+///
+/// The caller must first require matching extension dtypes, ignoring nullability.
 fn extension_storage(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
-    if !array.dtype().is_extension() {
-        return Ok(array.clone());
-    }
     if let Some(constant) = array.as_opt::<Constant>() {
         return Ok(ConstantArray::new(
             constant.scalar().as_extension().to_storage_scalar(),
@@ -238,7 +223,12 @@ fn extension_storage(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<A
     Ok(ext.storage_array().clone())
 }
 
+/// Compares scalars with matching dtypes, ignoring nullability.
+///
+/// Returns null if either operand is null, and an error if the operand dtypes differ.
 pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: CompareOperator) -> VortexResult<Scalar> {
+    validate_comparison_dtypes(lhs.dtype(), rhs.dtype())?;
+
     if lhs.is_null() | rhs.is_null() {
         return Ok(Scalar::null(DType::Bool(Nullability::Nullable)));
     }
@@ -266,6 +256,16 @@ pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: CompareOperator) -> Vort
     Ok(Scalar::bool(b, nullability))
 }
 
+/// Checks the shared dtype requirement before comparison fast paths or encoding kernels run.
+pub(super) fn validate_comparison_dtypes(lhs: &DType, rhs: &DType) -> VortexResult<()> {
+    vortex_ensure!(
+        lhs.eq_ignore_nullability(rhs),
+        "Comparison requires equal dtypes ignoring nullability, got {lhs} and {rhs}"
+    );
+
+    Ok(())
+}
+
 /// Returns the predicate `Ordering -> bool` for a comparison operator.
 #[inline]
 pub(super) fn ordering_predicate(op: CompareOperator) -> fn(Ordering) -> bool {