Support borrowed RowFn input decoding

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

@@ -49,7 +49,8 @@ pub unsafe trait InputElement: 'static {
     /// null rows to the row closure.
     const DENSE_SAFE: bool;
 
-    /// Whether [`decode`](Self::decode) is infallible for _legal_ input data.
+    /// Whether [`decode`](Self::decode) and [`decode_ref`](Self::decode_ref) are infallible for
+    /// _legal_ input data.
     ///
     /// This excludes infrastructural failures such as IO or allocation.
     /// It is independent of
@@ -67,6 +68,18 @@ pub unsafe trait InputElement: 'static {
     /// and other invocation-invariant work into this method.
     fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column>;
 
+    /// Decode a borrowed `array` into its column representation.
+    ///
+    /// The default clones `array` for [`decode`](Self::decode). Override this when a common
+    /// representation decodes without taking ownership of the array.
+    ///
+    /// The result must preserve the valid-row values and row count of [`decode`](Self::decode),
+    /// and must retain its storage after `array` is dropped. [`DENSE_SAFE`](Self::DENSE_SAFE) and
+    /// [`DECODE_INFALLIBLE`](Self::DECODE_INFALLIBLE) apply to this method too.
+    fn decode_ref(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column> {
+        Self::decode(array.clone(), ctx)
+    }
+
     /// Decode one non-null batch-constant input.
     ///
     /// `array` retains its logical batch length. Implementations can extract one value directly

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

@@ -14,6 +14,7 @@ use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::Constant;
+use crate::arrays::Primitive;
 use crate::arrays::PrimitiveArray;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
@@ -45,9 +46,23 @@ unsafe impl<T: NativePType> InputElement for T {
     }
 
     fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column> {
+        // A canonical input shares its buffer without executing the array.
+        if let Some(primitive) = array.as_opt::<Primitive>() {
+            return Ok(primitive.to_buffer::<T>());
+        }
+
         Ok(array.execute::<PrimitiveArray>(ctx)?.into_buffer::<T>())
     }
 
+    #[inline]
+    fn decode_ref(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column> {
+        if let Some(primitive) = array.as_opt::<Primitive>() {
+            return Ok(primitive.to_buffer::<T>());
+        }
+
+        Self::decode(array.clone(), ctx)
+    }
+
     fn decode_constant(array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult<Self::Constant> {
         let Some(constant) = array.as_opt::<Constant>() else {
             vortex_bail!(

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

@@ -4,6 +4,7 @@
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 
+use rstest::rstest;
 use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -16,6 +17,7 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::VortexSessionExecute;
 use crate::array_session;
+use crate::arrays::BoolArray;
 use crate::arrays::Constant;
 use crate::arrays::ConstantArray;
 use crate::arrays::ExtensionArray;
@@ -105,6 +107,32 @@ i64_test_element!(
 );
 i64_test_element!(DenseFallible, false);
 
+#[rstest]
+#[case::canonical(PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(), vec![1, 2, 3])]
+#[case::encoded(ConstantArray::new(7_i64, 3).into_array(), vec![7, 7, 7])]
+fn test_borrowed_primitive_decode_retains_storage(
+    #[case] array: ArrayRef,
+    #[case] expected: Vec<i64>,
+) -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let column = <i64 as InputElement>::decode_ref(&array, &mut ctx)?;
+    drop(array);
+
+    assert_eq!(column.as_slice(), expected);
+    Ok(())
+}
+
+#[test]
+fn test_default_borrowed_decode_retains_storage() -> VortexResult<()> {
+    let array = BoolArray::from_iter([true, false, true]).into_array();
+    let mut ctx = array_session().create_execution_ctx();
+    let column = <bool as InputElement>::decode_ref(&array, &mut ctx)?;
+    drop(array);
+
+    assert_eq!(column.iter().collect::<Vec<_>>(), [true, false, true]);
+    Ok(())
+}
+
 #[test]
 fn test_view_lens_match_checks_each_view() {
     let first: &[i64] = &[1, 2];