Support borrowed RowFn input decoding

Keep the owned decoder as the required extension point and default borrowed decoding to it. Native primitive inputs share their canonical buffers directly. Retain decoded storage after the input handle is dropped.

Extracted from #9979.

Signed-off-by: Robert Kruszewski <github@robertk.io>
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: robert3005

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];