Add ColumnarValue::to_array_variant method (follow-up to #22784)

datafusion/expr-common/src/columnar_value.rs

@@ -226,6 +226,26 @@ impl ColumnarValue {
         }
     }
 
+    /// Converts this value into the [`Self::Array`] variant, returning a
+    /// [`ColumnarValue`] that is guaranteed to be [`Self::Array`].
+    ///
+    /// A [`Self::Scalar`] is expanded into an array of `num_rows` by repeating
+    /// the value (which is less efficient than handling the scalar directly); a
+    /// [`Self::Array`] is returned unchanged, without validating its length.
+    ///
+    /// Unlike [`Self::to_array`], which returns the bare [`ArrayRef`], this keeps
+    /// the value wrapped as a [`ColumnarValue`]. This is useful when later logic
+    /// still threads a [`ColumnarValue`] but needs to assume a non-scalar (array)
+    /// input.
+    ///
+    /// # Errors
+    ///
+    /// Errors if `self` is a Scalar that fails to be converted into an array of
+    /// the requested size.
+    pub fn to_array_variant(&self, num_rows: usize) -> Result<ColumnarValue> {
+        Ok(ColumnarValue::Array(self.to_array(num_rows)?))
+    }
+
     /// Null columnar values are implemented as a null array in order to pass batch
     /// num_rows
     pub fn create_null_array(num_rows: usize) -> Self {
@@ -497,6 +517,36 @@ mod tests {
         );
     }
 
+    #[test]
+    fn to_array_variant() {
+        // Scalar is expanded into an array of the requested length.
+        let scalar = ColumnarValue::Scalar(ScalarValue::Int32(Some(42)));
+        match scalar.to_array_variant(3).unwrap() {
+            ColumnarValue::Array(arr) => assert_eq!(&arr, &make_array(42, 3)),
+            ColumnarValue::Scalar(_) => panic!("expected the Array variant"),
+        }
+
+        // Array is returned unchanged.
+        let arr = make_array(1, 3);
+        match ColumnarValue::Array(Arc::clone(&arr))
+            .to_array_variant(3)
+            .unwrap()
+        {
+            ColumnarValue::Array(out) => assert_eq!(&out, &arr),
+            ColumnarValue::Scalar(_) => panic!("expected the Array variant"),
+        }
+
+        // Unlike `to_array_of_size`, an existing array is not length-validated:
+        // `num_rows` is ignored when the value is already an array.
+        match ColumnarValue::Array(Arc::clone(&arr))
+            .to_array_variant(5)
+            .unwrap()
+        {
+            ColumnarValue::Array(out) => assert_eq!(&out, &arr),
+            ColumnarValue::Scalar(_) => panic!("expected the Array variant"),
+        }
+    }
+
     #[test]
     fn values_to_arrays() {
         // (input, expected)

datafusion/functions-nested/src/map.rs

@@ -63,13 +63,6 @@ fn can_evaluate_to_const(args: &[ColumnarValue]) -> bool {
         .all(|arg| matches!(arg, ColumnarValue::Scalar(_)))
 }
 
-fn expand_if_scalar(arg: &ColumnarValue, rows: usize) -> Result<ColumnarValue> {
-    match arg {
-        ColumnarValue::Scalar(s) => Ok(ColumnarValue::Array(s.to_array_of_size(rows)?)),
-        ColumnarValue::Array(a) => Ok(ColumnarValue::Array(Arc::clone(a))),
-    }
-}
-
 fn make_map_batch(args: &[ColumnarValue], number_rows: usize) -> Result<ColumnarValue> {
     let [keys_arg, values_arg] = take_function_args("make_map", args)?;
 
@@ -79,8 +72,8 @@ fn make_map_batch(args: &[ColumnarValue], number_rows: usize) -> Result<Columnar
     // are expanded to arrays which following logic expects
     let (keys_arg, values_arg) = if !can_evaluate_to_const {
         (
-            expand_if_scalar(keys_arg, number_rows)?,
-            expand_if_scalar(values_arg, number_rows)?,
+            keys_arg.to_array_variant(number_rows)?,
+            values_arg.to_array_variant(number_rows)?,
         )
     } else {
         (keys_arg.clone(), values_arg.clone())