Implement Spark-compatible `map_zip_with` (higher-order function over two maps)

未关闭
#252 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 1 天内回复

还没有人认领这个 Issue。

评估

难度
4/5
预计耗时
3-5 天
新手友好度
35/100
Issue 类型
功能
描述清晰度
描述清楚
活跃度
活跃
技术栈
rust, spark
领域
databases

调研方向

Start in crates/spark/src/function/lambda/mod.rs, currently an empty ScalarUDF placeholder, and see how other functions register with register_higher_order_function / SessionStateBuilder::higher_order_functions. Read DataFusion's HigherOrderUDFImpl in datafusion/expr/src/higher_order_function.rs and LambdaArgument::evaluate/spread_captures, then model the behavior on Spark's MapZipWith in higherOrderFunctions.scala. Reuse the helpers in crates/spark/src/function/map/utils.rs for MapArray rebuilding, and add crates/sqllogictest/test_files/lambda/map_zip_with.slt; done means the examples and edge cases in the issue pass.

由索引模型根据 Issue 内容生成。

描述

enhancement

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

What the function does

map_zip_with(map1, map2, func) merges two maps into one by calling func(k, v1, v2) for every key in the union of
their keys. For a key present in only one map, NULL is passed for the missing side. If an input map contains
duplicated keys, only the first entry of each key is used. The result type is
MAP<keyType, func.dataType> with valueContainsNull = func.nullable. The result is NULL if either map is NULL.
Spark evaluates map1 first and skips map2 if map1 is NULL. The lambda parameters are k (non-null) and v1,
v2 (always nullable). Non-lambda func arguments are wrapped in a hidden lambda. Available since Spark 3.0.0. No
aliases.

Key types. If the two key types differ, the MapZipWithCoercion rule casts both maps to the wider key type
(findWiderTypeForTwo), unless the cast would force nullability. If no common type exists, analysis fails with
DATATYPE_MISMATCH.MAP_ZIP_WITH_DIFF_TYPES, for example INT against MAP<INT, INT>. The key type must be
orderable.

Output order. Keys from map1 come first in their original order, followed by keys that appear only in map2, in
map2 order. Spark builds a LinkedHashMap for key types with proper equality, and does a linear search with
interpreted ordering otherwise (binary, struct and array keys). No dedup policy is applied to the result.

Examples
SELECT map_zip_with(map(1, 'a', 2, 'b'), map(1, 'x', 2, 'y'), (k, v1, v2) -> concat(v1, v2));
-- {1:"ax",2:"by"}
SELECT map_zip_with(map('a', 1, 'b', 2), map('b', 3, 'c', 4), (k, v1, v2) -> coalesce(v1, 0) + coalesce(v2, 0));
-- {"a":1,"b":5,"c":4}
SELECT map_zip_with(map('a', 1), map('a', 10), 100);
-- {"a":100}
SELECT map_zip_with(map('a', 1), map('a', 10), struct('left', 1, 'right', 10));
-- {"a":{"col1":"left","col2":1,"col3":"right","col4":10}}

These come from the ExpressionDescription and higher-order-functions.sql.out. HigherOrderFunctionsSuite, with
(k, v1, v2) -> k * v1 * v2, adds:

map1 map2 result
{1:10, 2:20, 3:30} {1:-1, 2:-2, 4:-4} {1:-10, 2:-80, 3:null, 4:null}
{1:10, 2:20, 3:30} {1:null, 2:-2, 3:null} {1:null, 2:-80, 3:null}
{1:10, 2:20, 3:30} {} {1:null, 2:null, 3:null}
{1:10, 2:20, 3:30} NULL NULL

With concat(k, v1, v2) on strings, {"d":"b","b":"d"} and {"c":null,"b":"t","a":null} give
{"d":null,"b":"bdt","c":null,"a":null}. The same test also covers binary keys, which use the linear-search path.

Edge cases to test
  • Either map NULL gives NULL. Both maps empty gives {}. One map empty gives every key of the other, with NULL for
    the missing side.
  • Key-only-in-one-side handling, NULL values that are actually present (indistinguishable from missing values inside
    the lambda), and the output order described above.
  • Duplicate keys inside an input map. Only the first occurrence is used, and later occurrences are ignored rather
    than raising DUPLICATED_MAP_KEY. Normal map construction forbids duplicates, but Arrow input can contain them.
  • Floating-point keys. With spark.sql.mapZipWithUsesJavaCollections=true (internal, default true, 4.1+), keys are
    normalized with NormalizeFloatingNumbers, so NaN matches NaN and -0.0 matches 0.0. 3.4 through 4.0 use a
    Scala LinkedHashMap with no normalization, and the config's doc says the Java path exists to avoid "NaN equality
    issues". Test NaN and -0.0 keys and expect the older versions to differ, at least for NaN.
  • Key type coercion: map(1, ...) (INT) against map(1L, ...) (BIGINT) gives BIGINT keys. Also test incompatible key
    types (MAP_ZIP_WITH_DIFF_TYPES) and non-orderable keys.
  • Different value types on the two sides, for example MAP<STRING, INT> and MAP<STRING, STRING>.
  • Lambda arity: exactly 3 parameters. (x, y) -> ... fails with INVALID_LAMBDA_FUNCTION_CALL.NUM_ARGS_MISMATCH.
  • Captured outer columns, nondeterministic bodies (once per output key), and ANSI errors inside the body.
  • Very large maps: mapSizeExceedArraySizeWhenZipMapError on the brute-force path.
Differences across Spark versions
  • 4.1.0 adds spark.sql.mapZipWithUsesJavaCollections (internal, default true). The fast path then uses a Java
    LinkedHashMap with float/double key normalization, which fixes NaN equality. 3.4, 3.5 and 4.0 use a Scala
    LinkedHashMap without normalization.
  • 4.0 changed the key-type comparison to DataTypeUtils.sameType. That is an internal refactor.
Implementation notes
  • Spark: MapZipWith in sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala
    (getKeysWithIndexesFastAsJava, getKeysWithIndexesFastUsingScala, getKeysWithIndexesBruteForce). Coercion
    is MapZipWithTypeCoercion in sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionHelper.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 should offer
    [key (non-null), v1 (nullable), v2 (nullable)]. Key coercion has to happen in coerce_value_types.
  • DataFusion core has no equivalent. Implementation outline:
    1. For each row, compute the key union and two index vectors into the value arrays, with NULL for a missing key.
      Use arrow_row or a hash on normalized keys.
    2. take the keys and both value arrays to get three aligned flattened columns.
    3. Evaluate the lambda once over those columns.
    4. Rebuild the MapArray.
  • 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. crates/spark/src/function/map/utils.rs has MapArray helpers.
  • Add crates/sqllogictest/test_files/lambda/map_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.

主要语言
Rust
星标
7
派生
15
平均合并
16 小时 43 分钟
30 天内合并 PR
26

环境准备

  • 没有 Dockerfile 或 Docker Compose 文件
  • 没有 Pull Request 模板
  • 阅读贡献指南

从这里开始

  1. 先读完整个 Issue,再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库,在一个分支上完成修改。
  4. 提交 Pull Request,并在描述里引用这个 Issue 编号。

apache/datafusion-spark 的其他 Issue

查看 apache/datafusion-spark 的全部 Issue

相似的 Issue

更多 Rust Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。