feature: Support Spark expression: arrays_zip by hsiang-c · Pull Request #3643 · apache/datafusion-comet

hsiang-c · 2026-03-07T00:20:37Z

Which issue does this PR close?

Rationale for this change

This PR supports Spark-compatible arrays_zip SQL function

What changes are included in this PR?

The Rust implementation copied and modified arrays_zip_inner implementation from DataFusion 53.0.0 to support custom field name in the resulting struct. It also included subsequent patches: feat: correct struct column names for arrays_zip return type datafusion#20886 and fix: arrays_zip/list_zip allow single array argument datafusion#21047.
At QueryPlan SerDe stage, if one of the arrays_zip arguments is NULL, we return NULL as Spark does. Examples of Spark's arrays_zip:

scala> spark.sql("SELECT arrays_zip(array(1, 2, 3), array(1), NULL)").show(100, false)
+------------------------------------------+
|arrays_zip(array(1, 2, 3), array(1), NULL)|
+------------------------------------------+
|NULL                                      |
+------------------------------------------+

scala> spark.sql("SELECT arrays_zip(NULL, array(1, 2, 3))").show(100, false)
+--------------------------------+
|arrays_zip(NULL, array(1, 2, 3))|
+--------------------------------+
|NULL                            |
+--------------------------------+

How are these changes tested?

By SQL File Tests, we covered cases such as single array argument, null arguments and custom field name in the resulting struct. Here is an example of custom file name:

scala> spark.sql("SELECT arrays_zip(b, a) FROM (SELECT array(1, 2, 3) as a, array(1, 2) as b)").show(100, false)
+---------------------------+
|arrays_zip(b, a)           |
+---------------------------+
|[{1, 1}, {2, 2}, {NULL, 3}]|
+---------------------------+


scala> spark.sql("SELECT arrays_zip(b, a) FROM (SELECT array(1, 2, 3) as a, array(1, 2) as b)").printSchema
warning: 1 deprecation (since 2.13.3); for details, enable `:setting -deprecation` or `:replay -deprecation`
root
 |-- arrays_zip(b, a): array (nullable = false)
 |    |-- element: struct (containsNull = false)
 |    |    |-- b: integer (nullable = true)
 |    |    |-- a: integer (nullable = true)

comphead · 2026-03-07T00:25:47Z

#3575

hsiang-c · 2026-03-07T00:36:35Z

Thanks @comphead

comphead

Thanks @hsiang-c I'm not sure this implementation actually uses the DataFusion

hsiang-c · 2026-04-14T19:46:44Z

@comphead Thanks for your review. This implementation doesn't use DataFusion for now b/c I need to pass names argument from Spark to arrays_zip_inner to parameterized the field key in the final struct.

parthchandra · 2026-04-15T00:27:12Z

spark/src/main/scala/org/apache/comet/serde/arrays.scala

+    // mimic Spark's ArraysZip behavior: returns NULL if any argument is NULL
+    val combinedNullCheck = expr.children.map(child => IsNotNull(child)).reduce(And)
+    val isNotNullExpr = exprToProtoInternal(combinedNullCheck, inputs, binding)
+    val nullLiteralProto = exprToProto(Literal(null, BooleanType), Seq.empty)


The null literal here uses BooleanType, but elsewhere in this file (e.g., CometArrayAppend at line 88) we use the return type of the expression. DF expects all arms of casewhen to have compatible types and this may cause an error.

parthchandra · 2026-04-15T00:33:06Z

spark/src/main/scala/org/apache/comet/serde/arrays.scala

+object CometArraysZip extends CometExpressionSerde[ArraysZip] {
+  override def getSupportLevel(expr: ArraysZip): SupportLevel = {
+    expr.dataType match {
+      case _: ArrayType => Compatible()


We should probably check the element type here. There have been issues noted in the past. See this for instance - #1308

parthchandra · 2026-04-15T00:36:07Z

native/spark-expr/src/array_funcs/arrays_zip.rs

+        let fields = self.fields(input_schema)?;
+        Ok(List(Arc::new(Field::new_list_field(
+            DataType::Struct(Fields::from(fields)),
+            true,


There is a slight mismatch here. Spark has this defined as non-nullable.

parthchandra · 2026-04-15T00:37:31Z

native/core/src/execution/planner.rs

                )))
            }
+            ExprStruct::ArraysZip(expr) => {
+                assert!(!expr.values.is_empty());


Better to return Err instead of asserting (which will cause a panic).

return Err(GeneralError("arrays_zip requires at least one argument".to_string()))

If you want to be extra safe, then you can also check

expr.values.len() == expr.names.len()

parthchandra · 2026-04-15T00:44:28Z

spark/src/test/resources/sql-tests/expressions/array/arrays_zip.sql

+SELECT arrays_zip(a) FROM (SELECT array(1, 2, 3) as a, null as b)
+
+query
+SELECT arrays_zip(b) FROM (SELECT array(1, 2, 3) as a, null as b)


nit: newline needed at end of file

parthchandra · 2026-04-15T00:44:45Z

native/proto/src/proto/expr.proto

+message ArraysZip {
+  repeated Expr values = 1;
+  repeated string names = 2;
+}


nit: newline needed at end of file

hsiang-c force-pushed the arrays_zip branch 2 times, most recently from 97d31a4 to be5dfce Compare March 20, 2026 18:33

hsiang-c mentioned this pull request Mar 26, 2026

tests: add SQL file tests for arrays_zip #3794

Draft

hsiang-c force-pushed the arrays_zip branch 2 times, most recently from a261445 to 3a21e25 Compare April 13, 2026 17:44

hsiang-c added 7 commits April 13, 2026 10:58

Define ArraysZip expr proto

9275466

Create ArraysZip SerDe

138f76b

Register SerDe to arrayExpressions

9fdba4d

Add SQL test

82e403f

Register expression to planner

40502c4

Rust wrapper around DF's arrays_zip

f6f66c2

Null checks

9af23d8

hsiang-c force-pushed the arrays_zip branch from 3a21e25 to 9af23d8 Compare April 13, 2026 17:58

Merge branch 'main' into arrays_zip

a1f1718

hsiang-c marked this pull request as ready for review April 13, 2026 18:21

hsiang-c mentioned this pull request Apr 13, 2026

feat: use arrays_zip from DataFusion #3575

Open

hsiang-c and others added 4 commits April 13, 2026 14:42

Fix clippy

c2e7f67

Update supported Spark expressions doc

74d6e57

Merge branch 'main' into arrays_zip

5ca87ce

Merge branch 'main' into arrays_zip

b35e028

comphead reviewed Apr 14, 2026

View reviewed changes

parthchandra reviewed Apr 15, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feature: Support Spark expression: arrays_zip#3643

feature: Support Spark expression: arrays_zip#3643
hsiang-c wants to merge 12 commits intoapache:mainfrom
hsiang-c:arrays_zip

hsiang-c commented Mar 7, 2026 •

edited

Loading

Uh oh!

comphead commented Mar 7, 2026

Uh oh!

hsiang-c commented Mar 7, 2026

Uh oh!

comphead left a comment

Uh oh!

hsiang-c commented Apr 14, 2026

Uh oh!

parthchandra Apr 15, 2026

Uh oh!

parthchandra Apr 15, 2026

Uh oh!

parthchandra Apr 15, 2026

Uh oh!

parthchandra Apr 15, 2026

Uh oh!

parthchandra Apr 15, 2026

Uh oh!

parthchandra Apr 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

hsiang-c commented Mar 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Uh oh!

comphead commented Mar 7, 2026

Uh oh!

hsiang-c commented Mar 7, 2026

Uh oh!

comphead left a comment

Choose a reason for hiding this comment

Uh oh!

hsiang-c commented Apr 14, 2026

Uh oh!

parthchandra Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hsiang-c commented Mar 7, 2026 •

edited

Loading