Introduce MaybeCompareFn trait to allow for partial compare specializations

encodings/byte-bool/src/compute/mod.rs

@@ -13,8 +13,8 @@ use vortex_scalar::Scalar;
 use super::ByteBoolArray;
 
 impl ArrayCompute for ByteBoolArray {
-    fn compare(&self) -> Option<&dyn CompareFn> {
-        Some(self)
+    fn compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        CompareFn::compare(self, array, operator)
     }
 
     fn fill_forward(&self) -> Option<&dyn FillForwardFn> {
@@ -100,12 +100,12 @@ impl TakeFn for ByteBoolArray {
 }
 
 impl CompareFn for ByteBoolArray {
-    fn compare(&self, other: &Array, op: Operator) -> VortexResult<Array> {
-        let canonical = other.clone().into_bool()?;
+    fn compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        let canonical = array.clone().into_bool()?;
         let lhs = BooleanBuffer::from(self.maybe_null_slice());
         let rhs = canonical.boolean_buffer();
 
-        let result_buf = match op {
+        let result_buf = match operator {
             Operator::Eq => lhs.bitxor(&rhs).not(),
             Operator::NotEq => lhs.bitxor(&rhs),
 

encodings/datetime-parts/src/array.rs

@@ -1,11 +1,14 @@
 use std::fmt::Debug;
 
 use serde::{Deserialize, Serialize};
+use vortex::array::StructArray;
 use vortex::stats::{ArrayStatisticsCompute, StatsSet};
 use vortex::validity::{ArrayValidity, LogicalValidity};
 use vortex::variants::{ArrayVariants, ExtensionArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
-use vortex::{impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, Canonical, IntoCanonical};
+use vortex::{
+    impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, Canonical, IntoArray, IntoCanonical,
+};
 use vortex_dtype::DType;
 use vortex_error::{vortex_bail, VortexResult};
 
@@ -89,7 +92,19 @@ impl ArrayVariants for DateTimePartsArray {
     }
 }
 
-impl ExtensionArrayTrait for DateTimePartsArray {}
+impl ExtensionArrayTrait for DateTimePartsArray {
+    fn storage_array(&self) -> Array {
+        // FIXME(ngates): this needs to be a tuple array so we can implement Compare
+        StructArray::try_new(
+            vec!["days".into(), "seconds".into(), "subseconds".into()].into(),
+            [self.days(), self.seconds(), self.subsecond()].into(),
+            self.len(),
+            self.logical_validity().into_validity(),
+        )
+        .expect("Failed to create struct array")
+        .into_array()
+    }
+}
 
 impl IntoCanonical for DateTimePartsArray {
     fn into_canonical(self) -> VortexResult<Canonical> {

vortex-array/src/array/bool/compute/compare.rs

@@ -8,14 +8,13 @@ use crate::{Array, IntoArray, IntoArrayVariant};
 
 impl CompareFn for BoolArray {
     // TODO(aduffy): replace these with Arrow compute kernels.
-    fn compare(&self, other: &Array, op: Operator) -> VortexResult<Array> {
-        let flattened = other.clone().into_bool()?;
+    fn compare(&self, array: &Array, operator: Operator) -> VortexResult<Array> {
+        let flattened = array.clone().into_bool()?;
         let lhs = self.boolean_buffer();
         let rhs = flattened.boolean_buffer();
-        let result_buf = match op {
+        let result_buf = match operator {
             Operator::Eq => lhs.bitxor(&rhs).not(),
             Operator::NotEq => lhs.bitxor(&rhs),
-
             Operator::Gt => lhs.bitand(&rhs.not()),
             Operator::Gte => lhs.bitor(&rhs.not()),
             Operator::Lt => lhs.not().bitand(&rhs),

vortex-array/src/array/bool/compute/mod.rs

@@ -1,6 +1,9 @@
+use vortex_error::VortexResult;
+
 use crate::array::BoolArray;
 use crate::compute::unary::{FillForwardFn, ScalarAtFn};
-use crate::compute::{AndFn, ArrayCompute, CompareFn, OrFn, SliceFn, TakeFn};
+use crate::compute::{AndFn, ArrayCompute, CompareFn, Operator, OrFn, SliceFn, TakeFn};
+use crate::Array;
 
 mod boolean;
 mod compare;
@@ -12,8 +15,8 @@ mod slice;
 mod take;
 
 impl ArrayCompute for BoolArray {
-    fn compare(&self) -> Option<&dyn CompareFn> {
-        Some(self)
+    fn compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        Some(CompareFn::compare(self, array, operator))
     }
 
     fn fill_forward(&self) -> Option<&dyn FillForwardFn> {

vortex-array/src/array/chunked/compute/mod.rs

@@ -17,8 +17,8 @@ impl ArrayCompute for ChunkedArray {
         Some(self)
     }
 
-    fn compare(&self) -> Option<&dyn CompareFn> {
-        Some(self)
+    fn compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        Some(CompareFn::compare(self, array, operator))
     }
 
     fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {

vortex-array/src/array/chunked/variants.rs

@@ -83,4 +83,13 @@ impl StructArrayTrait for ChunkedArray {
 
 impl ListArrayTrait for ChunkedArray {}
 
-impl ExtensionArrayTrait for ChunkedArray {}
+impl ExtensionArrayTrait for ChunkedArray {
+    fn storage_array(&self) -> Array {
+        ChunkedArray::from_iter(
+            self.chunks()
+                .into_iter()
+                .map(|chunk| chunk.with_dyn(|a| a.as_extension_array_unchecked().storage_array())),
+        )
+        .into_array()
+    }
+}

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

@@ -1,25 +1,21 @@
 use std::cmp::Ordering;
-use std::sync::Arc;
 
-use arrow_array::Datum;
-use arrow_ord::cmp;
 use vortex_dtype::Nullability;
 use vortex_error::{vortex_bail, vortex_err, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::array::constant::ConstantArray;
-use crate::arrow::FromArrowArray;
 use crate::compute::unary::{scalar_at, ScalarAtFn};
 use crate::compute::{
-    scalar_cmp, AndFn, ArrayCompute, CompareFn, FilterFn, Operator, OrFn, SearchResult,
+    scalar_cmp, AndFn, ArrayCompute, FilterFn, MaybeCompareFn, Operator, OrFn, SearchResult,
     SearchSortedFn, SearchSortedSide, SliceFn, TakeFn,
 };
 use crate::stats::{ArrayStatistics, Stat};
-use crate::{Array, ArrayDType, AsArray, IntoArray, IntoCanonical};
+use crate::{Array, ArrayDType, AsArray, IntoArray};
 
 impl ArrayCompute for ConstantArray {
-    fn compare(&self) -> Option<&dyn CompareFn> {
-        Some(self)
+    fn compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        MaybeCompareFn::maybe_compare(self, array, operator)
     }
 
     fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
@@ -99,30 +95,17 @@ impl SearchSortedFn for ConstantArray {
     }
 }
 
-impl CompareFn for ConstantArray {
-    fn compare(&self, rhs: &Array, operator: Operator) -> VortexResult<Array> {
-        if let Some(true) = rhs.statistics().get_as::<bool>(Stat::IsConstant) {
-            let lhs = self.scalar();
-            let rhs = scalar_at(rhs, 0)?;
-
-            let scalar = scalar_cmp(lhs, &rhs, operator);
-
-            Ok(ConstantArray::new(scalar, self.len()).into_array())
+impl MaybeCompareFn for ConstantArray {
+    fn maybe_compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        if let Some(true) = array.statistics().get_as::<bool>(Stat::IsConstant) {
+            Some({
+                let lhs = self.scalar();
+                let rhs = scalar_at(array, 0).expect("Expected scalar");
+                let scalar = scalar_cmp(lhs, &rhs, operator);
+                Ok(ConstantArray::new(scalar, self.len()).into_array())
+            })
         } else {
-            let datum = Arc::<dyn Datum>::from(self.scalar());
-            let rhs = rhs.clone().into_canonical()?.into_arrow();
-            let rhs = rhs.as_ref();
-
-            let boolean_array = match operator {
-                Operator::Eq => cmp::eq(datum.as_ref(), &rhs)?,
-                Operator::NotEq => cmp::neq(datum.as_ref(), &rhs)?,
-                Operator::Gt => cmp::gt(datum.as_ref(), &rhs)?,
-                Operator::Gte => cmp::gt_eq(datum.as_ref(), &rhs)?,
-                Operator::Lt => cmp::lt(datum.as_ref(), &rhs)?,
-                Operator::Lte => cmp::lt_eq(datum.as_ref(), &rhs)?,
-            };
-
-            Ok(Array::from_arrow(&boolean_array, true))
+            None
         }
     }
 }

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

@@ -3,7 +3,7 @@ use std::sync::Arc;
 
 use vortex_dtype::{DType, PType};
 use vortex_error::VortexError;
-use vortex_scalar::{Scalar, StructScalar};
+use vortex_scalar::{ExtScalar, Scalar, ScalarValue, StructScalar};
 
 use crate::array::constant::ConstantArray;
 use crate::iter::{Accessor, AccessorRef};
@@ -190,7 +190,28 @@ impl StructArrayTrait for ConstantArray {
 
 impl ListArrayTrait for ConstantArray {}
 
-impl ExtensionArrayTrait for ConstantArray {}
+impl ExtensionArrayTrait for ConstantArray {
+    fn storage_array(&self) -> Array {
+        let scalar_ext = ExtScalar::try_from(self.scalar()).expect("Expected an extension scalar");
+
+        // FIXME(ngates): there's not enough information to get the storage array.
+        let n = self.dtype().nullability();
+        let storage_dtype = match scalar_ext.value() {
+            ScalarValue::Bool(_) => DType::Binary(n),
+            ScalarValue::Primitive(pvalue) => DType::Primitive(pvalue.ptype(), n),
+            ScalarValue::Buffer(_) => DType::Binary(n),
+            ScalarValue::BufferString(_) => DType::Utf8(n),
+            ScalarValue::List(_) => panic!("List not supported"),
+            ScalarValue::Null => DType::Null,
+        };
+
+        ConstantArray::new(
+            Scalar::new(storage_dtype, scalar_ext.value().clone()),
+            self.len(),
+        )
+        .into_array()
+    }
+}
 
 #[cfg(test)]
 mod test {

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

@@ -1,10 +1,13 @@
 use vortex_error::VortexResult;
-use vortex_scalar::Scalar;
+use vortex_scalar::{ExtScalar, Scalar};
 
 use crate::array::extension::ExtensionArray;
+use crate::array::ConstantArray;
 use crate::compute::unary::{scalar_at, scalar_at_unchecked, CastFn, ScalarAtFn};
-use crate::compute::{slice, take, ArrayCompute, SliceFn, TakeFn};
-use crate::{Array, IntoArray};
+use crate::compute::{
+    compare, slice, take, ArrayCompute, MaybeCompareFn, Operator, SliceFn, TakeFn,
+};
+use crate::{Array, ArrayDType, IntoArray};
 
 impl ArrayCompute for ExtensionArray {
     fn cast(&self) -> Option<&dyn CastFn> {
@@ -14,6 +17,10 @@ impl ArrayCompute for ExtensionArray {
         None
     }
 
+    fn compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        MaybeCompareFn::maybe_compare(self, array, operator)
+    }
+
     fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
         Some(self)
     }
@@ -27,6 +34,29 @@ impl ArrayCompute for ExtensionArray {
     }
 }
 
+impl MaybeCompareFn for ExtensionArray {
+    fn maybe_compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        if let Ok(const_ext) = ConstantArray::try_from(array) {
+            return Some({
+                let scalar_ext =
+                    ExtScalar::try_from(const_ext.scalar()).expect("Expected ExtScalar");
+                let const_storage = ConstantArray::new(
+                    Scalar::new(self.storage().dtype().clone(), scalar_ext.value().clone()),
+                    const_ext.len(),
+                );
+                compare(&self.storage(), const_storage.array(), operator)
+            });
+        }
+
+        // FIXME(ngates): this is not necessarily true, any other encoding could be an extension
+        if let Ok(rhs_ext) = ExtensionArray::try_from(array) {
+            return Some(compare(&self.storage(), &rhs_ext.storage(), operator));
+        }
+
+        None
+    }
+}
+
 impl ScalarAtFn for ExtensionArray {
     fn scalar_at(&self, index: usize) -> VortexResult<Scalar> {
         Ok(Scalar::extension(

vortex-array/src/array/extension/mod.rs

@@ -60,7 +60,11 @@ impl ArrayVariants for ExtensionArray {
     }
 }
 
-impl ExtensionArrayTrait for ExtensionArray {}
+impl ExtensionArrayTrait for ExtensionArray {
+    fn storage_array(&self) -> Array {
+        self.storage()
+    }
+}
 
 impl IntoCanonical for ExtensionArray {
     fn into_canonical(self) -> VortexResult<Canonical> {

vortex-array/src/array/primitive/compute/compare.rs

@@ -9,8 +9,8 @@ use crate::compute::{CompareFn, Operator};
 use crate::{Array, IntoArray, IntoArrayVariant};
 
 impl CompareFn for PrimitiveArray {
-    fn compare(&self, other: &Array, operator: Operator) -> VortexResult<Array> {
-        let other = other.clone().into_primitive()?;
+    fn compare(&self, array: &Array, operator: Operator) -> VortexResult<Array> {
+        let other = array.clone().into_primitive()?;
 
         let match_mask = match_each_native_ptype!(self.ptype(), |$T| {
             apply_predicate(self.maybe_null_slice::<$T>(), other.maybe_null_slice::<$T>(), operator.to_fn::<$T>())

vortex-array/src/array/primitive/compute/mod.rs

@@ -1,6 +1,9 @@
+use vortex_error::VortexResult;
+
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::unary::{CastFn, FillForwardFn, ScalarAtFn, SubtractScalarFn};
-use crate::compute::{ArrayCompute, CompareFn, SearchSortedFn, SliceFn, TakeFn};
+use crate::compute::{ArrayCompute, CompareFn, Operator, SearchSortedFn, SliceFn, TakeFn};
+use crate::Array;
 
 mod cast;
 mod compare;
@@ -17,8 +20,8 @@ impl ArrayCompute for PrimitiveArray {
         Some(self)
     }
 
-    fn compare(&self) -> Option<&dyn CompareFn> {
-        Some(self)
+    fn compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        Some(CompareFn::compare(self, array, operator))
     }
 
     fn fill_forward(&self) -> Option<&dyn FillForwardFn> {

vortex-array/src/array/sparse/variants.rs

@@ -86,4 +86,17 @@ impl StructArrayTrait for SparseArray {
 
 impl ListArrayTrait for SparseArray {}
 
-impl ExtensionArrayTrait for SparseArray {}
+impl ExtensionArrayTrait for SparseArray {
+    fn storage_array(&self) -> Array {
+        SparseArray::try_new_with_offset(
+            self.indices().clone(),
+            self.values()
+                .with_dyn(|a| a.as_extension_array_unchecked().storage_array()),
+            self.len(),
+            self.indices_offset(),
+            self.fill_value().clone(),
+        )
+        .expect("Failed to create new sparse array")
+        .into_array()
+    }
+}