diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index fc9c5584c9e..5688cf8a587 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -45,9 +45,14 @@ ## 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)). 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). `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 @@ -55,6 +60,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`, 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`. ## map_keys @@ -79,7 +86,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 faf4d5dfda0..1d4afc5f431 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -59,8 +59,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; @@ -115,7 +113,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 crate::parquet::parquet_support::CometObjectStoreRegistry; @@ -761,6 +759,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 @@ -802,7 +809,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())); @@ -810,7 +816,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 8f030da455b..8196698e95c 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3666,6 +3666,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, @@ -3792,7 +3800,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 8fe19f0aad5..bd281a7e283 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -32,8 +32,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}; @@ -341,9 +341,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/lib.rs b/native/spark-expr/src/lib.rs index 0d7675a64e0..dd1e05f32a8 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -61,7 +61,9 @@ pub mod jvm_udf; mod conditional_funcs; mod conversion_funcs; mod map_funcs; -pub use map_funcs::{spark_map_sort, SparkMapExtract}; +pub use map_funcs::{ + spark_map_sort, SparkMapExtract, 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..57e922fc15b --- /dev/null +++ b/native/spark-expr/src/map_funcs/map_builders.rs @@ -0,0 +1,976 @@ +// 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 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. + +use crate::SparkError; +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::{exec_err, DataFusionError, HashSet, Result, ScalarValue}; +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 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(&args))? + } + 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 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(&args))? + } + 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) +} + +/// 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 +/// 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` 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(()); + } + 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); + 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 + // 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 { + check_keys_in_order(&flat_keys, start, end, nulls, last_value_wins, &mut seen)?; + } + } + Ok(()) +} + +/// 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, last_value_wins: bool) -> 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(); + + 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, end) = (offsets[row], offsets[row + 1]); + if element_nulls.is_some_and(|nulls| nulls.slice(start, end - start).null_count() > 0) { + continue; + } + 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(()) +} + +/// 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, Int32Type}; + 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_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( + 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 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")])); + 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]); + // `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 + /// 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" + ); + } + + /// 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!( + 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 644466d0320..1db58917350 100644 --- a/native/spark-expr/src/map_funcs/mod.rs +++ b/native/spark-expr/src/map_funcs/mod.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. +mod map_builders; mod map_extract; mod map_sort; +pub use map_builders::{SparkMapFromArrays, SparkMapFromEntries, SparkStrToMap}; pub use map_extract::SparkMapExtract; 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 4da95af18d0..8d888dc09bc 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -388,6 +388,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`. 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, + SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY).toString) + builder.build().toByteArray } 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 51fa428b543..7ed235b7a1d 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,41 +133,103 @@ 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 { + + /** + * 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 = + "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 + * 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 = + 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(MapKeyDedupPolicySupport.incompatibleReason) + Seq(MapBuilderSupport.collationKeyReason) + + override def getUnsupportedReasons(): Seq[String] = + Seq(NullGuardSupport.nondeterministicReason) 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 = + 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 + * 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. + * + * 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, inputs: Seq[Attribute], @@ -177,35 +240,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) - mapFromArraysExprProto <- scalarFunctionExprToProto("map", keysExpr, valuesExpr) + 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 @@ -217,20 +274,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, MapBuilderSupport.collationKeyReason) 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_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 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..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 @@ -58,4 +58,28 @@ 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 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)) + +-- and in the other direction +query expect_error(must have the same length) +SELECT map_from_arrays(array('a'), array(1, 2)) 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..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 @@ -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,52 @@ 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('a', 'b', 'a'), array(1, 2, 3)), + (array('a', 'a', 'b'), array(1, NULL, 3)), + (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)) + +-- 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)) + +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/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..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 @@ -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,39 @@ 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(struct('a', 1), struct('b', 2), struct('a', 3))), + (array(struct('a', 1), struct('a', CAST(NULL AS INT)), struct('b', 3))), + (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))) + +-- 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/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 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..74abb046baa --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/str_to_map_dedup_policy.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. + +-- 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') + +-- `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 diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index ebbdce406a3..13a19599f79 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 @@ -126,6 +127,265 @@ 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")) + } + } + + // 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. 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") { + // 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")) + } + } + } + + // 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( + 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 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") + } + } + + 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)) + } + } + } + + // 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 + // 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") { + val df = mapPolicyQuery(path, padding) + assertDuplicateMapKey(df) + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + checkAnswer(df.select("a", "e", "s"), mapPolicyLastWin) + } + } + } + } + } + } + } + + // 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 do not fix the dedup policy when the plan is materialized") { + 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) { + // 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 + if (cometEnabled == "true") { + checkCometOperators(stripAQEPlan(plan)) + } + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "LAST_WIN") { + checkAnswer(df.select("a", "e", "s"), mapPolicyLastWin) + } + } + // 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 + withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> "EXCEPTION") { + assertDuplicateMapKey(df) + } + } + } + } + } + } + } + + /** + * 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"))) + + 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( + 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") {