From 8d7bd78ea98a0fa3b73be2e08913f2562e006041 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 10 Sep 2026 22:07:55 +0000 Subject: [PATCH] fix: enforce null-key rejection and mapKeyDedupPolicy in native map construction `map_from_arrays` and `map_from_entries` built their maps without the entry checks Spark's `ArrayBasedMapBuilder` performs, so a `NULL` key inside the keys array produced a map with a `NULL` key instead of raising `NULL_MAP_KEY`, and `spark.sql.mapKeyDedupPolicy=LAST_WIN` fell the whole expression back to Spark. DataFusion 55 added `datafusion.spark.map_key_dedup_policy` and taught the `datafusion-spark` map kernels to follow it, which is the missing half. Forward Spark's `spark.sql.mapKeyDedupPolicy` to it across JNI, and pass the session's `ConfigOptions` into `ScalarFunctionExpr` so a kernel that reads a setting sees the session's value rather than DataFusion's defaults. New `SparkMapFromArrays` / `SparkMapFromEntries` / `SparkStrToMap` wrappers add the checks the upstream kernels do not perform and restate their errors as the Spark error classes `SparkErrorConverter` turns back into `QueryExecutionErrors`: a `NULL` key raises `NULL_MAP_KEY` ahead of any duplicate-key check, key and value arrays of different lengths raise `MAP_KEY_VALUE_DIFF_SIZES`, and a duplicate key under `EXCEPTION` raises `DUPLICATED_MAP_KEY` naming the key. `CometMapFromArrays` now emits `map_from_arrays`, which is null intolerant like Spark's, so the `CaseWhen` guard against NULL input arrays is no longer needed. A floating-point map key stays a documented difference: Spark normalizes `-0.0` to `+0.0` and canonicalizes `NaN` before storing a key, while the native builders compare the raw Arrow values. `spark.comet.exec.strictFloatingPoint` declines those key types. Closes #4680 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01REK8xCiYKTcqw1NGQniHXj --- .../expression-audits/map_funcs.md | 9 +- native/core/src/execution/jni_api.rs | 15 +- native/core/src/execution/planner.rs | 10 +- native/core/src/execution/spark_config.rs | 2 + native/spark-expr/src/comet_scalar_funcs.rs | 6 +- native/spark-expr/src/lib.rs | 2 +- .../spark-expr/src/map_funcs/map_builders.rs | 651 ++++++++++++++++++ native/spark-expr/src/map_funcs/mod.rs | 2 + .../org/apache/comet/CometExecIterator.scala | 7 + .../scala/org/apache/comet/serde/maps.scala | 106 ++- .../expressions/map/map_from_arrays.sql | 21 +- .../map/map_from_arrays_dedup_policy.sql | 29 +- .../expressions/map/map_from_entries.sql | 14 + .../map/map_from_entries_dedup_policy.sql | 32 +- .../sql-tests/expressions/map/str_to_map.sql | 8 +- .../map/str_to_map_dedup_policy.sql | 42 ++ .../comet/CometMapExpressionSuite.scala | 89 +++ .../org/apache/spark/sql/CometTestBase.scala | 20 +- 18 files changed, 958 insertions(+), 107 deletions(-) create mode 100644 native/spark-expr/src/map_funcs/map_builders.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/map/str_to_map_dedup_policy.sql diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index ea13e6ab130..779e307dd3d 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -45,9 +45,12 @@ ## map_from_arrays - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `MapFromArrays(left, right) extends BinaryExpression with NullIntolerant`; Spark uses `ArrayBasedMapBuilder` to detect duplicate keys (subject to `spark.sql.mapKeyDedupPolicy`) and rejects null keys with `RuntimeException("Cannot use null as map key")`. Comet `CometMapFromArrays` wraps the inputs in `CaseWhen(IsNotNull(left) AND IsNotNull(right), map(left, right), null)` so NULL-array inputs return NULL rather than triggering the previously reported native crash ([#3327](https://github.com/apache/datafusion-comet/issues/3327)). +- Spark 3.5.8 (audited 2026-05-27): baseline. `MapFromArrays(left, right) extends BinaryExpression with NullIntolerant`; Spark uses `ArrayBasedMapBuilder` to detect duplicate keys (subject to `spark.sql.mapKeyDedupPolicy`) and rejects null keys with `RuntimeException("Cannot use null as map key")`. Comet `CometMapFromArrays` wires the native `map_from_arrays` from `datafusion-spark`, which is null intolerant the same way, so NULL-array inputs return NULL rather than triggering the previously reported native crash ([#3327](https://github.com/apache/datafusion-comet/issues/3327)). - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; `NullIntolerant` trait replaced by `nullIntolerant: Boolean`. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. +- `ArrayBasedMapBuilder` semantics, reproduced natively rather than falling back ([#4680](https://github.com/apache/datafusion-comet/issues/4680)): a `NULL` key element raises `NULL_MAP_KEY`, ahead of any duplicate-key check, matching the order Spark applies them in; a duplicate key follows `spark.sql.mapKeyDedupPolicy`, which `CometExecIterator` forwards to the native session as `datafusion.spark.map_key_dedup_policy` (`EXCEPTION` raises `DUPLICATED_MAP_KEY` naming the key, `LAST_WIN` keeps the last value for the key). +- Known limitation: `ArrayBasedMapBuilder` normalizes a floating-point key before storing it (`-0.0` becomes `+0.0`, every `NaN` collapses to one), while the native builder compares the raw Arrow values, so a map built from both `-0.0` and `+0.0` keeps two entries where Spark reports a duplicate key. Gated only under `spark.comet.exec.strictFloatingPoint`, which marks the expression `Incompatible` for a floating-point key type. +- Spark raises `MAP_KEY_VALUE_DIFF_SIZES` when a row's key and value arrays differ in length; the native path raises the same error. ## map_from_entries @@ -55,6 +58,8 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `MapFromEntries(child) extends UnaryExpression with NullIntolerant`; expects an array of structs and produces a map. Wired as `CometScalarFunction("map_from_entries")`. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; trait refactor. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. +- `ArrayBasedMapBuilder` semantics, reproduced natively rather than falling back ([#4680](https://github.com/apache/datafusion-comet/issues/4680)): a `NULL` key element raises `NULL_MAP_KEY`, ahead of any duplicate-key check, matching the order Spark applies them in; a duplicate key follows `spark.sql.mapKeyDedupPolicy`, which `CometExecIterator` forwards to the native session as `datafusion.spark.map_key_dedup_policy` (`EXCEPTION` raises `DUPLICATED_MAP_KEY` naming the key, `LAST_WIN` keeps the last value for the key). +- Known limitation: `ArrayBasedMapBuilder` normalizes a floating-point key before storing it (`-0.0` becomes `+0.0`, every `NaN` collapses to one), while the native builder compares the raw Arrow values, so a map built from both `-0.0` and `+0.0` keeps two entries where Spark reports a duplicate key. Gated only under `spark.comet.exec.strictFloatingPoint`, which marks the expression `Incompatible` for a floating-point key type. - Known limitation: input arrays where the struct's key or value type contains `BinaryType` are marked `Incompatible` and fall back unless `spark.comet.expression.MapFromEntries.allowIncompatible=true`. ## map_keys @@ -74,7 +79,7 @@ ## str_to_map - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `StringToMap(text, pairDelim, keyValueDelim) extends TernaryExpression`; splits `text` on `pairDelim`, then each pair on `keyValueDelim` (default `","` and `":"`). Uses `ArrayBasedMapBuilder` for duplicate-key handling. Wired as `CometScalarFunction("str_to_map")`. +- Spark 3.5.8 (audited 2026-05-27): baseline. `StringToMap(text, pairDelim, keyValueDelim) extends TernaryExpression`; splits `text` on `pairDelim`, then each pair on `keyValueDelim` (default `","` and `":"`). Uses `ArrayBasedMapBuilder` for duplicate-key handling. Wired as `CometScalarFunction("str_to_map")`. The native `str_to_map` reads the duplicate-key policy from `datafusion.spark.map_key_dedup_policy`, which `CometExecIterator` forwards from `spark.sql.mapKeyDedupPolicy`. - Spark 4.0.1 (audited 2026-05-27): `inputTypes` widened to `StringTypeNonCSAICollation`; uses `CollationAwareUTF8String.splitSQL` with a `collationId`. Runtime unchanged for `UTF8_BINARY`. - Spark 4.1.1 (audited 2026-05-27): adds the `legacySplitTruncate` flag (driven by `spark.sql.legacy.truncateForEmptyRegexSplit`) to both `splitSQL` calls. The Comet native impl always behaves as if the flag were false, so `CometStrToMap` reads the config by string key and reports `Incompatible` when it is enabled; the `CodegenDispatchFallback` trait then routes the expression through the JVM codegen dispatcher rather than falling the whole projection back to Spark. Non-UTF8_BINARY collations on the input or the delimiters are handled the same way. diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 65a2d68ec18..2a80c488631 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -57,8 +57,6 @@ use datafusion_spark::function::datetime::to_utc_timestamp::SparkToUtcTimestamp; use datafusion_spark::function::hash::crc32::SparkCrc32; use datafusion_spark::function::hash::sha1::SparkSha1; use datafusion_spark::function::hash::sha2::SparkSha2; -use datafusion_spark::function::map::map_from_entries::MapFromEntries; -use datafusion_spark::function::map::str_to_map::SparkStrToMap; use datafusion_spark::function::math::expm1::SparkExpm1; use datafusion_spark::function::math::factorial::SparkFactorial; use datafusion_spark::function::math::hex::SparkHex; @@ -112,7 +110,7 @@ use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, - COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, SPARK_MAP_KEY_DEDUP_POLICY, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use datafusion_comet_proto::spark_operator::operator::OpStruct; @@ -715,6 +713,15 @@ fn prepare_datafusion_session_context( session_config.set_str("datafusion.execution.parquet.reorder_filters", "true"); } + // `map_from_arrays`, `map_from_entries` and `str_to_map` build their maps with the + // duplicate-key policy Spark's `ArrayBasedMapBuilder` uses. DataFusion spells the same + // setting `datafusion.spark.map_key_dedup_policy` and takes the same `EXCEPTION` / + // `LAST_WIN` values. Set before the `spark.comet.datafusion.*` testing escape hatch + // pass-through below, so an explicit override of the DataFusion key still wins. + if let Some(policy) = spark_config.get(SPARK_MAP_KEY_DEDUP_POLICY) { + session_config = session_config.set_str("datafusion.spark.map_key_dedup_policy", policy); + } + // Pass through DataFusion configs from Spark. // e.g: spark-shell --conf spark.comet.datafusion.sql_parser.parse_float_as_decimal=true // becomes datafusion.sql_parser.parse_float_as_decimal=true @@ -754,7 +761,6 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) { session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseNot::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkHex::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkWidthBucket::default())); - session_ctx.register_udf(ScalarUDF::new_from_impl(MapFromEntries::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkCrc32::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkLuhnCheck::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSpace::default())); @@ -762,7 +768,6 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) { session_ctx.register_udf(ScalarUDF::new_from_impl(SparkArrayContains::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkArrayRepeat::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBin::default())); - session_ctx.register_udf(ScalarUDF::new_from_impl(SparkStrToMap::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkUrlDecode::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkUrlEncode::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkTryUrlDecode::default())); diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 37d5e744415..e42855a421d 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3657,6 +3657,14 @@ impl PhysicalPlanner { } } + /// The session's `ConfigOptions`, so a kernel that reads one sees what + /// `prepare_datafusion_session_context` set rather than DataFusion's defaults. The map + /// builders read `datafusion.spark.map_key_dedup_policy` this way, which Comet forwards from + /// `spark.sql.mapKeyDedupPolicy`. + fn session_config_options(&self) -> Arc { + Arc::clone(self.session_ctx.copied_config().options()) + } + fn create_scalar_function_expr( &self, expr: &ScalarFunc, @@ -3783,7 +3791,7 @@ impl PhysicalPlanner { fun_expr, args.to_vec(), Arc::new(Field::new(fun_name, data_type.clone(), true)), - Arc::new(ConfigOptions::default()), + self.session_config_options(), )); // DF53 changed some UDFs (e.g. md5) to return StringViewArray at execution diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 4c2811cb5de..573e1e9544f 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -25,6 +25,8 @@ pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory"; pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; +/// Spark's duplicate map key policy, forwarded to `datafusion.spark.map_key_dedup_policy`. +pub(crate) const SPARK_MAP_KEY_DEDUP_POLICY: &str = "spark.sql.mapKeyDedupPolicy"; pub(crate) trait SparkConfig { fn get_bool(&self, name: &str) -> bool; diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index b5820144ea6..6091bfcc2a7 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -31,7 +31,8 @@ use crate::{ EvalMode, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, - SparkMakeTime, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, + SparkMakeTime, SparkMapFromArrays, SparkMapFromEntries, SparkNextDay, SparkSecondsToTimestamp, + SparkSizeFunc, SparkStrToMap, }; use arrow::datatypes::DataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; @@ -321,9 +322,12 @@ fn all_scalar_functions() -> Vec> { )), Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::default())), Arc::new(ScalarUDF::new_from_impl(SparkMakeTime::default())), + Arc::new(ScalarUDF::new_from_impl(SparkMapFromArrays::default())), + Arc::new(ScalarUDF::new_from_impl(SparkMapFromEntries::default())), Arc::new(ScalarUDF::new_from_impl(SparkNextDay::default())), Arc::new(ScalarUDF::new_from_impl(SparkSecondsToTimestamp::default())), Arc::new(ScalarUDF::new_from_impl(SparkSizeFunc::default())), + Arc::new(ScalarUDF::new_from_impl(SparkStrToMap::default())), Arc::new(ScalarUDF::new_from_impl(JsonArrayLength::default())), ] } diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 758f8ee3c90..026cb0b9a67 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -61,7 +61,7 @@ pub mod jvm_udf; mod conditional_funcs; mod conversion_funcs; mod map_funcs; -pub use map_funcs::spark_map_sort; +pub use map_funcs::{spark_map_sort, SparkMapFromArrays, SparkMapFromEntries, SparkStrToMap}; mod math_funcs; mod nondetermenistic_funcs; pub mod url_funcs; diff --git a/native/spark-expr/src/map_funcs/map_builders.rs b/native/spark-expr/src/map_funcs/map_builders.rs new file mode 100644 index 00000000000..0e3c881b609 --- /dev/null +++ b/native/spark-expr/src/map_funcs/map_builders.rs @@ -0,0 +1,651 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible `map_from_arrays`, `map_from_entries` and `str_to_map`. +//! +//! The `datafusion-spark` kernels build the `MapArray` and already follow Spark's +//! `spark.sql.mapKeyDedupPolicy`, which Comet forwards as +//! `datafusion.spark.map_key_dedup_policy`. These wrappers add the checks Spark's +//! `ArrayBasedMapBuilder` performs before inserting an entry, and restate the upstream errors +//! as the Spark error classes `SparkErrorConverter` turns back into `QueryExecutionErrors`: +//! +//! - a `NULL` key element raises `[NULL_MAP_KEY]`, ahead of any duplicate-key check, because +//! Spark rejects the `NULL` before it reaches the dedup map; +//! - a key array and value array of different lengths raise `[MAP_KEY_VALUE_DIFF_SIZES]`; +//! - a duplicate key under `EXCEPTION` raises `[DUPLICATED_MAP_KEY]` naming the key. +//! +//! `str_to_map` builds its keys by splitting a string, so it needs only the duplicate-key +//! restatement. + +use crate::SparkError; +use arrow::array::{Array, ArrayRef, AsArray, StructArray}; +use arrow::buffer::NullBuffer; +use arrow::datatypes::{DataType, FieldRef}; +use datafusion::common::{exec_err, DataFusionError, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, +}; +use datafusion_spark::function::map::map_from_arrays::MapFromArrays as DataFusionMapFromArrays; +use datafusion_spark::function::map::map_from_entries::MapFromEntries as DataFusionMapFromEntries; +use datafusion_spark::function::map::str_to_map::SparkStrToMap as DataFusionStrToMap; +use std::sync::Arc; + +/// Spark-compatible `map_from_arrays(keys, values)`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkMapFromArrays { + inner: DataFusionMapFromArrays, +} + +impl Default for SparkMapFromArrays { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapFromArrays { + pub fn new() -> Self { + Self { + inner: DataFusionMapFromArrays::new(), + } + } +} + +impl ScalarUDFImpl for SparkMapFromArrays { + fn name(&self) -> &str { + self.inner.name() + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + self.inner.return_type(arg_types) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + self.inner.return_field_from_args(args) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = expand_scalars(args)?; + match args.args.as_slice() { + [ColumnarValue::Array(keys), ColumnarValue::Array(values)] => { + validate_map_from_arrays(keys, values)? + } + other => return exec_err!("map_from_arrays expects 2 arguments, got {}", other.len()), + } + self.inner + .invoke_with_args(args) + .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Bare)) + } +} + +/// Spark-compatible `map_from_entries(entries)`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkMapFromEntries { + inner: DataFusionMapFromEntries, +} + +impl Default for SparkMapFromEntries { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapFromEntries { + pub fn new() -> Self { + Self { + inner: DataFusionMapFromEntries::new(), + } + } +} + +impl ScalarUDFImpl for SparkMapFromEntries { + fn name(&self) -> &str { + self.inner.name() + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + self.inner.return_type(arg_types) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + self.inner.return_field_from_args(args) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = expand_scalars(args)?; + match args.args.as_slice() { + [ColumnarValue::Array(entries)] => validate_map_from_entries(entries)?, + other => return exec_err!("map_from_entries expects 1 argument, got {}", other.len()), + } + self.inner + .invoke_with_args(args) + .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Bare)) + } +} + +/// Spark-compatible `str_to_map(text[, pair_delim[, key_value_delim]])`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkStrToMap { + inner: DataFusionStrToMap, +} + +impl Default for SparkStrToMap { + fn default() -> Self { + Self::new() + } +} + +impl SparkStrToMap { + pub fn new() -> Self { + Self { + inner: DataFusionStrToMap::new(), + } + } +} + +impl ScalarUDFImpl for SparkStrToMap { + fn name(&self) -> &str { + self.inner.name() + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + self.inner.return_type(arg_types) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + self.inner.return_field_from_args(args) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + // Splitting a string cannot produce a NULL key, so only the duplicate-key error needs + // restating here. + self.inner + .invoke_with_args(args) + .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Quoted)) + } +} + +/// Materializes scalar arguments so the validation below indexes rows the same way the kernel +/// does. `make_scalar_function` inside the kernel expands them anyway, so this only moves that +/// work earlier. +fn expand_scalars(mut args: ScalarFunctionArgs) -> Result { + let number_rows = args.number_rows; + for arg in args.args.iter_mut() { + if let ColumnarValue::Scalar(scalar) = arg { + *arg = ColumnarValue::Array(scalar.to_array_of_size(number_rows)?); + } + } + Ok(args) +} + +/// Rejects the inputs Spark's `MapFromArrays` rejects before building the map: a row whose key +/// and value arrays differ in length, and a `NULL` key element. +fn validate_map_from_arrays(keys: &ArrayRef, values: &ArrayRef) -> Result<()> { + // A `NULL`-typed argument makes every row a NULL map, which never reaches the builder. + if matches!(keys.data_type(), DataType::Null) || matches!(values.data_type(), DataType::Null) { + return Ok(()); + } + let (flat_keys, key_offsets) = list_values_and_offsets(keys)?; + let (_, value_offsets) = list_values_and_offsets(values)?; + if key_offsets.len() != value_offsets.len() { + return exec_err!("map_from_arrays: keys and values must have the same number of rows"); + } + let key_nulls = element_validity(&flat_keys); + + for row in 0..key_offsets.len().saturating_sub(1) { + // `MapFromArrays` is null intolerant, so a NULL input array yields a NULL map without + // evaluating the builder. + if !keys.is_valid(row) || !values.is_valid(row) { + continue; + } + let (start, end) = (key_offsets[row], key_offsets[row + 1]); + if end - start != value_offsets[row + 1] - value_offsets[row] { + return Err(SparkError::MapKeyValueDiffSizes.into()); + } + if let Some(nulls) = &key_nulls { + if nulls.slice(start, end - start).null_count() > 0 { + return Err(SparkError::NullMapKey.into()); + } + } + } + Ok(()) +} + +/// Rejects a `NULL` key element in the rows `map_from_entries` actually builds a map from. A row +/// is skipped when its entries array is NULL or holds a NULL `struct` element, since Spark +/// returns a NULL map for both without inserting any entry. +fn validate_map_from_entries(entries: &ArrayRef) -> Result<()> { + if matches!(entries.data_type(), DataType::Null) { + return Ok(()); + } + let (elements, offsets) = list_values_and_offsets(entries)?; + let Some(structs) = elements.as_any().downcast_ref::() else { + return exec_err!( + "map_from_entries: expected array>, got {:?}", + elements.data_type() + ); + }; + let Some(key_nulls) = element_validity(structs.column(0)) else { + return Ok(()); + }; + let element_nulls = structs.nulls(); + + for row in 0..offsets.len().saturating_sub(1) { + if !entries.is_valid(row) { + continue; + } + let (start, len) = (offsets[row], offsets[row + 1] - offsets[row]); + if element_nulls.is_some_and(|nulls| nulls.slice(start, len).null_count() > 0) { + continue; + } + if key_nulls.slice(start, len).null_count() > 0 { + return Err(SparkError::NullMapKey.into()); + } + } + Ok(()) +} + +/// The flattened element array of a list argument together with its per-row offsets. The offsets +/// index into the returned array, which a slice of the list does not itself narrow. +fn list_values_and_offsets(array: &ArrayRef) -> Result<(ArrayRef, Vec)> { + match array.data_type() { + DataType::List(_) => { + let list = array.as_list::(); + let offsets = list.offsets().iter().map(|o| *o as usize).collect(); + Ok((Arc::clone(list.values()), offsets)) + } + DataType::LargeList(_) => { + let list = array.as_list::(); + let offsets = list.offsets().iter().map(|o| *o as usize).collect(); + Ok((Arc::clone(list.values()), offsets)) + } + DataType::FixedSizeList(_, size) => { + let list = array.as_fixed_size_list(); + let size = *size as usize; + let offsets = (0..=list.len()).map(|row| row * size).collect(); + Ok((Arc::clone(list.values()), offsets)) + } + other => exec_err!("expected list, large_list or fixed_size_list, got {other:?}"), + } +} + +/// The per-element validity of a map key array, or `None` when no element is NULL. A `NullArray` +/// carries no null buffer even though all of its elements are NULL, so report one for it. +fn element_validity(array: &ArrayRef) -> Option { + if matches!(array.data_type(), DataType::Null) { + return Some(NullBuffer::new_null(array.len())); + } + array + .nulls() + .filter(|nulls| nulls.null_count() > 0) + .cloned() +} + +/// How the upstream kernel renders the offending key in its duplicate-key message. +#[derive(Clone, Copy)] +enum DuplicateKeyFormat { + /// The map builders write the key as-is, which is what Spark's `key.toString` produces. + Bare, + /// `str_to_map` single-quotes it. + Quoted, +} + +/// Restates the upstream duplicate-key error as `SparkError::DuplicatedMapKey` so the JVM side +/// raises Spark's `DUPLICATED_MAP_KEY` naming the same key. Any other error is passed through. +fn as_spark_error(error: DataFusionError, key_format: DuplicateKeyFormat) -> DataFusionError { + match duplicate_map_key(&error.to_string(), key_format) { + Some(key) => SparkError::DuplicatedMapKey { key }.into(), + None => error, + } +} + +/// The key named by `datafusion-spark`'s duplicate-key message. The +/// `*_reports_the_duplicate_key` tests pin the wordings this parses against the kernels +/// themselves, so an upstream rewording fails there rather than silently downgrading the error +/// to a generic execution failure. +fn duplicate_map_key(message: &str, key_format: DuplicateKeyFormat) -> Option { + let (open, close) = match key_format { + DuplicateKeyFormat::Bare => ("[DUPLICATED_MAP_KEY] Duplicate map key ", " was found"), + DuplicateKeyFormat::Quoted => ("[DUPLICATED_MAP_KEY] Duplicate map key '", "' was found"), + }; + let (_, tail) = message.split_once(open)?; + let (key, _) = tail.rsplit_once(close)?; + Some(key.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int32Array, ListArray, MapArray, StringArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Fields}; + use datafusion::common::config::{ConfigOptions, MapKeyDedupPolicy}; + use datafusion::common::ScalarValue; + + /// `[[1, 2], [3]]`-shaped keys, with `nulls` marking whole rows NULL. + fn int_list(values: Int32Array, offsets: &[i32], nulls: Option) -> ArrayRef { + let field = Arc::new(Field::new("item", DataType::Int32, true)); + Arc::new(ListArray::new( + field, + OffsetBuffer::new(offsets.to_vec().into()), + Arc::new(values), + nulls, + )) + } + + fn string_list(values: StringArray, offsets: &[i32], nulls: Option) -> ArrayRef { + let field = Arc::new(Field::new("item", DataType::Utf8, true)); + Arc::new(ListArray::new( + field, + OffsetBuffer::new(offsets.to_vec().into()), + Arc::new(values), + nulls, + )) + } + + /// `array>`, with `element_nulls` marking NULL entries. + fn entry_list( + keys: Int32Array, + values: StringArray, + offsets: &[i32], + element_nulls: Option, + ) -> ArrayRef { + let fields = Fields::from(vec![ + Field::new("key", DataType::Int32, true), + Field::new("value", DataType::Utf8, true), + ]); + let structs = StructArray::new( + fields.clone(), + vec![Arc::new(keys), Arc::new(values)], + element_nulls, + ); + let field = Arc::new(Field::new("item", DataType::Struct(fields), true)); + Arc::new(ListArray::new( + field, + OffsetBuffer::new(offsets.to_vec().into()), + Arc::new(structs), + None, + )) + } + + fn invoke( + udf: &dyn ScalarUDFImpl, + args: Vec, + policy: MapKeyDedupPolicy, + ) -> Result { + let arg_fields: Vec = args + .iter() + .enumerate() + .map(|(i, arg)| Arc::new(Field::new(format!("arg{i}"), arg.data_type().clone(), true))) + .collect(); + let scalar_arguments: Vec> = vec![None; args.len()]; + let return_field = udf.return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalar_arguments, + })?; + let mut config = ConfigOptions::default(); + config.spark.map_key_dedup_policy = policy; + let number_rows = args.first().map(|arg| arg.len()).unwrap_or(0); + udf.invoke_with_args(ScalarFunctionArgs { + args: args.into_iter().map(ColumnarValue::Array).collect(), + arg_fields, + number_rows, + return_field, + config_options: Arc::new(config), + }) + } + + fn map_result(value: ColumnarValue) -> MapArray { + match value { + ColumnarValue::Array(array) => array.as_map().clone(), + ColumnarValue::Scalar(scalar) => { + scalar.to_array().expect("scalar to array").as_map().clone() + } + } + } + + #[test] + fn map_from_arrays_rejects_null_key() { + let keys = int_list(Int32Array::from(vec![Some(1), None]), &[0, 2], None); + let values = string_list(StringArray::from(vec![Some("a"), Some("b")]), &[0, 2], None); + let err = invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[NULL_MAP_KEY]"), "{err}"); + } + + #[test] + fn map_from_arrays_ignores_null_key_in_a_null_row() { + // Row 0's keys array is NULL, so Spark returns a NULL map without inspecting its keys. + let keys = int_list( + Int32Array::from(vec![None, Some(1)]), + &[0, 1, 2], + Some(NullBuffer::from(vec![false, true])), + ); + let values = string_list( + StringArray::from(vec![Some("a"), Some("b")]), + &[0, 1, 2], + None, + ); + let result = map_result( + invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::Exception, + ) + .unwrap(), + ); + assert!(result.is_null(0)); + assert_eq!(result.value_offsets(), &[0, 0, 1]); + } + + #[test] + fn map_from_arrays_rejects_key_value_length_mismatch() { + let keys = int_list(Int32Array::from(vec![1, 2]), &[0, 2], None); + let values = string_list(StringArray::from(vec![Some("a")]), &[0, 1], None); + let err = invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[MAP_KEY_VALUE_DIFF_SIZES]"), "{err}"); + } + + /// Pins the upstream message `duplicate_map_key` parses: a wording change upstream fails here + /// rather than silently downgrading the error to a generic execution failure. + #[test] + fn map_from_arrays_reports_the_duplicate_key() { + let keys = int_list(Int32Array::from(vec![7, 7]), &[0, 2], None); + let values = string_list(StringArray::from(vec![Some("a"), Some("b")]), &[0, 2], None); + let err = invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("[DUPLICATED_MAP_KEY] Cannot create map with duplicate keys: 7."), + "{err}" + ); + } + + /// Spark's `duplicateMapKeyFoundError` reports `key.toString`, so a string key carries no + /// quotes. `str_to_map` quotes its key and `map_from_arrays` does not, which is why the two + /// go through different `DuplicateKeyFormat`s. + #[test] + fn map_from_arrays_reports_a_string_duplicate_key_unquoted() { + let field = Arc::new(Field::new("item", DataType::Utf8, true)); + let keys: ArrayRef = Arc::new(ListArray::new( + field, + OffsetBuffer::new(vec![0i32, 2].into()), + Arc::new(StringArray::from(vec![Some("a"), Some("a")])), + None, + )); + let values = string_list(StringArray::from(vec![Some("1"), Some("2")]), &[0, 2], None); + let err = invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("[DUPLICATED_MAP_KEY] Cannot create map with duplicate keys: a."), + "{err}" + ); + } + + #[test] + fn map_from_arrays_honours_last_win() { + let keys = int_list(Int32Array::from(vec![7, 7]), &[0, 2], None); + let values = string_list(StringArray::from(vec![Some("a"), Some("b")]), &[0, 2], None); + let result = map_result( + invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::LastWin, + ) + .unwrap(), + ); + assert_eq!(result.value_offsets(), &[0, 1]); + let values = result.entries().column(1).as_string::().clone(); + assert_eq!(values.value(0), "b"); + } + + #[test] + fn map_from_entries_rejects_null_key() { + let entries = entry_list( + Int32Array::from(vec![Some(1), None]), + StringArray::from(vec![Some("a"), Some("b")]), + &[0, 2], + None, + ); + let err = invoke( + &SparkMapFromEntries::default(), + vec![entries], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[NULL_MAP_KEY]"), "{err}"); + } + + #[test] + fn map_from_entries_ignores_a_null_entry() { + // A NULL struct element makes the whole row a NULL map, so its NULL key is never a key. + let entries = entry_list( + Int32Array::from(vec![None, Some(2)]), + StringArray::from(vec![None, Some("b")]), + &[0, 1, 2], + Some(NullBuffer::from(vec![false, true])), + ); + let result = map_result( + invoke( + &SparkMapFromEntries::default(), + vec![entries], + MapKeyDedupPolicy::Exception, + ) + .unwrap(), + ); + assert!(result.is_null(0)); + assert_eq!(result.value_offsets(), &[0, 0, 1]); + } + + #[test] + fn map_from_entries_honours_last_win() { + let entries = entry_list( + Int32Array::from(vec![7, 7]), + StringArray::from(vec![Some("a"), Some("b")]), + &[0, 2], + None, + ); + let result = map_result( + invoke( + &SparkMapFromEntries::default(), + vec![entries], + MapKeyDedupPolicy::LastWin, + ) + .unwrap(), + ); + assert_eq!(result.value_offsets(), &[0, 1]); + let values = result.entries().column(1).as_string::().clone(); + assert_eq!(values.value(0), "b"); + } + + #[test] + fn str_to_map_reports_the_duplicate_key() { + let text: ArrayRef = Arc::new(StringArray::from(vec![Some("a:1,b:2,a:3")])); + let err = invoke( + &SparkStrToMap::default(), + vec![text], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("[DUPLICATED_MAP_KEY] Cannot create map with duplicate keys: a."), + "{err}" + ); + } + + #[test] + fn str_to_map_honours_last_win() { + let text: ArrayRef = Arc::new(StringArray::from(vec![Some("a:1,b:2,a:3")])); + let result = map_result( + invoke( + &SparkStrToMap::default(), + vec![text], + MapKeyDedupPolicy::LastWin, + ) + .unwrap(), + ); + assert_eq!(result.value_offsets(), &[0, 2]); + } + + #[test] + fn duplicate_map_key_ignores_unrelated_errors() { + assert_eq!( + duplicate_map_key("Execution error: something else", DuplicateKeyFormat::Bare), + None + ); + assert_eq!( + duplicate_map_key( + "Execution error: something else", + DuplicateKeyFormat::Quoted + ), + None + ); + } +} diff --git a/native/spark-expr/src/map_funcs/mod.rs b/native/spark-expr/src/map_funcs/mod.rs index 7288b847a83..99fdc6eeda2 100644 --- a/native/spark-expr/src/map_funcs/mod.rs +++ b/native/spark-expr/src/map_funcs/mod.rs @@ -15,5 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod map_builders; mod map_sort; +pub use map_builders::{SparkMapFromArrays, SparkMapFromEntries, SparkStrToMap}; pub use map_sort::spark_map_sort; diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index e2c132904d5..386ea69c192 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -358,6 +358,13 @@ object CometExecIterator extends Logging { CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(SQLConf.get).toString) + // The native map constructors (map_from_arrays, map_from_entries, str_to_map) resolve + // duplicate keys with this policy, which the native side reads as + // `datafusion.spark.map_key_dedup_policy`. + builder.putEntries( + SQLConf.MAP_KEY_DEDUP_POLICY.key, + SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString) + builder.build().toByteArray } diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala index 51fa428b543..e6bf2a62975 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -23,8 +23,9 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ +import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.serde.QueryPlanSerde.{createBinaryExpr, exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto} +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto} import org.apache.comet.shims.CometTypeShim /** @@ -132,40 +133,44 @@ object CometMapExtract extends CometExpressionSerde[GetMapValue] { } } -private object MapKeyDedupPolicySupport { - val incompatibleReason: String = - s"`${SQLConf.MAP_KEY_DEDUP_POLICY.key}` is set to " + - s"`${SQLConf.MapKeyDedupPolicy.LAST_WIN}`; Comet's native map construction " + - "does not implement LAST_WIN dedup semantics." - - val nullKeyReason: String = - "Spark rejects a `NULL` element inside the keys array with a `RuntimeException`" + - " (`Cannot use null as map key`); Comet's native `map_from_arrays` / `map_from_entries`" + - " does not detect a per-element `NULL` key and produces a map with a `NULL` key instead" + - " ([#4680](https://github.com/apache/datafusion-comet/issues/4680))." - - def isLastWin: Boolean = - SQLConf.get - .getConf(SQLConf.MAP_KEY_DEDUP_POLICY) - .toString - .equalsIgnoreCase(SQLConf.MapKeyDedupPolicy.LAST_WIN.toString) +/** + * Shared gate for the native map constructors (`map_from_arrays`, `map_from_entries`), which + * reproduce Spark's `ArrayBasedMapBuilder`: they reject a `NULL` key with `NULL_MAP_KEY` and + * follow `spark.sql.mapKeyDedupPolicy`, whose value Comet forwards to the native session as + * `datafusion.spark.map_key_dedup_policy`. + */ +private object MapBuilderSupport { + + /** + * `ArrayBasedMapBuilder` normalizes a floating-point key before storing it, so a `-0.0` key is + * stored as `+0.0` and every `NaN` collapses to one canonical `NaN`. The native builders + * compare the raw Arrow values, so a map built from both `-0.0` and `+0.0` keeps two entries + * where Spark reports a duplicate key. This is a note rather than a decline because a map keyed + * on `-0.0` or `NaN` is rare; `spark.comet.exec.strictFloatingPoint` declines it for users who + * want the guarantee. + */ + val floatingPointKeyNote: String = + "Spark normalizes a floating-point map key, so a `-0.0` key is stored as `+0.0` and all " + + "`NaN` keys collapse into one. Comet's native map construction compares the raw Arrow " + + "values, so `-0.0` and `+0.0` stay distinct keys rather than a duplicate key. Set " + + s"`${COMET_EXEC_STRICT_FLOATING_POINT.key}=true` to fall back to Spark for a " + + "floating-point map key." + + /** The support level for a map constructor whose result has key type `keyType`. */ + def keySupport(keyType: DataType): SupportLevel = + SupportLevel + .strictFloatingPointReason(keyType, "Map construction on a floating-point key") + .map(reason => Incompatible(Some(reason))) + .getOrElse(Compatible(None)) } object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { - override def getIncompatibleReasons(): Seq[String] = - Seq(MapKeyDedupPolicySupport.incompatibleReason) - override def getCompatibleNotes(): Seq[String] = - Seq(MapKeyDedupPolicySupport.nullKeyReason) + Seq(MapBuilderSupport.floatingPointKeyNote) - override def getSupportLevel(expr: MapFromArrays): SupportLevel = { - if (MapKeyDedupPolicySupport.isLastWin) { - Incompatible(Some(MapKeyDedupPolicySupport.incompatibleReason)) - } else { - Compatible(None) - } - } + override def getSupportLevel(expr: MapFromArrays): SupportLevel = + MapBuilderSupport.keySupport(expr.dataType.keyType) override def convert( expr: MapFromArrays, @@ -173,38 +178,9 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { binding: Boolean): Option[ExprOuterClass.Expr] = { val keysExpr = exprToProtoInternal(expr.left, inputs, binding) val valuesExpr = exprToProtoInternal(expr.right, inputs, binding) - val keyType = expr.left.dataType.asInstanceOf[ArrayType].elementType - val valueType = expr.right.dataType.asInstanceOf[ArrayType].elementType - val returnType = MapType(keyType = keyType, valueType = valueType) - for { - andBinaryExprProto <- createAndBinaryExpr(expr, inputs, binding) - mapFromArraysExprProto <- scalarFunctionExprToProto("map", keysExpr, valuesExpr) - nullLiteralExprProto <- exprToProtoInternal(Literal(null, returnType), inputs, binding) - } yield { - val caseWhenExprProto = ExprOuterClass.CaseWhen - .newBuilder() - .addWhen(andBinaryExprProto) - .addThen(mapFromArraysExprProto) - .setElseExpr(nullLiteralExprProto) - .build() - ExprOuterClass.Expr - .newBuilder() - .setCaseWhen(caseWhenExprProto) - .build() - } - } - - private def createAndBinaryExpr( - expr: MapFromArrays, - inputs: Seq[Attribute], - binding: Boolean): Option[ExprOuterClass.Expr] = { - createBinaryExpr( - expr, - IsNotNull(expr.left), - IsNotNull(expr.right), - inputs, - binding, - (builder, binaryExpr) => builder.setAnd(binaryExpr)) + // Native `map_from_arrays` is null intolerant like Spark's: a NULL keys or values array + // yields a NULL map for that row, so no CaseWhen guard is needed here. + scalarFunctionExprToProto("map_from_arrays", keysExpr, valuesExpr) } } @@ -217,20 +193,18 @@ object CometMapFromEntries "`BinaryType` is not supported as a map value in `map_from_entries`" override def getIncompatibleReasons(): Seq[String] = - Seq(keyUnsupportedReason, valueUnsupportedReason, MapKeyDedupPolicySupport.incompatibleReason) + Seq(keyUnsupportedReason, valueUnsupportedReason) override def getCompatibleNotes(): Seq[String] = - Seq(MapKeyDedupPolicySupport.nullKeyReason) + Seq(MapBuilderSupport.floatingPointKeyNote) override def getSupportLevel(expr: MapFromEntries): SupportLevel = { if (SupportLevel.containsType(expr.dataType.keyType, classOf[BinaryType])) { Incompatible(Some(keyUnsupportedReason)) } else if (SupportLevel.containsType(expr.dataType.valueType, classOf[BinaryType])) { Incompatible(Some(valueUnsupportedReason)) - } else if (MapKeyDedupPolicySupport.isLastWin) { - Incompatible(Some(MapKeyDedupPolicySupport.incompatibleReason)) } else { - Compatible(None) + MapBuilderSupport.keySupport(expr.dataType.keyType) } } } diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql index 178c07f432a..6ff24fd85ec 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays.sql @@ -58,4 +58,23 @@ query SELECT map_from_arrays(array('a'), NULL) query -SELECT map_from_arrays(NULL, NULL) \ No newline at end of file +SELECT map_from_arrays(NULL, NULL) + +-- Spark's ArrayBasedMapBuilder rejects a NULL key element outright, ahead of the duplicate-key +-- check, and resolves duplicates by the default `spark.sql.mapKeyDedupPolicy` = `EXCEPTION`. +-- `map_from_arrays_dedup_policy.sql` covers `LAST_WIN`. + +query expect_error(NULL_MAP_KEY) +SELECT map_from_arrays(array('a', NULL), array(1, 2)) + +-- a NULL key is reported as such even when it repeats, which a duplicate check would see first +query expect_error(NULL_MAP_KEY) +SELECT map_from_arrays(array(CAST(NULL AS STRING), NULL), array(1, 2)) + +query expect_error(DUPLICATED_MAP_KEY) +SELECT map_from_arrays(array('a', 'a'), array(1, 2)) + +-- key and value arrays of different lengths. Spark reports this through a `_LEGACY_ERROR_TEMP_*` +-- condition whose number moves between Spark versions, so match on the message instead. +query expect_error(must have the same length) +SELECT map_from_arrays(array('a', 'b'), array(1)) diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql index fffaf5f9a92..70517905b28 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_dedup_policy.sql @@ -15,10 +15,10 @@ -- specific language governing permissions and limitations -- under the License. --- Verifies that `map_from_arrays` falls back to Spark when `spark.sql.mapKeyDedupPolicy` is set --- to `LAST_WIN`. Spark's ArrayBasedMapBuilder keeps the last occurrence of each duplicate key; --- Comet's native `map` scalar has no LAST_WIN path, so it must fall back. The default `EXCEPTION` --- mode agrees with Comet and is covered by `map_from_arrays.sql`. +-- Verifies that `map_from_arrays` follows `spark.sql.mapKeyDedupPolicy` = `LAST_WIN`, keeping +-- the last value for each duplicate key. Comet forwards the policy to the native builder as +-- `datafusion.spark.map_key_dedup_policy`, so the query stays native rather than falling back. +-- The default `EXCEPTION` mode is covered by `map_from_arrays.sql`. -- Config: spark.sql.mapKeyDedupPolicy=LAST_WIN @@ -29,13 +29,22 @@ statement INSERT INTO test_map_from_arrays_dedup VALUES (array('a', 'b', 'c'), array(1, 2, 3)), (array('a', 'a', 'b'), array(1, 2, 3)), - (array('x', 'x'), array(10, 20)) + (array('x', 'x'), array(10, 20)), + (array(), array()), + (NULL, array(99)) --- literal duplicate keys under LAST_WIN: Spark keeps the last value; Comet must fall back. -query expect_fallback(mapKeyDedupPolicy) +-- literal duplicate keys: the last value wins +query SELECT map_from_arrays(array('a', 'a', 'b'), array(1, 2, 3)) --- column input falls back the same way; the incompat branch is triggered by the SQLConf value, --- not per-row content. -query expect_fallback(mapKeyDedupPolicy) +-- three occurrences of the same key collapse to the last one +query +SELECT map_from_arrays(array('a', 'a', 'a'), array(1, 2, 3)) + +-- column input, including rows without duplicates and a NULL row +query SELECT map_from_arrays(k, v) FROM test_map_from_arrays_dedup + +-- LAST_WIN does not weaken the NULL key check +query expect_error(NULL_MAP_KEY) +SELECT map_from_arrays(array('a', NULL), array(1, 2)) diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_entries.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_entries.sql index 74723509334..cdbdba4e2bb 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_from_entries.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_entries.sql @@ -35,3 +35,17 @@ SELECT map_from_entries(array(struct(10, cast('x' as binary)))) -- literal arguments query spark_answer_only SELECT map_from_entries(array(struct('x', 10), struct('y', 20), struct('z', 30))) + +-- Spark's ArrayBasedMapBuilder rejects a NULL key element outright, ahead of the duplicate-key +-- check, and resolves duplicates by the default `spark.sql.mapKeyDedupPolicy` = `EXCEPTION`. +-- `map_from_entries_dedup_policy.sql` covers `LAST_WIN`. + +query expect_error(NULL_MAP_KEY) +SELECT map_from_entries(array(struct(CAST(NULL AS STRING), 1), struct('b', 2))) + +query expect_error(DUPLICATED_MAP_KEY) +SELECT map_from_entries(array(struct('a', 1), struct('a', 2))) + +-- a NULL entry makes the whole map NULL, so its NULL key is never inserted +query +SELECT map_from_entries(array(CAST(NULL AS struct), struct('b' AS key, 2 AS value))) diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_entries_dedup_policy.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_entries_dedup_policy.sql index feba7951933..c344e583e19 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/map_from_entries_dedup_policy.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_entries_dedup_policy.sql @@ -15,15 +15,12 @@ -- specific language governing permissions and limitations -- under the License. --- Verifies that `map_from_entries` falls back to Spark when `spark.sql.mapKeyDedupPolicy` is set --- to `LAST_WIN`. `CometMapFromEntries` mixes in `CodegenDispatchFallback`, so its native --- `Incompatible` normally routes through the JVM codegen dispatcher; we disable the dispatcher --- here so the incompat branch surfaces as a genuine Spark fallback rather than in-pipeline --- codegen. The default `EXCEPTION` mode agrees with Comet and is covered by --- `map_from_entries.sql`. +-- Verifies that `map_from_entries` follows `spark.sql.mapKeyDedupPolicy` = `LAST_WIN`, keeping +-- the last value for each duplicate key. Comet forwards the policy to the native builder as +-- `datafusion.spark.map_key_dedup_policy`, so the query stays native rather than routing through +-- the JVM codegen dispatcher. The default `EXCEPTION` mode is covered by `map_from_entries.sql`. -- Config: spark.sql.mapKeyDedupPolicy=LAST_WIN --- Config: spark.comet.exec.scalaUDF.codegen.enabled=false statement CREATE TABLE test_map_from_entries_dedup(entries array>) USING parquet @@ -32,13 +29,22 @@ statement INSERT INTO test_map_from_entries_dedup VALUES (array(struct('a', 1), struct('b', 2), struct('c', 3))), (array(struct('a', 1), struct('a', 2), struct('b', 3))), - (array(struct('x', 10), struct('x', 20))) + (array(struct('x', 10), struct('x', 20))), + (array()), + (NULL) --- literal duplicate keys under LAST_WIN: Spark keeps the last value; Comet must fall back. -query expect_fallback(mapKeyDedupPolicy) +-- literal duplicate keys: the last value wins +query SELECT map_from_entries(array(struct('a', 1), struct('a', 2), struct('b', 3))) --- column input falls back the same way; the incompat branch is triggered by the SQLConf value, --- not per-row content. -query expect_fallback(mapKeyDedupPolicy) +-- three occurrences of the same key collapse to the last one +query +SELECT map_from_entries(array(struct('a', 1), struct('a', 2), struct('a', 3))) + +-- column input, including rows without duplicates and a NULL row +query SELECT map_from_entries(entries) FROM test_map_from_entries_dedup + +-- LAST_WIN does not weaken the NULL key check +query expect_error(NULL_MAP_KEY) +SELECT map_from_entries(array(struct(CAST(NULL AS STRING), 1), struct('b', 2))) diff --git a/spark/src/test/resources/sql-tests/expressions/map/str_to_map.sql b/spark/src/test/resources/sql-tests/expressions/map/str_to_map.sql index 7db1242fd4e..1642c68f4c7 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/str_to_map.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/str_to_map.sql @@ -70,10 +70,10 @@ SELECT str_to_map('a') query SELECT str_to_map('a=1&b=2&c=3', '&', '=') --- Duplicate keys: EXCEPTION policy (Spark 3.0+ default) --- TODO: Add LAST_WIN policy tests when spark.sql.mapKeyDedupPolicy config is supported --- query --- SELECT str_to_map('a:1,b:2,a:3') +-- Duplicate keys under the default EXCEPTION policy; `str_to_map_dedup_policy.sql` covers +-- LAST_WIN. +query expect_error(DUPLICATED_MAP_KEY) +SELECT str_to_map('a:1,b:2,a:3') -- NULL input returns NULL query diff --git a/spark/src/test/resources/sql-tests/expressions/map/str_to_map_dedup_policy.sql b/spark/src/test/resources/sql-tests/expressions/map/str_to_map_dedup_policy.sql new file mode 100644 index 00000000000..f3aab4eb8af --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/str_to_map_dedup_policy.sql @@ -0,0 +1,42 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Verifies that `str_to_map` follows `spark.sql.mapKeyDedupPolicy` = `LAST_WIN`, keeping the +-- last value for each duplicate key. Comet forwards the policy to the native kernel as +-- `datafusion.spark.map_key_dedup_policy`. The default `EXCEPTION` mode is covered by +-- `str_to_map.sql`. + +-- Config: spark.sql.mapKeyDedupPolicy=LAST_WIN + +statement +CREATE TABLE test_str_to_map_dedup(s string) USING parquet + +statement +INSERT INTO test_str_to_map_dedup VALUES + ('a:1,b:2,a:3'), + ('a:1,b:2,c:3'), + ('x:1,x:2,x:3'), + (NULL) + +query +SELECT str_to_map('a:1,b:2,a:3') + +query +SELECT str_to_map(s) FROM test_str_to_map_dedup + +query +SELECT str_to_map(s, ',', ':') FROM test_str_to_map_dedup diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index f4a559b872b..9d4302be0be 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -126,6 +126,95 @@ class CometMapExpressionSuite extends CometTestBase { } } + // Spark builds both `map_from_arrays` and `map_from_entries` through `ArrayBasedMapBuilder`, + // which rejects a NULL key outright and resolves duplicate keys by + // `spark.sql.mapKeyDedupPolicy`. Comet forwards that policy to the native builders as + // `datafusion.spark.map_key_dedup_policy`, so both engines must agree on the answer and on the + // error. Each query reads a column so constant folding cannot evaluate it on the driver, which + // would take the native builders out of the picture. + // https://github.com/apache/datafusion-comet/issues/4680 + private def withMapBuilderTable(f: String => Unit): Unit = { + val table = "map_builder_input" + withTable(table) { + sql(s"CREATE TABLE $table(k INT, v STRING) USING parquet") + sql(s"INSERT INTO $table VALUES (1, 'a'), (2, 'b'), (3, 'c')") + f(table) + } + } + + test("map_from_arrays - null key is rejected") { + withMapBuilderTable { table => + val exception = checkSparkError( + sql(s"SELECT map_from_arrays(array(k, CAST(NULL AS INT)), array(v, v)) FROM $table"), + "NULL_MAP_KEY") + assert(exception.getMessage.contains("Cannot use null as map key")) + } + } + + test("map_from_arrays - a null input array gives a null map") { + withMapBuilderTable { table => + checkSparkAnswerAndOperator( + sql(s"""SELECT map_from_arrays(CASE WHEN k > 1 THEN array(k) END, array(v)), + | map_from_arrays(array(k), CASE WHEN k > 2 THEN array(v) END) + |FROM $table""".stripMargin)) + } + } + + test("map_from_arrays - key and value arrays of different lengths are rejected") { + withMapBuilderTable { table => + // Spark reports this through a `_LEGACY_ERROR_TEMP_*` condition whose number moves between + // Spark versions, so hold the two engines to each other rather than naming the condition. + checkSparkErrorParity(sql(s"SELECT map_from_arrays(array(k, k + 1), array(v)) FROM $table")) + } + } + + test("map_from_arrays - duplicate key follows spark.sql.mapKeyDedupPolicy") { + withMapBuilderTable { table => + val query = s"SELECT map_from_arrays(array(k, k), array(v, concat(v, 'x'))) FROM $table" + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + // One row, so both engines name the same offending key. + val exception = checkSparkError(sql(s"$query WHERE k = 2"), "DUPLICATED_MAP_KEY") + assert(exception.getMessage.contains("Duplicate map key 2 was found")) + } + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + checkSparkAnswerAndOperator(sql(query)) + } + } + } + + test("map_from_entries - null key is rejected") { + withMapBuilderTable { table => + val exception = checkSparkError( + sql(s"SELECT map_from_entries(array(struct(CAST(NULL AS INT), v))) FROM $table"), + "NULL_MAP_KEY") + assert(exception.getMessage.contains("Cannot use null as map key")) + } + } + + test("map_from_entries - a null entry gives a null map") { + withMapBuilderTable { table => + checkSparkAnswerAndOperator( + sql(s"""SELECT map_from_entries(array(CASE WHEN k > 1 THEN struct(k, v) END)) + |FROM $table""".stripMargin)) + } + } + + test("map_from_entries - duplicate key follows spark.sql.mapKeyDedupPolicy") { + withMapBuilderTable { table => + // `struct` names a column argument after the column, so both entries need explicit field + // names for `array` to see one struct type. + val query = "SELECT map_from_entries(array(struct(k AS key, v AS value), " + + s"struct(k AS key, concat(v, 'x') AS value))) FROM $table" + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + val exception = checkSparkError(sql(s"$query WHERE k = 2"), "DUPLICATED_MAP_KEY") + assert(exception.getMessage.contains("Duplicate map key 2 was found")) + } + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + checkSparkAnswerAndOperator(sql(query)) + } + } + } + test("size with map input") { withTempDir { dir => withTempView("t1") { diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index a2bbe415cf5..78fad80df0e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -448,13 +448,27 @@ abstract class CometTestBase protected def checkSparkError( df: DataFrame, errorClass: String): SparkThrowable with Throwable = { + val actual = checkSparkErrorParity(df, Some(errorClass)) + assert(actual.getErrorClass == errorClass) + actual + } + + /** + * Checks native execution and that both engines fail with the same exception type, error class + * and SQLSTATE. Use this rather than `checkSparkError` for an error Spark still reports through + * a `_LEGACY_ERROR_TEMP_*` condition, whose number moves between Spark versions. + */ + protected def checkSparkErrorParity( + df: DataFrame, + errorClass: Option[String] = None): SparkThrowable with Throwable = { checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) def structuredError( error: Option[Throwable], engine: String): SparkThrowable with Throwable = { - val failure = error.getOrElse(fail(s"$engine did not fail with $errorClass")) + val expectation = errorClass.map(c => s" with $c").getOrElse("") + val failure = error.getOrElse(fail(s"$engine did not fail$expectation")) val chain = causeChain(failure) assert(!chain.exists(_.isInstanceOf[CometNativeException]), s"$engine: $failure") chain.collect { case e: SparkThrowable with Throwable => e }.lastOption.getOrElse { @@ -464,9 +478,9 @@ abstract class CometTestBase val expected = structuredError(sparkError, "Spark") val actual = structuredError(cometError, "Comet") - assert(expected.getErrorClass == errorClass) + errorClass.foreach(c => assert(expected.getErrorClass == c)) assert(actual.getClass == expected.getClass) - assert(actual.getErrorClass == errorClass) + assert(actual.getErrorClass == expected.getErrorClass) assert(actual.getSqlState == expected.getSqlState) actual }