Implement Spark-compatible `transform_keys` (higher-order map function)

オープン
#253 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
42/100
issue の種類
機能追加
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
rust, spark
領域
database

調査の方向性

Start by reading datafusion/expr/src/higher_order_function.rs to see how HigherOrderUDFImpl and LambdaArgument are evaluated, then inspect crates/spark/src/function/map/utils.rs for the existing dedup and NULL_MAP_KEY handling. Registration runs through crates/spark/src/function/lambda/mod.rs, currently an empty placeholder, so add the HigherOrderUDF list and register_higher_order_function wiring. Done when crates/sqllogictest/test_files/lambda/transform_keys.slt covers the examples and edge cases and the function has #[user_doc].

索引モデルが issue の本文から書いたものです。

説明

enhancement

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

What the function does

transform_keys(expr, func) returns a map whose keys are func(k, v) for each entry of the map expr. The values
are unchanged. The result type is MAP<func.dataType, valueType> with the input's valueContainsNull. A NULL map
gives NULL, and the result is nullable if the input is. The lambda takes exactly 2 parameters: k (non-null) and
v (nullable when valueContainsNull is true).

The new key type must be a valid map key, so checkInputDataTypes uses TypeUtils.checkForMapKeyType. A key type
that is or contains a MAP fails with DATATYPE_MISMATCH.INVALID_MAP_KEY_TYPE. expr must be a MAP. A
non-lambda func is wrapped in a hidden lambda. Available since Spark 3.0.0. No aliases.

The result is built with ArrayBasedMapBuilder, which means:

  • a NULL produced key fails with NULL_MAP_KEY;
  • duplicate produced keys follow spark.sql.mapKeyDedupPolicy. The default EXCEPTION fails with
    DUPLICATED_MAP_KEY (key, mapKeyDedupPolicy parameters). Under LAST_WIN the key keeps the position of its
    first occurrence and takes the value of its last occurrence;
  • from 4.0, float and double keys are normalized (-0.0 becomes 0.0, and NaN values are unified) unless
    spark.sql.legacy.disableMapKeyNormalization=true.
Examples
SELECT transform_keys(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), (k, v) -> k + 1);
-- {2:1,3:2,4:3}
SELECT transform_keys(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), (k, v) -> k + v);
-- {2:1,4:2,6:3}
SELECT transform_keys(map('a', 1, 'b', 2), 42);
-- error DUPLICATED_MAP_KEY
SELECT transform_keys(map('a', 1, 'b', 2), map('new', 'key'));
-- error DATATYPE_MISMATCH.INVALID_MAP_KEY_TYPE
SELECT transform_keys(map('a', null), (k, v) -> v);
-- error NULL_MAP_KEY

These come from the ExpressionDescription, higher-order-functions.sql.out and DataFrameFunctionsSuite.
HigherOrderFunctionsSuite adds:

map lambda result
{1:1, 2:2, 3:3, 4:4} k % 3 DUPLICATED_MAP_KEY (key 1)
{1:1, 2:2, 3:3, 4:4} k % 3, LAST_WIN {1:4, 2:2, 0:3}
{1:1, 2:null, 3:3} k + 1 {2:1, 3:null, 4:3}
{"a":"xy", "bb":"yz", "ccc":null} length(k) + 1 {2:"xy", 3:"yz", 4:null}
{} any {}
NULL any NULL
Edge cases to test
  • A NULL map gives NULL. An empty map gives {}, and no error is possible.
  • A lambda that returns NULL for any entry gives NULL_MAP_KEY, even under LAST_WIN.
  • Duplicate produced keys under both spark.sql.mapKeyDedupPolicy values. Check the error parameters
    ("key" -> "1", "mapKeyDedupPolicy" -> "\"spark.sql.mapKeyDedupPolicy\""). Under LAST_WIN, check both the
    position and the value: {1:4, 2:2, 0:3} above. The datafusion-spark config is
    datafusion.spark.map_key_dedup_policy.
  • Float and double produced keys, for example (k, v) -> if(k = 1, 0.0D, -0.0D) on a two-entry map. By code
    reading, 4.0+ normalization makes 0.0 and -0.0 collide, so the dedup policy applies (DUPLICATED_MAP_KEY by
    default). In 3.4 and 3.5 they are distinct keys. Also test NaN.
  • Chained calls: transform_keys(transform_keys(m, (k, v) -> k + 1), (k, v) -> k + v) gives {3:1, 5:2, 7:3, 9:4}
    for {1:1, 2:2, 3:3, 4:4}.
  • Key-type change (string to int, as in the table above). Struct and array keys are allowed. Map keys are rejected
    at analysis time.
  • Lambda arity: k -> k fails with INVALID_LAMBDA_FUNCTION_CALL.NUM_ARGS_MISMATCH (expected 1, actual 2), and so
    does (k, v, x) -> .... DataFusion's planner allows a lambda to declare fewer parameters than offered, so this check
    has to be explicit.
  • Captured outer columns, nondeterministic key expressions, and ANSI errors inside the lambda. With ANSI on,
    k + 1 can overflow and raise ARITHMETIC_OVERFLOW.
  • A non-map argument (transform_keys(array(...), ...)) fails with DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE.
Differences across Spark versions
  • 4.0 adds float and double key normalization to ArrayBasedMapBuilder (spark.sql.legacy.disableMapKeyNormalization,
    default false).
  • 3.5 marks TransformKeys as stateful, because the map builder is reused. This is internal only.
  • ANSI mode is on by default from 4.0.
Implementation notes
  • Spark: TransformKeys in sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala.
    The builder is sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ArrayBasedMapBuilder.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, with config_options available in HigherOrderFunctionArgs) is on branch-56, which this crate uses.
  • DataFusion core has no higher-order map functions. Evaluate the lambda once over the flattened keys and values of
    the MapArray entries, then rebuild the map with deduplication.
  • In this crate, crates/spark/src/function/map/utils.rs already implements Spark's dedup policy
    (map_from_keys_values_offsets_nulls / map_deduplicate_keys, which raise [DUPLICATED_MAP_KEY] or keep the last
    value). It is wired to datafusion.spark.map_key_dedup_policy and used by map_from_arrays,
    map_from_entries and str_to_map. Reuse it, and check that it also enforces NULL_MAP_KEY and the 4.0+ float
    normalization.
  • 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/transform_keys.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分
マージ済み PR(30日)
26

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

apache/datafusion-spark のほかの issue

apache/datafusion-spark の issue をすべて見る

似ている issue

Rust の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。