Optional fallback in new is_constant compute function (#2485)

Author: AdamGS

encodings/alp/src/alp/compute/mod.rs

@@ -5,7 +5,7 @@ use std::fmt::Debug;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::compute::{
     between, filter, scalar_at, slice, take, BetweenFn, BetweenOptions, CompareFn, FilterFn,
-    IsConstantFn, ScalarAtFn, SliceFn, StrictComparison, TakeFn,
+    ScalarAtFn, SliceFn, StrictComparison, TakeFn,
 };
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::vtable::ComputeVTable;
@@ -30,10 +30,6 @@ impl ComputeVTable for ALPEncoding {
         Some(self)
     }
 
-    fn is_constant_fn(&self) -> Option<&dyn IsConstantFn<&dyn Array>> {
-        Some(self)
-    }
-
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
         Some(self)
     }
@@ -121,12 +117,6 @@ impl FilterFn<&ALPArray> for ALPEncoding {
     }
 }
 
-impl IsConstantFn<&ALPArray> for ALPEncoding {
-    fn is_constant(&self, _array: &ALPArray) -> VortexResult<Option<bool>> {
-        Ok(None)
-    }
-}
-
 impl BetweenFn<&ALPArray> for ALPEncoding {
     fn between(
         &self,

encodings/alp/src/alp_rd/compute/is_constant.rs

@@ -1,21 +1,16 @@
-use vortex_array::compute::{FilterFn, IsConstantFn, MaskFn, ScalarAtFn, SliceFn, TakeFn};
+use vortex_array::compute::{FilterFn, MaskFn, ScalarAtFn, SliceFn, TakeFn};
 use vortex_array::vtable::ComputeVTable;
 use vortex_array::Array;
 
 use crate::ALPRDEncoding;
 
 mod filter;
-mod is_constant;
 mod mask;
 mod scalar_at;
 mod slice;
 mod take;
 
 impl ComputeVTable for ALPRDEncoding {
-    fn is_constant_fn(&self) -> Option<&dyn IsConstantFn<&dyn Array>> {
-        Some(self)
-    }
-
     fn filter_fn(&self) -> Option<&dyn FilterFn<&dyn Array>> {
         Some(self)
     }

encodings/alp/src/alp_rd/compute/mod.rs

@@ -1,6 +1,5 @@
 use vortex_array::compute::{
-    between, BetweenFn, BetweenOptions, FilterFn, IsConstantFn, ScalarAtFn, SearchSortedFn,
-    SliceFn, TakeFn,
+    between, BetweenFn, BetweenOptions, FilterFn, ScalarAtFn, SearchSortedFn, SliceFn, TakeFn,
 };
 use vortex_array::vtable::ComputeVTable;
 use vortex_array::{Array, ArrayRef, IntoArray};
@@ -9,7 +8,6 @@ use vortex_error::VortexResult;
 use crate::{BitPackedArray, BitPackedEncoding};
 
 mod filter;
-mod is_constant;
 mod scalar_at;
 mod search_sorted;
 mod slice;
@@ -20,10 +18,6 @@ impl ComputeVTable for BitPackedEncoding {
         Some(self)
     }
 
-    fn is_constant_fn(&self) -> Option<&dyn IsConstantFn<&dyn Array>> {
-        Some(self)
-    }
-
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
         Some(self)
     }

encodings/fastlanes/src/bitpacking/compute/is_constant.rs

@@ -3,7 +3,7 @@ use std::sync::RwLock;
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_scalar::{Scalar, ScalarValue};
 
-use crate::compute::{is_constant, min_max, scalar_at, sum, MinMaxResult};
+use crate::compute::{is_constant_opts, min_max, scalar_at, sum, IsConstantOpts, MinMaxResult};
 use crate::stats::{Precision, Stat, Statistics, StatsSet};
 use crate::{Array, ArrayImpl};
 
@@ -86,7 +86,15 @@ impl<A: Array + ArrayImpl> Statistics for A {
                 if self.is_empty() {
                     None
                 } else {
-                    Some(is_constant(self)?.into())
+                    Some(
+                        is_constant_opts(
+                            self,
+                            &IsConstantOpts {
+                                fallback_to_canonicalize: true,
+                            },
+                        )?
+                        .into(),
+                    )
                 }
             }
             _ => {

encodings/fastlanes/src/bitpacking/compute/mod.rs

@@ -84,6 +84,7 @@ impl MinMaxFn<&ListArray> for ListEncoding {
 
 impl IsConstantFn<&ListArray> for ListEncoding {
     fn is_constant(&self, _array: &ListArray) -> VortexResult<Option<bool>> {
+        // TODO(adam): Do we want to fallback to arrow here?
         Ok(None)
     }
 }

vortex-array/src/array/statistics.rs

@@ -7,6 +7,7 @@ use vortex_dtype::DType;
 use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 
 use crate::arrow::{FromArrowArray, IntoArrowArray};
+use crate::compute::is_constant;
 use crate::encoding::Encoding;
 use crate::{Array, ArrayRef};
 
@@ -85,13 +86,15 @@ pub fn binary_boolean(
         vortex_bail!("Boolean operations are only supported on boolean arrays")
     }
 
+    let rhs_is_constant = is_constant(rhs)?;
+
     // If LHS is constant, then we make sure it's on the RHS.
-    if lhs.is_constant() && !rhs.is_constant() {
+    if is_constant(lhs)? && !rhs_is_constant {
         return binary_boolean(rhs, lhs, op);
     }
 
     // If the RHS is constant and the LHS is Arrow, we can't do any better than arrow_compare.
-    if lhs.is_arrow() && (rhs.is_arrow() || rhs.is_constant()) {
+    if lhs.is_arrow() && (rhs.is_arrow() || rhs_is_constant) {
         return arrow_boolean(lhs.to_array(), rhs.to_array(), op);
     }
 

vortex-array/src/arrays/list/compute/mod.rs

@@ -7,6 +7,7 @@ use vortex_dtype::{DType, NativePType, Nullability};
 use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 use vortex_scalar::Scalar;
 
+use super::is_constant;
 use crate::arrays::ConstantArray;
 use crate::arrow::{from_arrow_array_with_len, Datum};
 use crate::encoding::Encoding;
@@ -124,8 +125,11 @@ pub fn compare(left: &dyn Array, right: &dyn Array, operator: Operator) -> Vorte
         return Ok(ConstantArray::new(Scalar::null(result_dtype), left.len()).into_array());
     }
 
+    let left_is_constant = is_constant(left)?;
+    let right_is_constant = is_constant(right)?;
+
     // Always try to put constants on the right-hand side so encodings can optimise themselves.
-    if left.is_constant() && !right.is_constant() {
+    if left_is_constant && !right_is_constant {
         return compare(right, left, operator.swap());
     }
 
@@ -151,7 +155,7 @@ pub fn compare(left: &dyn Array, right: &dyn Array, operator: Operator) -> Vorte
 
     // Only log missing compare implementation if there's possibly better one than arrow,
     // i.e. lhs isn't arrow or rhs isn't arrow or constant
-    if !(left.is_arrow() && (right.is_arrow() || right.is_constant())) {
+    if !(left.is_arrow() && (right.is_arrow() || right_is_constant)) {
         log::debug!(
             "No compare implementation found for LHS {}, RHS {}, and operator {} (or inverse)",
             right.encoding(),

vortex-array/src/compute/boolean.rs

@@ -27,6 +27,16 @@ where
     }
 }
 
+#[derive(Default, Clone)]
+pub struct IsConstantOpts {
+    pub fallback_to_canonicalize: bool,
+}
+
+pub fn is_constant(array: &dyn Array) -> VortexResult<bool> {
+    let opts = IsConstantOpts::default();
+    is_constant_opts(array, &opts)
+}
+
 /// Computes whether an array has constant values.
 /// An array is constant IFF at least one of the following conditions apply:
 /// 1. It has one elements.
@@ -37,7 +47,7 @@ where
 ///
 /// If the array has some null values but is not all null, it'll never be constant.
 /// **Please note:** Might return false negatives if a specific encoding couldn't make a determination.
-pub fn is_constant(array: &dyn Array) -> VortexResult<bool> {
+pub fn is_constant_opts(array: &dyn Array, opts: &IsConstantOpts) -> VortexResult<bool> {
     match array.len() {
         // Our current semantics are that we can always get a value out of a constant array. We might want to change that in the future.
         0 => return Ok(false),
@@ -96,15 +106,19 @@ pub fn is_constant(array: &dyn Array) -> VortexResult<bool> {
             array.encoding()
         );
 
-        let array = array.to_canonical()?;
-
-        if let Some(is_constant_fn) = array.as_ref().vtable().is_constant_fn() {
-            is_constant_fn.is_constant(array.as_ref())?
+        if opts.fallback_to_canonicalize {
+            let array = array.to_canonical()?;
+
+            if let Some(is_constant_fn) = array.as_ref().vtable().is_constant_fn() {
+                is_constant_fn.is_constant(array.as_ref())?
+            } else {
+                vortex_bail!(
+                    "No is_constant function for canonical array: {}",
+                    array.as_ref().encoding(),
+                )
+            }
         } else {
-            vortex_bail!(
-                "No is_constant function for canonical array: {}",
-                array.as_ref().encoding(),
-            )
+            None
         }
     };