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

Abierto
  • ¿Algo no está bien?
#253 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 1 día

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
42/100
Tipo de issue
Nueva funcionalidad
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
rust, spark
Área
database

Línea de trabajo

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].

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Rust
Estrellas
7
Forks
15
Merge medio
16 h 43 min
PR fusionados (30 d)
26

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de apache/datafusion-spark

Todos los issues de apache/datafusion-spark

Issues similares

Más issues de Rust

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.