Skip to content

[Feature] Support Spark expression: arrays_zip #3151

@andygrove

Description

@andygrove

What is the problem the feature request solves?

Note: This issue was generated with AI assistance. The specification details have been extracted from Spark documentation and may need verification.

Comet does not currently support the Spark arrays_zip function, causing queries using this function to fall back to Spark's JVM execution instead of running natively on DataFusion.

The ArraysZip expression combines multiple arrays into a single array of structs by transposing elements at corresponding positions. Each resulting struct contains fields named "0", "1", "2", etc., with values from the input arrays at the same index position.

Supporting this expression would allow more Spark workloads to benefit from Comet's native acceleration.

Describe the potential solution

Spark Specification

Syntax:

arrays_zip(array1, array2, ...)

Arguments:

Argument Type Description
children Seq[Expression] Variable number of array expressions to be zipped together
names Seq[Expression] Field names for the resulting struct fields (typically auto-generated as "0", "1", "2", etc.)

Return Type: Array of structs, where each struct contains fields corresponding to elements from input arrays at the same position.

Supported Data Types:
All data types are supported for array elements, including:

  • Numeric types (byte, short, int, long, float, double, decimal)
  • String and binary types
  • Boolean type
  • Date and timestamp types
  • Complex types (arrays, maps, structs)
  • Null values

Edge Cases:

  • Null arrays: If an input array is null, the corresponding field in all output structs will be null
  • Empty arrays: Empty input arrays contribute null values to all positions in the output
  • Mismatched lengths: Shorter arrays are padded with nulls; longer arrays determine the output length
  • All empty inputs: Results in an empty array
  • Single array input: Creates array of single-field structs

Examples:

-- Basic usage with arrays of same length
SELECT arrays_zip(array(1, 2), array(2, 3), array(3, 4));
-- Result: [{"0":1,"1":2,"2":3},{"0":2,"1":3,"2":4}]

-- Arrays with different lengths
SELECT arrays_zip(array(1, 2, 3), array('a', 'b'));
-- Result: [{"0":1,"1":"a"},{"0":2,"1":"b"},{"0":3,"1":null}]

-- With null values
SELECT arrays_zip(array(1, null, 3), array('x', 'y', 'z'));
-- Result: [{"0":1,"1":"x"},{"0":null,"1":"y"},{"0":3,"1":"z"}]
// DataFrame API usage
import org.apache.spark.sql.functions._

df.select(arrays_zip(col("array1"), col("array2"), col("array3")))

// Using with explode to create rows
df.select(explode(arrays_zip(col("array1"), col("array2"))))

Implementation Approach

See the Comet guide on adding new expressions for detailed instructions.

  1. Scala Serde: Add expression handler in spark/src/main/scala/org/apache/comet/serde/
  2. Register: Add to appropriate map in QueryPlanSerde.scala
  3. Protobuf: Add message type in native/proto/src/proto/expr.proto if needed
  4. Rust: Implement in native/spark-expr/src/ (check if DataFusion has built-in support first)

Additional context

Difficulty: Medium
Spark Expression Class: org.apache.spark.sql.catalyst.expressions.ArraysZip

Related:

  • explode() - Often used with arrays_zip to create rows from zipped arrays
  • array() - Creates arrays that can be used as input
  • struct() - Creates individual struct values
  • zip_with() - Alternative for element-wise array operations with custom logic

This issue was auto-generated from Spark reference documentation.

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions