Make null-tolerant decoding opt in

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

Author: connortsui20

vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs

@@ -37,6 +37,10 @@ unsafe impl InputElement for bool {
         Ok(array.execute::<BoolArray>(ctx)?.into_bit_buffer())
     }
 
+    fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
+        Ok(true)
+    }
+
     fn get(column: &Self::Column, index: usize) -> bool {
         column.value(index)
     }

vortex-array/src/scalar_fn/unstable/row/types/element/input.rs

@@ -57,8 +57,17 @@ pub unsafe trait InputElement: 'static {
     /// invocation-invariant work into this method.
     fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column>;
 
-    /// Decode `array` _without_ assuming every row is valid, or `Ok(None)` when this element
-    /// cannot for this particular array.
+    /// Whether [`decode_null_tolerant`](Self::decode_null_tolerant) can decode this array.
+    ///
+    /// The conservative default declines. An implementation whose ordinary decode is safe and
+    /// infallible over null payloads can return `true`. Other implementations can inspect `array`
+    /// and opt in only for supported representations.
+    fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
+        Ok(false)
+    }
+
+    /// Decode `array` _without_ assuming every row is valid, or return `Ok(None)` when this element
+    /// cannot decode this particular array.
     ///
     /// Override this for a non-dense-safe representation that can still place safe placeholders in
     /// null slots. The skip-invalid executor never reads those slots.
@@ -69,7 +78,7 @@ pub unsafe trait InputElement: 'static {
         array: ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<Self::Column>> {
-        if Self::DENSE_SAFE {
+        if Self::can_decode_null_tolerant(&array)? {
             Self::decode(array, ctx).map(Some)
         } else {
             Ok(None)

vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs

@@ -44,6 +44,10 @@ unsafe impl<T: NativePType> InputElement for T {
         Ok(array.execute::<PrimitiveArray>(ctx)?.into_buffer::<T>())
     }
 
+    fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
+        Ok(true)
+    }
+
     fn get(column: &Self::Column, index: usize) -> T {
         column[index]
     }

vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs

@@ -63,6 +63,16 @@ impl<T: InputElement> ArgColumn<T> {
             .map(Self))
     }
 
+    fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult<bool> {
+        // Batch execution short-circuits null constants before selecting this path, so a
+        // non-empty constant can always use the ordinary decode.
+        if batch_constant(array).is_some() && !array.is_empty() {
+            return Ok(true);
+        }
+
+        T::can_decode_null_tolerant(array)
+    }
+
     fn get(&self, index: usize) -> T::Elem<'_> {
         match &self.0 {
             ArgColumnKind::PerRow(column) => T::get(column, index),
@@ -153,6 +163,12 @@ pub trait ElementTuple: 'static + private::Sealed {
     /// A dense deferred-error retry starts another invocation over filtered valid rows.
     fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult<Self::Columns>;
 
+    /// Whether every input can be decoded without assuming that all rows are valid.
+    ///
+    /// The tuple checks this before decoding any column, so a decline does not discard work from
+    /// earlier arguments.
+    fn can_decode_null_tolerant(args: &dyn ExecutionArgs) -> VortexResult<bool>;
+
     /// Decode every input column once while tolerating null rows.
     ///
     /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid
@@ -225,6 +241,10 @@ impl ElementTuple for () {
         Ok(())
     }
 
+    fn can_decode_null_tolerant(_args: &dyn ExecutionArgs) -> VortexResult<bool> {
+        Ok(true)
+    }
+
     fn decode_null_tolerant(
         _args: &dyn ExecutionArgs,
         _ctx: &mut ExecutionCtx,
@@ -291,10 +311,21 @@ macro_rules! element_tuple {
                 Ok(($(ArgColumn::<$t>::decode(args.get($idx)?, ctx)?,)+))
             }
 
+            fn can_decode_null_tolerant(args: &dyn ExecutionArgs) -> VortexResult<bool> {
+                Ok($({
+                    let array = args.get($idx)?;
+                    ArgColumn::<$t>::can_decode_null_tolerant(&array)?
+                } &&)+ true)
+            }
+
             fn decode_null_tolerant(
                 args: &dyn ExecutionArgs,
                 ctx: &mut ExecutionCtx,
             ) -> VortexResult<Option<Self::Columns>> {
+                if !Self::can_decode_null_tolerant(args)? {
+                    return Ok(None);
+                }
+
                 Ok(Some((
                     $(match ArgColumn::<$t>::decode_null_tolerant(args.get($idx)?, ctx)? {
                         Some(column) => column,

vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs

@@ -1,21 +1,114 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_mask::Mask;
 
+use super::ElementTuple;
 use super::element_tuple::batch_constant;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
+use crate::VortexSessionExecute;
+use crate::array_session;
 use crate::arrays::Constant;
 use crate::arrays::ConstantArray;
 use crate::arrays::ExtensionArray;
 use crate::arrays::MaskedArray;
+use crate::arrays::PrimitiveArray;
+use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::extension::datetime::TimeUnit;
 use crate::extension::datetime::Timestamp;
+use crate::scalar_fn::VecExecutionArgs;
+use crate::scalar_fn::unstable::row::InputElement;
 use crate::validity::Validity;
 
+static DECODE_CALLS: AtomicUsize = AtomicUsize::new(0);
+
+macro_rules! i64_test_element {
+    ($element:ident, $decode_fallible:literal $(, $can_decode:item)?) => {
+        struct $element;
+
+        // SAFETY: the view and unchecked access delegate to the `i64` implementation.
+        unsafe impl InputElement for $element {
+            type Column = Buffer<i64>;
+            type View<'a> = &'a [i64];
+            type Elem<'a> = i64;
+
+            const DENSE_SAFE: bool = true;
+            const DECODE_FALLIBLE: bool = $decode_fallible;
+
+            fn validate(dtype: &DType) -> VortexResult<()> {
+                <i64 as InputElement>::validate(dtype)
+            }
+
+            fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column> {
+                DECODE_CALLS.fetch_add(1, Ordering::Relaxed);
+                <i64 as InputElement>::decode(array, ctx)
+            }
+
+            $($can_decode)?
+
+            fn get(column: &Self::Column, index: usize) -> i64 {
+                <i64 as InputElement>::get(column, index)
+            }
+
+            fn view(column: &Self::Column) -> Self::View<'_> {
+                <i64 as InputElement>::view(column)
+            }
+
+            fn view_len(view: &Self::View<'_>) -> usize {
+                <i64 as InputElement>::view_len(view)
+            }
+
+            fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> i64
+            where
+                Self: 'a,
+            {
+                <i64 as InputElement>::get_from_view(view, index)
+            }
+
+            unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> i64
+            where
+                Self: 'a,
+            {
+                // SAFETY: forwarded from this method's contract.
+                unsafe { <i64 as InputElement>::get_from_view_unchecked(view, index) }
+            }
+        }
+    };
+}
+
+i64_test_element!(
+    DecodeProbe,
+    false,
+    fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
+        Ok(true)
+    }
+);
+i64_test_element!(DenseFallible, true);
+
+#[test]
+fn test_null_tolerant_decline_precedes_decoding() -> VortexResult<()> {
+    DECODE_CALLS.store(0, Ordering::Relaxed);
+    let first = PrimitiveArray::from_iter([1_i64, 2]).into_array();
+    let second = PrimitiveArray::from_iter([3_i64, 4]).into_array();
+    let args = VecExecutionArgs::new(vec![first, second], 2);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let columns = <(DecodeProbe, DenseFallible)>::decode_null_tolerant(&args, &mut ctx)?;
+
+    assert!(columns.is_none());
+    assert_eq!(DECODE_CALLS.load(Ordering::Relaxed), 0);
+    Ok(())
+}
+
 #[test]
 fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> {
     let child = ConstantArray::new(7_i64, 3).into_array();