feat: support Spark-compatible arrays_zip function

datafusion/functions-nested/src/arrays_zip.rs

@@ -17,17 +17,14 @@
 
 //! [`ScalarUDFImpl`] definitions for arrays_zip function.
 
-use crate::utils::make_scalar_function;
-use arrow::array::{
-    Array, ArrayRef, Capacities, ListArray, MutableArrayData, StructArray, new_null_array,
-};
+use arrow::array::{Array, ArrayRef, Capacities, ListArray, MutableArrayData, StringArray, StructArray, new_null_array, Int32Array};
 use arrow::buffer::{NullBuffer, OffsetBuffer};
 use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null};
 use arrow::datatypes::{DataType, Field, Fields};
 use datafusion_common::cast::{
     as_fixed_size_list_array, as_large_list_array, as_list_array,
 };
-use datafusion_common::{Result, exec_err};
+use datafusion_common::{Result, ScalarValue, exec_err};
 use datafusion_expr::{
     ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
 };
@@ -62,13 +59,13 @@ make_udf_expr_and_func!(
 +---------------------------------------------------+
 | arrays_zip([1, 2, 3], ['a', 'b', 'c'])             |
 +---------------------------------------------------+
-| [{c0: 1, c1: a}, {c0: 2, c1: b}, {c0: 3, c1: c}] |
+| [{1: 1, 2: a}, {1: 2, 2: b}, {1: 3, 2: c}] |
 +---------------------------------------------------+
 > select arrays_zip([1, 2], [3, 4, 5]);
 +---------------------------------------------------+
 | arrays_zip([1, 2], [3, 4, 5])                       |
 +---------------------------------------------------+
-| [{c0: 1, c1: 3}, {c0: 2, c1: 4}, {c0: , c1: 5}]  |
+| [{1: 1, 2: 3}, {1: 2, 2: 4}, {1: NULL, 2: 5}]  |
 +---------------------------------------------------+
 ```"#,
     argument(name = "array1", description = "First array expression."),
@@ -138,7 +135,33 @@ impl ScalarUDFImpl for ArraysZip {
         &self,
         args: datafusion_expr::ScalarFunctionArgs,
     ) -> Result<ColumnarValue> {
-        make_scalar_function(arrays_zip_inner)(&args.args)
+        let args = &args.args;
+        let strings_vec: Vec<String> = (1..args.len() + 1)
+            .into_iter()
+            .map(|i| i.to_string())
+            .collect();
+        let names = vec![Arc::new(StringArray::from(strings_vec)) as ArrayRef];
+
+        let len = args
+            .iter()
+            .fold(Option::<usize>::None, |acc, arg| match arg {
+                ColumnarValue::Scalar(_) => acc,
+                ColumnarValue::Array(a) => Some(a.len()),
+            });
+
+        let is_scalar = len.is_none();
+
+        let args = ColumnarValue::values_to_arrays(args)?;
+
+        let result = (arrays_zip_inner)(&args, &names);
+
+        if is_scalar {
+            // If all inputs are scalar, keeps output as scalar
+            let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
+            result.map(ColumnarValue::Scalar)
+        } else {
+            result.map(ColumnarValue::Array)
+        }
     }
 
     fn aliases(&self) -> &[String] {
@@ -156,11 +179,19 @@ impl ScalarUDFImpl for ArraysZip {
 /// has one field per input array. If arrays within a row have different
 /// lengths, shorter arrays are padded with NULLs.
 /// Supports List, LargeList, and Null input types.
-fn arrays_zip_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
+pub fn arrays_zip_inner(args: &[ArrayRef], names: &[ArrayRef]) -> Result<ArrayRef> {
     if args.len() < 2 {
         return exec_err!("arrays_zip requires at least two arguments");
     }
 
+    // if args.len() != names.len() {
+    //     return exec_err!(
+    //         "The numbers of zipped arrays: {} and field names: {} should be the same",
+    //         args.len(),
+    //         names.len()
+    //     );
+    // }
+
     let num_rows = args[0].len();
 
     // Build a type-erased ListColumnView for each argument.
@@ -224,10 +255,20 @@ fn arrays_zip_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
         .map(|v| v.as_ref().map(|view| view.values.to_data()))
         .collect();
 
+    let n: Vec<Vec<&str>>  = names.iter().map(|child| {
+        let values = child.as_any().downcast_ref::<StringArray>().unwrap();
+        values.iter().map(|v: Option<&str>| v.unwrap()).collect()
+    }).collect();
+
+    // dbg!("{}", &n[0]);
+
     let struct_fields: Fields = element_types
         .iter()
         .enumerate()
-        .map(|(i, dt)| Field::new(format!("{}", i + 1), dt.clone(), true))
+        .map(|(i, dt)| {
+            println!("{}", &n[0][i]);
+            Field::new(format!("{}", &n[0][i]), dt.clone(), true)
+        })
         .collect::<Vec<_>>()
         .into();
 

datafusion/functions-nested/src/utils.rs

@@ -50,7 +50,7 @@ pub(crate) fn check_datatypes(name: &str, args: &[&ArrayRef]) -> Result<()> {
 }
 
 /// array function wrapper that differentiates between scalar (length 1) and array.
-pub(crate) fn make_scalar_function<F>(
+pub fn make_scalar_function<F>(
     inner: F,
 ) -> impl Fn(&[ColumnarValue]) -> Result<ColumnarValue>
 where

datafusion/spark/src/function/array/arrays_zip.rs

@@ -0,0 +1,128 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null};
+use arrow::datatypes::{DataType, Field, Fields};
+
+use datafusion_common::{Result, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
+
+use datafusion_functions_nested::arrays_zip::arrays_zip_inner;
+use std::any::Any;
+use std::sync::Arc;
+use arrow::array::{ArrayRef, StringArray};
+
+/// Spark-compatible `arrays_zip` function.
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkArraysZip {
+    signature: Signature,
+    aliases: Vec<String>,
+}
+
+impl Default for SparkArraysZip {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SparkArraysZip {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::variadic_any(Volatility::Immutable),
+            aliases: vec![String::from("list_zip")],
+        }
+    }
+}
+
+impl ScalarUDFImpl for SparkArraysZip {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "arrays_zip"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        if arg_types.is_empty() {
+            return exec_err!("arrays_zip requires at least two arguments");
+        }
+
+        let mut fields = Vec::with_capacity(arg_types.len());
+        for (i, arg_type) in arg_types.iter().enumerate() {
+            let element_type = match arg_type {
+                List(field) | LargeList(field) | FixedSizeList(field, _) => {
+                    field.data_type().clone()
+                }
+                Null => Null,
+                dt => {
+                    return exec_err!("arrays_zip expects array arguments, got {dt}");
+                }
+            };
+            fields.push(Field::new(format!("{}", i), element_type, true));
+        }
+
+        Ok(List(Arc::new(Field::new_list_field(
+            DataType::Struct(Fields::from(fields)),
+            true,
+        ))))
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: datafusion_expr::ScalarFunctionArgs,
+    ) -> Result<ColumnarValue> {
+        let args = &args.args;
+
+        // TODO: make configurable: zero-based, one-based
+        // &args.config_options.execution.enable_ansi_mode;
+        let strings_vec: Vec<String> = (0..args.len())
+            .into_iter()
+            .map(|i| i.to_string())
+            .collect();
+        let names = vec![Arc::new(StringArray::from(strings_vec)) as ArrayRef];
+
+        let len = args
+            .iter()
+            .fold(Option::<usize>::None, |acc, arg| match arg {
+                ColumnarValue::Scalar(_) => acc,
+                ColumnarValue::Array(a) => Some(a.len()),
+            });
+
+        let is_scalar = len.is_none();
+
+        let args = ColumnarValue::values_to_arrays(args)?;
+
+        let result = arrays_zip_inner(&args, &names);
+
+        if is_scalar {
+            // If all inputs are scalar, keeps output as scalar
+            let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
+            result.map(ColumnarValue::Scalar)
+        } else {
+            result.map(ColumnarValue::Array)
+        }
+    }
+
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+}

datafusion/spark/src/function/array/mod.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 pub mod array_contains;
+pub mod arrays_zip;
 pub mod repeat;
 pub mod shuffle;
 pub mod slice;
@@ -30,6 +31,7 @@ make_udf_function!(spark_array::SparkArray, array);
 make_udf_function!(shuffle::SparkShuffle, shuffle);
 make_udf_function!(repeat::SparkArrayRepeat, array_repeat);
 make_udf_function!(slice::SparkSlice, slice);
+make_udf_function!(arrays_zip::SparkArraysZip, arrays_zip);
 
 pub mod expr_fn {
     use datafusion_functions::export_functions;
@@ -55,6 +57,11 @@ pub mod expr_fn {
         "Returns a slice of the array from the start index with the given length.",
         array start length
     ));
+    export_functions!((
+        arrays_zip,
+        "Returns an array of structs created by combining the elements of each input array at the same index. If the arrays have different lengths, shorter arrays are padded with NULLs.",
+        args
+    ));
 }
 
 pub fn functions() -> Vec<Arc<ScalarUDF>> {
@@ -64,5 +71,6 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
         shuffle(),
         array_repeat(),
         slice(),
+        arrays_zip(),
     ]
 }