Implement Spark-compatible `zip_with` (higher-order function over two arrays)

Aperta
#250 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
40/100
Tipo di issue
Funzionalità
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
rust, sql
Ambito
databases

Direzione di ricerca

Start in crates/spark/src/function/lambda/mod.rs, an empty ScalarUDF placeholder, and read HigherOrderUDFImpl in datafusion/expr/src/higher_order_function.rs on branch-56. Implement zip_with with lambda_parameters returning two nullable element-type fields, then register it via register_higher_order_function/SessionStateBuilder::higher_order_functions. Add crates/sqllogictest/test_files/lambda/zip_with.slt covering the Spark examples plus padding, NULL and empty-array cases; done means those pass and the function has #[user_doc].

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

enhancement

Part of #176 and apache/datafusion-comet#7300.

What the function does

zip_with(left, right, func) merges two arrays element by element with a two-argument lambda (x, y) -> .... If
one array is shorter, it is padded with NULLs at the end to the length of the longer one before func is applied.
The result type is ARRAY<func.dataType> with containsNull = func.nullable. The result is NULL if either input
array is NULL, and nullable if either input is nullable. Both lambda parameters are always declared nullable,
because of the padding. Both arguments must be ARRAYs, and their element types can differ. There is no index
argument. Non-lambda func arguments are wrapped in a hidden lambda. Available since Spark 2.4.0. No aliases.

Examples
SELECT zip_with(array(1, 2, 3), array('a', 'b', 'c'), (x, y) -> (y, x));
-- [{"y":"a","x":1},{"y":"b","x":2},{"y":"c","x":3}]
SELECT zip_with(array(1, 2), array(3, 4), (x, y) -> x + y);
-- [4,6]
SELECT zip_with(array('a', 'b', 'c'), array('d', 'e', 'f'), (x, y) -> concat(x, y));
-- ["ad","be","cf"]
SELECT zip_with(array('a'), array('d', null, 'f'), (x, y) -> coalesce(x, y));
-- ["a",null,"f"]
SELECT zip_with(array(1, 2, 3), array(4, 5, 6), 100);
-- [100,100,100]
-- nested view (x, ys, zs) = (1, [32,97], [[12,99],[123,42],[1]]), (2, [77,-76], [[6,96,65],[-1,-2]]), (3, [12], [[17]])
SELECT zip_with(ys, zs, (a, b) -> a + size(b)) FROM nested;
-- [34,99,null] / [80,-74] / [13]

These come from the ExpressionDescription and higher-order-functions.sql.out. The last example is listed in view
order, and it shows padding: row 1 has 2 ys and 3 zs, so the third a is NULL.

Edge cases to test
  • Arrays of different lengths, in both directions: zip_with([1,2,3], [1,2,3,4], (x, y) -> x + y) gives
    [2,4,6,null], and zip_with([1,null], [1,null,3], add) and zip_with([1,null,3], [1,null], add) both give
    [2,null,null] (HigherOrderFunctionsSuite). The padding must be NULL even when the input element type is
    non-nullable.
  • The lambda still runs on padded positions, so a body that ignores NULLs produces values there:
    coalesce(x, y) above gives "f".
  • Either array NULL gives NULL. Spark evaluates left first and does not evaluate right if left is NULL.
  • Both arrays empty gives []. One empty array gives a result as long as the other, computed with all-NULL values on
    one side.
  • A NULL inner array within an array of arrays:
    zip_with(aai1, aai1, (a1, a2) -> cast(transform(a1, x -> x + 1) as string)) on [[1,2,3],null,[4,5]] gives
    ["[2, 3, 4]",null,"[5, 6]"].
  • Element types that differ between left and right (int and string), and struct results built with (y, x).
    Spark names the struct fields after the lambda variables ({"y":...,"x":...}).
  • Captured outer columns, nested lambdas, and nondeterministic bodies evaluated once per output position.
  • Errors inside the lambda under ANSI mode, for example x / y with a zero divisor. They are raised only for
    positions that are actually evaluated, which includes padded positions.
  • Lambda arity: exactly 2 parameters, otherwise INVALID_LAMBDA_FUNCTION_CALL.NUM_ARGS_MISMATCH. A non-array
    argument fails with DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE.
Differences across Spark versions
  • None found in ZipWith between 3.4 and 4.2. ANSI mode is on by default from 4.0.
Implementation notes
  • Spark: ZipWith in sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala.
  • Infrastructure. DataFusion's higher-order function framework (HigherOrderUDFImpl in
    datafusion/expr/src/higher_order_function.rs, and LambdaArgument::evaluate with spread_captures for captured
    columns) is on branch-56, which this crate uses. lambda_parameters must return two parameter fields, both
    nullable, with the two element types.
  • DataFusion core has no two-array higher-order function. arrays_zip is not a higher-order function. To evaluate
    vectorized, build per row the max(len_l, len_r) positions, then take both value arrays with NULL indices for the
    padding, which gives two aligned flattened columns. Evaluate the body once over them, and spread captures to the
    same positions. The result offsets are the per-row maximum lengths.
  • This crate's crates/spark/src/function/lambda/mod.rs is an empty ScalarUDF placeholder. It needs a
    HigherOrderUDF list and registration through register_higher_order_function /
    SessionStateBuilder::higher_order_functions.
  • Add crates/sqllogictest/test_files/lambda/zip_with.slt and document the function with #[user_doc].

This issue was generated by an LLM (Claude) from the Apache Spark source code and tests. Examples come from Spark's ExpressionDescription examples, test suites and .sql.out golden files, but please verify them against Spark before relying on them.

Lingua principale
Rust
Stelle
7
Fork
15
Merge medio
16h 43m
PR unite (30g)
26

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di apache/datafusion-spark

Tutte le issue di apache/datafusion-spark

Issue simili

Altre issue su Rust

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.