Decode RowFn constants directly (#9680)

## Summary

- Tracking Issue: #9128
- Depends on #9626

Adds a constant-specific representation to `InputElement`, so batch
constants no longer become one-row decoded columns.

## Changes

`decode_constant` retains the logical batch length while primitive and
Boolean inputs extract one native value. Complex inputs can keep
ordinary column decoding behind the same hook.

## API Changes

Adds `InputElement::Constant`, `decode_constant`, and `get_constant`.

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

Author: connortsui20

vortex-array/src/scalar_fn/unstable/row/batch/tests.rs

@@ -98,6 +98,7 @@ struct DenseRetryI64;
 // SAFETY: the view is a slice, and its reported length is the buffer length.
 unsafe impl InputElement for ValidOnlyI64 {
     type Column = Buffer<i64>;
+    type Constant = i64;
     type View<'a> = &'a [i64];
     type Elem<'a> = i64;
 
@@ -112,6 +113,10 @@ unsafe impl InputElement for ValidOnlyI64 {
         <i64 as InputElement>::decode(array, ctx)
     }
 
+    fn decode_constant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Constant> {
+        <i64 as InputElement>::decode_constant(array, ctx)
+    }
+
     fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
         Ok(true)
     }
@@ -120,6 +125,10 @@ unsafe impl InputElement for ValidOnlyI64 {
         column[index]
     }
 
+    fn get_constant(constant: &Self::Constant) -> Self::Elem<'_> {
+        *constant
+    }
+
     fn view(column: &Self::Column) -> Self::View<'_> {
         column.as_slice()
     }
@@ -132,6 +141,7 @@ unsafe impl InputElement for ValidOnlyI64 {
 // SAFETY: the view is a slice, and its reported length is the buffer length.
 unsafe impl InputElement for FilterOnlyI64 {
     type Column = Buffer<i64>;
+    type Constant = i64;
     type View<'a> = &'a [i64];
     type Elem<'a> = i64;
 
@@ -152,10 +162,21 @@ unsafe impl InputElement for FilterOnlyI64 {
         Ok(values)
     }
 
+    fn decode_constant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Constant> {
+        let value = <i64 as InputElement>::decode_constant(array, ctx)?;
+        vortex_ensure!(value != i64::MIN, "test input contains an invalid payload",);
+
+        Ok(value)
+    }
+
     fn get(column: &Self::Column, index: usize) -> Self::Elem<'_> {
         column[index]
     }
 
+    fn get_constant(constant: &Self::Constant) -> Self::Elem<'_> {
+        *constant
+    }
+
     fn view(column: &Self::Column) -> Self::View<'_> {
         column.as_slice()
     }
@@ -168,6 +189,7 @@ unsafe impl InputElement for FilterOnlyI64 {
 // SAFETY: the view is a slice, and its reported length is the buffer length.
 unsafe impl InputElement for DenseRetryI64 {
     type Column = Buffer<i64>;
+    type Constant = i64;
     type View<'a> = &'a [i64];
     type Elem<'a> = i64;
 
@@ -182,10 +204,18 @@ unsafe impl InputElement for DenseRetryI64 {
         <i64 as InputElement>::decode(array, ctx)
     }
 
+    fn decode_constant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Constant> {
+        <i64 as InputElement>::decode_constant(array, ctx)
+    }
+
     fn get(column: &Self::Column, index: usize) -> Self::Elem<'_> {
         column[index]
     }
 
+    fn get_constant(constant: &Self::Constant) -> Self::Elem<'_> {
+        *constant
+    }
+
     fn view(column: &Self::Column) -> Self::View<'_> {
         column.as_slice()
     }

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

@@ -4,21 +4,25 @@
 use vortex_buffer::BitBuffer;
 use vortex_compute::lane_kernels::IndexedSource;
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 
 use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::BoolArray;
+use crate::arrays::Constant;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
+use crate::scalar::ScalarValue;
 use crate::scalar_fn::unstable::row::InputElement;
 use crate::scalar_fn::unstable::row::OutputElement;
 use crate::validity::Validity;
 
 // SAFETY: the view is a bit buffer, and its reported length is the buffer length.
 unsafe impl InputElement for bool {
     type Column = BitBuffer;
+    type Constant = bool;
     type View<'a> = &'a BitBuffer;
     type Elem<'a> = bool;
 
@@ -38,6 +42,21 @@ unsafe impl InputElement for bool {
         Ok(array.execute::<BoolArray>(ctx)?.into_bit_buffer())
     }
 
+    fn decode_constant(array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult<Self::Constant> {
+        let Some(constant) = array.as_opt::<Constant>() else {
+            vortex_bail!(
+                "a Boolean batch constant must use the Constant encoding, got {}",
+                array.encoding_id()
+            );
+        };
+        let scalar = constant.scalar();
+        let Some(ScalarValue::Bool(value)) = scalar.value() else {
+            vortex_bail!("a Boolean batch constant must contain a non-null value, got {scalar}");
+        };
+
+        Ok(*value)
+    }
+
     fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
         Ok(true)
     }
@@ -46,6 +65,10 @@ unsafe impl InputElement for bool {
         column.value(index)
     }
 
+    fn get_constant(constant: &Self::Constant) -> bool {
+        *constant
+    }
+
     fn view(column: &Self::Column) -> Self::View<'_> {
         column
     }

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

@@ -3,8 +3,8 @@
 
 //! Typed decoding and row access for one input column.
 //!
-//! [`InputElement`] separates invocation-wide decoding from the checked and unchecked access paths
-//! used by row kernels.
+//! [`InputElement`] separates invocation-wide column and batch-constant decoding from the checked
+//! and unchecked access paths used by row kernels.
 
 use vortex_error::VortexResult;
 
@@ -27,6 +27,11 @@ pub unsafe trait InputElement: 'static {
     /// The decoded column representation supporting `O(1)` row access.
     type Column;
 
+    /// The decoded representation of one non-null batch-constant input.
+    ///
+    /// This representation stores one logical value regardless of the input array's batch length.
+    type Constant;
+
     /// The row-loop view of a decoded column.
     ///
     /// This can borrow a cheaper representation than [`Column`](Self::Column). Primitive elements,
@@ -62,6 +67,12 @@ pub unsafe trait InputElement: 'static {
     /// and other invocation-invariant work into this method.
     fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column>;
 
+    /// Decode one non-null batch-constant input.
+    ///
+    /// `array` retains its logical batch length. Implementations can extract one value directly
+    /// without slicing or materializing a one-row column.
+    fn decode_constant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Constant>;
+
     /// Whether [`decode_null_tolerant`](Self::decode_null_tolerant) can decode this array.
     ///
     /// The conservative default declines. An implementation whose ordinary decode is safe and
@@ -90,6 +101,9 @@ pub unsafe trait InputElement: 'static {
     /// Read one row without repeating batch-constant work from [`decode`](Self::decode).
     fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>;
 
+    /// Read the element stored in a decoded batch constant.
+    fn get_constant(constant: &Self::Constant) -> Self::Elem<'_>;
+
     /// Borrow the representation used inside the row loop.
     ///
     /// Executors call this before the hot loop. For every index below the returned view's length,

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

@@ -9,17 +9,20 @@ use vortex_error::vortex_ensure_eq;
 use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
+use crate::arrays::Constant;
 use crate::arrays::PrimitiveArray;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
 use crate::dtype::Nullability;
+use crate::scalar::ScalarValue;
 use crate::scalar_fn::unstable::row::InputElement;
 use crate::scalar_fn::unstable::row::OutputElement;
 use crate::validity::Validity;
 
 // SAFETY: the view is a native slice, and its reported length is the slice length.
 unsafe impl<T: NativePType> InputElement for T {
     type Column = Buffer<T>;
+    type Constant = T;
     type View<'a> = &'a [T];
     type Elem<'a> = T;
 
@@ -44,6 +47,24 @@ unsafe impl<T: NativePType> InputElement for T {
         Ok(array.execute::<PrimitiveArray>(ctx)?.into_buffer::<T>())
     }
 
+    fn decode_constant(array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult<Self::Constant> {
+        let Some(constant) = array.as_opt::<Constant>() else {
+            vortex_bail!(
+                "a primitive batch constant must use the Constant encoding, got {}",
+                array.encoding_id()
+            );
+        };
+        let scalar = constant.scalar();
+        let Some(ScalarValue::Primitive(value)) = scalar.value() else {
+            vortex_bail!(
+                "a primitive batch constant must contain a non-null {} value, got {scalar}",
+                T::PTYPE
+            );
+        };
+
+        value.cast::<T>()
+    }
+
     fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
         Ok(true)
     }
@@ -52,6 +73,10 @@ unsafe impl<T: NativePType> InputElement for T {
         column[index]
     }
 
+    fn get_constant(constant: &Self::Constant) -> T {
+        *constant
+    }
+
     fn view(column: &Self::Column) -> Self::View<'_> {
         column.as_slice()
     }

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

@@ -21,7 +21,7 @@ use crate::scalar_fn::ExecutionArgs;
 use crate::scalar_fn::unstable::row::InputElement;
 use crate::scalar_fn::unstable::row::ViewLen;
 
-/// One decoded input, collapsed to a single row when it is constant for the batch.
+/// One decoded input, with batch constants represented by a single value.
 pub struct ArgColumn<T: InputElement>(
     /// The decoded argument, classified by how the row loop addresses it.
     pub(super) ArgColumnKind<T>,
@@ -31,40 +31,34 @@ pub(super) enum ArgColumnKind<T: InputElement> {
     /// A decoded column covering the full batch; executors validate its length before traversal.
     Column(T::Column),
 
-    /// Exactly one decoded row, established by [`ArgColumn::try_from_const`].
-    Const(T::Column),
+    /// One decoded batch-constant value.
+    Const(T::Constant),
 }
 
 impl<T: InputElement> ArgColumn<T> {
-    fn try_from_const(column: T::Column) -> VortexResult<Self> {
-        let decoded_len = T::view(&column).len();
-        vortex_ensure_eq!(
-            decoded_len,
-            1,
-            "a decoded batch-constant input must contain exactly 1 row, got {decoded_len}",
-        );
-
-        Ok(Self(ArgColumnKind::Const(column)))
-    }
-
     fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
-        // An empty input has no row 0 to slice, and its row loop runs zero times either way.
+        // An empty input has no value to decode, and its row loop runs zero times either way.
         if let Some(const_array) = batch_const(&array)
             && !array.is_empty()
         {
-            return Self::try_from_const(T::decode(const_array.slice(0..1)?, ctx)?);
+            return Ok(Self(ArgColumnKind::Const(T::decode_constant(
+                const_array,
+                ctx,
+            )?)));
         }
 
         Ok(Self(ArgColumnKind::Column(T::decode(array, ctx)?)))
     }
 
     fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Option<Self>> {
-        // Batch execution short-circuits null constants before selecting a strategy, so a
-        // constant reaching this path is non-null and can use the ordinary decode.
+        // Batch execution short-circuits null constants before selecting a strategy.
         if let Some(const_array) = batch_const(&array)
             && !array.is_empty()
         {
-            return Self::try_from_const(T::decode(const_array.slice(0..1)?, ctx)?).map(Some);
+            return T::decode_constant(const_array, ctx)
+                .map(ArgColumnKind::Const)
+                .map(Self)
+                .map(Some);
         }
 
         Ok(T::decode_null_tolerant(array, ctx)?
@@ -85,7 +79,7 @@ impl<T: InputElement> ArgColumn<T> {
     fn get(&self, index: usize) -> T::Elem<'_> {
         match &self.0 {
             ArgColumnKind::Column(column) => T::get(column, index),
-            ArgColumnKind::Const(column) => T::get(column, 0),
+            ArgColumnKind::Const(constant) => T::get_constant(constant),
         }
     }
 
@@ -97,7 +91,7 @@ impl<T: InputElement> ArgColumn<T> {
     }
 
     fn addresses_rows(&self, row_count: usize) -> bool {
-        // A constant is validated when constructed and is always read at index zero.
+        // A constant represents one value for every logical row.
         match &self.0 {
             ArgColumnKind::Column(column) => T::view(column).len() == row_count,
             ArgColumnKind::Const(_) => true,
@@ -107,7 +101,7 @@ impl<T: InputElement> ArgColumn<T> {
     fn const_value(&self) -> Option<T::Elem<'_>> {
         match &self.0 {
             ArgColumnKind::Column(_) => None,
-            ArgColumnKind::Const(column) => Some(T::get(column, 0)),
+            ArgColumnKind::Const(constant) => Some(T::get_constant(constant)),
         }
     }
 }
@@ -215,8 +209,8 @@ pub trait ElementTuple: 'static + private::Sealed {
     /// Whether every argument decoded at full batch length contains exactly `row_count` rows.
     ///
     /// `Columns` cannot implement [`ViewLen`] because a batch-constant `ArgColumn` stores one
-    /// decoded row while logically addressing the full batch. The batch-constant constructor
-    /// validates that one-row representation, so this method checks only non-constant columns.
+    /// decoded value while logically addressing the full batch. This method therefore checks only
+    /// non-constant columns.
     fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool;
 
     /// Read one row from borrowed views.

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

@@ -52,9 +52,9 @@ pub(in crate::scalar_fn::unstable::row) fn decoded_source<'a, Args: IndexedEleme
 enum ArgColumnSource<'a, T: InputElement> {
     Rows(T::View<'a>),
 
-    /// A validated one-row view that logically addresses `row_count` rows.
+    /// One decoded value that logically addresses `row_count` rows.
     Constant {
-        view: T::View<'a>,
+        constant: &'a T::Constant,
         row_count: usize,
     },
 }
@@ -66,10 +66,10 @@ impl<'a, T: InputElement> ArgColumnSource<'a, T> {
                 let view = T::view(column);
                 (view.len() == row_count).then_some(Self::Rows(view))
             }
-            ArgColumnKind::Const(column) => {
-                let view = T::view(column);
-                (view.len() == 1).then_some(Self::Constant { view, row_count })
-            }
+            ArgColumnKind::Const(constant) => Some(Self::Constant {
+                constant,
+                row_count,
+            }),
         }
     }
 }
@@ -91,10 +91,7 @@ impl<'a, T: InputElement> IndexedSource for ArgColumnSource<'a, T> {
                 // caller guarantees that `index` is below the source length.
                 unsafe { T::get_from_view_unchecked(view, index) }
             }
-            Self::Constant { view, .. } => {
-                // SAFETY: `try_new` checked that this exact retained view contains row zero.
-                unsafe { T::get_from_view_unchecked(view, 0) }
-            }
+            Self::Constant { constant, .. } => T::get_constant(constant),
         }
     }
 }