feat: Support Spark expression: arrays_zip (#3643)

* Define ArraysZip expr proto

* Create ArraysZip SerDe

* Register SerDe to arrayExpressions

* Add SQL test

* Register expression to planner

* Rust wrapper around DF's arrays_zip

* Null checks

* Update supported Spark expressions doc

Author: hsiang-c

docs/spark_expressions_support.md

@@ -98,7 +98,7 @@
 - [x] array_repeat
 - [x] array_union
 - [x] arrays_overlap
-- [ ] arrays_zip
+- [x] arrays_zip
 - [x] element_at
 - [ ] flatten
 - [x] get

native/core/src/execution/planner.rs

@@ -68,7 +68,8 @@ use datafusion::{
 };
 use datafusion_comet_spark_expr::{
     create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle,
-    BloomFilterAgg, BloomFilterMightContain, CsvWriteOptions, EvalMode, SumInteger, ToCsv,
+    BloomFilterAgg, BloomFilterMightContain, CsvWriteOptions, EvalMode, SparkArraysZipFunc,
+    SumInteger, ToCsv,
 };
 use datafusion_spark::function::aggregate::collect::SparkCollectSet;
 use iceberg::expr::Bind;
@@ -94,7 +95,6 @@ use datafusion::physical_expr::window::WindowExpr;
 use datafusion::physical_expr::LexOrdering;
 
 use crate::parquet::parquet_exec::init_datasource_exec;
-
 use arrow::array::{
     new_empty_array, Array, ArrayRef, BinaryBuilder, BooleanArray, Date32Array, Decimal128Array,
     Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray,
@@ -683,6 +683,24 @@ impl PhysicalPlanner {
                     csv_write_options,
                 )))
             }
+            ExprStruct::ArraysZip(expr) => {
+                if expr.values.is_empty() {
+                    return Err(GeneralError(
+                        "arrays_zip requires at least one argument".to_string(),
+                    ));
+                }
+
+                let children = expr
+                    .values
+                    .iter()
+                    .map(|child| self.create_expr(child, Arc::clone(&input_schema)))
+                    .collect::<Result<Vec<_>, _>>()?;
+
+                Ok(Arc::new(SparkArraysZipFunc::new(
+                    children,
+                    expr.names.clone(),
+                )))
+            }
             expr => Err(GeneralError(format!("Not implemented: {expr:?}"))),
         }
     }

native/core/src/execution/planner/expression_registry.rs

@@ -103,6 +103,7 @@ pub enum ExpressionType {
     Randn,
     SparkPartitionId,
     MonotonicallyIncreasingId,
+    ArraysZip,
 
     // Time functions
     Hour,
@@ -381,6 +382,7 @@ impl ExpressionRegistry {
             Some(ExprStruct::MonotonicallyIncreasingId(_)) => {
                 Ok(ExpressionType::MonotonicallyIncreasingId)
             }
+            Some(ExprStruct::ArraysZip(_)) => Ok(ExpressionType::ArraysZip),
 
             Some(ExprStruct::Hour(_)) => Ok(ExpressionType::Hour),
             Some(ExprStruct::Minute(_)) => Ok(ExpressionType::Minute),

native/proto/src/proto/expr.proto

@@ -89,6 +89,7 @@ message Expr {
     FromJson from_json = 66;
     ToCsv to_csv = 67;
     HoursTransform hours_transform = 68;
+    ArraysZip arrays_zip = 69;
   }
 
   // Optional QueryContext for error reporting (contains SQL text and position)
@@ -495,3 +496,10 @@ message ArrayJoin {
 message Rand {
   int64 seed = 1;
 }
+
+// Spark's ArraysZip takes children: Seq[Expression] and names: Seq[Expression]
+// https://github.com/apache/spark/blob/branch-4.1/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala#L296
+message ArraysZip {
+  repeated Expr values = 1;
+  repeated string names = 2;
+}