From 32bf30b9a69ba267188e54df03b138a990b96028 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Thu, 10 Sep 2026 22:07:55 +0000 Subject: [PATCH 01/16] 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 --- .../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 } From e4cd899d40073d62ca36b3e9695b13535e4ca059 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Tue, 15 Sep 2026 08:46:14 +0000 Subject: [PATCH 02/16] fix: read the right row when a map builder gets a sliced list argument The upstream `datafusion-spark` map kernels read each row's entries at its own offset but build the mask selecting the surviving keys from zero, then apply that mask to the list's whole values array. Arrow's `filter` accepts a predicate shorter than the array it filters, so on a sliced argument the mismatch silently returns keys belonging to earlier rows rather than raising: keys `[[10], [20]]` and values `[[100], [200]]`, both sliced to the second row, built `{10: 200}` instead of `{20: 200}`. A `LIMIT` above a projection produces such an argument. Compact any list argument whose values hold more than its offsets address before validating or delegating, so the kernels see the layout they assume. `map_from_entries` reached the same helper before this branch, so the bug is not new to `map_from_arrays`; a fix belongs upstream as well. Reported by @rich7420. --- .../spark-expr/src/map_funcs/map_builders.rs | 117 +++++++++++++++++- 1 file changed, 113 insertions(+), 4 deletions(-) diff --git a/native/spark-expr/src/map_funcs/map_builders.rs b/native/spark-expr/src/map_funcs/map_builders.rs index 0e3c881b609..531c878e077 100644 --- a/native/spark-expr/src/map_funcs/map_builders.rs +++ b/native/spark-expr/src/map_funcs/map_builders.rs @@ -32,8 +32,9 @@ //! restatement. use crate::SparkError; -use arrow::array::{Array, ArrayRef, AsArray, StructArray}; +use arrow::array::{Array, ArrayRef, AsArray, StructArray, UInt32Array}; use arrow::buffer::NullBuffer; +use arrow::compute::take; use arrow::datatypes::{DataType, FieldRef}; use datafusion::common::{exec_err, DataFusionError, Result}; use datafusion::logical_expr::{ @@ -82,7 +83,8 @@ impl ScalarUDFImpl for SparkMapFromArrays { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let args = expand_scalars(args)?; + let mut args = expand_scalars(args)?; + compact_list_arguments(&mut args)?; match args.args.as_slice() { [ColumnarValue::Array(keys), ColumnarValue::Array(values)] => { validate_map_from_arrays(keys, values)? @@ -133,7 +135,8 @@ impl ScalarUDFImpl for SparkMapFromEntries { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let args = expand_scalars(args)?; + let mut args = expand_scalars(args)?; + compact_list_arguments(&mut 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()), @@ -203,6 +206,47 @@ fn expand_scalars(mut args: ScalarFunctionArgs) -> Result { Ok(args) } +/// Rebuilds any list argument whose entries do not start at offset zero. +/// +/// The upstream kernels read each row's entries at its own offset but build the mask that selects +/// the surviving keys from zero, then apply that mask to the list's whole values array. Arrow's +/// `filter` accepts a predicate shorter than the array it filters, so on a sliced argument the +/// mismatch silently selects keys belonging to earlier rows instead of raising. A `LIMIT` above a +/// projection is enough to produce one, so bring the argument back to offset zero first. +fn compact_list_arguments(args: &mut ScalarFunctionArgs) -> Result<()> { + for arg in args.args.iter_mut() { + if let ColumnarValue::Array(array) = arg { + if !entries_start_at_zero(array) { + let indices = UInt32Array::from_iter_values(0..array.len() as u32); + *arg = ColumnarValue::Array(take(array.as_ref(), &indices, None)?); + } + } + } + Ok(()) +} + +/// Whether a list argument's values hold exactly the entries its offsets address, which is what +/// the upstream kernels assume. Any other array type is left alone. +fn entries_start_at_zero(array: &ArrayRef) -> bool { + match array.data_type() { + DataType::List(_) => { + let list = array.as_list::(); + let offsets = list.offsets(); + offsets[0] == 0 && offsets[offsets.len() - 1] as usize == list.values().len() + } + DataType::LargeList(_) => { + let list = array.as_list::(); + let offsets = list.offsets(); + offsets[0] == 0 && offsets[offsets.len() - 1] as usize == list.values().len() + } + DataType::FixedSizeList(_, size) => { + let list = array.as_fixed_size_list(); + list.values().len() == list.len() * *size as usize + } + _ => true, + } +} + /// 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<()> { @@ -343,7 +387,7 @@ mod tests { use super::*; use arrow::array::{Int32Array, ListArray, MapArray, StringArray}; use arrow::buffer::OffsetBuffer; - use arrow::datatypes::{Field, Fields}; + use arrow::datatypes::{Field, Fields, Int32Type}; use datafusion::common::config::{ConfigOptions, MapKeyDedupPolicy}; use datafusion::common::ScalarValue; @@ -634,6 +678,71 @@ mod tests { assert_eq!(result.value_offsets(), &[0, 2]); } + /// A `LIMIT` above a projection hands the kernel a sliced list. The mask the upstream helper + /// builds is zero-based while it reads entries at each row's own offset, so without + /// `compact_list_arguments` this reads a preceding row's key instead of raising. + #[test] + fn map_from_arrays_reads_the_right_row_of_a_sliced_list() { + let keys = int_list(Int32Array::from(vec![10, 20]), &[0, 1, 2], None); + let values = string_list( + StringArray::from(vec![Some("100"), Some("200")]), + &[0, 1, 2], + None, + ); + let result = map_result( + invoke( + &SparkMapFromArrays::default(), + vec![keys.slice(1, 1), values.slice(1, 1)], + MapKeyDedupPolicy::Exception, + ) + .unwrap(), + ); + assert_eq!(result.len(), 1); + assert_eq!( + result + .entries() + .column(0) + .as_primitive::() + .value(0), + 20 + ); + assert_eq!( + result.entries().column(1).as_string::().value(0), + "200" + ); + } + + #[test] + fn map_from_entries_reads_the_right_row_of_a_sliced_list() { + let entries = entry_list( + Int32Array::from(vec![10, 20]), + StringArray::from(vec![Some("100"), Some("200")]), + &[0, 1, 2], + None, + ); + let result = map_result( + invoke( + &SparkMapFromEntries::default(), + vec![entries.slice(1, 1)], + MapKeyDedupPolicy::Exception, + ) + .unwrap(), + ); + assert_eq!(result.len(), 1); + assert_eq!( + result + .entries() + .column(0) + .as_primitive::() + .value(0), + 20 + ); + assert_eq!( + result.entries().column(1).as_string::().value(0), + "200" + ); + } + #[test] fn duplicate_map_key_ignores_unrelated_errors() { assert_eq!( From 0a37af9d2f683df729be845f22c7dc81374bbf7e Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Tue, 15 Sep 2026 08:49:31 +0000 Subject: [PATCH 03/16] fix: report whichever of a null or duplicate map key comes first Spark's `ArrayBasedMapBuilder` inserts entries one at a time, so for keys `[1, 1, NULL]` under `EXCEPTION` it raises `DUPLICATED_MAP_KEY` on the second entry and never reaches the null. The validation pre-scanned a whole row for null keys before delegating, so it reported `NULL_MAP_KEY` instead, and its comments described the precedence as categorical rather than positional. Walk each row's keys in insertion order and raise on the first offending entry, so the two errors order the way Spark orders them, across rows as well as within one. The walk runs only when the keys carry a `NULL`: without one the kernel's own duplicate check already names the key Spark would. Under `LAST_WIN` a duplicate overwrites rather than raising, so only the null check applies. Reported by @rich7420. --- .../spark-expr/src/map_funcs/map_builders.rs | 194 ++++++++++++++++-- 1 file changed, 176 insertions(+), 18 deletions(-) diff --git a/native/spark-expr/src/map_funcs/map_builders.rs b/native/spark-expr/src/map_funcs/map_builders.rs index 531c878e077..394eba45928 100644 --- a/native/spark-expr/src/map_funcs/map_builders.rs +++ b/native/spark-expr/src/map_funcs/map_builders.rs @@ -23,10 +23,11 @@ //! `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. +//! - a key array and value array of different lengths raise `[MAP_KEY_VALUE_DIFF_SIZES]`, which +//! Spark checks before it builds anything; +//! - a `NULL` key raises `[NULL_MAP_KEY]` and, under `EXCEPTION`, a duplicate key raises +//! `[DUPLICATED_MAP_KEY]` naming the key. Spark inserts entries one at a time, so whichever +//! comes first in the row decides which of the two it reports. //! //! `str_to_map` builds its keys by splitting a string, so it needs only the duplicate-key //! restatement. @@ -36,7 +37,8 @@ use arrow::array::{Array, ArrayRef, AsArray, StructArray, UInt32Array}; use arrow::buffer::NullBuffer; use arrow::compute::take; use arrow::datatypes::{DataType, FieldRef}; -use datafusion::common::{exec_err, DataFusionError, Result}; +use datafusion::common::config::MapKeyDedupPolicy; +use datafusion::common::{exec_err, DataFusionError, HashSet, Result, ScalarValue}; use datafusion::logical_expr::{ ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, }; @@ -87,7 +89,7 @@ impl ScalarUDFImpl for SparkMapFromArrays { compact_list_arguments(&mut args)?; match args.args.as_slice() { [ColumnarValue::Array(keys), ColumnarValue::Array(values)] => { - validate_map_from_arrays(keys, values)? + validate_map_from_arrays(keys, values, last_value_wins(&args))? } other => return exec_err!("map_from_arrays expects 2 arguments, got {}", other.len()), } @@ -138,7 +140,9 @@ impl ScalarUDFImpl for SparkMapFromEntries { let mut args = expand_scalars(args)?; compact_list_arguments(&mut args)?; match args.args.as_slice() { - [ColumnarValue::Array(entries)] => validate_map_from_entries(entries)?, + [ColumnarValue::Array(entries)] => { + validate_map_from_entries(entries, last_value_wins(&args))? + } other => return exec_err!("map_from_entries expects 1 argument, got {}", other.len()), } self.inner @@ -206,6 +210,11 @@ fn expand_scalars(mut args: ScalarFunctionArgs) -> Result { Ok(args) } +/// Whether the session asks for Spark's `LAST_WIN` duplicate key policy. +fn last_value_wins(args: &ScalarFunctionArgs) -> bool { + args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin +} + /// Rebuilds any list argument whose entries do not start at offset zero. /// /// The upstream kernels read each row's entries at its own offset but build the mask that selects @@ -248,8 +257,12 @@ fn entries_start_at_zero(array: &ArrayRef) -> bool { } /// 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<()> { +/// and value arrays differ in length, and a `NULL` or duplicate key. +fn validate_map_from_arrays( + keys: &ArrayRef, + values: &ArrayRef, + last_value_wins: bool, +) -> 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(()); @@ -260,6 +273,7 @@ fn validate_map_from_arrays(keys: &ArrayRef, values: &ArrayRef) -> Result<()> { return exec_err!("map_from_arrays: keys and values must have the same number of rows"); } let key_nulls = element_validity(&flat_keys); + let mut seen = HashSet::new(); for row in 0..key_offsets.len().saturating_sub(1) { // `MapFromArrays` is null intolerant, so a NULL input array yields a NULL map without @@ -272,18 +286,16 @@ fn validate_map_from_arrays(keys: &ArrayRef, values: &ArrayRef) -> Result<()> { 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()); - } + check_keys_in_order(&flat_keys, start, end, nulls, last_value_wins, &mut seen)?; } } 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 +/// Rejects a `NULL` or duplicate key 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<()> { +fn validate_map_from_entries(entries: &ArrayRef, last_value_wins: bool) -> Result<()> { if matches!(entries.data_type(), DataType::Null) { return Ok(()); } @@ -299,17 +311,52 @@ fn validate_map_from_entries(entries: &ArrayRef) -> Result<()> { }; let element_nulls = structs.nulls(); + let keys = structs.column(0); + let mut seen = HashSet::new(); + 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) { + let (start, end) = (offsets[row], offsets[row + 1]); + if element_nulls.is_some_and(|nulls| nulls.slice(start, end - start).null_count() > 0) { continue; } - if key_nulls.slice(start, len).null_count() > 0 { + check_keys_in_order(keys, start, end, &key_nulls, last_value_wins, &mut seen)?; + } + Ok(()) +} + +/// Walks one row's keys in the order Spark's `ArrayBasedMapBuilder` inserts them, so whichever of +/// a `NULL` key and a duplicate key comes first is the one reported, as Spark reports it. Only +/// reached when the keys carry a `NULL` somewhere: without one, the kernel's own duplicate check +/// already names the same key Spark would. +#[allow(clippy::allow_attributes, clippy::mutable_key_type)] // ScalarValue is used as a hash key +fn check_keys_in_order( + flat_keys: &ArrayRef, + start: usize, + end: usize, + key_nulls: &NullBuffer, + last_value_wins: bool, + seen: &mut HashSet, +) -> Result<()> { + seen.clear(); + for index in start..end { + if key_nulls.is_null(index) { return Err(SparkError::NullMapKey.into()); } + // `LAST_WIN` overwrites a duplicate rather than raising, so only the `NULL` check is + // left to do in that mode. + if last_value_wins { + continue; + } + let key = ScalarValue::try_from_array(flat_keys, index)?.compacted(); + if !seen.insert(key.clone()) { + return Err(SparkError::DuplicatedMapKey { + key: key.to_string(), + } + .into()); + } } Ok(()) } @@ -743,6 +790,117 @@ mod tests { ); } + /// Spark inserts entries one at a time, so a duplicate at an earlier index is reported even + /// though a `NULL` key follows it. + #[test] + fn map_from_arrays_reports_a_duplicate_before_a_later_null_key() { + let keys = int_list( + Int32Array::from(vec![Some(1), Some(1), None]), + &[0, 3], + None, + ); + let values = string_list( + StringArray::from(vec![Some("a"), Some("b"), Some("c")]), + &[0, 3], + None, + ); + let err = invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[DUPLICATED_MAP_KEY]"), "{err}"); + } + + /// The mirror case: the `NULL` comes first, so it is the one reported. + #[test] + fn map_from_arrays_reports_a_null_key_before_a_later_duplicate() { + let keys = int_list( + Int32Array::from(vec![None, Some(1), Some(1)]), + &[0, 3], + None, + ); + let values = string_list( + StringArray::from(vec![Some("a"), Some("b"), Some("c")]), + &[0, 3], + None, + ); + let err = invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[NULL_MAP_KEY]"), "{err}"); + } + + /// A duplicate in an earlier row wins over a `NULL` key in a later one. + #[test] + fn map_from_arrays_reports_the_first_offending_row() { + let keys = int_list( + Int32Array::from(vec![Some(1), Some(1), None, Some(2)]), + &[0, 2, 4], + None, + ); + let values = string_list( + StringArray::from(vec![Some("a"), Some("b"), Some("c"), Some("d")]), + &[0, 2, 4], + None, + ); + let err = invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[DUPLICATED_MAP_KEY]"), "{err}"); + } + + #[test] + fn map_from_entries_reports_a_duplicate_before_a_later_null_key() { + let entries = entry_list( + Int32Array::from(vec![Some(1), Some(1), None]), + StringArray::from(vec![Some("a"), Some("b"), Some("c")]), + &[0, 3], + None, + ); + let err = invoke( + &SparkMapFromEntries::default(), + vec![entries], + MapKeyDedupPolicy::Exception, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[DUPLICATED_MAP_KEY]"), "{err}"); + } + + /// Under `LAST_WIN` a duplicate is not an error, so a `NULL` key is still reported. + #[test] + fn last_win_still_rejects_a_null_key_after_a_duplicate() { + let keys = int_list( + Int32Array::from(vec![Some(1), Some(1), None]), + &[0, 3], + None, + ); + let values = string_list( + StringArray::from(vec![Some("a"), Some("b"), Some("c")]), + &[0, 3], + None, + ); + let err = invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::LastWin, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[NULL_MAP_KEY]"), "{err}"); + } + #[test] fn duplicate_map_key_ignores_unrelated_errors() { assert_eq!( From 846819af0be698925c0815d3081da94026912139 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Tue, 15 Sep 2026 08:52:34 +0000 Subject: [PATCH 04/16] feat: decline a collated key type in the native map constructors `ArrayBasedMapBuilder` keys its dedup map on `TypeUtils.getInterpretedOrdering` once the key type contains a string, so under `UTF8_LCASE` the keys 'a' and 'A' are one key. The native builders compare the raw Arrow bytes and would keep both, missing the duplicate Spark reports or the overwrite Spark performs under `LAST_WIN`. `MapKeySupport` already declines a collated key for `map_extract` for the same reason; `MapBuilderSupport` only gated floating-point keys. Report `Incompatible` for a collated key type in both constructors. `CometMapFromArrays` falls back to Spark, while `CometMapFromEntries` mixes in `CodegenDispatchFallback` and stays in the Comet pipeline running Spark's own generated code. The new fixture pins both routes. Reported by @andygrove. --- .../scala/org/apache/comet/serde/maps.scala | 28 ++++++++-- .../map/map_builders_collation.sql | 52 +++++++++++++++++++ 2 files changed, 75 insertions(+), 5 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/map/map_builders_collation.sql 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 e6bf2a62975..daa27bb10bd 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -156,16 +156,34 @@ private object MapBuilderSupport { s"`${COMET_EXEC_STRICT_FLOATING_POINT.key}=true` to fall back to Spark for a " + "floating-point map key." + /** + * `ArrayBasedMapBuilder` keys its dedup map on `TypeUtils.getInterpretedOrdering` once the key + * type contains a string, so under `UTF8_LCASE` the keys `'a'` and `'A'` are one key. The + * native builders compare the raw Arrow bytes and would keep both, missing the duplicate that + * Spark reports (or, under `LAST_WIN`, the overwrite Spark performs). `MapKeySupport` declines + * a collated key for `map_extract` for the same reason. + */ + val collationKeyReason: String = + "Comet's native map construction compares string keys as `UTF8_BINARY`, so it cannot honour " + + "a non-default collation when it looks for a duplicate 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)) + if (hasNonDefaultStringCollation(keyType)) { + Incompatible(Some(collationKeyReason)) + } else { + 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(MapBuilderSupport.collationKeyReason) + override def getCompatibleNotes(): Seq[String] = Seq(MapBuilderSupport.floatingPointKeyNote) @@ -193,7 +211,7 @@ object CometMapFromEntries "`BinaryType` is not supported as a map value in `map_from_entries`" override def getIncompatibleReasons(): Seq[String] = - Seq(keyUnsupportedReason, valueUnsupportedReason) + Seq(keyUnsupportedReason, valueUnsupportedReason, MapBuilderSupport.collationKeyReason) override def getCompatibleNotes(): Seq[String] = Seq(MapBuilderSupport.floatingPointKeyNote) diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_builders_collation.sql b/spark/src/test/resources/sql-tests/expressions/map/map_builders_collation.sql new file mode 100644 index 00000000000..3673b854b98 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/map_builders_collation.sql @@ -0,0 +1,52 @@ +-- 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. + +-- MinSparkVersion: 4.0 + +-- Spark 4.0+ supports string collations. `ArrayBasedMapBuilder` keys its dedup map on +-- `TypeUtils.getInterpretedOrdering` once the key type contains a string, so under `UTF8_LCASE` +-- the keys 'a' and 'A' are one key and Spark raises `DUPLICATED_MAP_KEY`. Comet's native +-- builders compare the raw Arrow bytes and would keep both, so both constructors decline a +-- collated key type outright, whether or not a given row actually collides. +-- +-- The keys below are distinct under `UTF8_LCASE` so both engines return a map and the queries +-- can check where the expression ran. `CometMapFromArrays` has no codegen dispatcher, so it +-- falls back to Spark; `CometMapFromEntries` mixes in `CodegenDispatchFallback`, so it stays in +-- the Comet pipeline running Spark's own generated code. +-- +-- `size` wraps each call so the projection's output type is an `int`. A map with a collated key +-- is not a supported Comet output type, and that check runs first: returning the map itself +-- takes the whole plan off Comet with no expression-level reason, testing nothing here. + +statement +CREATE TABLE test_map_builders_collation(k string) USING parquet + +statement +INSERT INTO test_map_builders_collation VALUES ('a'), ('b') + +query expect_fallback(cannot honour a non-default collation) +SELECT size(map_from_arrays( + array(CAST(k AS STRING COLLATE UTF8_LCASE), + CAST(concat(k, 'z') AS STRING COLLATE UTF8_LCASE)), + array(1, 2))) +FROM test_map_builders_collation + +query expect_dispatch(map_from_entries) +SELECT size(map_from_entries(array( + struct(CAST(k AS STRING COLLATE UTF8_LCASE) AS key, 1 AS value), + struct(CAST(concat(k, 'z') AS STRING COLLATE UTF8_LCASE) AS key, 2 AS value)))) +FROM test_map_builders_collation From f7f87f3c656b1ff98cc1bcdafaee3960345415da Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Tue, 15 Sep 2026 08:52:44 +0000 Subject: [PATCH 05/16] docs: scope the floating-point map key note to what Spark actually does The note claimed Spark normalizes a floating-point map key before storing it, full stop. Two corrections, both checked against Spark's sources: `ArrayBasedMapBuilder` gained `keyNormalizer` in 4.0, alongside `spark.sql.legacy.disableMapKeyNormalization`. The 3.5 builder has no normalizer and no reference to `NormalizeFloatingNumbers`, so on 3.4 and 3.5 the native builders already match Spark and there is nothing to warn about. On 4.0+ the two functions differ. `MapFromArrays` calls `ArrayBasedMapBuilder.from`, which returns the input arrays untouched when no key repeated, so a lone `-0.0` key stays `-0.0` in Spark as it does natively; only duplicate detection diverges. `MapFromEntries` puts entries one at a time and always calls `build()`, so Spark stores the normalized key and returns `+0.0` where Comet returns `-0.0`. The gate stays unconditional. Declining on 3.4 and 3.5 costs only a fallback that `spark.comet.exec.strictFloatingPoint` users opted into. Reported by @andygrove. --- .../expression-audits/map_funcs.md | 4 +-- .../scala/org/apache/comet/serde/maps.scala | 36 +++++++++++++------ 2 files changed, 27 insertions(+), 13 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index 779e307dd3d..309d6e0f53d 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -49,7 +49,7 @@ - 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. +- Known limitation: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. `from` returns the input arrays untouched when no key repeated, so the stored keys match Spark either way and only duplicate detection diverges. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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 @@ -59,7 +59,7 @@ - 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: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. Unlike `map_from_arrays`, this expression always calls `build()`, so Spark stores the normalized key and returns `+0.0` for a `-0.0` key where Comet returns `-0.0`. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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 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 daa27bb10bd..20b80ee1f58 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -142,19 +142,33 @@ object CometMapExtract extends CometExpressionSerde[GetMapValue] { 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. + * Floating-point keys differ from Spark only on 4.0 and later, and differently per function. + * `ArrayBasedMapBuilder` gained `keyNormalizer` in 4.0 (with + * `spark.sql.legacy.disableMapKeyNormalization` to turn it off); 3.4 and 3.5 do not normalize + * at all, so the native builders already match there. + * + * On 4.0+ the normalized key decides duplicates for both functions, so a map built from both + * `-0.0` and `+0.0` is one key in Spark and two natively. What each function stores then + * diverges: `MapFromArrays` calls `ArrayBasedMapBuilder.from`, which returns the input arrays + * untouched when no key repeated, so a lone `-0.0` key stays `-0.0` in Spark too; while + * `MapFromEntries` puts entries one at a time and always calls `build()`, which emits the + * normalized keys, so a lone `-0.0` key comes back as `+0.0` in Spark and as `-0.0` natively. + * + * A note rather than a decline, because a map keyed on `-0.0` or `NaN` is rare; + * `spark.comet.exec.strictFloatingPoint` declines it for anyone who wants the guarantee. That + * gate is not conditioned on the Spark version: declining on 3.4 and 3.5 costs those users + * nothing beyond a fallback they opted into. */ 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." + "On Spark 4.0 and later, `ArrayBasedMapBuilder` normalizes a floating-point map key before " + + "comparing it, so `-0.0` counts as the same key as `+0.0` and all `NaN`s count as one " + + "key. Comet's native map construction 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. " + + "`map_from_entries` also stores the normalized key, so Spark returns `+0.0` for a `-0.0` " + + "key where Comet returns `-0.0`; `map_from_arrays` keeps the original keys in both " + + "engines when nothing repeated. Spark 3.4 and 3.5 do not normalize at all, so they match " + + s"Comet already. Set `${COMET_EXEC_STRICT_FLOATING_POINT.key}=true` to fall back to Spark " + + "for a floating-point map key." /** * `ArrayBasedMapBuilder` keys its dedup map on `TypeUtils.getInterpretedOrdering` once the key From 0021e46e2a3d3c5bebc3b3cb2301e00ed2df5629 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Tue, 15 Sep 2026 08:53:28 +0000 Subject: [PATCH 06/16] test: pin the map length mismatch error class instead of comparing engines The length mismatch test avoided naming Spark's condition because I assumed the `_LEGACY_ERROR_TEMP_*` number moved between Spark versions, and added `checkSparkErrorParity` to `CometTestBase` to work around it. The assumption was never checked and is wrong: `mapDataKeyArrayLengthDiffersFromValueArrayLengthError` raises `_LEGACY_ERROR_TEMP_2128` in 3.4.3, 3.5.8 and 4.1.3 alike. Name the condition in the test and drop the helper, which leaves `CometTestBase` untouched by this branch. Reported by @andygrove. --- .../expressions/map/map_from_arrays.sql | 5 +++-- .../comet/CometMapExpressionSuite.scala | 8 +++++--- .../org/apache/spark/sql/CometTestBase.scala | 20 +++---------------- 3 files changed, 11 insertions(+), 22 deletions(-) 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 6ff24fd85ec..be29663cc34 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 @@ -74,7 +74,8 @@ 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. +-- key and value arrays of different lengths. Spark reports this through a legacy condition, +-- `_LEGACY_ERROR_TEMP_2128` in every version Comet supports; matching on the message keeps the +-- fixture readable. query expect_error(must have the same length) SELECT map_from_arrays(array('a', 'b'), array(1)) diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index cd578573491..66f3085ec20 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -162,9 +162,11 @@ class CometMapExpressionSuite extends CometTestBase { 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")) + // Spark reports this through a legacy condition rather than a named one, but the number is + // the same in every version Comet supports (checked in 3.4.3, 3.5.8 and 4.1.3). + checkSparkError( + sql(s"SELECT map_from_arrays(array(k, k + 1), array(v)) FROM $table"), + "_LEGACY_ERROR_TEMP_2128") } } 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 78fad80df0e..a2bbe415cf5 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -448,27 +448,13 @@ 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 expectation = errorClass.map(c => s" with $c").getOrElse("") - val failure = error.getOrElse(fail(s"$engine did not fail$expectation")) + val failure = error.getOrElse(fail(s"$engine did not fail with $errorClass")) val chain = causeChain(failure) assert(!chain.exists(_.isInstanceOf[CometNativeException]), s"$engine: $failure") chain.collect { case e: SparkThrowable with Throwable => e }.lastOption.getOrElse { @@ -478,9 +464,9 @@ abstract class CometTestBase val expected = structuredError(sparkError, "Spark") val actual = structuredError(cometError, "Comet") - errorClass.foreach(c => assert(expected.getErrorClass == c)) + assert(expected.getErrorClass == errorClass) assert(actual.getClass == expected.getClass) - assert(actual.getErrorClass == expected.getErrorClass) + assert(actual.getErrorClass == errorClass) assert(actual.getSqlState == expected.getSqlState) actual } From 308570296f7aa8619e339dd9e8c857f22381a9d5 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Tue, 15 Sep 2026 14:10:13 +0000 Subject: [PATCH 07/16] test: map_from_entries stays native under LAST_WIN in the routing fixtures `routing_map_legacy_disabled.sql` and `routing_map_legacy_enabled.sql` arrived with #5918 and pin how `map_from_entries` routes under `spark.sql.mapKeyDedupPolicy=LAST_WIN`. They encode the behavior this branch removes: `MapFromEntries` reported `Incompatible` under `LAST_WIN`, so `spark.comet.exec.scalaUDF.codegen.enabled` decided whether it fell back to Spark or ran through the JVM codegen dispatcher. The native builder now reads the policy from `datafusion.spark.map_key_dedup_policy`, so the expression is `Compatible` and stays native under either setting of that flag. Expect native in both fixtures. No routing coverage is lost. `map_from_entries` is still `Incompatible` for a `BinaryType` key or value, and `routing_maps_disabled.sql` and `routing_maps_enabled.sql` exercise its fallback and dispatch routes that way. `str_to_map` keeps its expectations in both fixtures: it declines for `spark.sql.legacy.truncateForEmptyRegexSplit`, which this branch does not touch. --- .../expressions/map/routing_map_legacy_disabled.sql | 6 +++++- .../expressions/map/routing_map_legacy_enabled.sql | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/spark/src/test/resources/sql-tests/expressions/map/routing_map_legacy_disabled.sql b/spark/src/test/resources/sql-tests/expressions/map/routing_map_legacy_disabled.sql index a390e18e091..4b6ac3de8a5 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/routing_map_legacy_disabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/routing_map_legacy_disabled.sql @@ -30,5 +30,9 @@ INSERT INTO routing_map_legacy VALUES ('a:1,b:2', array(named_struct('key', 'a', query expect_fallback(str_to_map: spark.comet.exec.scalaUDF.codegen.enabled=false) SELECT str_to_map(s) FROM routing_map_legacy -query expect_fallback(map_from_entries: spark.comet.exec.scalaUDF.codegen.enabled=false) +-- `MapFromEntries` no longer declines under `LAST_WIN`: the native builder reads the policy from +-- `datafusion.spark.map_key_dedup_policy`, so it stays native whatever the codegen flag says. Its +-- dispatch and fallback routes are still covered by the `BinaryType` queries in +-- `routing_maps_enabled.sql` and `routing_maps_disabled.sql`. +query expect_native(map_from_entries) SELECT map_from_entries(e) FROM routing_map_legacy diff --git a/spark/src/test/resources/sql-tests/expressions/map/routing_map_legacy_enabled.sql b/spark/src/test/resources/sql-tests/expressions/map/routing_map_legacy_enabled.sql index afbfc95dba4..71e7d8de272 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/routing_map_legacy_enabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/routing_map_legacy_enabled.sql @@ -30,5 +30,9 @@ INSERT INTO routing_map_legacy VALUES ('a:1,b:2', array(named_struct('key', 'a', query expect_dispatch(str_to_map) SELECT str_to_map(s) FROM routing_map_legacy -query expect_dispatch(map_from_entries) +-- `MapFromEntries` no longer declines under `LAST_WIN`: the native builder reads the policy from +-- `datafusion.spark.map_key_dedup_policy`, so it stays native whatever the codegen flag says. Its +-- dispatch and fallback routes are still covered by the `BinaryType` queries in +-- `routing_maps_enabled.sql` and `routing_maps_disabled.sql`. +query expect_native(map_from_entries) SELECT map_from_entries(e) FROM routing_map_legacy From 8f54bdceb95fbd8da5554b07e74d36db59ec3e81 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Wed, 16 Sep 2026 09:13:33 +0000 Subject: [PATCH 08/16] fix: keep the map_from_arrays null-array guard so ANSI casts short-circuit Spark's `BinaryExpression.eval` returns NULL as soon as the left input is NULL and never evaluates the right one, so a failing cast in the values argument does not run for a row whose keys array is NULL. Comet evaluates both argument subtrees, so removing the `CaseWhen` guard made `map_from_arrays(k, array(CAST(v AS INT)))` raise `CAST_INVALID_INPUT` under ANSI where Spark returns NULL. Restore the guard. Its `AND` over `IsNotNull` on both arguments lets the native side skip the values expression once the keys array is known NULL, which matches Spark's evaluation order. The earlier commit removed it on the grounds that native `map_from_arrays` is null intolerant; that is true of the result but says nothing about which subtrees get evaluated. Add a regression test that fails with `CAST_INVALID_INPUT` without the guard and passes with it. --- .../scala/org/apache/comet/serde/maps.scala | 48 +++++++++++++++++-- .../comet/CometMapExpressionSuite.scala | 16 +++++++ 2 files changed, 60 insertions(+), 4 deletions(-) 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 20b80ee1f58..c3de6785342 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -25,7 +25,7 @@ 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.{exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto} +import org.apache.comet.serde.QueryPlanSerde.{createBinaryExpr, exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto} import org.apache.comet.shims.CometTypeShim /** @@ -204,15 +204,55 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { override def getSupportLevel(expr: MapFromArrays): SupportLevel = MapBuilderSupport.keySupport(expr.dataType.keyType) + /** + * Native `map_from_arrays` already returns a NULL map for a NULL input array, so the `CaseWhen` + * below guards evaluation order rather than the result. `BinaryExpression.eval` returns as soon + * as the left input is NULL and never evaluates the right one, so under ANSI a failing cast in + * the values argument never runs for a row whose keys array is NULL. Comet evaluates both + * argument subtrees, so without the guard that cast raises where Spark returns NULL. Wrapping + * the call lets the `AND` short-circuit skip the values expression. + * + * @see + * https://github.com/apache/datafusion-comet/pull/5854#discussion_r4016898751 + */ override def convert( expr: MapFromArrays, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val keysExpr = exprToProtoInternal(expr.left, inputs, binding) val valuesExpr = exprToProtoInternal(expr.right, inputs, binding) - // 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) + 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_from_arrays", 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)) } } diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 66f3085ec20..5f9e6d9ac87 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -151,6 +151,22 @@ class CometMapExpressionSuite extends CometTestBase { } } + // Spark's `BinaryExpression.eval` returns NULL the moment the left input is NULL and never + // evaluates the right one, so a failing cast in the values argument does not run for a row whose + // keys array is NULL. Comet evaluates both argument subtrees, so the serde keeps a `CaseWhen` + // guard: its `AND` lets the native side skip the values expression once the keys are known NULL. + // https://github.com/apache/datafusion-comet/pull/5854#discussion_r4016898751 + test("map_from_arrays - a null keys array skips the values expression under ANSI") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTable("map_short_circuit") { + sql("CREATE TABLE map_short_circuit(k ARRAY, v STRING) USING parquet") + sql("INSERT INTO map_short_circuit VALUES (NULL, 'bad')") + checkSparkAnswerAndOperator( + sql("SELECT map_from_arrays(k, array(CAST(v AS INT))) FROM map_short_circuit")) + } + } + } + test("map_from_arrays - a null input array gives a null map") { withMapBuilderTable { table => checkSparkAnswerAndOperator( From 6d1d327b2a882cbc711c9e0da702bc4818f77369 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Fri, 18 Sep 2026 00:44:44 +0000 Subject: [PATCH 09/16] fix: nest the map_from_arrays null guards so a NULL keys array skips the values expression The previous commit restored `CASE WHEN keys IS NOT NULL AND values IS NOT NULL` around `map_from_arrays`, and its regression test passed only because the table held a single row. DataFusion's `AND` skips its right side when the left side is false on every row of the batch, or on at most a fifth of them; in a batch where most rows do have keys it evaluates `values IS NOT NULL` on the whole batch, so a failing cast in the values array still runs for the NULL-keys row and raises under ANSI where Spark returns NULL. Nest one `CaseWhen` per argument instead, as #5846 does. DataFusion evaluates a THEN branch only on the rows its WHEN selected, so the values expression is never evaluated for a row whose keys array is NULL, which is what `BinaryExpression.eval` does. Rewrite the regression test as a five-row table written in one partition so every row shares a batch and most rows have keys. It fails with `CAST_INVALID_INPUT` against the `AND` guard and passes with the nested guards; the three `map_from_arrays` tests from #5846 pass as well. --- .../expression-audits/map_funcs.md | 2 +- .../scala/org/apache/comet/serde/maps.scala | 46 +++++++++---------- .../comet/CometMapExpressionSuite.scala | 20 ++++++-- 3 files changed, 39 insertions(+), 29 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index 09f3d9cb0b8..928965db7ca 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -45,7 +45,7 @@ ## 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` 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 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)). The serde still nests `CASE WHEN left IS NOT NULL THEN (CASE WHEN right IS NOT NULL THEN map_from_arrays(left, right) END) END` around the call: `BinaryExpression.eval` never evaluates `right` for a row whose `left` is NULL, and DataFusion evaluates a THEN branch only on the rows its WHEN selected, so a failing cast in the values array does not run for such a row. A single `left IS NOT NULL AND right IS NOT NULL` guard does not give that, since DataFusion's `AND` evaluates its right side on the whole batch unless the left side is false on all or most rows. - 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). 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 c3de6785342..a4b63bafc90 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -25,7 +25,7 @@ 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 /** @@ -205,12 +205,16 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { MapBuilderSupport.keySupport(expr.dataType.keyType) /** - * Native `map_from_arrays` already returns a NULL map for a NULL input array, so the `CaseWhen` - * below guards evaluation order rather than the result. `BinaryExpression.eval` returns as soon - * as the left input is NULL and never evaluates the right one, so under ANSI a failing cast in - * the values argument never runs for a row whose keys array is NULL. Comet evaluates both - * argument subtrees, so without the guard that cast raises where Spark returns NULL. Wrapping - * the call lets the `AND` short-circuit skip the values expression. + * Native `map_from_arrays` already returns a NULL map for a NULL input array, so the guards + * below are about evaluation order rather than the result. `BinaryExpression.eval` returns as + * soon as the left input is NULL and never evaluates the right one, so under ANSI a failing + * cast in the values argument never runs for a row whose keys array is NULL. Nesting one + * `CaseWhen` per argument reproduces that: DataFusion evaluates a THEN branch only on the rows + * its WHEN selected, so the values expression is never evaluated for a row whose keys array is + * NULL. A single `keys IS NOT NULL AND values IS NOT NULL` guard is not enough, because + * DataFusion's `AND` skips its right side only when the left side is false on every row of the + * batch, or on most of them; a batch where most rows do have keys evaluates the values + * expression on all of them, the NULL-keys rows included. * * @see * https://github.com/apache/datafusion-comet/pull/5854#discussion_r4016898751 @@ -225,35 +229,29 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { val valueType = expr.right.dataType.asInstanceOf[ArrayType].elementType val returnType = MapType(keyType = keyType, valueType = valueType) for { - andBinaryExprProto <- createAndBinaryExpr(expr, inputs, binding) + keysNotNullExprProto <- exprToProtoInternal(IsNotNull(expr.left), inputs, binding) + valuesNotNullExprProto <- exprToProtoInternal(IsNotNull(expr.right), inputs, binding) mapFromArraysExprProto <- scalarFunctionExprToProto("map_from_arrays", keysExpr, valuesExpr) nullLiteralExprProto <- exprToProtoInternal(Literal(null, returnType), inputs, binding) } yield { - val caseWhenExprProto = ExprOuterClass.CaseWhen + val valuesGuardProto = ExprOuterClass.CaseWhen .newBuilder() - .addWhen(andBinaryExprProto) + .addWhen(valuesNotNullExprProto) .addThen(mapFromArraysExprProto) .setElseExpr(nullLiteralExprProto) .build() + val keysGuardProto = ExprOuterClass.CaseWhen + .newBuilder() + .addWhen(keysNotNullExprProto) + .addThen(ExprOuterClass.Expr.newBuilder().setCaseWhen(valuesGuardProto).build()) + .setElseExpr(nullLiteralExprProto) + .build() ExprOuterClass.Expr .newBuilder() - .setCaseWhen(caseWhenExprProto) + .setCaseWhen(keysGuardProto) .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)) - } } object CometMapFromEntries diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 5f9e6d9ac87..46ff6ff1023 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -153,14 +153,26 @@ class CometMapExpressionSuite extends CometTestBase { // Spark's `BinaryExpression.eval` returns NULL the moment the left input is NULL and never // evaluates the right one, so a failing cast in the values argument does not run for a row whose - // keys array is NULL. Comet evaluates both argument subtrees, so the serde keeps a `CaseWhen` - // guard: its `AND` lets the native side skip the values expression once the keys are known NULL. + // keys array is NULL. The serde nests one `CaseWhen` per argument so the native side evaluates + // the values expression only on rows whose keys array is not NULL. The rows with keys outnumber + // the row without on purpose, and all of them sit in one batch: a single `AND` guard skips its + // right side only when the left side is false on every row of the batch, or on most of them, so + // this batch would evaluate the cast on the NULL-keys row as well. // https://github.com/apache/datafusion-comet/pull/5854#discussion_r4016898751 test("map_from_arrays - a null keys array skips the values expression under ANSI") { withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { withTable("map_short_circuit") { - sql("CREATE TABLE map_short_circuit(k ARRAY, v STRING) USING parquet") - sql("INSERT INTO map_short_circuit VALUES (NULL, 'bad')") + // One partition, so every row lands in the same file and the same batch. + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0, 5, 1, 1) + .selectExpr( + "IF(id = 0, CAST(NULL AS ARRAY), array(CAST(id AS INT))) AS k", + "IF(id = 0, 'bad', CAST(id AS STRING)) AS v") + .write + .format("parquet") + .saveAsTable("map_short_circuit") + } checkSparkAnswerAndOperator( sql("SELECT map_from_arrays(k, array(CAST(v AS INT))) FROM map_short_circuit")) } From 855b0392bc2391c46b73fb91903637f388aa6ea3 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Fri, 18 Sep 2026 00:44:45 +0000 Subject: [PATCH 10/16] test: pin the entry order of a duplicate key under LAST_WIN `ArrayBasedMapBuilder` fixes a key's slot at its first occurrence and only replaces its value, so `['a', 'b', 'a']` with `[1, 2, 3]` becomes `{a -> 3, b -> 2}` rather than `{b -> 2, a -> 3}`. The datafusion-spark kernels mirror that, but no test here told the two apart: the existing LAST_WIN cases repeat one key, or repeat a key only in adjacent slots. Add Rust unit tests for `map_from_arrays`, `map_from_entries` and `str_to_map` that assert the keys and values in order, and fixture rows for the same case. A map compares equal in any entry order, so the fixtures pin the order through `map_keys` and `map_values`. Also cover a NULL as the value that wins. --- .../spark-expr/src/map_funcs/map_builders.rs | 58 +++++++++++++++++++ .../map/map_from_arrays_dedup_policy.sql | 17 ++++++ .../map/map_from_entries_dedup_policy.sql | 17 ++++++ .../map/str_to_map_dedup_policy.sql | 10 ++++ 4 files changed, 102 insertions(+) diff --git a/native/spark-expr/src/map_funcs/map_builders.rs b/native/spark-expr/src/map_funcs/map_builders.rs index 394eba45928..57e922fc15b 100644 --- a/native/spark-expr/src/map_funcs/map_builders.rs +++ b/native/spark-expr/src/map_funcs/map_builders.rs @@ -635,6 +635,34 @@ mod tests { assert_eq!(values.value(0), "b"); } + #[test] + fn map_from_arrays_keeps_a_duplicate_key_in_its_first_position_under_last_win() { + // `ArrayBasedMapBuilder` fixes a key's slot at its first occurrence and only replaces + // the value, so `[1, 2, 1]` with `[a, b, c]` is `{1 -> c, 2 -> b}`, never + // `{2 -> b, 1 -> c}`. + let keys = int_list(Int32Array::from(vec![1, 2, 1]), &[0, 3], None); + let values = string_list( + StringArray::from(vec![Some("a"), Some("b"), Some("c")]), + &[0, 3], + None, + ); + let result = map_result( + invoke( + &SparkMapFromArrays::default(), + vec![keys, values], + MapKeyDedupPolicy::LastWin, + ) + .unwrap(), + ); + assert_eq!(result.value_offsets(), &[0, 2]); + assert_eq!( + result.keys().as_primitive::().values().as_ref(), + &[1, 2] + ); + let values = result.values().as_string::(); + assert_eq!((values.value(0), values.value(1)), ("c", "b")); + } + #[test] fn map_from_entries_rejects_null_key() { let entries = entry_list( @@ -695,6 +723,31 @@ mod tests { assert_eq!(values.value(0), "b"); } + #[test] + fn map_from_entries_keeps_a_duplicate_key_in_its_first_position_under_last_win() { + let entries = entry_list( + Int32Array::from(vec![1, 2, 1]), + StringArray::from(vec![Some("a"), Some("b"), Some("c")]), + &[0, 3], + None, + ); + let result = map_result( + invoke( + &SparkMapFromEntries::default(), + vec![entries], + MapKeyDedupPolicy::LastWin, + ) + .unwrap(), + ); + assert_eq!(result.value_offsets(), &[0, 2]); + assert_eq!( + result.keys().as_primitive::().values().as_ref(), + &[1, 2] + ); + let values = result.values().as_string::(); + assert_eq!((values.value(0), values.value(1)), ("c", "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")])); @@ -723,6 +776,11 @@ mod tests { .unwrap(), ); assert_eq!(result.value_offsets(), &[0, 2]); + // `a` keeps the slot of its first occurrence and takes its last value. + let keys = result.keys().as_string::(); + let values = result.values().as_string::(); + assert_eq!((keys.value(0), values.value(0)), ("a", "3")); + assert_eq!((keys.value(1), values.value(1)), ("b", "2")); } /// A `LIMIT` above a projection hands the kernel a sliced list. The mask the upstream helper 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 70517905b28..cd8ecf10056 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 @@ -30,6 +30,8 @@ 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('a', 'b', 'a'), array(1, 2, 3)), + (array('a', 'a', 'b'), array(1, NULL, 3)), (array(), array()), (NULL, array(99)) @@ -41,10 +43,25 @@ SELECT map_from_arrays(array('a', 'a', 'b'), array(1, 2, 3)) query SELECT map_from_arrays(array('a', 'a', 'a'), array(1, 2, 3)) +-- a repeated key keeps the position of its first occurrence and takes its last value, as +-- `ArrayBasedMapBuilder` does: {a -> 3, b -> 2}. Maps compare equal in any entry order, so +-- `map_keys` and `map_values` pin the order. +query +SELECT map_keys(map_from_arrays(array('a', 'b', 'a'), array(1, 2, 3))), + map_values(map_from_arrays(array('a', 'b', 'a'), array(1, 2, 3))) + +-- a NULL can be the value that wins +query +SELECT map_from_arrays(array('a', 'a', 'b'), array(1, NULL, 3)) + -- column input, including rows without duplicates and a NULL row query SELECT map_from_arrays(k, v) FROM test_map_from_arrays_dedup +-- the same rows with their entry order pinned +query +SELECT map_keys(map_from_arrays(k, v)), map_values(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_dedup_policy.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_entries_dedup_policy.sql index c344e583e19..3d3b46c4f41 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 @@ -30,6 +30,8 @@ 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('a', 1), struct('b', 2), struct('a', 3))), + (array(struct('a', 1), struct('a', CAST(NULL AS INT)), struct('b', 3))), (array()), (NULL) @@ -41,10 +43,25 @@ SELECT map_from_entries(array(struct('a', 1), struct('a', 2), struct('b', 3))) query SELECT map_from_entries(array(struct('a', 1), struct('a', 2), struct('a', 3))) +-- a repeated key keeps the position of its first occurrence and takes its last value, as +-- `ArrayBasedMapBuilder` does: {a -> 3, b -> 2}. Maps compare equal in any entry order, so +-- `map_keys` and `map_values` pin the order. +query +SELECT map_keys(map_from_entries(array(struct('a', 1), struct('b', 2), struct('a', 3)))), + map_values(map_from_entries(array(struct('a', 1), struct('b', 2), struct('a', 3)))) + +-- a NULL can be the value that wins +query +SELECT map_from_entries(array(struct('a', 1), struct('a', CAST(NULL AS INT)), struct('b', 3))) + -- column input, including rows without duplicates and a NULL row query SELECT map_from_entries(entries) FROM test_map_from_entries_dedup +-- the same rows with their entry order pinned +query +SELECT map_keys(map_from_entries(entries)), map_values(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_dedup_policy.sql b/spark/src/test/resources/sql-tests/expressions/map/str_to_map_dedup_policy.sql index f3aab4eb8af..74abb046baa 100644 --- 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 @@ -35,8 +35,18 @@ INSERT INTO test_str_to_map_dedup VALUES query SELECT str_to_map('a:1,b:2,a:3') +-- `a` keeps the position of its first occurrence and takes its last value, as +-- `ArrayBasedMapBuilder` does: {a -> 3, b -> 2}. Maps compare equal in any entry order, so +-- `map_keys` and `map_values` pin the order. +query +SELECT map_keys(str_to_map('a:1,b:2,a:3')), map_values(str_to_map('a:1,b:2,a:3')) + query SELECT str_to_map(s) FROM test_str_to_map_dedup +-- the same rows with their entry order pinned +query +SELECT map_keys(str_to_map(s)), map_values(str_to_map(s)) FROM test_str_to_map_dedup + query SELECT str_to_map(s, ',', ':') FROM test_str_to_map_dedup From 101b2da0a9272bacd4fe0bf6a719740423002be0 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Fri, 18 Sep 2026 00:47:11 +0000 Subject: [PATCH 11/16] test: cover a values array longer than the keys array in map_from_arrays --- .../resources/sql-tests/expressions/map/map_from_arrays.sql | 4 ++++ 1 file changed, 4 insertions(+) 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 be29663cc34..0606163ef7c 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 @@ -79,3 +79,7 @@ SELECT map_from_arrays(array('a', 'a'), array(1, 2)) -- fixture readable. query expect_error(must have the same length) SELECT map_from_arrays(array('a', 'b'), array(1)) + +-- and in the other direction +query expect_error(must have the same length) +SELECT map_from_arrays(array('a'), array(1, 2)) From 8c9ebdf43dd4b0d3f768a9101ad3994072d4e244 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Fri, 18 Sep 2026 10:02:39 +0000 Subject: [PATCH 12/16] fix: decline a nondeterministic child of map_from_arrays before it reaches the null guards The nested null guards serialize each child a second time inside the `map_from_arrays` call, so a stateful child advances independently in each copy: the guard's copy sees every row while the constructor's copy sees only the rows the guard selected. With map_from_arrays(IF(monotonically_increasing_id() % 2 != 0, array(1), NULL), array(2)) over sixteen rows in one partition, Spark returns eight maps and Comet returned four (#5781). Under LAST_WIN this case used to fall back for the policy alone, so running the policy natively exposed it there. Port `NullGuardSupport` from #5867 unchanged in name, reason and position, so that PR rebases by dropping the hunk, and decline a nondeterministic child in `CometMapFromArrays.getSupportLevel` as `Unsupported`; the projection falls back to Spark, which evaluates the child once. #5867 still routes the same decline through the JVM codegen dispatcher and applies it to `size`, `array_append` and `arrays_zip`. Cover it with the query above as a Scala test on a one-partition table under LAST_WIN, the same query in `map_from_arrays_dedup_policy.sql`, and `map_from_arrays_nondeterministic_child.sql` for the default policy, which mirrors the fixture in #5867 with `expect_fallback` in place of `expect_dispatch`. --- .../expression-audits/map_funcs.md | 1 + .../scala/org/apache/comet/serde/arrays.scala | 21 ++++++++ .../scala/org/apache/comet/serde/maps.scala | 13 ++++- .../map/map_from_arrays_dedup_policy.sql | 13 +++++ ...map_from_arrays_nondeterministic_child.sql | 53 +++++++++++++++++++ .../comet/CometMapExpressionSuite.scala | 23 ++++++++ 6 files changed, 123 insertions(+), 1 deletion(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_nondeterministic_child.sql diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index 928965db7ca..cd19e494aac 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -51,6 +51,7 @@ - `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: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. `from` returns the input arrays untouched when no key repeated, so the stored keys match Spark either way and only duplicate detection diverges. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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. +- Known limitation: the two null guards serialize each child a second time inside the `map_from_arrays` call, so a nondeterministic child such as `monotonically_increasing_id()` would advance independently in each copy and the result would drift from Spark ([#5781](https://github.com/apache/datafusion-comet/issues/5781)). `CometMapFromArrays` declines such a child as `Unsupported` through `NullGuardSupport` and the projection falls back to Spark; [#5867](https://github.com/apache/datafusion-comet/pull/5867) routes the same decline through the JVM codegen dispatcher and applies it to `size`, `array_append` and `arrays_zip` as well. ## map_from_entries diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index cca9f63f8bf..e14872f4cf3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -51,6 +51,27 @@ object CometArrayRemove } } +/** + * Shared gate for serdes whose native NULL guard (`CASE WHEN child IS NOT NULL`) serializes the + * child twice: a stateful child drifts between the two copies, so it is declined and Spark + * evaluates it once, through the JVM codegen dispatcher where the serde mixes in + * `CodegenDispatchFallback` and through a fallback otherwise. Nullability is not consulted: a + * non-nullable stateful child only stays in step because DataFusion skips the filter when the + * guard matches every row, which is not a contract to lean on. + */ +private[serde] object NullGuardSupport { + + val nondeterministicReason: String = + "a nondeterministic operand: the native NULL guard serializes the operand twice, " + + "and the two copies of a stateful operand drift apart" + + /** `Unsupported` when any of `children` is nondeterministic, otherwise `None`. */ + def nondeterministicChild(children: Seq[Expression]): Option[SupportLevel] = + children + .find(child => !child.deterministic) + .map(_ => Unsupported(Some(nondeterministicReason))) +} + object CometArrayAppend extends CometExpressionSerde[ArrayAppend] with ArraysBase { override def convert( 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 a4b63bafc90..7ed235b7a1d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -198,11 +198,16 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { override def getIncompatibleReasons(): Seq[String] = Seq(MapBuilderSupport.collationKeyReason) + override def getUnsupportedReasons(): Seq[String] = + Seq(NullGuardSupport.nondeterministicReason) + override def getCompatibleNotes(): Seq[String] = Seq(MapBuilderSupport.floatingPointKeyNote) override def getSupportLevel(expr: MapFromArrays): SupportLevel = - MapBuilderSupport.keySupport(expr.dataType.keyType) + NullGuardSupport + .nondeterministicChild(expr.children) + .getOrElse(MapBuilderSupport.keySupport(expr.dataType.keyType)) /** * Native `map_from_arrays` already returns a NULL map for a NULL input array, so the guards @@ -216,8 +221,14 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { * batch, or on most of them; a batch where most rows do have keys evaluates the values * expression on all of them, the NULL-keys rows included. * + * Each guard serializes its child a second time inside the `map_from_arrays` call, so a + * stateful child would advance independently in each copy and the result would drift from + * Spark; `getSupportLevel` declines a nondeterministic child for that reason. + * * @see * https://github.com/apache/datafusion-comet/pull/5854#discussion_r4016898751 + * @see + * https://github.com/apache/datafusion-comet/pull/5854#discussion_r4043896247 */ override def convert( expr: MapFromArrays, 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 cd8ecf10056..72d516e940c 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 @@ -65,3 +65,16 @@ SELECT map_keys(map_from_arrays(k, v)), map_values(map_from_arrays(k, v)) FROM t -- 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)) + +statement +CREATE TABLE test_map_from_arrays_dedup_nondet(id bigint) USING parquet + +statement +INSERT INTO test_map_from_arrays_dedup_nondet SELECT id FROM range(0, 16) + +-- A nondeterministic child used to fall back for the policy alone. The serde's null guards +-- serialize each child twice, so a stateful child would drift between the two copies; it is +-- declined and the projection falls back to Spark, which evaluates it once. +-- `map_from_arrays_nondeterministic_child.sql` has the default-policy cases. +query expect_fallback(nondeterministic operand) +SELECT id, map_from_arrays(IF(monotonically_increasing_id() % 2 != 0, array(1), NULL), array(2)) FROM test_map_from_arrays_dedup_nondet diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_nondeterministic_child.sql b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_nondeterministic_child.sql new file mode 100644 index 00000000000..6be4efa0dcd --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/map_from_arrays_nondeterministic_child.sql @@ -0,0 +1,53 @@ +-- 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. + +-- `CometMapFromArrays` reproduces Spark's NULL propagation and evaluation order with nested +-- `CASE WHEN keys IS NOT NULL` / `CASE WHEN values IS NOT NULL` guards that serialize each child a +-- second time inside the `map_from_arrays` call. A stateful child advances each copy +-- independently: the guard's copy sees every row while the constructor's copy sees only the rows +-- the guard selected, so the result silently drifts from Spark (#5781). The serde declines a +-- nondeterministic child and the projection falls back to Spark, which evaluates it once. A +-- deterministic nullable child keeps the native guards. `map_from_arrays_dedup_policy.sql` covers +-- the same decline under `LAST_WIN`. + +statement +CREATE TABLE test_map_from_arrays_nondet(_1 int, k array, v array) USING parquet + +statement +INSERT INTO test_map_from_arrays_nondet +SELECT id, IF(id % 4 = 3, NULL, array(id, id + 100)), array(id * 2, id * 2 + 1) FROM range(0, 16) + +-- Spark returns {1 -> 2} on every row whose keys array is non-NULL and NULL on the rest. +query expect_fallback(nondeterministic operand) +SELECT _1, map_from_arrays(IF(monotonically_increasing_id() % 2 = 0, array(1), CAST(NULL AS ARRAY)), array(2)) AS m +FROM test_map_from_arrays_nondet + +-- The guards cover both children, so a stateful values array is declined the same way. +query expect_fallback(nondeterministic operand) +SELECT _1, map_from_arrays(array(1), IF(monotonically_increasing_id() % 2 = 0, array(2), CAST(NULL AS ARRAY))) AS m +FROM test_map_from_arrays_nondet + +-- A non-nullable stateful child is declined too, rather than relying on the guard matching every +-- row. +query expect_fallback(nondeterministic operand) +SELECT _1, map_from_arrays(array(monotonically_increasing_id()), array(2)) AS m +FROM test_map_from_arrays_nondet + +-- A deterministic nullable child stays on the native guarded path. +query expect_native(map_from_arrays) +SELECT _1, map_from_arrays(IF(_1 % 2 = 0, array(1), CAST(NULL AS ARRAY)), array(2)) AS m +FROM test_map_from_arrays_nondet diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 46ff6ff1023..609b1c81050 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -179,6 +179,29 @@ class CometMapExpressionSuite extends CometTestBase { } } + // Both null guards serialize their child a second time inside the `map_from_arrays` call, so a + // stateful child advances independently in each copy: with `monotonically_increasing_id()` + // deciding which rows have keys, the guard's copy sees every row while the constructor's copy + // sees only the rows the guard selected, and half of the expected maps come back NULL (#5781). + // Such a child is declined, so the projection runs in Spark, which evaluates it once. Under + // LAST_WIN this case used to fall back for the policy alone; the decline keeps it correct now + // that the policy runs natively. + // https://github.com/apache/datafusion-comet/pull/5854#discussion_r4043896247 + test("map_from_arrays - a nondeterministic child falls back under LAST_WIN") { + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + withTable("map_nondeterministic") { + // One partition, so both copies of the child would see the same sixteen-row batch. + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 16, 1, 1).write.format("parquet").saveAsTable("map_nondeterministic") + } + checkSparkAnswerAndFallbackReason( + "SELECT id, map_from_arrays(IF(monotonically_increasing_id() % 2 != 0, array(1), NULL), " + + "array(2)) FROM map_nondeterministic", + "nondeterministic operand") + } + } + } + test("map_from_arrays - a null input array gives a null map") { withMapBuilderTable { table => checkSparkAnswerAndOperator( From 5ba49238b385b44809a3a913308244254eba06ef Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Sat, 19 Sep 2026 09:53:31 +0000 Subject: [PATCH 13/16] fix: capture the map dedup policy with the expression rather than per native iterator Spark's `ArrayBasedMapBuilder` reads `spark.sql.mapKeyDedupPolicy` when an expression is first evaluated and keeps that builder, so a Dataset executed again after the session setting changed still builds its maps under the policy it started with. `CometExecIterator` forwarded the setting from the task SQLConf on every new native iterator instead, so the same Dataset switched behavior: executed under LAST_WIN, it raised DUPLICATED_MAP_KEY on its next action once EXCEPTION was set, where Spark kept returning the LAST_WIN map. Carry the policy with the expression. New `MapFromArrays`, `MapFromEntries` and `StrToMap` proto messages hold a `map_key_dedup_policy` that the serde reads when it converts the plan, which the Dataset reuses across actions, the way `Hour` carries its timezone. The planner builds the native wrappers with that policy, and they hand it to the datafusion-spark kernels through the session options those kernels read, so whatever the session holds at execution time does not apply. Drop the per-iterator forwarding and the session-level `datafusion.spark.map_key_dedup_policy` setting, and the name-based registrations the dedicated messages replace. Cover it with a test that executes one Dataset twice across a policy change, for all three constructors and in both directions, with Comet disabled and enabled; it failed on the second `collect()` before this change. A native unit test checks that the wrapper's policy wins over the session option. --- .../expression-audits/map_funcs.md | 8 +- native/core/src/execution/jni_api.rs | 11 +- native/core/src/execution/planner.rs | 96 ++++++++- native/core/src/execution/spark_config.rs | 2 - native/proto/src/proto/expr.proto | 30 +++ native/spark-expr/src/comet_scalar_funcs.rs | 7 +- .../spark-expr/src/map_funcs/map_builders.rs | 202 ++++++++++-------- .../org/apache/comet/CometExecIterator.scala | 7 - .../scala/org/apache/comet/serde/maps.scala | 67 +++++- .../comet/CometMapExpressionSuite.scala | 54 ++++- 10 files changed, 359 insertions(+), 125 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index cd19e494aac..f723631162d 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -48,7 +48,7 @@ - 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)). The serde still nests `CASE WHEN left IS NOT NULL THEN (CASE WHEN right IS NOT NULL THEN map_from_arrays(left, right) END) END` around the call: `BinaryExpression.eval` never evaluates `right` for a row whose `left` is NULL, and DataFusion evaluates a THEN branch only on the rows its WHEN selected, so a failing cast in the values array does not run for such a row. A single `left IS NOT NULL AND right IS NOT NULL` guard does not give that, since DataFusion's `AND` evaluates its right side on the whole batch unless the left side is false on all or most rows. - 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). +- `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 the serde reads when the plan is converted and carries with the expression (`map_key_dedup_policy` on the expression's proto message), so an executed Dataset keeps its policy the way Spark's builder does (`EXCEPTION` raises `DUPLICATED_MAP_KEY` naming the key, `LAST_WIN` keeps the last value for the key). - Known limitation: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. `from` returns the input arrays untouched when no key repeated, so the stored keys match Spark either way and only duplicate detection diverges. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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. - Known limitation: the two null guards serialize each child a second time inside the `map_from_arrays` call, so a nondeterministic child such as `monotonically_increasing_id()` would advance independently in each copy and the result would drift from Spark ([#5781](https://github.com/apache/datafusion-comet/issues/5781)). `CometMapFromArrays` declines such a child as `Unsupported` through `NullGuardSupport` and the projection falls back to Spark; [#5867](https://github.com/apache/datafusion-comet/pull/5867) routes the same decline through the JVM codegen dispatcher and applies it to `size`, `array_append` and `arrays_zip` as well. @@ -56,10 +56,10 @@ ## map_from_entries - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.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 3.5.8 (audited 2026-05-27): baseline. `MapFromEntries(child) extends UnaryExpression with NullIntolerant`; expects an array of structs and produces a map. Wired through the `MapFromEntries` message, which carries the dedup policy captured when the plan is converted. - 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). +- `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 the serde reads when the plan is converted and carries with the expression (`map_key_dedup_policy` on the expression's proto message), so an executed Dataset keeps its policy the way Spark's builder does (`EXCEPTION` raises `DUPLICATED_MAP_KEY` naming the key, `LAST_WIN` keeps the last value for the key). - Known limitation: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. Unlike `map_from_arrays`, this expression always calls `build()`, so Spark stores the normalized key and returns `+0.0` for a `-0.0` key where Comet returns `-0.0`. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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`. @@ -85,7 +85,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")`. 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 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 through the `StrToMap` message, which carries the `spark.sql.mapKeyDedupPolicy` captured when the plan is converted; the native wrapper hands that policy to the kernel. - 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 59c4a9c1bad..2fcb0e749f5 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -111,7 +111,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, SPARK_MAP_KEY_DEDUP_POLICY, + COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use datafusion_comet_proto::spark_operator::operator::OpStruct; @@ -742,15 +742,6 @@ 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 diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index fcd9768b1c0..2bef77849ff 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -102,6 +102,7 @@ use crate::execution::shuffle::{CometPartitioning, CompressionCodec}; use crate::execution::spark_plan::SparkPlan; use crate::parquet::objectstore::s3_blob_fs_support::normalize_object_store_url; use crate::parquet::parquet_support::prepare_object_store_with_configs; +use datafusion::common::config::MapKeyDedupPolicy; use datafusion::common::scalar::ScalarStructBuilder; use datafusion::common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter}, @@ -110,6 +111,7 @@ use datafusion::common::{ use datafusion::datasource::listing::PartitionedFile; use datafusion::logical_expr::type_coercion::functions::fields_with_udf; use datafusion::logical_expr::type_coercion::other::get_coerce_type_for_case_expression; +use datafusion::logical_expr::ScalarUDFImpl; use datafusion::logical_expr::{ AggregateUDF, ReturnFieldArgs, ScalarUDF, TypeSignature, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, @@ -151,8 +153,8 @@ use datafusion_comet_spark_expr::{ jvm_udf::JvmScalarUdfExpr, ApproxPercentile, ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, DecimalRescaleCheckOverflow, GetArrayStructFields, GetStructField, HllPlusPlus, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, Regr, - RegrType, SparkCastOptions, Stddev, SumDecimal, ToJson, UnboundColumn, Variance, - WideDecimalBinaryExpr, WideDecimalOp, + RegrType, SparkCastOptions, SparkMapFromArrays, SparkMapFromEntries, SparkStrToMap, Stddev, + SumDecimal, ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp, }; use itertools::Itertools; use jni::objects::{Global, JObject}; @@ -706,6 +708,43 @@ impl PhysicalPlanner { query_context, ))) } + ExprStruct::MapFromArrays(expr) => { + let keys = + self.create_expr(expr.keys.as_ref().unwrap(), Arc::clone(&input_schema))?; + let values = + self.create_expr(expr.values.as_ref().unwrap(), Arc::clone(&input_schema))?; + self.create_map_builder_expr( + SparkMapFromArrays::new(map_key_dedup_policy(&expr.map_key_dedup_policy)?), + vec![keys, values], + &input_schema, + ) + } + ExprStruct::MapFromEntries(expr) => { + let entries = + self.create_expr(expr.entries.as_ref().unwrap(), Arc::clone(&input_schema))?; + self.create_map_builder_expr( + SparkMapFromEntries::new(map_key_dedup_policy(&expr.map_key_dedup_policy)?), + vec![entries], + &input_schema, + ) + } + ExprStruct::StrToMap(expr) => { + let text = + self.create_expr(expr.text.as_ref().unwrap(), Arc::clone(&input_schema))?; + let pair_delimiter = self.create_expr( + expr.pair_delimiter.as_ref().unwrap(), + Arc::clone(&input_schema), + )?; + let key_value_delimiter = self.create_expr( + expr.key_value_delimiter.as_ref().unwrap(), + Arc::clone(&input_schema), + )?; + self.create_map_builder_expr( + SparkStrToMap::new(map_key_dedup_policy(&expr.map_key_dedup_policy)?), + vec![text, pair_delimiter, key_value_delimiter], + &input_schema, + ) + } ExprStruct::ScalarFunc(expr) => { let func = self.create_scalar_function_expr(expr, input_schema); match expr.func.as_ref() { @@ -3668,13 +3707,53 @@ 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`. + /// `prepare_datafusion_session_context` set rather than DataFusion's defaults. fn session_config_options(&self) -> Arc { Arc::clone(self.session_ctx.copied_config().options()) } + /// Builds one of the map constructors. Its `ScalarUDF` carries the + /// `spark.sql.mapKeyDedupPolicy` captured when the plan was converted, so the policy stays + /// fixed for the plan's lifetime the way Spark's `ArrayBasedMapBuilder` keeps the one it was + /// created with. + fn create_map_builder_expr( + &self, + udf: impl ScalarUDFImpl + 'static, + args: Vec>, + input_schema: &Schema, + ) -> Result, ExecutionError> { + let udf = ScalarUDF::new_from_impl(udf); + let name = udf.name().to_string(); + let input_expr_types = args + .iter() + .map(|arg| arg.data_type(input_schema)) + .collect::, _>>()?; + let arg_fields: Vec<_> = input_expr_types + .into_iter() + .enumerate() + .map(|(i, data_type)| Arc::new(Field::new(format!("arg{i}"), data_type, true))) + .collect(); + let scalar_arguments = args + .iter() + .map(|arg| { + arg.as_ref() + .downcast_ref::() + .map(|lit| lit.value()) + }) + .collect::>(); + let return_field = udf.return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalar_arguments, + })?; + Ok(Arc::new(ScalarFunctionExpr::new( + &name, + Arc::new(udf), + args, + return_field, + self.session_config_options(), + ))) + } + fn create_scalar_function_expr( &self, expr: &ScalarFunc, @@ -5099,6 +5178,13 @@ fn needs_fields_coercion(sig: &TypeSignature) -> bool { } } +/// The `spark.sql.mapKeyDedupPolicy` a map constructor captured when the plan was converted. +fn map_key_dedup_policy(policy: &str) -> Result { + policy + .parse::() + .map_err(|error| GeneralError(error.to_string())) +} + #[cfg(test)] mod tests { use futures::{poll, StreamExt}; diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 573e1e9544f..4c2811cb5de 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -25,8 +25,6 @@ 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/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 34d502daac4..f385fbc91a9 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -94,6 +94,9 @@ message Expr { Shuffle shuffle = 72; RandStr rand_str = 73; Uuid uuid = 74; + MapFromArrays map_from_arrays = 75; + MapFromEntries map_from_entries = 76; + StrToMap str_to_map = 77; } reserved 20; @@ -527,6 +530,33 @@ message ScalarFunc { bool fail_on_error = 4; } +// The map constructors carry the `spark.sql.mapKeyDedupPolicy` captured when the plan was +// converted. Spark's `ArrayBasedMapBuilder` reads the policy when the expression is first +// evaluated and keeps it, so a Dataset executed again after the session setting changed still +// builds its maps under the policy it started with; the converted plan is reused across actions +// the same way, so the policy travels with the expression rather than being read again by each +// native iterator. +message MapFromArrays { + Expr keys = 1; + Expr values = 2; + // `EXCEPTION` or `LAST_WIN`. + string map_key_dedup_policy = 3; +} + +message MapFromEntries { + Expr entries = 1; + // `EXCEPTION` or `LAST_WIN`. + string map_key_dedup_policy = 2; +} + +message StrToMap { + Expr text = 1; + Expr pair_delimiter = 2; + Expr key_value_delimiter = 3; + // `EXCEPTION` or `LAST_WIN`. + string map_key_dedup_policy = 4; +} + message CaseWhen { // The expr field is added to be consistent with CaseExpr definition in DataFusion. // This field is not really used. When constructing a CaseExpr, this expr field diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index e44e5b07564..c9b50a02a03 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -31,8 +31,8 @@ use crate::{ EvalMode, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek, SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, - SparkMakeInterval, SparkMakeTime, SparkMapExtract, SparkMapFromArrays, SparkMapFromEntries, - SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, SparkStrToMap, SparkWeekDay, + SparkMakeInterval, SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, + SparkSizeFunc, SparkWeekDay, }; use arrow::datatypes::DataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; @@ -336,12 +336,9 @@ fn all_scalar_functions() -> Vec> { // returns the value itself rather than a one-element list (#5795). It carries the same // `element_at` alias so both registry entries the override replaces point here. Arc::new(ScalarUDF::new_from_impl(SparkMapExtract::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/map_funcs/map_builders.rs b/native/spark-expr/src/map_funcs/map_builders.rs index 57e922fc15b..30f0e3710f8 100644 --- a/native/spark-expr/src/map_funcs/map_builders.rs +++ b/native/spark-expr/src/map_funcs/map_builders.rs @@ -18,10 +18,13 @@ //! 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`: +//! `spark.sql.mapKeyDedupPolicy`, which they read as `datafusion.spark.map_key_dedup_policy` +//! from the session options. Each wrapper carries the policy the plan captured when it was +//! converted instead, so an executed plan keeps its policy the way Spark's `ArrayBasedMapBuilder` +//! keeps the one it was created with, and hands the kernel session options that say so. The +//! wrappers also add the checks `ArrayBasedMapBuilder` performs before inserting an entry, and +//! restate the upstream errors as the Spark error classes `SparkErrorConverter` turns back into +//! `QueryExecutionErrors`: //! //! - a key array and value array of different lengths raise `[MAP_KEY_VALUE_DIFF_SIZES]`, which //! Spark checks before it builds anything; @@ -37,7 +40,7 @@ use arrow::array::{Array, ArrayRef, AsArray, StructArray, UInt32Array}; use arrow::buffer::NullBuffer; use arrow::compute::take; use arrow::datatypes::{DataType, FieldRef}; -use datafusion::common::config::MapKeyDedupPolicy; +use datafusion::common::config::{ConfigOptions, MapKeyDedupPolicy}; use datafusion::common::{exec_err, DataFusionError, HashSet, Result, ScalarValue}; use datafusion::logical_expr::{ ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -45,28 +48,33 @@ use datafusion::logical_expr::{ 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::hash::{Hash, Hasher}; use std::sync::Arc; /// Spark-compatible `map_from_arrays(keys, values)`. -#[derive(Debug, PartialEq, Eq, Hash)] +#[derive(Debug, PartialEq, Eq)] pub struct SparkMapFromArrays { inner: DataFusionMapFromArrays, -} - -impl Default for SparkMapFromArrays { - fn default() -> Self { - Self::new() - } + policy: MapKeyDedupPolicy, } impl SparkMapFromArrays { - pub fn new() -> Self { + /// `policy` is the `spark.sql.mapKeyDedupPolicy` captured when the plan was converted. + pub fn new(policy: MapKeyDedupPolicy) -> Self { Self { inner: DataFusionMapFromArrays::new(), + policy, } } } +impl Hash for SparkMapFromArrays { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + last_value_wins(self.policy).hash(state); + } +} + impl ScalarUDFImpl for SparkMapFromArrays { fn name(&self) -> &str { self.inner.name() @@ -85,11 +93,11 @@ impl ScalarUDFImpl for SparkMapFromArrays { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let mut args = expand_scalars(args)?; + let mut args = expand_scalars(with_policy(args, self.policy))?; compact_list_arguments(&mut args)?; match args.args.as_slice() { [ColumnarValue::Array(keys), ColumnarValue::Array(values)] => { - validate_map_from_arrays(keys, values, last_value_wins(&args))? + validate_map_from_arrays(keys, values, last_value_wins(self.policy))? } other => return exec_err!("map_from_arrays expects 2 arguments, got {}", other.len()), } @@ -100,25 +108,29 @@ impl ScalarUDFImpl for SparkMapFromArrays { } /// Spark-compatible `map_from_entries(entries)`. -#[derive(Debug, PartialEq, Eq, Hash)] +#[derive(Debug, PartialEq, Eq)] pub struct SparkMapFromEntries { inner: DataFusionMapFromEntries, -} - -impl Default for SparkMapFromEntries { - fn default() -> Self { - Self::new() - } + policy: MapKeyDedupPolicy, } impl SparkMapFromEntries { - pub fn new() -> Self { + /// `policy` is the `spark.sql.mapKeyDedupPolicy` captured when the plan was converted. + pub fn new(policy: MapKeyDedupPolicy) -> Self { Self { inner: DataFusionMapFromEntries::new(), + policy, } } } +impl Hash for SparkMapFromEntries { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + last_value_wins(self.policy).hash(state); + } +} + impl ScalarUDFImpl for SparkMapFromEntries { fn name(&self) -> &str { self.inner.name() @@ -137,11 +149,11 @@ impl ScalarUDFImpl for SparkMapFromEntries { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let mut args = expand_scalars(args)?; + let mut args = expand_scalars(with_policy(args, self.policy))?; compact_list_arguments(&mut args)?; match args.args.as_slice() { [ColumnarValue::Array(entries)] => { - validate_map_from_entries(entries, last_value_wins(&args))? + validate_map_from_entries(entries, last_value_wins(self.policy))? } other => return exec_err!("map_from_entries expects 1 argument, got {}", other.len()), } @@ -152,25 +164,29 @@ impl ScalarUDFImpl for SparkMapFromEntries { } /// Spark-compatible `str_to_map(text[, pair_delim[, key_value_delim]])`. -#[derive(Debug, PartialEq, Eq, Hash)] +#[derive(Debug, PartialEq, Eq)] pub struct SparkStrToMap { inner: DataFusionStrToMap, -} - -impl Default for SparkStrToMap { - fn default() -> Self { - Self::new() - } + policy: MapKeyDedupPolicy, } impl SparkStrToMap { - pub fn new() -> Self { + /// `policy` is the `spark.sql.mapKeyDedupPolicy` captured when the plan was converted. + pub fn new(policy: MapKeyDedupPolicy) -> Self { Self { inner: DataFusionStrToMap::new(), + policy, } } } +impl Hash for SparkStrToMap { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + last_value_wins(self.policy).hash(state); + } +} + impl ScalarUDFImpl for SparkStrToMap { fn name(&self) -> &str { self.inner.name() @@ -192,7 +208,7 @@ impl ScalarUDFImpl for SparkStrToMap { // Splitting a string cannot produce a NULL key, so only the duplicate-key error needs // restating here. self.inner - .invoke_with_args(args) + .invoke_with_args(with_policy(args, self.policy)) .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Quoted)) } } @@ -210,9 +226,20 @@ fn expand_scalars(mut args: ScalarFunctionArgs) -> Result { Ok(args) } -/// Whether the session asks for Spark's `LAST_WIN` duplicate key policy. -fn last_value_wins(args: &ScalarFunctionArgs) -> bool { - args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin +/// Whether `policy` is Spark's `LAST_WIN` duplicate key policy. +fn last_value_wins(policy: MapKeyDedupPolicy) -> bool { + policy == MapKeyDedupPolicy::LastWin +} + +/// Hands the kernel the policy the plan captured. The kernels read it from the session options +/// as `datafusion.spark.map_key_dedup_policy`, so those are replaced when they say otherwise. +fn with_policy(mut args: ScalarFunctionArgs, policy: MapKeyDedupPolicy) -> ScalarFunctionArgs { + if args.config_options.spark.map_key_dedup_policy != policy { + let mut options = ConfigOptions::clone(&args.config_options); + options.spark.map_key_dedup_policy = policy; + args.config_options = Arc::new(options); + } + args } /// Rebuilds any list argument whose entries do not start at offset zero. @@ -484,10 +511,16 @@ mod tests { )) } - fn invoke( + fn invoke(udf: &dyn ScalarUDFImpl, args: Vec) -> Result { + invoke_with_session_policy(udf, args, MapKeyDedupPolicy::default()) + } + + /// Invokes `udf` in a session whose `datafusion.spark.map_key_dedup_policy` is + /// `session_policy`, which the wrapper's own policy must take precedence over. + fn invoke_with_session_policy( udf: &dyn ScalarUDFImpl, args: Vec, - policy: MapKeyDedupPolicy, + session_policy: MapKeyDedupPolicy, ) -> Result { let arg_fields: Vec = args .iter() @@ -500,7 +533,7 @@ mod tests { scalar_arguments: &scalar_arguments, })?; let mut config = ConfigOptions::default(); - config.spark.map_key_dedup_policy = policy; + config.spark.map_key_dedup_policy = session_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(), @@ -525,9 +558,8 @@ mod tests { 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(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys, values], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -549,9 +581,8 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys, values], - MapKeyDedupPolicy::Exception, ) .unwrap(), ); @@ -564,9 +595,8 @@ mod tests { 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(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys, values], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -580,9 +610,8 @@ mod tests { 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(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys, values], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -606,9 +635,8 @@ mod tests { )); let values = string_list(StringArray::from(vec![Some("1"), Some("2")]), &[0, 2], None); let err = invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys, values], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -624,9 +652,8 @@ mod tests { let values = string_list(StringArray::from(vec![Some("a"), Some("b")]), &[0, 2], None); let result = map_result( invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::LastWin), vec![keys, values], - MapKeyDedupPolicy::LastWin, ) .unwrap(), ); @@ -648,9 +675,8 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::LastWin), vec![keys, values], - MapKeyDedupPolicy::LastWin, ) .unwrap(), ); @@ -663,6 +689,31 @@ mod tests { assert_eq!((values.value(0), values.value(1)), ("c", "b")); } + #[test] + fn the_plan_policy_wins_over_the_session_option() { + // The plan captured its policy when it was converted; whatever the native session holds + // at execution time must not override it. + 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_with_session_policy( + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), + vec![keys(), values()], + MapKeyDedupPolicy::LastWin, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("[DUPLICATED_MAP_KEY]"), "{err}"); + let result = map_result( + invoke_with_session_policy( + &SparkMapFromArrays::new(MapKeyDedupPolicy::LastWin), + vec![keys(), values()], + MapKeyDedupPolicy::Exception, + ) + .unwrap(), + ); + assert_eq!(result.value_offsets(), &[0, 1]); + } + #[test] fn map_from_entries_rejects_null_key() { let entries = entry_list( @@ -672,9 +723,8 @@ mod tests { None, ); let err = invoke( - &SparkMapFromEntries::default(), + &SparkMapFromEntries::new(MapKeyDedupPolicy::Exception), vec![entries], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -692,9 +742,8 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromEntries::default(), + &SparkMapFromEntries::new(MapKeyDedupPolicy::Exception), vec![entries], - MapKeyDedupPolicy::Exception, ) .unwrap(), ); @@ -712,9 +761,8 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromEntries::default(), + &SparkMapFromEntries::new(MapKeyDedupPolicy::LastWin), vec![entries], - MapKeyDedupPolicy::LastWin, ) .unwrap(), ); @@ -733,9 +781,8 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromEntries::default(), + &SparkMapFromEntries::new(MapKeyDedupPolicy::LastWin), vec![entries], - MapKeyDedupPolicy::LastWin, ) .unwrap(), ); @@ -752,9 +799,8 @@ mod tests { 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(), + &SparkStrToMap::new(MapKeyDedupPolicy::Exception), vec![text], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -768,12 +814,7 @@ mod tests { 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(), + invoke(&SparkStrToMap::new(MapKeyDedupPolicy::LastWin), vec![text]).unwrap(), ); assert_eq!(result.value_offsets(), &[0, 2]); // `a` keeps the slot of its first occurrence and takes its last value. @@ -796,9 +837,8 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys.slice(1, 1), values.slice(1, 1)], - MapKeyDedupPolicy::Exception, ) .unwrap(), ); @@ -827,9 +867,8 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromEntries::default(), + &SparkMapFromEntries::new(MapKeyDedupPolicy::Exception), vec![entries.slice(1, 1)], - MapKeyDedupPolicy::Exception, ) .unwrap(), ); @@ -863,9 +902,8 @@ mod tests { None, ); let err = invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys, values], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -886,9 +924,8 @@ mod tests { None, ); let err = invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys, values], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -909,9 +946,8 @@ mod tests { None, ); let err = invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), vec![keys, values], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -927,9 +963,8 @@ mod tests { None, ); let err = invoke( - &SparkMapFromEntries::default(), + &SparkMapFromEntries::new(MapKeyDedupPolicy::Exception), vec![entries], - MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -950,9 +985,8 @@ mod tests { None, ); let err = invoke( - &SparkMapFromArrays::default(), + &SparkMapFromArrays::new(MapKeyDedupPolicy::LastWin), vec![keys, values], - MapKeyDedupPolicy::LastWin, ) .unwrap_err() .to_string(); diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 386ea69c192..e2c132904d5 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -358,13 +358,6 @@ 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 7ed235b7a1d..7f946733f41 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -136,11 +136,22 @@ object CometMapExtract extends CometExpressionSerde[GetMapValue] { /** * 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`. + * follow `spark.sql.mapKeyDedupPolicy`, which every constructor carries with its expression. */ private object MapBuilderSupport { + /** + * The `spark.sql.mapKeyDedupPolicy` a map constructor carries into the native plan. + * + * Spark's `ArrayBasedMapBuilder` reads the policy when the expression is first evaluated and + * the expression keeps that builder, so a Dataset executed again after the session setting + * changed still builds its maps under the policy it started with. Reading the setting here, + * when the plan is converted, gives the native plan the same lifetime: the converted plan is + * reused across actions, so the policy travels with the expression rather than being read again + * by each native iterator. + */ + def dedupPolicy: String = SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString + /** * Floating-point keys differ from Spark only on 4.0 and later, and differently per function. * `ArrayBasedMapBuilder` gained `keyNormalizer` in 4.0 (with @@ -234,17 +245,25 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { expr: MapFromArrays, inputs: Seq[Attribute], 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 { keysNotNullExprProto <- exprToProtoInternal(IsNotNull(expr.left), inputs, binding) valuesNotNullExprProto <- exprToProtoInternal(IsNotNull(expr.right), inputs, binding) - mapFromArraysExprProto <- scalarFunctionExprToProto("map_from_arrays", keysExpr, valuesExpr) + keysExprProto <- exprToProtoInternal(expr.left, inputs, binding) + valuesExprProto <- exprToProtoInternal(expr.right, inputs, binding) nullLiteralExprProto <- exprToProtoInternal(Literal(null, returnType), inputs, binding) } yield { + val mapFromArraysExprProto = ExprOuterClass.Expr + .newBuilder() + .setMapFromArrays( + ExprOuterClass.MapFromArrays + .newBuilder() + .setKeys(keysExprProto) + .setValues(valuesExprProto) + .setMapKeyDedupPolicy(MapBuilderSupport.dedupPolicy)) + .build() val valuesGuardProto = ExprOuterClass.CaseWhen .newBuilder() .addWhen(valuesNotNullExprProto) @@ -266,7 +285,7 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { } object CometMapFromEntries - extends CometScalarFunction[MapFromEntries]("map_from_entries") + extends CometExpressionSerde[MapFromEntries] with CodegenDispatchFallback { val keyUnsupportedReason = "`BinaryType` is not supported as a map key in `map_from_entries`" @@ -288,10 +307,25 @@ object CometMapFromEntries MapBuilderSupport.keySupport(expr.dataType.keyType) } } + + override def convert( + expr: MapFromEntries, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + exprToProtoInternal(expr.child, inputs, binding).map { entriesExprProto => + ExprOuterClass.Expr + .newBuilder() + .setMapFromEntries( + ExprOuterClass.MapFromEntries + .newBuilder() + .setEntries(entriesExprProto) + .setMapKeyDedupPolicy(MapBuilderSupport.dedupPolicy)) + .build() + } } object CometStrToMap - extends CometScalarFunction[StringToMap]("str_to_map") + extends CometExpressionSerde[StringToMap] with CometTypeShim with CodegenDispatchFallback { @@ -321,6 +355,25 @@ object CometStrToMap Compatible(None) } } + + override def convert( + expr: StringToMap, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + for { + textExprProto <- exprToProtoInternal(expr.text, inputs, binding) + pairDelimExprProto <- exprToProtoInternal(expr.pairDelim, inputs, binding) + keyValueDelimExprProto <- exprToProtoInternal(expr.keyValueDelim, inputs, binding) + } yield ExprOuterClass.Expr + .newBuilder() + .setStrToMap( + ExprOuterClass.StrToMap + .newBuilder() + .setText(textExprProto) + .setPairDelimiter(pairDelimExprProto) + .setKeyValueDelimiter(keyValueDelimExprProto) + .setMapKeyDedupPolicy(MapBuilderSupport.dedupPolicy)) + .build() } object CometCreateMap extends CometCodegenDispatch[CreateMap] diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 609b1c81050..6e76f09cd1e 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -22,7 +22,8 @@ package org.apache.comet import scala.util.Random import org.apache.hadoop.fs.Path -import org.apache.spark.sql.CometTestBase +import org.apache.spark.SparkThrowable +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.ArrayContains import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -235,6 +236,57 @@ class CometMapExpressionSuite extends CometTestBase { } } + // Spark's `ArrayBasedMapBuilder` reads `spark.sql.mapKeyDedupPolicy` when the expression is + // first evaluated and keeps that builder, so a Dataset executed again after the setting changed + // still builds its maps under the policy it started with. Comet captures the policy when it + // converts the plan, which the Dataset reuses across actions, so both engines keep it; reading + // the setting again for every native iterator would apply the new one instead. + // https://github.com/apache/datafusion-comet/pull/5854#discussion_r4049790875 + test("map constructors keep the dedup policy of an executed Dataset") { + withTempPath { dir => + val path = dir.getCanonicalPath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 1, 1, 1).write.parquet(path) + } + def query(): DataFrame = spark.read + .parquet(path) + .selectExpr( + "map_from_arrays(array(id, id), array(1, 2)) AS a", + "map_from_entries(array(struct(id, 1), struct(id, 2))) AS e", + "str_to_map(concat(CAST(id AS STRING), ':1,', CAST(id AS STRING), ':2')) AS s") + val lastWin = Seq(Row(Map(0L -> 2), Map(0L -> 2), Map("0" -> "2"))) + for (cometEnabled <- Seq("false", "true")) { + withSQLConf(CometConf.COMET_ENABLED.key -> cometEnabled) { + // Executed under LAST_WIN, the Dataset keeps that policy once EXCEPTION is set. + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + val df = query() + checkAnswer(df, lastWin) + if (cometEnabled == "true") { + checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + } + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + checkAnswer(df, lastWin) + } + } + // Executed under EXCEPTION, the Dataset keeps raising once LAST_WIN is set. + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + val df = query() + assertDuplicateMapKey(df) + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + assertDuplicateMapKey(df) + } + } + } + } + } + } + + private def assertDuplicateMapKey(df: DataFrame): Unit = { + val error = intercept[Throwable](df.collect()) + val sparkError = causeChain(error).collect { case e: SparkThrowable => e }.lastOption + assert(sparkError.exists(_.getErrorClass == "DUPLICATED_MAP_KEY"), s"$error") + } + test("map_from_entries - null key is rejected") { withMapBuilderTable { table => val exception = checkSparkError( From 69620536200fb3b19284d0512aa83214818a7883 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Sat, 19 Sep 2026 18:52:24 +0000 Subject: [PATCH 14/16] fix: read the map dedup policy when a plan is first executed, and keep it Spark's `ArrayBasedMapBuilder` is a lazy field of the map expression, so `spark.sql.mapKeyDedupPolicy` is read the first time the expression is evaluated and the expression keeps that builder afterwards. The previous commit read the setting in the serde, which runs when the plan is converted, and `explain()` converts a plan without evaluating anything. A Dataset explained under EXCEPTION and then first collected under LAST_WIN therefore raised DUPLICATED_MAP_KEY where Spark returns the map, and the reverse direction returned a map where Spark raises. Capture the policy in a `lazy val` on `CometNativeExec` instead. It is forced by the first `doExecuteColumnar`, which `explain` does not reach, and the node lives in the cached `executedPlan`, so every later action reuses the value. The captured value reaches native through the plan's config map, so `serializeCometSQLConfs` no longer reads the task SQLConf for this setting; the shuffle-writer and write paths pass it the same way. That makes the proto messages the previous commit added unnecessary: the problem was never where the policy travels but when it is read. Revert them, along with the planner arms, the serde overrides and the wrappers' constructor policy, so the constructors are wired as they were and the kernels read the session option again. Cover the new case with a test that materializes the plan under one policy and first executes it under the other, in both directions and with Comet disabled and enabled; it fails both ways without this change. --- .../expression-audits/map_funcs.md | 8 +- native/core/src/execution/jni_api.rs | 11 +- native/core/src/execution/planner.rs | 96 +-------- native/core/src/execution/spark_config.rs | 2 + native/proto/src/proto/expr.proto | 30 --- native/spark-expr/src/comet_scalar_funcs.rs | 7 +- .../spark-expr/src/map_funcs/map_builders.rs | 202 ++++++++---------- .../org/apache/comet/CometExecIterator.scala | 22 +- .../scala/org/apache/comet/serde/maps.scala | 67 +----- .../apache/spark/sql/comet/CometExecRDD.scala | 6 +- .../sql/comet/CometIcebergWriteExec.scala | 4 +- .../sql/comet/CometNativeWriteExec.scala | 4 +- .../shuffle/CometNativeShuffleWriter.scala | 3 +- .../apache/spark/sql/comet/operators.scala | 26 ++- .../comet/CometMapExpressionSuite.scala | 65 +++++- 15 files changed, 229 insertions(+), 324 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index f723631162d..a53401b7a78 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -48,7 +48,7 @@ - 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)). The serde still nests `CASE WHEN left IS NOT NULL THEN (CASE WHEN right IS NOT NULL THEN map_from_arrays(left, right) END) END` around the call: `BinaryExpression.eval` never evaluates `right` for a row whose `left` is NULL, and DataFusion evaluates a THEN branch only on the rows its WHEN selected, so a failing cast in the values array does not run for such a row. A single `left IS NOT NULL AND right IS NOT NULL` guard does not give that, since DataFusion's `AND` evaluates its right side on the whole batch unless the left side is false on all or most rows. - 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 the serde reads when the plan is converted and carries with the expression (`map_key_dedup_policy` on the expression's proto message), so an executed Dataset keeps its policy the way Spark's builder does (`EXCEPTION` raises `DUPLICATED_MAP_KEY` naming the key, `LAST_WIN` keeps the last value for the key). +- `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`, forwarded 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). `CometNativeExec.mapKeyDedupPolicy` reads the setting once, when the plan is first executed, and every native iterator for that plan is given the value it captured; Spark's `ArrayBasedMapBuilder` is a lazy field of the expression and captures the policy at the same moment, so a plan that is explained or executed again across a change to the setting builds its maps the same way in both engines. - Known limitation: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. `from` returns the input arrays untouched when no key repeated, so the stored keys match Spark either way and only duplicate detection diverges. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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. - Known limitation: the two null guards serialize each child a second time inside the `map_from_arrays` call, so a nondeterministic child such as `monotonically_increasing_id()` would advance independently in each copy and the result would drift from Spark ([#5781](https://github.com/apache/datafusion-comet/issues/5781)). `CometMapFromArrays` declines such a child as `Unsupported` through `NullGuardSupport` and the projection falls back to Spark; [#5867](https://github.com/apache/datafusion-comet/pull/5867) routes the same decline through the JVM codegen dispatcher and applies it to `size`, `array_append` and `arrays_zip` as well. @@ -56,10 +56,10 @@ ## map_from_entries - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.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 through the `MapFromEntries` message, which carries the dedup policy captured when the plan is converted. +- 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 the serde reads when the plan is converted and carries with the expression (`map_key_dedup_policy` on the expression's proto message), so an executed Dataset keeps its policy the way Spark's builder does (`EXCEPTION` raises `DUPLICATED_MAP_KEY` naming the key, `LAST_WIN` keeps the last value for the key). +- `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`, forwarded 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). `CometNativeExec.mapKeyDedupPolicy` reads the setting once, when the plan is first executed, and every native iterator for that plan is given the value it captured; Spark's `ArrayBasedMapBuilder` is a lazy field of the expression and captures the policy at the same moment, so a plan that is explained or executed again across a change to the setting builds its maps the same way in both engines. - Known limitation: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. Unlike `map_from_arrays`, this expression always calls `build()`, so Spark stores the normalized key and returns `+0.0` for a `-0.0` key where Comet returns `-0.0`. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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`. @@ -85,7 +85,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 through the `StrToMap` message, which carries the `spark.sql.mapKeyDedupPolicy` captured when the plan is converted; the native wrapper hands that policy to the kernel. +- 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 2fcb0e749f5..59c4a9c1bad 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -111,7 +111,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; @@ -742,6 +742,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 diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 2bef77849ff..fcd9768b1c0 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -102,7 +102,6 @@ use crate::execution::shuffle::{CometPartitioning, CompressionCodec}; use crate::execution::spark_plan::SparkPlan; use crate::parquet::objectstore::s3_blob_fs_support::normalize_object_store_url; use crate::parquet::parquet_support::prepare_object_store_with_configs; -use datafusion::common::config::MapKeyDedupPolicy; use datafusion::common::scalar::ScalarStructBuilder; use datafusion::common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter}, @@ -111,7 +110,6 @@ use datafusion::common::{ use datafusion::datasource::listing::PartitionedFile; use datafusion::logical_expr::type_coercion::functions::fields_with_udf; use datafusion::logical_expr::type_coercion::other::get_coerce_type_for_case_expression; -use datafusion::logical_expr::ScalarUDFImpl; use datafusion::logical_expr::{ AggregateUDF, ReturnFieldArgs, ScalarUDF, TypeSignature, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, @@ -153,8 +151,8 @@ use datafusion_comet_spark_expr::{ jvm_udf::JvmScalarUdfExpr, ApproxPercentile, ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, DecimalRescaleCheckOverflow, GetArrayStructFields, GetStructField, HllPlusPlus, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, Regr, - RegrType, SparkCastOptions, SparkMapFromArrays, SparkMapFromEntries, SparkStrToMap, Stddev, - SumDecimal, ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp, + RegrType, SparkCastOptions, Stddev, SumDecimal, ToJson, UnboundColumn, Variance, + WideDecimalBinaryExpr, WideDecimalOp, }; use itertools::Itertools; use jni::objects::{Global, JObject}; @@ -708,43 +706,6 @@ impl PhysicalPlanner { query_context, ))) } - ExprStruct::MapFromArrays(expr) => { - let keys = - self.create_expr(expr.keys.as_ref().unwrap(), Arc::clone(&input_schema))?; - let values = - self.create_expr(expr.values.as_ref().unwrap(), Arc::clone(&input_schema))?; - self.create_map_builder_expr( - SparkMapFromArrays::new(map_key_dedup_policy(&expr.map_key_dedup_policy)?), - vec![keys, values], - &input_schema, - ) - } - ExprStruct::MapFromEntries(expr) => { - let entries = - self.create_expr(expr.entries.as_ref().unwrap(), Arc::clone(&input_schema))?; - self.create_map_builder_expr( - SparkMapFromEntries::new(map_key_dedup_policy(&expr.map_key_dedup_policy)?), - vec![entries], - &input_schema, - ) - } - ExprStruct::StrToMap(expr) => { - let text = - self.create_expr(expr.text.as_ref().unwrap(), Arc::clone(&input_schema))?; - let pair_delimiter = self.create_expr( - expr.pair_delimiter.as_ref().unwrap(), - Arc::clone(&input_schema), - )?; - let key_value_delimiter = self.create_expr( - expr.key_value_delimiter.as_ref().unwrap(), - Arc::clone(&input_schema), - )?; - self.create_map_builder_expr( - SparkStrToMap::new(map_key_dedup_policy(&expr.map_key_dedup_policy)?), - vec![text, pair_delimiter, key_value_delimiter], - &input_schema, - ) - } ExprStruct::ScalarFunc(expr) => { let func = self.create_scalar_function_expr(expr, input_schema); match expr.func.as_ref() { @@ -3707,53 +3668,13 @@ impl PhysicalPlanner { } /// The session's `ConfigOptions`, so a kernel that reads one sees what - /// `prepare_datafusion_session_context` set rather than DataFusion's defaults. + /// `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()) } - /// Builds one of the map constructors. Its `ScalarUDF` carries the - /// `spark.sql.mapKeyDedupPolicy` captured when the plan was converted, so the policy stays - /// fixed for the plan's lifetime the way Spark's `ArrayBasedMapBuilder` keeps the one it was - /// created with. - fn create_map_builder_expr( - &self, - udf: impl ScalarUDFImpl + 'static, - args: Vec>, - input_schema: &Schema, - ) -> Result, ExecutionError> { - let udf = ScalarUDF::new_from_impl(udf); - let name = udf.name().to_string(); - let input_expr_types = args - .iter() - .map(|arg| arg.data_type(input_schema)) - .collect::, _>>()?; - let arg_fields: Vec<_> = input_expr_types - .into_iter() - .enumerate() - .map(|(i, data_type)| Arc::new(Field::new(format!("arg{i}"), data_type, true))) - .collect(); - let scalar_arguments = args - .iter() - .map(|arg| { - arg.as_ref() - .downcast_ref::() - .map(|lit| lit.value()) - }) - .collect::>(); - let return_field = udf.return_field_from_args(ReturnFieldArgs { - arg_fields: &arg_fields, - scalar_arguments: &scalar_arguments, - })?; - Ok(Arc::new(ScalarFunctionExpr::new( - &name, - Arc::new(udf), - args, - return_field, - self.session_config_options(), - ))) - } - fn create_scalar_function_expr( &self, expr: &ScalarFunc, @@ -5178,13 +5099,6 @@ fn needs_fields_coercion(sig: &TypeSignature) -> bool { } } -/// The `spark.sql.mapKeyDedupPolicy` a map constructor captured when the plan was converted. -fn map_key_dedup_policy(policy: &str) -> Result { - policy - .parse::() - .map_err(|error| GeneralError(error.to_string())) -} - #[cfg(test)] mod tests { use futures::{poll, StreamExt}; 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/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index f385fbc91a9..34d502daac4 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -94,9 +94,6 @@ message Expr { Shuffle shuffle = 72; RandStr rand_str = 73; Uuid uuid = 74; - MapFromArrays map_from_arrays = 75; - MapFromEntries map_from_entries = 76; - StrToMap str_to_map = 77; } reserved 20; @@ -530,33 +527,6 @@ message ScalarFunc { bool fail_on_error = 4; } -// The map constructors carry the `spark.sql.mapKeyDedupPolicy` captured when the plan was -// converted. Spark's `ArrayBasedMapBuilder` reads the policy when the expression is first -// evaluated and keeps it, so a Dataset executed again after the session setting changed still -// builds its maps under the policy it started with; the converted plan is reused across actions -// the same way, so the policy travels with the expression rather than being read again by each -// native iterator. -message MapFromArrays { - Expr keys = 1; - Expr values = 2; - // `EXCEPTION` or `LAST_WIN`. - string map_key_dedup_policy = 3; -} - -message MapFromEntries { - Expr entries = 1; - // `EXCEPTION` or `LAST_WIN`. - string map_key_dedup_policy = 2; -} - -message StrToMap { - Expr text = 1; - Expr pair_delimiter = 2; - Expr key_value_delimiter = 3; - // `EXCEPTION` or `LAST_WIN`. - string map_key_dedup_policy = 4; -} - message CaseWhen { // The expr field is added to be consistent with CaseExpr definition in DataFusion. // This field is not really used. When constructing a CaseExpr, this expr field diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index c9b50a02a03..e44e5b07564 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -31,8 +31,8 @@ use crate::{ EvalMode, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek, SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, - SparkMakeInterval, SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, - SparkSizeFunc, SparkWeekDay, + SparkMakeInterval, SparkMakeTime, SparkMapExtract, SparkMapFromArrays, SparkMapFromEntries, + SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, SparkStrToMap, SparkWeekDay, }; use arrow::datatypes::DataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; @@ -336,9 +336,12 @@ fn all_scalar_functions() -> Vec> { // returns the value itself rather than a one-element list (#5795). It carries the same // `element_at` alias so both registry entries the override replaces point here. Arc::new(ScalarUDF::new_from_impl(SparkMapExtract::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/map_funcs/map_builders.rs b/native/spark-expr/src/map_funcs/map_builders.rs index 30f0e3710f8..57e922fc15b 100644 --- a/native/spark-expr/src/map_funcs/map_builders.rs +++ b/native/spark-expr/src/map_funcs/map_builders.rs @@ -18,13 +18,10 @@ //! 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 they read as `datafusion.spark.map_key_dedup_policy` -//! from the session options. Each wrapper carries the policy the plan captured when it was -//! converted instead, so an executed plan keeps its policy the way Spark's `ArrayBasedMapBuilder` -//! keeps the one it was created with, and hands the kernel session options that say so. The -//! wrappers also add the checks `ArrayBasedMapBuilder` performs before inserting an entry, and -//! restate the upstream errors as the Spark error classes `SparkErrorConverter` turns back into -//! `QueryExecutionErrors`: +//! `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 key array and value array of different lengths raise `[MAP_KEY_VALUE_DIFF_SIZES]`, which //! Spark checks before it builds anything; @@ -40,7 +37,7 @@ use arrow::array::{Array, ArrayRef, AsArray, StructArray, UInt32Array}; use arrow::buffer::NullBuffer; use arrow::compute::take; use arrow::datatypes::{DataType, FieldRef}; -use datafusion::common::config::{ConfigOptions, MapKeyDedupPolicy}; +use datafusion::common::config::MapKeyDedupPolicy; use datafusion::common::{exec_err, DataFusionError, HashSet, Result, ScalarValue}; use datafusion::logical_expr::{ ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -48,33 +45,28 @@ use datafusion::logical_expr::{ 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::hash::{Hash, Hasher}; use std::sync::Arc; /// Spark-compatible `map_from_arrays(keys, values)`. -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkMapFromArrays { inner: DataFusionMapFromArrays, - policy: MapKeyDedupPolicy, +} + +impl Default for SparkMapFromArrays { + fn default() -> Self { + Self::new() + } } impl SparkMapFromArrays { - /// `policy` is the `spark.sql.mapKeyDedupPolicy` captured when the plan was converted. - pub fn new(policy: MapKeyDedupPolicy) -> Self { + pub fn new() -> Self { Self { inner: DataFusionMapFromArrays::new(), - policy, } } } -impl Hash for SparkMapFromArrays { - fn hash(&self, state: &mut H) { - self.inner.hash(state); - last_value_wins(self.policy).hash(state); - } -} - impl ScalarUDFImpl for SparkMapFromArrays { fn name(&self) -> &str { self.inner.name() @@ -93,11 +85,11 @@ impl ScalarUDFImpl for SparkMapFromArrays { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let mut args = expand_scalars(with_policy(args, self.policy))?; + let mut args = expand_scalars(args)?; compact_list_arguments(&mut args)?; match args.args.as_slice() { [ColumnarValue::Array(keys), ColumnarValue::Array(values)] => { - validate_map_from_arrays(keys, values, last_value_wins(self.policy))? + validate_map_from_arrays(keys, values, last_value_wins(&args))? } other => return exec_err!("map_from_arrays expects 2 arguments, got {}", other.len()), } @@ -108,29 +100,25 @@ impl ScalarUDFImpl for SparkMapFromArrays { } /// Spark-compatible `map_from_entries(entries)`. -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkMapFromEntries { inner: DataFusionMapFromEntries, - policy: MapKeyDedupPolicy, +} + +impl Default for SparkMapFromEntries { + fn default() -> Self { + Self::new() + } } impl SparkMapFromEntries { - /// `policy` is the `spark.sql.mapKeyDedupPolicy` captured when the plan was converted. - pub fn new(policy: MapKeyDedupPolicy) -> Self { + pub fn new() -> Self { Self { inner: DataFusionMapFromEntries::new(), - policy, } } } -impl Hash for SparkMapFromEntries { - fn hash(&self, state: &mut H) { - self.inner.hash(state); - last_value_wins(self.policy).hash(state); - } -} - impl ScalarUDFImpl for SparkMapFromEntries { fn name(&self) -> &str { self.inner.name() @@ -149,11 +137,11 @@ impl ScalarUDFImpl for SparkMapFromEntries { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let mut args = expand_scalars(with_policy(args, self.policy))?; + let mut args = expand_scalars(args)?; compact_list_arguments(&mut args)?; match args.args.as_slice() { [ColumnarValue::Array(entries)] => { - validate_map_from_entries(entries, last_value_wins(self.policy))? + validate_map_from_entries(entries, last_value_wins(&args))? } other => return exec_err!("map_from_entries expects 1 argument, got {}", other.len()), } @@ -164,29 +152,25 @@ impl ScalarUDFImpl for SparkMapFromEntries { } /// Spark-compatible `str_to_map(text[, pair_delim[, key_value_delim]])`. -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkStrToMap { inner: DataFusionStrToMap, - policy: MapKeyDedupPolicy, +} + +impl Default for SparkStrToMap { + fn default() -> Self { + Self::new() + } } impl SparkStrToMap { - /// `policy` is the `spark.sql.mapKeyDedupPolicy` captured when the plan was converted. - pub fn new(policy: MapKeyDedupPolicy) -> Self { + pub fn new() -> Self { Self { inner: DataFusionStrToMap::new(), - policy, } } } -impl Hash for SparkStrToMap { - fn hash(&self, state: &mut H) { - self.inner.hash(state); - last_value_wins(self.policy).hash(state); - } -} - impl ScalarUDFImpl for SparkStrToMap { fn name(&self) -> &str { self.inner.name() @@ -208,7 +192,7 @@ impl ScalarUDFImpl for SparkStrToMap { // Splitting a string cannot produce a NULL key, so only the duplicate-key error needs // restating here. self.inner - .invoke_with_args(with_policy(args, self.policy)) + .invoke_with_args(args) .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Quoted)) } } @@ -226,20 +210,9 @@ fn expand_scalars(mut args: ScalarFunctionArgs) -> Result { Ok(args) } -/// Whether `policy` is Spark's `LAST_WIN` duplicate key policy. -fn last_value_wins(policy: MapKeyDedupPolicy) -> bool { - policy == MapKeyDedupPolicy::LastWin -} - -/// Hands the kernel the policy the plan captured. The kernels read it from the session options -/// as `datafusion.spark.map_key_dedup_policy`, so those are replaced when they say otherwise. -fn with_policy(mut args: ScalarFunctionArgs, policy: MapKeyDedupPolicy) -> ScalarFunctionArgs { - if args.config_options.spark.map_key_dedup_policy != policy { - let mut options = ConfigOptions::clone(&args.config_options); - options.spark.map_key_dedup_policy = policy; - args.config_options = Arc::new(options); - } - args +/// Whether the session asks for Spark's `LAST_WIN` duplicate key policy. +fn last_value_wins(args: &ScalarFunctionArgs) -> bool { + args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin } /// Rebuilds any list argument whose entries do not start at offset zero. @@ -511,16 +484,10 @@ mod tests { )) } - fn invoke(udf: &dyn ScalarUDFImpl, args: Vec) -> Result { - invoke_with_session_policy(udf, args, MapKeyDedupPolicy::default()) - } - - /// Invokes `udf` in a session whose `datafusion.spark.map_key_dedup_policy` is - /// `session_policy`, which the wrapper's own policy must take precedence over. - fn invoke_with_session_policy( + fn invoke( udf: &dyn ScalarUDFImpl, args: Vec, - session_policy: MapKeyDedupPolicy, + policy: MapKeyDedupPolicy, ) -> Result { let arg_fields: Vec = args .iter() @@ -533,7 +500,7 @@ mod tests { scalar_arguments: &scalar_arguments, })?; let mut config = ConfigOptions::default(); - config.spark.map_key_dedup_policy = session_policy; + 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(), @@ -558,8 +525,9 @@ mod tests { 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::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -581,8 +549,9 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::Exception, ) .unwrap(), ); @@ -595,8 +564,9 @@ mod tests { 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::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -610,8 +580,9 @@ mod tests { 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::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -635,8 +606,9 @@ mod tests { )); let values = string_list(StringArray::from(vec![Some("1"), Some("2")]), &[0, 2], None); let err = invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -652,8 +624,9 @@ mod tests { let values = string_list(StringArray::from(vec![Some("a"), Some("b")]), &[0, 2], None); let result = map_result( invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::LastWin), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::LastWin, ) .unwrap(), ); @@ -675,8 +648,9 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::LastWin), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::LastWin, ) .unwrap(), ); @@ -689,31 +663,6 @@ mod tests { assert_eq!((values.value(0), values.value(1)), ("c", "b")); } - #[test] - fn the_plan_policy_wins_over_the_session_option() { - // The plan captured its policy when it was converted; whatever the native session holds - // at execution time must not override it. - 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_with_session_policy( - &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), - vec![keys(), values()], - MapKeyDedupPolicy::LastWin, - ) - .unwrap_err() - .to_string(); - assert!(err.contains("[DUPLICATED_MAP_KEY]"), "{err}"); - let result = map_result( - invoke_with_session_policy( - &SparkMapFromArrays::new(MapKeyDedupPolicy::LastWin), - vec![keys(), values()], - MapKeyDedupPolicy::Exception, - ) - .unwrap(), - ); - assert_eq!(result.value_offsets(), &[0, 1]); - } - #[test] fn map_from_entries_rejects_null_key() { let entries = entry_list( @@ -723,8 +672,9 @@ mod tests { None, ); let err = invoke( - &SparkMapFromEntries::new(MapKeyDedupPolicy::Exception), + &SparkMapFromEntries::default(), vec![entries], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -742,8 +692,9 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromEntries::new(MapKeyDedupPolicy::Exception), + &SparkMapFromEntries::default(), vec![entries], + MapKeyDedupPolicy::Exception, ) .unwrap(), ); @@ -761,8 +712,9 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromEntries::new(MapKeyDedupPolicy::LastWin), + &SparkMapFromEntries::default(), vec![entries], + MapKeyDedupPolicy::LastWin, ) .unwrap(), ); @@ -781,8 +733,9 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromEntries::new(MapKeyDedupPolicy::LastWin), + &SparkMapFromEntries::default(), vec![entries], + MapKeyDedupPolicy::LastWin, ) .unwrap(), ); @@ -799,8 +752,9 @@ mod tests { 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::new(MapKeyDedupPolicy::Exception), + &SparkStrToMap::default(), vec![text], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -814,7 +768,12 @@ mod tests { 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::new(MapKeyDedupPolicy::LastWin), vec![text]).unwrap(), + invoke( + &SparkStrToMap::default(), + vec![text], + MapKeyDedupPolicy::LastWin, + ) + .unwrap(), ); assert_eq!(result.value_offsets(), &[0, 2]); // `a` keeps the slot of its first occurrence and takes its last value. @@ -837,8 +796,9 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys.slice(1, 1), values.slice(1, 1)], + MapKeyDedupPolicy::Exception, ) .unwrap(), ); @@ -867,8 +827,9 @@ mod tests { ); let result = map_result( invoke( - &SparkMapFromEntries::new(MapKeyDedupPolicy::Exception), + &SparkMapFromEntries::default(), vec![entries.slice(1, 1)], + MapKeyDedupPolicy::Exception, ) .unwrap(), ); @@ -902,8 +863,9 @@ mod tests { None, ); let err = invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -924,8 +886,9 @@ mod tests { None, ); let err = invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -946,8 +909,9 @@ mod tests { None, ); let err = invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::Exception), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -963,8 +927,9 @@ mod tests { None, ); let err = invoke( - &SparkMapFromEntries::new(MapKeyDedupPolicy::Exception), + &SparkMapFromEntries::default(), vec![entries], + MapKeyDedupPolicy::Exception, ) .unwrap_err() .to_string(); @@ -985,8 +950,9 @@ mod tests { None, ); let err = invoke( - &SparkMapFromArrays::new(MapKeyDedupPolicy::LastWin), + &SparkMapFromArrays::default(), vec![keys, values], + MapKeyDedupPolicy::LastWin, ) .unwrap_err() .to_string(); diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index e2c132904d5..de61af3b495 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -77,7 +77,8 @@ class CometExecIterator( encryptedFilePaths: Seq[String] = Seq.empty, shuffleBlockIterators: Map[Int, CometShuffleBlockIterator] = Map.empty, taskFilePaths: Seq[String] = Seq.empty, - shufflePartitionPusher: Option[ShufflePartitionPusher] = None) + shufflePartitionPusher: Option[ShufflePartitionPusher] = None, + mapKeyDedupPolicy: Option[String] = None) extends Iterator[ColumnarBatch] with Logging { @@ -94,7 +95,7 @@ class CometExecIterator( val localDiskDirs = SparkEnv.get.blockManager.getLocalDiskDirs // serialize Comet related Spark configs in protobuf format - val protobufSparkConfigs = CometExecIterator.serializeCometSQLConfs() + val protobufSparkConfigs = CometExecIterator.serializeCometSQLConfs(mapKeyDedupPolicy) // Create keyUnwrapper if encryption is enabled val keyUnwrapper = if (encryptedFilePaths.nonEmpty) { @@ -335,7 +336,13 @@ object CometExecIterator extends Logging { private def cometSqlConfs: Map[String, String] = SQLConf.get.getAllConfs.filter(_._1.startsWith(CometConf.COMET_PREFIX)) - def serializeCometSQLConfs(): Array[Byte] = { + /** + * @param mapKeyDedupPolicy + * the `spark.sql.mapKeyDedupPolicy` the plan captured when it was first executed, or `None` + * to read the current value. See + * [[org.apache.spark.sql.comet.CometNativeExec.mapKeyDedupPolicy]]. + */ + def serializeCometSQLConfs(mapKeyDedupPolicy: Option[String] = None): Array[Byte] = { val builder = ConfigMap.newBuilder() cometSqlConfs.foreach { case (k, v) => if (k.startsWith(s"${CometConf.COMET_PREFIX}.datafusion.")) { @@ -358,6 +365,15 @@ 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`. Spark's `ArrayBasedMapBuilder` reads it once, when + // the expression is first evaluated, so the plan captures it then and passes it in here rather + // than letting every native iterator read whatever the session holds at the time. + builder.putEntries( + SQLConf.MAP_KEY_DEDUP_POLICY.key, + mapKeyDedupPolicy.getOrElse(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 7f946733f41..7ed235b7a1d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -136,22 +136,11 @@ object CometMapExtract extends CometExpressionSerde[GetMapValue] { /** * 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`, which every constructor carries with its expression. + * follow `spark.sql.mapKeyDedupPolicy`, whose value Comet forwards to the native session as + * `datafusion.spark.map_key_dedup_policy`. */ private object MapBuilderSupport { - /** - * The `spark.sql.mapKeyDedupPolicy` a map constructor carries into the native plan. - * - * Spark's `ArrayBasedMapBuilder` reads the policy when the expression is first evaluated and - * the expression keeps that builder, so a Dataset executed again after the session setting - * changed still builds its maps under the policy it started with. Reading the setting here, - * when the plan is converted, gives the native plan the same lifetime: the converted plan is - * reused across actions, so the policy travels with the expression rather than being read again - * by each native iterator. - */ - def dedupPolicy: String = SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString - /** * Floating-point keys differ from Spark only on 4.0 and later, and differently per function. * `ArrayBasedMapBuilder` gained `keyNormalizer` in 4.0 (with @@ -245,25 +234,17 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { expr: MapFromArrays, inputs: Seq[Attribute], 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 { keysNotNullExprProto <- exprToProtoInternal(IsNotNull(expr.left), inputs, binding) valuesNotNullExprProto <- exprToProtoInternal(IsNotNull(expr.right), inputs, binding) - keysExprProto <- exprToProtoInternal(expr.left, inputs, binding) - valuesExprProto <- exprToProtoInternal(expr.right, inputs, binding) + mapFromArraysExprProto <- scalarFunctionExprToProto("map_from_arrays", keysExpr, valuesExpr) nullLiteralExprProto <- exprToProtoInternal(Literal(null, returnType), inputs, binding) } yield { - val mapFromArraysExprProto = ExprOuterClass.Expr - .newBuilder() - .setMapFromArrays( - ExprOuterClass.MapFromArrays - .newBuilder() - .setKeys(keysExprProto) - .setValues(valuesExprProto) - .setMapKeyDedupPolicy(MapBuilderSupport.dedupPolicy)) - .build() val valuesGuardProto = ExprOuterClass.CaseWhen .newBuilder() .addWhen(valuesNotNullExprProto) @@ -285,7 +266,7 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { } object CometMapFromEntries - extends CometExpressionSerde[MapFromEntries] + extends CometScalarFunction[MapFromEntries]("map_from_entries") with CodegenDispatchFallback { val keyUnsupportedReason = "`BinaryType` is not supported as a map key in `map_from_entries`" @@ -307,25 +288,10 @@ object CometMapFromEntries MapBuilderSupport.keySupport(expr.dataType.keyType) } } - - override def convert( - expr: MapFromEntries, - inputs: Seq[Attribute], - binding: Boolean): Option[ExprOuterClass.Expr] = - exprToProtoInternal(expr.child, inputs, binding).map { entriesExprProto => - ExprOuterClass.Expr - .newBuilder() - .setMapFromEntries( - ExprOuterClass.MapFromEntries - .newBuilder() - .setEntries(entriesExprProto) - .setMapKeyDedupPolicy(MapBuilderSupport.dedupPolicy)) - .build() - } } object CometStrToMap - extends CometExpressionSerde[StringToMap] + extends CometScalarFunction[StringToMap]("str_to_map") with CometTypeShim with CodegenDispatchFallback { @@ -355,25 +321,6 @@ object CometStrToMap Compatible(None) } } - - override def convert( - expr: StringToMap, - inputs: Seq[Attribute], - binding: Boolean): Option[ExprOuterClass.Expr] = - for { - textExprProto <- exprToProtoInternal(expr.text, inputs, binding) - pairDelimExprProto <- exprToProtoInternal(expr.pairDelim, inputs, binding) - keyValueDelimExprProto <- exprToProtoInternal(expr.keyValueDelim, inputs, binding) - } yield ExprOuterClass.Expr - .newBuilder() - .setStrToMap( - ExprOuterClass.StrToMap - .newBuilder() - .setText(textExprProto) - .setPairDelimiter(pairDelimExprProto) - .setKeyValueDelimiter(keyValueDelimExprProto) - .setMapKeyDedupPolicy(MapBuilderSupport.dedupPolicy)) - .build() } object CometCreateMap extends CometCodegenDispatch[CreateMap] diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index d9e0bf3a4c9..4dbceca889c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -67,7 +67,8 @@ private[spark] class CometExecRDD( broadcastedHadoopConfForEncryption: Option[Broadcast[SerializableConfiguration]] = None, encryptedFilePaths: Seq[String] = Seq.empty, shuffleScanIndices: Set[Int] = Set.empty, - @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty) + @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty, + mapKeyDedupPolicy: Option[String] = None) extends RDD[ColumnarBatch](sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { // Determine partition count: from inputs if available, otherwise from parameter @@ -133,7 +134,8 @@ private[spark] class CometExecRDD( broadcastedHadoopConfForEncryption, encryptedFilePaths, shuffleBlockIters, - taskFilePaths = partition.filePaths) + taskFilePaths = partition.filePaths, + mapKeyDedupPolicy = mapKeyDedupPolicy) // Register ScalarSubqueries so native code can look them up subqueries.foreach(sub => CometScalarSubquery.setSubquery(it.id, sub)) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala index 00a5047625f..56500216108 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala @@ -238,6 +238,7 @@ case class CometIcebergWriteExec( // column each (see `build_output_schema` in `iceberg_write.rs`). val numOutputCols = 2 val capturedNativeOp = nativeOp + val capturedMapKeyDedupPolicy = mapKeyDedupPolicy childRDD.mapPartitionsInternal { iter => val partitionId = TaskContext.getPartitionId() @@ -277,7 +278,8 @@ case class CometIcebergWriteExec( numPartitions, partitionId, None, - Seq.empty) + Seq.empty, + mapKeyDedupPolicy = Some(capturedMapKeyDedupPolicy)) execIterator } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala index f0d10b17667..424604416a0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala @@ -180,6 +180,7 @@ case class CometNativeWriteExec( val capturedCommitter = committer val capturedJobTrackerID = jobTrackerID val capturedNativeOp = nativeOp + val capturedMapKeyDedupPolicy = mapKeyDedupPolicy val capturedAccumulator = taskCommitMessagesAccum // Capture accumulator for use in tasks // Execute native write operation with task-level commit protocol @@ -242,7 +243,8 @@ case class CometNativeWriteExec( numPartitions, partitionId, None, - Seq.empty) + Seq.empty, + mapKeyDedupPolicy = Some(capturedMapKeyDedupPolicy)) // Wrap the iterator to handle task commit/abort and capture TaskCommitMessage new Iterator[ColumnarBatch] { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index 21dd7686302..48618f7e5e1 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -198,7 +198,8 @@ class CometNativeShuffleWriter[K, V]( ctx.broadcastedHadoopConfForEncryption, ctx.encryptedFilePaths, shuffleBlockIters, - shufflePartitionPusher = remoteDestination.map(_.callback)) + shufflePartitionPusher = remoteDestination.map(_.callback), + mapKeyDedupPolicy = Some(ctx.mapKeyDedupPolicy)) // Register subqueries against the iterator id so native callbacks resolve them to values. ctx.subqueries.foreach { sub => diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 0e160faf47a..bae17cde4a2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -44,6 +44,7 @@ import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregat import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.SerializableConfiguration @@ -546,7 +547,8 @@ private[comet] case class NativeExecContext( // binary when this context rides on the non-transient CometShuffleDependency.nativeShuffleSpec. @transient perPartitionByKey: Map[String, Array[Array[Byte]]], shuffleScanIndices: Set[Int], - hasScanInput: Boolean) { + hasScanInput: Boolean, + mapKeyDedupPolicy: String = SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString) { // Catch shape divergence (e.g. broadcast scans with different partition counts after DPP // filtering) at construction so consumers don't trip ArrayIndexOutOfBoundsException at // partition idx access time. @@ -570,6 +572,22 @@ abstract class CometNativeExec extends CometExec { /** The Comet native operator */ def nativeOp: Operator + /** + * The `spark.sql.mapKeyDedupPolicy` the native map constructors in this plan build their maps + * with, read the first time the plan is executed and kept from then on. + * + * Spark's `ArrayBasedMapBuilder` is a lazy field of the map expression, so it reads the policy + * when the expression is first evaluated and the expression keeps that builder for every later + * action. Neither materializing nor explaining a plan evaluates anything, so a Dataset + * explained under one policy and then executed under another builds its maps under the second + * one, and a Dataset executed twice across a change keeps the first one. A `lazy val` on this + * node gives the same two properties: it is forced by the first `doExecuteColumnar`, which + * `explain` does not reach, and the node lives in the cached `executedPlan`, so later actions + * reuse the value. + */ + private[comet] lazy val mapKeyDedupPolicy: String = + SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString + override protected def doPrepare(): Unit = prepareSubqueries(this) override lazy val metrics: Map[String, SQLMetric] = @@ -624,7 +642,8 @@ abstract class CometNativeExec extends CometExec { ctx.subqueries, ctx.broadcastedHadoopConfForEncryption, ctx.encryptedFilePaths, - ctx.shuffleScanIndices) { + ctx.shuffleScanIndices, + mapKeyDedupPolicy = Some(ctx.mapKeyDedupPolicy)) { override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { val res = super.compute(split, context) if (ctx.hasScanInput) { @@ -825,7 +844,8 @@ abstract class CometNativeExec extends CometExec { commonByKey = commonByKey, perPartitionByKey = perPartitionByKey, shuffleScanIndices = shuffleScanIndices, - hasScanInput = sparkPlans.exists(_.isInstanceOf[CometNativeScanExec])) + hasScanInput = sparkPlans.exists(_.isInstanceOf[CometNativeScanExec]), + mapKeyDedupPolicy = mapKeyDedupPolicy) } /** diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 6e76f09cd1e..b3f23fca356 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -248,13 +248,8 @@ class CometMapExpressionSuite extends CometTestBase { withSQLConf(CometConf.COMET_ENABLED.key -> "false") { spark.range(0, 1, 1, 1).write.parquet(path) } - def query(): DataFrame = spark.read - .parquet(path) - .selectExpr( - "map_from_arrays(array(id, id), array(1, 2)) AS a", - "map_from_entries(array(struct(id, 1), struct(id, 2))) AS e", - "str_to_map(concat(CAST(id AS STRING), ':1,', CAST(id AS STRING), ':2')) AS s") - val lastWin = Seq(Row(Map(0L -> 2), Map(0L -> 2), Map("0" -> "2"))) + def query(): DataFrame = mapPolicyQuery(path) + val lastWin = mapPolicyLastWin for (cometEnabled <- Seq("false", "true")) { withSQLConf(CometConf.COMET_ENABLED.key -> cometEnabled) { // Executed under LAST_WIN, the Dataset keeps that policy once EXCEPTION is set. @@ -281,6 +276,62 @@ class CometMapExpressionSuite extends CometTestBase { } } + // Spark initializes `ArrayBasedMapBuilder`, and with it reads the policy, when the expression is + // first evaluated. Materializing the plan does not evaluate anything, so a Dataset explained + // under one policy and then executed under another builds its maps under the second one. Comet + // converts its plan when `executedPlan` is materialized, which `explain()` also triggers, so the + // policy cannot be read there. + // https://github.com/apache/datafusion-comet/pull/5854#issuecomment-5744116922 + test("map constructors capture the dedup policy at first execution, not at planning") { + // AQE converts each query stage as it runs, which hides the difference between the two moments. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { dir => + val path = dir.getCanonicalPath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 1, 1, 1).write.parquet(path) + } + for (cometEnabled <- Seq("false", "true")) { + withSQLConf(CometConf.COMET_ENABLED.key -> cometEnabled) { + // Explained under EXCEPTION, first executed under LAST_WIN: LAST_WIN builds the maps, + // and keeps doing so once the setting goes back. + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + val df = mapPolicyQuery(path) + val plan = df.queryExecution.executedPlan + if (cometEnabled == "true") { + checkCometOperators(stripAQEPlan(plan)) + } + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + checkAnswer(df, mapPolicyLastWin) + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + checkAnswer(df, mapPolicyLastWin) + } + } + } + // Explained under LAST_WIN, first executed under EXCEPTION: EXCEPTION rejects the + // duplicate. + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + val df = mapPolicyQuery(path) + df.queryExecution.executedPlan + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + assertDuplicateMapKey(df) + } + } + } + } + } + } + } + + /** One row whose three map constructors each see the duplicate key `0`. */ + private def mapPolicyQuery(path: String): DataFrame = spark.read + .parquet(path) + .selectExpr( + "map_from_arrays(array(id, id), array(1, 2)) AS a", + "map_from_entries(array(struct(id, 1), struct(id, 2))) AS e", + "str_to_map(concat(CAST(id AS STRING), ':1,', CAST(id AS STRING), ':2')) AS s") + + private def mapPolicyLastWin: Seq[Row] = Seq(Row(Map(0L -> 2), Map(0L -> 2), Map("0" -> "2"))) + private def assertDuplicateMapKey(df: DataFrame): Unit = { val error = intercept[Throwable](df.collect()) val sparkError = causeChain(error).collect { case e: SparkThrowable => e }.lastOption From 362f288ffc5054b079d26d61710fe57630927811 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Sun, 20 Sep 2026 12:02:27 +0000 Subject: [PATCH 15/16] fix: read the map dedup policy per native plan, matching Spark outside whole-stage codegen Spark reads `spark.sql.mapKeyDedupPolicy` into `ArrayBasedMapBuilder`, a lazy field of the map expression, so when it reads it depends on how the projection runs. Outside whole-stage codegen the projection is rebuilt in every task and the setting is read again on each action; inside it the builder is created once on the driver, in the first action, and kept. The previous commit froze the policy on the first execution, which matches the whole-stage case. Measuring both engines across the two codegen paths, the two directions of a policy change and both the repeated-action and materialize-then-execute scenarios, that freeze matched Spark in five of eight comparisons; reading the setting when each native plan is built matches seven of eight. The freeze also diverges by returning a map where Spark raises `DUPLICATED_MAP_KEY`, while reading per plan diverges by raising where Spark returns a map, so the remaining difference is loud rather than silent. Drop the freeze and the parameter it threaded through `NativeExecContext`, `CometExecRDD`, `CometExecIterator`, the shuffle writer and both write execs. The one case that cannot also be matched, a Dataset executed more than once across a change to the setting with the projection inside whole-stage codegen, is recorded in the map_funcs expression audit: Comet replaces the operator before `CollapseCodegenStages` runs, so the plan it sees carries no record of which path Spark would have taken. Cover both scenarios: one test runs a Dataset twice across a change in each direction, in both of the configurations that leave whole-stage codegen, and one materializes the plan under one policy and first executes it under the other. --- .../expression-audits/map_funcs.md | 5 +- .../org/apache/comet/CometExecIterator.scala | 19 +--- .../apache/spark/sql/comet/CometExecRDD.scala | 6 +- .../sql/comet/CometIcebergWriteExec.scala | 4 +- .../sql/comet/CometNativeWriteExec.scala | 4 +- .../shuffle/CometNativeShuffleWriter.scala | 1 - .../apache/spark/sql/comet/operators.scala | 26 +---- .../comet/CometMapExpressionSuite.scala | 107 ++++++++++-------- 8 files changed, 77 insertions(+), 95 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index a53401b7a78..840a34758a9 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -48,9 +48,10 @@ - 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)). The serde still nests `CASE WHEN left IS NOT NULL THEN (CASE WHEN right IS NOT NULL THEN map_from_arrays(left, right) END) END` around the call: `BinaryExpression.eval` never evaluates `right` for a row whose `left` is NULL, and DataFusion evaluates a THEN branch only on the rows its WHEN selected, so a failing cast in the values array does not run for such a row. A single `left IS NOT NULL AND right IS NOT NULL` guard does not give that, since DataFusion's `AND` evaluates its right side on the whole batch unless the left side is false on all or most rows. - 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`, forwarded 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). `CometNativeExec.mapKeyDedupPolicy` reads the setting once, when the plan is first executed, and every native iterator for that plan is given the value it captured; Spark's `ArrayBasedMapBuilder` is a lazy field of the expression and captures the policy at the same moment, so a plan that is explained or executed again across a change to the setting builds its maps the same way in both engines. +- `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`, forwarded 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). `CometExecIterator.serializeCometSQLConfs` reads the setting when it builds the native plan for a task, so materializing or explaining a plan does not fix it and a Dataset re-executed after a change to the setting uses the new value. - Known limitation: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. `from` returns the input arrays untouched when no key repeated, so the stored keys match Spark either way and only duplicate detection diverges. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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. +- Known limitation: Spark reads `spark.sql.mapKeyDedupPolicy` into `ArrayBasedMapBuilder`, a lazy field of the map expression, so *when* it reads it depends on how the projection runs. Outside whole-stage codegen (the flag off, or a projection wider than `spark.sql.codegen.maxFields`) the projection is rebuilt in every task and the setting is read again on each action, which is what Comet does. Inside whole-stage codegen Spark creates the builder once on the driver, in the first action, and keeps it, so a Dataset re-executed after a change to the setting still builds its maps under the policy it started with, where Comet uses the new one. Comet cannot tell the two apart: it replaces the operator before `CollapseCodegenStages` runs, so the plan it sees carries no record of which path Spark would have taken. Matching the whole-stage case instead would mean returning a map where Spark raises `DUPLICATED_MAP_KEY` in the other three configurations, so the loud divergence is preferred over the silent one. Only a Dataset that is executed more than once across a change to the setting is affected. - Known limitation: the two null guards serialize each child a second time inside the `map_from_arrays` call, so a nondeterministic child such as `monotonically_increasing_id()` would advance independently in each copy and the result would drift from Spark ([#5781](https://github.com/apache/datafusion-comet/issues/5781)). `CometMapFromArrays` declines such a child as `Unsupported` through `NullGuardSupport` and the projection falls back to Spark; [#5867](https://github.com/apache/datafusion-comet/pull/5867) routes the same decline through the JVM codegen dispatcher and applies it to `size`, `array_append` and `arrays_zip` as well. ## map_from_entries @@ -59,7 +60,7 @@ - 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`, forwarded 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). `CometNativeExec.mapKeyDedupPolicy` reads the setting once, when the plan is first executed, and every native iterator for that plan is given the value it captured; Spark's `ArrayBasedMapBuilder` is a lazy field of the expression and captures the policy at the same moment, so a plan that is explained or executed again across a change to the setting builds its maps the same way in both engines. +- `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`, forwarded 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). `CometExecIterator.serializeCometSQLConfs` reads the setting when it builds the native plan for a task, so materializing or explaining a plan does not fix it and a Dataset re-executed after a change to the setting uses the new value. - Known limitation: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. Unlike `map_from_arrays`, this expression always calls `build()`, so Spark stores the normalized key and returns `+0.0` for a `-0.0` key where Comet returns `-0.0`. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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`. diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index e61f1ac5fb6..8d888dc09bc 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -82,7 +82,6 @@ class CometExecIterator( shuffleBlockIterators: Map[Int, CometShuffleBlockIterator] = Map.empty, taskFilePaths: Seq[String] = Seq.empty, shufflePartitionPusher: Option[ShufflePartitionPusher] = None, - mapKeyDedupPolicy: Option[String] = None, capturePartitionOffsets: Boolean = false) extends Iterator[ColumnarBatch] with Logging { @@ -100,7 +99,7 @@ class CometExecIterator( val localDiskDirs = SparkEnv.get.blockManager.getLocalDiskDirs // serialize Comet related Spark configs in protobuf format - val protobufSparkConfigs = CometExecIterator.serializeCometSQLConfs(mapKeyDedupPolicy) + val protobufSparkConfigs = CometExecIterator.serializeCometSQLConfs() // Create keyUnwrapper if encryption is enabled val keyUnwrapper = if (encryptedFilePaths.nonEmpty) { @@ -366,13 +365,7 @@ object CometExecIterator extends Logging { private def cometSqlConfs: Map[String, String] = SQLConf.get.getAllConfs.filter(_._1.startsWith(CometConf.COMET_PREFIX)) - /** - * @param mapKeyDedupPolicy - * the `spark.sql.mapKeyDedupPolicy` the plan captured when it was first executed, or `None` - * to read the current value. See - * [[org.apache.spark.sql.comet.CometNativeExec.mapKeyDedupPolicy]]. - */ - def serializeCometSQLConfs(mapKeyDedupPolicy: Option[String] = None): Array[Byte] = { + def serializeCometSQLConfs(): Array[Byte] = { val builder = ConfigMap.newBuilder() cometSqlConfs.foreach { case (k, v) => if (k.startsWith(s"${CometConf.COMET_PREFIX}.datafusion.")) { @@ -397,12 +390,12 @@ object CometExecIterator extends Logging { // 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`. Spark's `ArrayBasedMapBuilder` reads it once, when - // the expression is first evaluated, so the plan captures it then and passes it in here rather - // than letting every native iterator read whatever the session holds at the time. + // `datafusion.spark.map_key_dedup_policy`. Read here, when the native plan for a task is + // built, which is where Spark's `ArrayBasedMapBuilder` reads it for a projection outside + // whole-stage codegen. See the note on `map_from_arrays` in the map_funcs expression audit. builder.putEntries( SQLConf.MAP_KEY_DEDUP_POLICY.key, - mapKeyDedupPolicy.getOrElse(SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString)) + SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString) builder.build().toByteArray } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index ad0322a65e2..1d876dfb83f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -67,8 +67,7 @@ private[spark] class CometExecRDD( broadcastedHadoopConfForEncryption: Option[Broadcast[SerializableConfiguration]] = None, encryptedFilePaths: Seq[String] = Seq.empty, shuffleScanIndices: Set[Int] = Set.empty, - @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty, - mapKeyDedupPolicy: Option[String] = None) + @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty) extends RDD[ColumnarBatch](sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { // Determine partition count: from inputs if available, otherwise from parameter @@ -137,8 +136,7 @@ private[spark] class CometExecRDD( broadcastedHadoopConfForEncryption, encryptedFilePaths, shuffleBlockIters, - taskFilePaths = partition.filePaths, - mapKeyDedupPolicy = mapKeyDedupPolicy) + taskFilePaths = partition.filePaths) // Register ScalarSubqueries so native code can look them up subqueries.foreach(sub => CometScalarSubquery.setSubquery(it.id, sub)) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala index 56500216108..00a5047625f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala @@ -238,7 +238,6 @@ case class CometIcebergWriteExec( // column each (see `build_output_schema` in `iceberg_write.rs`). val numOutputCols = 2 val capturedNativeOp = nativeOp - val capturedMapKeyDedupPolicy = mapKeyDedupPolicy childRDD.mapPartitionsInternal { iter => val partitionId = TaskContext.getPartitionId() @@ -278,8 +277,7 @@ case class CometIcebergWriteExec( numPartitions, partitionId, None, - Seq.empty, - mapKeyDedupPolicy = Some(capturedMapKeyDedupPolicy)) + Seq.empty) execIterator } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala index 424604416a0..f0d10b17667 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala @@ -180,7 +180,6 @@ case class CometNativeWriteExec( val capturedCommitter = committer val capturedJobTrackerID = jobTrackerID val capturedNativeOp = nativeOp - val capturedMapKeyDedupPolicy = mapKeyDedupPolicy val capturedAccumulator = taskCommitMessagesAccum // Capture accumulator for use in tasks // Execute native write operation with task-level commit protocol @@ -243,8 +242,7 @@ case class CometNativeWriteExec( numPartitions, partitionId, None, - Seq.empty, - mapKeyDedupPolicy = Some(capturedMapKeyDedupPolicy)) + Seq.empty) // Wrap the iterator to handle task commit/abort and capture TaskCommitMessage new Iterator[ColumnarBatch] { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index 0503716b88a..fce4291deb6 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -196,7 +196,6 @@ class CometNativeShuffleWriter[K, V]( ctx.encryptedFilePaths, shuffleBlockIters, shufflePartitionPusher = remoteDestination.map(_.callback), - mapKeyDedupPolicy = Some(ctx.mapKeyDedupPolicy), // Only a local destination publishes partition offsets; RSS reports lengths through its // pusher instead. capturePartitionOffsets = localOutput.isDefined) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 824ba25d473..fe1a2e637ad 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -44,7 +44,6 @@ import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregat import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.unsafe.Platform @@ -768,8 +767,7 @@ private[comet] case class NativeExecContext( // binary when this context rides on the non-transient CometShuffleDependency.nativeShuffleSpec. @transient perPartitionByKey: Map[String, Array[Array[Byte]]], shuffleScanIndices: Set[Int], - hasScanInput: Boolean, - mapKeyDedupPolicy: String = SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString) { + hasScanInput: Boolean) { // Catch shape divergence (e.g. broadcast scans with different partition counts after DPP // filtering) at construction so consumers don't trip ArrayIndexOutOfBoundsException at // partition idx access time. @@ -793,22 +791,6 @@ abstract class CometNativeExec extends CometExec { /** The Comet native operator */ def nativeOp: Operator - /** - * The `spark.sql.mapKeyDedupPolicy` the native map constructors in this plan build their maps - * with, read the first time the plan is executed and kept from then on. - * - * Spark's `ArrayBasedMapBuilder` is a lazy field of the map expression, so it reads the policy - * when the expression is first evaluated and the expression keeps that builder for every later - * action. Neither materializing nor explaining a plan evaluates anything, so a Dataset - * explained under one policy and then executed under another builds its maps under the second - * one, and a Dataset executed twice across a change keeps the first one. A `lazy val` on this - * node gives the same two properties: it is forced by the first `doExecuteColumnar`, which - * `explain` does not reach, and the node lives in the cached `executedPlan`, so later actions - * reuse the value. - */ - private[comet] lazy val mapKeyDedupPolicy: String = - SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString - override protected def doPrepare(): Unit = prepareSubqueries(this) override lazy val metrics: Map[String, SQLMetric] = @@ -864,8 +846,7 @@ abstract class CometNativeExec extends CometExec { ctx.subqueries, ctx.broadcastedHadoopConfForEncryption, ctx.encryptedFilePaths, - ctx.shuffleScanIndices, - mapKeyDedupPolicy = Some(ctx.mapKeyDedupPolicy)) { + ctx.shuffleScanIndices) { override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { val res = super.compute(split, context) if (ctx.hasScanInput) { @@ -1066,8 +1047,7 @@ abstract class CometNativeExec extends CometExec { commonByKey = commonByKey, perPartitionByKey = perPartitionByKey, shuffleScanIndices = shuffleScanIndices, - hasScanInput = sparkPlans.exists(_.isInstanceOf[CometNativeScanExec]), - mapKeyDedupPolicy = mapKeyDedupPolicy) + hasScanInput = sparkPlans.exists(_.isInstanceOf[CometNativeScanExec])) } /** diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index b3f23fca356..13a19599f79 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -242,33 +242,50 @@ class CometMapExpressionSuite extends CometTestBase { // converts the plan, which the Dataset reuses across actions, so both engines keep it; reading // the setting again for every native iterator would apply the new one instead. // https://github.com/apache/datafusion-comet/pull/5854#discussion_r4049790875 - test("map constructors keep the dedup policy of an executed Dataset") { - withTempPath { dir => - val path = dir.getCanonicalPath - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - spark.range(0, 1, 1, 1).write.parquet(path) - } - def query(): DataFrame = mapPolicyQuery(path) - val lastWin = mapPolicyLastWin - for (cometEnabled <- Seq("false", "true")) { - withSQLConf(CometConf.COMET_ENABLED.key -> cometEnabled) { - // Executed under LAST_WIN, the Dataset keeps that policy once EXCEPTION is set. - withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { - val df = query() - checkAnswer(df, lastWin) - if (cometEnabled == "true") { - checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + // Spark's `ArrayBasedMapBuilder` is a lazy field of the map expression, so it reads + // `spark.sql.mapKeyDedupPolicy` the first time the expression is evaluated. Outside whole-stage + // codegen the projection is rebuilt in every task, so that happens again on each action and a + // Dataset re-executed after the setting changed builds its maps under the new policy. Comet + // reads the setting when it builds the native plan for a task, which lands in the same place. + // Inside whole-stage codegen Spark instead creates the builder once, on the driver, and keeps + // it; Comet cannot tell the two apart, because it replaces the operator before + // `CollapseCodegenStages` runs. That one divergence is recorded in the map_funcs expression + // audit. + // https://github.com/apache/datafusion-comet/pull/5854#issuecomment-5745846643 + test("map constructors follow a dedup policy change between actions") { + // AQE converts each query stage as it runs, which hides when the setting is read. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { dir => + val path = dir.getCanonicalPath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 1, 1, 1).write.parquet(path) + } + // The two ways a projection runs outside whole-stage codegen: the flag is off, or the + // projection is wider than `spark.sql.codegen.maxFields`. + val outsideWholeStageCodegen = Seq( + (Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false"), Seq.empty[String]), + (Seq.empty[(String, String)], (1 to 100).map(i => s"id + $i AS c$i"))) + for ((codegenConf, padding) <- outsideWholeStageCodegen; + cometEnabled <- Seq("false", "true")) { + withSQLConf((codegenConf :+ (CometConf.COMET_ENABLED.key -> cometEnabled)): _*) { + // Run under LAST_WIN, then again under EXCEPTION: the duplicate is rejected. + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + val df = mapPolicyQuery(path, padding) + checkAnswer(df.select("a", "e", "s"), mapPolicyLastWin) + if (cometEnabled == "true") { + checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + } + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + assertDuplicateMapKey(df) + } } + // Run under EXCEPTION, then again under LAST_WIN: the last value wins. withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { - checkAnswer(df, lastWin) - } - } - // Executed under EXCEPTION, the Dataset keeps raising once LAST_WIN is set. - withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { - val df = query() - assertDuplicateMapKey(df) - withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + val df = mapPolicyQuery(path, padding) assertDuplicateMapKey(df) + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + checkAnswer(df.select("a", "e", "s"), mapPolicyLastWin) + } } } } @@ -276,14 +293,12 @@ class CometMapExpressionSuite extends CometTestBase { } } - // Spark initializes `ArrayBasedMapBuilder`, and with it reads the policy, when the expression is - // first evaluated. Materializing the plan does not evaluate anything, so a Dataset explained - // under one policy and then executed under another builds its maps under the second one. Comet - // converts its plan when `executedPlan` is materialized, which `explain()` also triggers, so the - // policy cannot be read there. + // Materializing a plan evaluates nothing, so it must not fix the policy in either engine: a + // Dataset explained under one policy and first executed under another builds its maps under the + // second one. Comet converts its plan when `executedPlan` is materialized, which `explain()` + // also triggers, so the setting cannot be read there. // https://github.com/apache/datafusion-comet/pull/5854#issuecomment-5744116922 - test("map constructors capture the dedup policy at first execution, not at planning") { - // AQE converts each query stage as it runs, which hides the difference between the two moments. + test("map constructors do not fix the dedup policy when the plan is materialized") { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { withTempPath { dir => val path = dir.getCanonicalPath @@ -292,8 +307,7 @@ class CometMapExpressionSuite extends CometTestBase { } for (cometEnabled <- Seq("false", "true")) { withSQLConf(CometConf.COMET_ENABLED.key -> cometEnabled) { - // Explained under EXCEPTION, first executed under LAST_WIN: LAST_WIN builds the maps, - // and keeps doing so once the setting goes back. + // Materialized under EXCEPTION, first executed under LAST_WIN: LAST_WIN builds them. withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { val df = mapPolicyQuery(path) val plan = df.queryExecution.executedPlan @@ -301,14 +315,10 @@ class CometMapExpressionSuite extends CometTestBase { checkCometOperators(stripAQEPlan(plan)) } withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { - checkAnswer(df, mapPolicyLastWin) - withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { - checkAnswer(df, mapPolicyLastWin) - } + checkAnswer(df.select("a", "e", "s"), mapPolicyLastWin) } } - // Explained under LAST_WIN, first executed under EXCEPTION: EXCEPTION rejects the - // duplicate. + // Materialized under LAST_WIN, first executed under EXCEPTION: EXCEPTION rejects it. withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { val df = mapPolicyQuery(path) df.queryExecution.executedPlan @@ -322,13 +332,18 @@ class CometMapExpressionSuite extends CometTestBase { } } - /** One row whose three map constructors each see the duplicate key `0`. */ - private def mapPolicyQuery(path: String): DataFrame = spark.read - .parquet(path) - .selectExpr( - "map_from_arrays(array(id, id), array(1, 2)) AS a", - "map_from_entries(array(struct(id, 1), struct(id, 2))) AS e", - "str_to_map(concat(CAST(id AS STRING), ':1,', CAST(id AS STRING), ':2')) AS s") + /** + * One row whose three map constructors each see the duplicate key `0`, with `padding` extra + * columns for callers that need the projection to exceed `spark.sql.codegen.maxFields`. + */ + private def mapPolicyQuery(path: String, padding: Seq[String] = Seq.empty): DataFrame = + spark.read + .parquet(path) + .selectExpr(Seq( + "map_from_arrays(array(id, id), array(1, 2)) AS a", + "map_from_entries(array(struct(id, 1), struct(id, 2))) AS e", + "str_to_map(concat(CAST(id AS STRING), ':1,', CAST(id AS STRING), ':2')) AS s") ++ + padding: _*) private def mapPolicyLastWin: Seq[Row] = Seq(Row(Map(0L -> 2), Map(0L -> 2), Map("0" -> "2"))) From 88c877b1501ddbefb770ab735e423721d8701fef Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Sun, 20 Sep 2026 14:34:52 +0000 Subject: [PATCH 16/16] docs: use prettier emphasis style in the map_funcs audit The Preflight job runs `prettier --check "**/*.md"`, which normalizes emphasis to underscores. --- docs/source/contributor-guide/expression-audits/map_funcs.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index 840a34758a9..5688cf8a587 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -51,7 +51,7 @@ - `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`, forwarded 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). `CometExecIterator.serializeCometSQLConfs` reads the setting when it builds the native plan for a task, so materializing or explaining a plan does not fix it and a Dataset re-executed after a change to the setting uses the new value. - Known limitation: on Spark 4.0+, `ArrayBasedMapBuilder` normalizes a floating-point key before comparing it (`keyNormalizer`, added in 4.0 with `spark.sql.legacy.disableMapKeyNormalization`), so `-0.0` and `+0.0` are one key and all `NaN`s are one key; the native builder compares the raw Arrow values and keeps them apart. `from` returns the input arrays untouched when no key repeated, so the stored keys match Spark either way and only duplicate detection diverges. Spark 3.4 and 3.5 do not normalize, so they already match. Gated 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. -- Known limitation: Spark reads `spark.sql.mapKeyDedupPolicy` into `ArrayBasedMapBuilder`, a lazy field of the map expression, so *when* it reads it depends on how the projection runs. Outside whole-stage codegen (the flag off, or a projection wider than `spark.sql.codegen.maxFields`) the projection is rebuilt in every task and the setting is read again on each action, which is what Comet does. Inside whole-stage codegen Spark creates the builder once on the driver, in the first action, and keeps it, so a Dataset re-executed after a change to the setting still builds its maps under the policy it started with, where Comet uses the new one. Comet cannot tell the two apart: it replaces the operator before `CollapseCodegenStages` runs, so the plan it sees carries no record of which path Spark would have taken. Matching the whole-stage case instead would mean returning a map where Spark raises `DUPLICATED_MAP_KEY` in the other three configurations, so the loud divergence is preferred over the silent one. Only a Dataset that is executed more than once across a change to the setting is affected. +- Known limitation: Spark reads `spark.sql.mapKeyDedupPolicy` into `ArrayBasedMapBuilder`, a lazy field of the map expression, so _when_ it reads it depends on how the projection runs. Outside whole-stage codegen (the flag off, or a projection wider than `spark.sql.codegen.maxFields`) the projection is rebuilt in every task and the setting is read again on each action, which is what Comet does. Inside whole-stage codegen Spark creates the builder once on the driver, in the first action, and keeps it, so a Dataset re-executed after a change to the setting still builds its maps under the policy it started with, where Comet uses the new one. Comet cannot tell the two apart: it replaces the operator before `CollapseCodegenStages` runs, so the plan it sees carries no record of which path Spark would have taken. Matching the whole-stage case instead would mean returning a map where Spark raises `DUPLICATED_MAP_KEY` in the other three configurations, so the loud divergence is preferred over the silent one. Only a Dataset that is executed more than once across a change to the setting is affected. - Known limitation: the two null guards serialize each child a second time inside the `map_from_arrays` call, so a nondeterministic child such as `monotonically_increasing_id()` would advance independently in each copy and the result would drift from Spark ([#5781](https://github.com/apache/datafusion-comet/issues/5781)). `CometMapFromArrays` declines such a child as `Unsupported` through `NullGuardSupport` and the projection falls back to Spark; [#5867](https://github.com/apache/datafusion-comet/pull/5867) routes the same decline through the JVM codegen dispatcher and applies it to `size`, `array_append` and `arrays_zip` as well. ## map_from_entries