From 4f6a1cee80e1491e5d38565e1e32789382967cae Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 9 Sep 2026 08:12:52 -0600 Subject: [PATCH 1/3] perf: vectorize the native map lookup behind element_at and GetMapValue `GetMapValue` and `element_at(, key)` both serialise to the DataFusion `map_extract` UDF, and the planner then unwrapped its one-element list with a second `ListExtract` pass. `general_map_extract_inner` re-slices the query key and every candidate key into a fresh `ArrayRef` per comparison and compares them through `dyn Array` equality, which made the lookup roughly 35x more expensive than any other Comet map kernel and slower than Spark itself. Add `SparkMapExtract`, registered under the same `map_extract` name so it overrides the DataFusion one. It runs a single Arrow `eq` over the batch's map entries, scans the resulting bitmask for each row's first match, and gathers the values with one `take`. It also returns the value directly rather than a one-element list, so the `ListExtract` wrapper in the planner goes away. Semantics are unchanged: first matching entry wins, and a missing key, a NULL map row and a NULL lookup key all yield NULL. Key types whose Spark equality the native lookup cannot reproduce are still declined by `MapKeySupport`. `eq` rejects nested key types, so an element-wise comparison remains as a backstop for those. Fixes #5795 --- native/core/src/execution/planner.rs | 11 - native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/map_extract.rs | 131 ++++ native/spark-expr/src/comet_scalar_funcs.rs | 5 +- native/spark-expr/src/lib.rs | 2 +- .../spark-expr/src/map_funcs/map_extract.rs | 566 ++++++++++++++++++ native/spark-expr/src/map_funcs/mod.rs | 2 + .../comet/CometMapExpressionSuite.scala | 28 + 8 files changed, 736 insertions(+), 13 deletions(-) create mode 100644 native/spark-expr/benches/map_extract.rs create mode 100644 native/spark-expr/src/map_funcs/map_extract.rs diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index fe01bdcfeec..67261cb82e5 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -708,17 +708,6 @@ impl PhysicalPlanner { ExprStruct::ScalarFunc(expr) => { let func = self.create_scalar_function_expr(expr, input_schema); match expr.func.as_ref() { - // DataFusion map_extract returns array of struct entries even if lookup by key - // Apache Spark wants a single value, so wrap the result into additional list extraction - "map_extract" => Ok(Arc::new(ListExtract::new( - func?, - Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), - None, - true, - false, - None, // No expr_id for internal map_extract wrapper - Arc::clone(&self.query_context_registry), - ))), // DataFusion 49 hardcodes return type for MD5 built in function as UTF8View // which is not yet supported in Comet // Converting forcibly to UTF8. To be removed after UTF8View supported diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 98e5a999044..29055b2ae19 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -215,6 +215,10 @@ harness = false name = "map_sort" harness = false +[[bench]] +name = "map_extract" +harness = false + [[bench]] name = "to_time" harness = false diff --git a/native/spark-expr/benches/map_extract.rs b/native/spark-expr/benches/map_extract.rs new file mode 100644 index 00000000000..8be11476375 --- /dev/null +++ b/native/spark-expr/benches/map_extract.rs @@ -0,0 +1,131 @@ +// 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. + +//! Benchmarks for the map lookup behind `GetMapValue` and `element_at(, key)`. +//! +//! Each shape is run against both Comet's `SparkMapExtract` and the +//! `datafusion-functions-nested` `map_extract` it overrides, so the gap that motivated +//! stays visible. + +use arrow::array::builder::{MapBuilder, StringBuilder}; +use arrow::array::{ArrayRef, MapFieldNames, StringArray}; +use arrow::datatypes::{DataType, Field}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::common::config::ConfigOptions; +use datafusion::common::ScalarValue; +use datafusion::functions_nested::map_extract::map_extract_udf; +use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_comet_spark_expr::SparkMapExtract; +use std::hint::black_box; +use std::sync::Arc; + +const BATCH_SIZE: usize = 8192; +/// Distinct keys per map column, as in the issue's `attrs map` dataset. +const DISTINCT_KEYS: usize = 60; + +/// `BATCH_SIZE` rows of `map`, every tenth row NULL, each non-null row holding +/// `entries_per_map` entries drawn from `DISTINCT_KEYS` keys. The stride is coprime with +/// `DISTINCT_KEYS` so a given lookup key lands at a different entry position in every row rather +/// than always being found (or missed) at the same depth. +fn string_map(entries_per_map: usize) -> ArrayRef { + let mut builder = MapBuilder::new( + Some(MapFieldNames { + entry: "entries".into(), + key: "key".into(), + value: "value".into(), + }), + StringBuilder::new(), + StringBuilder::new(), + ); + for row in 0..BATCH_SIZE { + if row % 10 == 0 { + builder.append(false).unwrap(); + continue; + } + for entry in 0..entries_per_map { + builder + .keys() + .append_value(format!("a{}", (row * 13 + entry * 7) % DISTINCT_KEYS)); + builder.values().append_value(format!("v{}", row % 400)); + } + builder.append(true).unwrap(); + } + Arc::new(builder.finish()) +} + +/// One lookup key per row, so the key cannot be hoisted out of the comparison. +fn per_row_keys() -> ArrayRef { + Arc::new(StringArray::from_iter_values( + (0..BATCH_SIZE).map(|row| format!("a{}", (row * 29 + 11) % DISTINCT_KEYS)), + )) +} + +fn call(udf: &dyn ScalarUDFImpl, args: &[ColumnarValue]) { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args.to_vec(), + arg_fields: vec![], + number_rows: BATCH_SIZE, + return_field: Arc::new(Field::new("result", DataType::Utf8, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap(), + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + let comet = SparkMapExtract::new(); + let datafusion = map_extract_udf(); + let mut group = c.benchmark_group("map_extract"); + + for entries in [2usize, 8, 32] { + let map = string_map(entries); + let cases: [(&str, Vec); 2] = [ + ( + "constant_key", + vec![ + ColumnarValue::Array(Arc::clone(&map)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("a1".to_string()))), + ], + ), + ( + "per_row_key", + vec![ + ColumnarValue::Array(Arc::clone(&map)), + ColumnarValue::Array(per_row_keys()), + ], + ), + ]; + for (case, args) in cases { + group.bench_with_input( + BenchmarkId::new(format!("comet/{case}"), entries), + &args, + |b, args| b.iter(|| call(&comet, args)), + ); + group.bench_with_input( + BenchmarkId::new(format!("datafusion/{case}"), entries), + &args, + |b, args| b.iter(|| call(datafusion.inner().as_ref(), args)), + ); + } + } + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index b5820144ea6..171a8f5c53f 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -31,7 +31,7 @@ use crate::{ EvalMode, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, - SparkMakeTime, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, + SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, }; use arrow::datatypes::DataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; @@ -321,6 +321,9 @@ fn all_scalar_functions() -> Vec> { )), Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::default())), Arc::new(ScalarUDF::new_from_impl(SparkMakeTime::default())), + // Overrides datafusion-functions-nested' `map_extract` with a vectorized lookup that + // returns the value itself rather than a one-element list (#5795). + Arc::new(ScalarUDF::new_from_impl(SparkMapExtract::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())), diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 758f8ee3c90..ef3dd58efb4 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -61,7 +61,7 @@ pub mod jvm_udf; mod conditional_funcs; mod conversion_funcs; mod map_funcs; -pub use map_funcs::spark_map_sort; +pub use map_funcs::{spark_map_extract, spark_map_sort, SparkMapExtract}; mod math_funcs; mod nondetermenistic_funcs; pub mod url_funcs; diff --git a/native/spark-expr/src/map_funcs/map_extract.rs b/native/spark-expr/src/map_funcs/map_extract.rs new file mode 100644 index 00000000000..26d3d522147 --- /dev/null +++ b/native/spark-expr/src/map_funcs/map_extract.rs @@ -0,0 +1,566 @@ +// 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. + +use arrow::array::{ + new_null_array, Array, ArrayRef, BooleanBufferBuilder, MapArray, NullBufferBuilder, Scalar, + UInt32Array, +}; +use arrow::buffer::BooleanBuffer; +use arrow::compute::kernels::cmp::eq; +use arrow::compute::take; +use arrow::datatypes::{DataType, FieldRef}; +use datafusion::common::utils::take_function_args; +use datafusion::common::{exec_err, Result as DataFusionResult}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use std::sync::Arc; + +/// Spark's map lookup: `GetMapValue` (`m[k]`) and `element_at(, k)`. +/// +/// Overrides DataFusion's `map_extract` under the same name, and differs from it in two ways: +/// +/// - it returns the matched **value** rather than a one-element list, so the planner does not +/// have to unwrap the list with a second `ListExtract` pass (see `planner.rs`); +/// - the lookup is vectorized. DataFusion's `general_map_extract_inner` re-slices the query key +/// and every candidate key into a fresh `ArrayRef` per comparison and compares them through +/// `dyn Array` equality, which made a constant-key lookup roughly 35x more expensive than any +/// other Comet map kernel and slower than Spark itself +/// ([#5795](https://github.com/apache/datafusion-comet/issues/5795)). Here a single Arrow +/// `eq` covers the whole batch of entries at once, the per-row work is a bit scan over the +/// resulting mask, and the values are gathered with one `take`. +/// +/// Spark's own lookup returns the first entry whose key compares equal, so the mask scan stops at +/// the first match too. A missing key, a `NULL` map row, and a `NULL` lookup key all produce +/// `NULL`, matching `GetMapValueUtil.getValueEval` and `ElementAt`'s map overload. +/// +/// Key types whose Spark equality this cannot reproduce (floating point, non-default collations, +/// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala` declines them so the +/// expression falls back to Spark. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkMapExtract { + signature: Signature, +} + +impl Default for SparkMapExtract { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapExtract { + pub fn new() -> Self { + Self { + // `user_defined` so `coerce_types` runs and casts the lookup key to the map's key + // type; Comet's planner applies that coercion to the argument expression. + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkMapExtract { + fn name(&self) -> &str { + "map_extract" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult { + let [map_type, _] = take_function_args(self.name(), arg_types)?; + Ok(map_entry_fields(map_type)?.1.data_type().clone()) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> DataFusionResult> { + let [map_type, _] = take_function_args(self.name(), arg_types)?; + Ok(vec![ + map_type.clone(), + map_entry_fields(map_type)?.0.data_type().clone(), + ]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult { + let [map_arg, key_arg] = take_function_args(self.name(), &args.args)?; + spark_map_extract(map_arg, key_arg, args.number_rows) + } +} + +/// The `(key, value)` fields of a `Map`'s entry struct. +fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef, &FieldRef)> { + match map_type { + DataType::Map(entries, _) => match entries.data_type() { + DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0], &fields[1])), + other => exec_err!("map_extract: map entries must be a two-field struct, got {other}"), + }, + other => exec_err!("map_extract: the first argument must be a map, got {other}"), + } +} + +/// Look up `key_arg` in each row of `map_arg`, returning the matched value or `NULL`. +pub fn spark_map_extract( + map_arg: &ColumnarValue, + key_arg: &ColumnarValue, + number_rows: usize, +) -> DataFusionResult { + let map_ref: ArrayRef = match map_arg { + ColumnarValue::Array(array) => Arc::clone(array), + ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?, + }; + let Some(map_array) = map_ref.as_any().downcast_ref::() else { + return exec_err!( + "map_extract: the first argument must be a map, got {}", + map_ref.data_type() + ); + }; + + let num_rows = map_array.len(); + let value_type = map_array.value_type(); + + // Arrow keeps a sliced `MapArray`'s entries child intact and slices only the offsets, so the + // offsets index the *unsliced* keys/values and the visible entries are the half-open range + // [entries_start, entries_end). Comparing only that window keeps a native OFFSET from paying + // for the entries it skipped. + let offsets = map_array.offsets(); + let entries_start = offsets[0] as usize; + let entries_end = offsets[num_rows] as usize; + if entries_start == entries_end { + // Every row is empty or NULL, so nothing can match. + return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows))); + } + let window_len = entries_end - entries_start; + let keys = map_array.keys().slice(entries_start, window_len); + + let matched = match key_arg { + ColumnarValue::Scalar(scalar) => { + if scalar.is_null() { + // Spark map keys are never NULL, so a NULL lookup key matches nothing. + return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows))); + } + let key = scalar.to_array_of_size(1)?; + key_match_mask(&keys, &key, true)? + } + ColumnarValue::Array(key_array) => { + if key_array.len() != num_rows { + return exec_err!( + "map_extract: expected {num_rows} lookup keys, got {}", + key_array.len() + ); + } + // One vectorized compare needs a lookup key per *entry*, not per row, so gather each + // row's key across that row's entries. Entries in a gap between two rows (offsets are + // only required to be monotonic) keep index 0; the per-row scan below never reads + // those positions. + let mut gather = vec![0u32; window_len]; + for row in 0..num_rows { + let start = offsets[row] as usize - entries_start; + let end = offsets[row + 1] as usize - entries_start; + gather[start..end].fill(row as u32); + } + let per_entry_key = take(key_array, &UInt32Array::from(gather), None)?; + key_match_mask(&keys, &per_entry_key, false)? + } + }; + + // Gather the first matching entry of each row. Map offsets are `i32`, so an entry index always + // fits in `u32`. + let mut indices = vec![0u32; num_rows]; + let mut nulls = NullBufferBuilder::new(num_rows); + for row in 0..num_rows { + let start = offsets[row] as usize - entries_start; + let end = offsets[row + 1] as usize - entries_start; + let found = (start..end).find(|&i| matched.value(i)); + if let Some(i) = found { + indices[row] = (i + entries_start) as u32; + } + nulls.append(found.is_some()); + } + let indices = UInt32Array::new(indices.into(), nulls.finish()); + + Ok(ColumnarValue::Array(take( + map_array.values(), + &indices, + None, + )?)) +} + +/// A bit per map entry: set where the stored key equals the lookup key. `lookup` is either a +/// length-1 array broadcast over every entry (constant key) or one key per entry. +fn key_match_mask( + keys: &ArrayRef, + lookup: &ArrayRef, + lookup_is_scalar: bool, +) -> DataFusionResult { + // The planner casts the lookup key to the map's declared key type, so a mismatch here means + // the runtime encoding is not the declared one (a dictionary-encoded key column, say). Reject + // it rather than comparing incomparable encodings and reporting every row as a miss. + if keys.data_type() != lookup.data_type() { + return exec_err!( + "map_extract: lookup key type {} does not match the map key type {}", + lookup.data_type(), + keys.data_type() + ); + } + let compared = if lookup_is_scalar { + eq(keys, &Scalar::new(Arc::clone(lookup))) + } else { + eq(keys, lookup) + }; + match compared { + Ok(mask) => { + // A NULL on either side compares as NULL, which is not a match. + let (values, nulls) = mask.into_parts(); + Ok(match nulls { + Some(nulls) if nulls.null_count() > 0 => &values & nulls.inner(), + _ => values, + }) + } + // `eq` rejects nested key types. `MapKeySupport` declines those before they reach the + // native lookup, but keep DataFusion's element-wise comparison as a backstop so this + // kernel is never less capable than the one it replaces. + Err(_) => Ok(elementwise_match_mask(keys, lookup, lookup_is_scalar)), + } +} + +fn elementwise_match_mask( + keys: &ArrayRef, + lookup: &ArrayRef, + lookup_is_scalar: bool, +) -> BooleanBuffer { + let mut builder = BooleanBufferBuilder::new(keys.len()); + for i in 0..keys.len() { + let lookup_row = if lookup_is_scalar { 0 } else { i }; + builder.append( + !lookup.is_null(lookup_row) + && keys.slice(i, 1).as_ref() == lookup.slice(lookup_row, 1).as_ref(), + ); + } + builder.finish() +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int32Array, StringArray, StructArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Fields}; + use datafusion::common::ScalarValue; + + /// One row of a test map: `None` for a NULL map, otherwise its `(key, value)` entries. + type MapRow<'a> = Option)>>; + + /// `{"a": 1, "b": 2}`, `{}`, `{"c": 3, "a": 30}`, NULL, `{"b": NULL}` + fn test_map() -> MapArray { + map_from(vec![ + Some(vec![("a", Some(1)), ("b", Some(2))]), + Some(vec![]), + Some(vec![("c", Some(3)), ("a", Some(30))]), + None, + Some(vec![("b", None)]), + ]) + } + + fn map_from(rows: Vec) -> MapArray { + let mut keys = Vec::new(); + let mut values = Vec::new(); + let mut offsets = vec![0i32]; + let mut nulls = NullBufferBuilder::new(rows.len()); + for row in &rows { + match row { + Some(entries) => { + for (k, v) in entries { + keys.push(*k); + values.push(*v); + } + nulls.append(true); + } + None => nulls.append(false), + } + offsets.push(keys.len() as i32); + } + + let key_field = Arc::new(Field::new("key", DataType::Utf8, false)); + let value_field = Arc::new(Field::new("value", DataType::Int32, true)); + let entries = StructArray::new( + Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), + vec![ + Arc::new(StringArray::from(keys)) as ArrayRef, + Arc::new(Int32Array::from(values)) as ArrayRef, + ], + None, + ); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![key_field, value_field])), + false, + )); + MapArray::try_new( + entries_field, + OffsetBuffer::new(offsets.into()), + entries, + nulls.finish(), + false, + ) + .unwrap() + } + + fn extract(map: MapArray, key: ColumnarValue) -> Vec> { + let num_rows = map.len(); + let result = spark_map_extract(&ColumnarValue::Array(Arc::new(map)), &key, num_rows) + .unwrap() + .into_array(num_rows) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + (0..result.len()) + .map(|i| (!result.is_null(i)).then(|| result.value(i))) + .collect() + } + + fn key(value: &str) -> ColumnarValue { + ColumnarValue::Scalar(ScalarValue::Utf8(Some(value.to_string()))) + } + + #[test] + fn constant_key_hit_and_miss() { + // A found key returns its value; an empty row, a NULL row, and a row without the key all + // return NULL, as does a row whose stored value is NULL. + assert_eq!( + extract(test_map(), key("a")), + vec![Some(1), None, Some(30), None, None] + ); + assert_eq!( + extract(test_map(), key("b")), + vec![Some(2), None, None, None, None] + ); + assert_eq!(extract(test_map(), key("zz")), vec![None; 5]); + } + + #[test] + fn duplicate_keys_return_the_first_match() { + // Spark's `GetMapValueUtil` scans entries in order and stops at the first equal key, so a + // map that kept duplicates (EXCEPTION dedup is a write-side check) resolves to the first. + let map = map_from(vec![Some(vec![("a", Some(1)), ("a", Some(2))])]); + assert_eq!(extract(map, key("a")), vec![Some(1)]); + } + + #[test] + fn null_lookup_key_matches_nothing() { + assert_eq!( + extract(test_map(), ColumnarValue::Scalar(ScalarValue::Utf8(None))), + vec![None; 5] + ); + } + + #[test] + fn per_row_lookup_key() { + // A different key per row, including a NULL key and a key looked up against a NULL map. + let keys: ArrayRef = Arc::new(StringArray::from(vec![ + Some("b"), + Some("a"), + Some("c"), + Some("a"), + None, + ])); + assert_eq!( + extract(test_map(), ColumnarValue::Array(keys)), + vec![Some(2), None, Some(3), None, None] + ); + } + + #[test] + fn sliced_map_keeps_original_entry_offsets() { + // Arrow slices a MapArray's offsets but not its entries, so the visible rows start part + // way into the keys/values children. Reading the entries from index 0 would look up the + // wrong rows. + let map = test_map().slice(2, 3); + assert_eq!(extract(map.clone(), key("a")), vec![Some(30), None, None]); + assert_eq!(extract(map.clone(), key("c")), vec![Some(3), None, None]); + assert_eq!(extract(map, key("b")), vec![None, None, None]); + + let keys: ArrayRef = Arc::new(StringArray::from(vec![Some("c"), Some("c"), Some("b")])); + assert_eq!( + extract(test_map().slice(2, 3), ColumnarValue::Array(keys)), + vec![Some(3), None, None] + ); + } + + #[test] + fn all_rows_empty_or_null() { + // The no-entries fast path still has to produce one NULL per row, typed as the value type. + let map = map_from(vec![Some(vec![]), None, Some(vec![])]); + assert_eq!(extract(map, key("a")), vec![None; 3]); + } + + #[test] + fn empty_input() { + let map = map_from(vec![]); + assert_eq!(extract(map, key("a")), Vec::>::new()); + } + + #[test] + fn non_string_keys() { + // Integer keys go down the same vectorized compare; pin that the gathered value lines up + // with the matching key rather than the key's position. + let key_field = Arc::new(Field::new("key", DataType::Int32, false)); + let value_field = Arc::new(Field::new("value", DataType::Utf8, true)); + let entries = StructArray::new( + Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, + Arc::new(StringArray::from(vec!["x", "y", "z"])) as ArrayRef, + ], + None, + ); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![key_field, value_field])), + false, + )); + let map = MapArray::try_new( + entries_field, + OffsetBuffer::new(vec![0i32, 2, 3].into()), + entries, + None, + false, + ) + .unwrap(); + + let result = spark_map_extract( + &ColumnarValue::Array(Arc::new(map)), + &ColumnarValue::Scalar(ScalarValue::Int32(Some(20))), + 2, + ) + .unwrap() + .into_array(2) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.value(0), "y"); + assert!(result.is_null(1)); + } + + #[test] + fn scalar_map_argument() { + // A constant map is expanded to the batch length; the lookup still runs per row. + let map = map_from(vec![Some(vec![("a", Some(7))])]); + let scalar = ColumnarValue::Scalar(ScalarValue::Map(Arc::new(map))); + let result = spark_map_extract(&scalar, &key("a"), 3) + .unwrap() + .into_array(3) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.len(), 3); + assert!((0..3).all(|i| result.value(i) == 7)); + } + + #[test] + fn return_and_coerce_types() { + let udf = SparkMapExtract::new(); + let map_type = test_map().data_type().clone(); + // The result is the map's value type, not a list of it. + assert_eq!( + udf.return_type(&[map_type.clone(), DataType::Utf8]) + .unwrap(), + DataType::Int32 + ); + // A wider lookup key is narrowed to the map's key type by the planner. + assert_eq!( + udf.coerce_types(&[map_type, DataType::LargeUtf8]).unwrap(), + vec![test_map().data_type().clone(), DataType::Utf8] + ); + } + + #[test] + fn non_map_first_argument_is_rejected() { + let udf = SparkMapExtract::new(); + assert!(udf + .return_type(&[DataType::Int32, DataType::Int32]) + .is_err()); + let err = spark_map_extract( + &ColumnarValue::Array(Arc::new(Int32Array::from(vec![1]))), + &ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + 1, + ) + .unwrap_err(); + assert!(err.to_string().contains("must be a map")); + } + + #[test] + fn mismatched_key_type_is_rejected() { + // A key the planner did not cast to the map's key type cannot be compared. Erroring keeps + // the mismatch visible instead of reporting every row as a miss. + let err = spark_map_extract( + &ColumnarValue::Array(Arc::new(test_map())), + &ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + 5, + ) + .unwrap_err(); + assert!(err.to_string().contains("does not match the map key type")); + } + + #[test] + fn nested_key_falls_back_to_elementwise_comparison() { + // `eq` refuses nested types. `MapKeySupport` keeps these on Spark, but the backstop has to + // still find the key rather than error. + let inner = Arc::new(Field::new("item", DataType::Int32, true)); + let key_values = arrow::array::ListArray::new( + Arc::clone(&inner), + OffsetBuffer::new(vec![0i32, 1, 2].into()), + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + None, + ); + let key_field = Arc::new(Field::new("key", key_values.data_type().clone(), false)); + let value_field = Arc::new(Field::new("value", DataType::Int32, true)); + let entries = StructArray::new( + Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), + vec![ + Arc::new(key_values) as ArrayRef, + Arc::new(Int32Array::from(vec![11, 22])) as ArrayRef, + ], + None, + ); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![key_field, value_field])), + false, + )); + let map = MapArray::try_new( + entries_field, + OffsetBuffer::new(vec![0i32, 2].into()), + entries, + None, + false, + ) + .unwrap(); + + let lookup: ArrayRef = Arc::new(arrow::array::ListArray::new( + inner, + OffsetBuffer::new(vec![0i32, 1].into()), + Arc::new(Int32Array::from(vec![2])) as ArrayRef, + None, + )); + let result = spark_map_extract( + &ColumnarValue::Array(Arc::new(map)), + &ColumnarValue::Array(lookup), + 1, + ) + .unwrap() + .into_array(1) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.value(0), 22); + } +} diff --git a/native/spark-expr/src/map_funcs/mod.rs b/native/spark-expr/src/map_funcs/mod.rs index 7288b847a83..d9de257d1ad 100644 --- a/native/spark-expr/src/map_funcs/mod.rs +++ b/native/spark-expr/src/map_funcs/mod.rs @@ -15,5 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod map_extract; mod map_sort; +pub use map_extract::{spark_map_extract, SparkMapExtract}; pub use map_sort::spark_map_sort; diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index f4a559b872b..58f17f2867d 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -460,6 +460,34 @@ class CometMapExpressionSuite extends CometTestBase { } } + // The native lookup compares a whole batch of map entries in one pass and then reads each row's + // window out of the resulting mask. A native OFFSET slices the batch, and Arrow keeps a sliced + // MapArray's original entry offsets, so the visible entries start part way into the keys child -- + // the same trap `mapsort` hit below. Reading the mask from index 0 would answer every row with + // some other row's entries. + test("element_at on a sliced map reads the visible entries") { + withParquetTable( + (0 until 20).map(i => (i, Map(s"a${i % 5}" -> i, "shared" -> (i * 10)))), + "t") { + checkSparkAnswerAndOperator( + "SELECT _1, element_at(_2, 'a3'), element_at(_2, 'shared') " + + "FROM (SELECT * FROM t ORDER BY _1 LIMIT 15 OFFSET 5)") + } + } + + // A lookup key that varies per row takes a different path than a constant key: the key has to be + // lined up against every entry of its own row. Rows whose key is missing, whose map is NULL, and + // whose key is NULL all have to come back NULL. + test("element_at on a map column with a per-row lookup key") { + val rows = (0 until 20).map { i => + val map = if (i % 7 == 0) null else Map(s"a${i % 5}" -> i, s"b${i % 3}" -> (i * 10)) + (if (i % 11 == 0) null else s"a${i % 6}", map) + } + withParquetTable(rows, "t") { + checkSparkAnswerAndOperator("SELECT _1, element_at(_2, _1), _2[_1] FROM t") + } + } + test("mapsort on a sliced map does not overrun the sorted entries") { // A native OFFSET slices the batch and Arrow keeps a sliced MapArray's original entry offsets, // so `mapsort` receives a map whose first entry offset is nonzero. `spark_map_sort` takes only From 7d527c148781bc07b2974d9d82a8e0659e7d69a0 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 9 Sep 2026 13:55:33 -0600 Subject: [PATCH 2/3] fix: read NULL from a NULL map row whose entries were retained The gather validity only recorded whether a match was found, so a NULL map row still returned the value of a matching entry inside its offset range. Arrow does not require a null row's range to be empty: adding a parent null mask over intact child buffers, which Comet's own struct-field helper does, leaves the entries in place, and the UDF layer supplies no validity mask. Spark returns NULL for a NULL map under both ANSI modes, and only element_at has a nullable-input guard upstream of this kernel, so GetMapValue was exposed directly. The row's map validity is now consulted before the mask scan, which also skips scanning a null row's entries. Two tests cover it, one per key path; both return Some(7) for the null row without the fix. --- .../spark-expr/src/map_funcs/map_extract.rs | 68 ++++++++++++++++++- 1 file changed, 67 insertions(+), 1 deletion(-) diff --git a/native/spark-expr/src/map_funcs/map_extract.rs b/native/spark-expr/src/map_funcs/map_extract.rs index 26d3d522147..1d95b22a154 100644 --- a/native/spark-expr/src/map_funcs/map_extract.rs +++ b/native/spark-expr/src/map_funcs/map_extract.rs @@ -178,9 +178,21 @@ pub fn spark_map_extract( // Gather the first matching entry of each row. Map offsets are `i32`, so an entry index always // fits in `u32`. + // + // A NULL map row reads NULL whatever its entries hold. Arrow does not require a null row's + // offset range to be empty, and neither Comet's struct-field helper (which adds a parent null + // mask while preserving the child buffers) nor the UDF execution layer clears those entries, so + // a null row can carry a live `a -> 7` that would otherwise match. Spark returns NULL for a + // NULL map under both ANSI modes, for `element_at` and for `GetMapValue` alike, and only + // `element_at` has a nullable-input guard upstream of this kernel. + let map_nulls = map_array.nulls(); let mut indices = vec![0u32; num_rows]; let mut nulls = NullBufferBuilder::new(num_rows); for row in 0..num_rows { + if map_nulls.is_some_and(|n| n.is_null(row)) { + nulls.append(false); + continue; + } let start = offsets[row] as usize - entries_start; let end = offsets[row + 1] as usize - entries_start; let found = (start..end).find(|&i| matched.value(i)); @@ -256,7 +268,7 @@ fn elementwise_match_mask( mod tests { use super::*; use arrow::array::{Int32Array, StringArray, StructArray}; - use arrow::buffer::OffsetBuffer; + use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{Field, Fields}; use datafusion::common::ScalarValue; @@ -365,6 +377,60 @@ mod tests { ); } + /// A NULL row whose entries were retained rather than dropped. `map_from` gives a NULL row an + /// empty offset range, which is the shape Arrow's builders produce, but nothing in the format + /// requires it: adding a parent null mask over intact child buffers leaves the entries in + /// place. Such a row must still read NULL, not the value its live entry holds. + fn null_row_with_retained_entries() -> MapArray { + let key_field = Arc::new(Field::new("key", DataType::Utf8, false)); + let value_field = Arc::new(Field::new("value", DataType::Int32, true)); + let entries = StructArray::new( + Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), + vec![ + Arc::new(StringArray::from(vec!["a", "a", "b"])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(7), Some(1), Some(2)])) as ArrayRef, + ], + None, + ); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![key_field, value_field])), + false, + )); + MapArray::try_new( + entries_field, + // Row 0 is NULL but still spans entry 0 (`a -> 7`); row 1 is a live `{a: 1, b: 2}`. + OffsetBuffer::new(vec![0i32, 1, 3].into()), + entries, + Some(NullBuffer::from(vec![false, true])), + false, + ) + .unwrap() + } + + #[test] + fn null_map_row_reads_null_even_with_live_entries() { + assert_eq!( + extract(null_row_with_retained_entries(), key("a")), + vec![None, Some(1)] + ); + assert_eq!( + extract(null_row_with_retained_entries(), key("b")), + vec![None, Some(2)] + ); + } + + #[test] + fn null_map_row_reads_null_with_a_per_row_lookup_key() { + // The per-row key path gathers a key per entry, so the NULL row's entry is compared + // against that row's own key. It must still be masked out. + let keys: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), Some("b")])); + assert_eq!( + extract(null_row_with_retained_entries(), ColumnarValue::Array(keys)), + vec![None, Some(2)] + ); + } + #[test] fn per_row_lookup_key() { // A different key per row, including a NULL key and a key looked up against a NULL map. From 68735f070c6b0565fa881aeee5b9ccb8cd6c1a43 Mon Sep 17 00:00:00 2001 From: test Date: Thu, 10 Sep 2026 11:15:22 -0600 Subject: [PATCH 3/3] refactor: address review feedback on the vectorized map lookup Argument validation no longer depends on the data. The lookup key's type and length are properties of the call, so they now run before the empty-entries fast path; previously a batch whose maps were all empty or all NULL answered NULLs while the next batch of the same query errored. The nested-key backstop dispatches on the key type rather than on an `eq` failure. With the length and type already checked, nesting is the only thing `eq` still rejects, so a catch-all `Err(_)` could only hide a genuine error behind a silent per-row throughput cliff. `SparkMapExtract` now declares the `element_at` alias DataFusion's `map_extract` declares. `register_udf` inserts one registry entry per alias, so without it the override was partial and `element_at` still resolved to the list-returning kernel. Nothing in Comet emits that name, so this is consistency, not a fix. `spark_map_extract` is no longer exported: only the UDF has a caller outside the module. The sliced-map Scala test now looks up through `_2[k]`. `element_at` on a nullable operand is wrapped in `CASE WHEN _2 IS NOT NULL` under ANSI, and DataFusion's CaseExpr runs the THEN branch through `filter_record_batch`, which compacts the entries child and resets the first offset to 0 as soon as any row is NULL -- so the old test only sliced by accident. `GetMapValue` has no such guard, and the fixture now carries NULL rows to keep that honest. Reverting the `entries_start` arithmetic still fails the test. `element_at_map.sql` covers the admitted key types the fixtures missed: boolean, tinyint, smallint, bigint, decimal, date, timestamp and timestamp_ntz. Arrow's `eq` is stricter about the exact Arrow type than the `ArrayData` equality it replaced, so a disagreement between it and `coerce_types` would surface as a query failure. A narrower decimal key is not expressible: Spark rejects it at analysis with MAP_FUNCTION_DIFF_TYPES. Also: one `map_of` helper for the four test fixtures, a note that the benchmark ratio measures pinned DataFusion 55.0.0 rather than a permanent gap, and comment corrections for the `DataType::Null` capability this drops and for the null-row-with-retained-entries shape, which no traced producer emits. --- native/spark-expr/benches/map_extract.rs | 7 + native/spark-expr/src/comet_scalar_funcs.rs | 3 +- native/spark-expr/src/lib.rs | 2 +- .../spark-expr/src/map_funcs/map_extract.rs | 281 ++++++++++-------- native/spark-expr/src/map_funcs/mod.rs | 2 +- .../expressions/map/element_at_map.sql | 54 ++++ .../comet/CometMapExpressionSuite.scala | 19 +- 7 files changed, 235 insertions(+), 133 deletions(-) diff --git a/native/spark-expr/benches/map_extract.rs b/native/spark-expr/benches/map_extract.rs index 8be11476375..1914c016f89 100644 --- a/native/spark-expr/benches/map_extract.rs +++ b/native/spark-expr/benches/map_extract.rs @@ -20,6 +20,13 @@ //! Each shape is run against both Comet's `SparkMapExtract` and the //! `datafusion-functions-nested` `map_extract` it overrides, so the gap that motivated //! stays visible. +//! +//! Read the ratio as a measurement of the pinned DataFusion 55.0.0, not as a permanent gap. +//! DataFusion main has since rewritten `general_map_extract_inner` around a single +//! `make_comparator` over the batch, so the per-comparison `ArrayRef` slicing that dominates the +//! baseline here is specific to the version Comet ships today, and the baseline arm will get much +//! faster at the next DataFusion bump. What survives that bump is the rest of the case for this +//! kernel: one `eq` plus one `take`, and the `ListExtract` unwrapping pass this removes. use arrow::array::builder::{MapBuilder, StringBuilder}; use arrow::array::{ArrayRef, MapFieldNames, StringArray}; diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index 171a8f5c53f..35aa5c6ab37 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -322,7 +322,8 @@ fn all_scalar_functions() -> Vec> { Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::default())), Arc::new(ScalarUDF::new_from_impl(SparkMakeTime::default())), // Overrides datafusion-functions-nested' `map_extract` with a vectorized lookup that - // returns the value itself rather than a one-element list (#5795). + // 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(SparkNextDay::default())), Arc::new(ScalarUDF::new_from_impl(SparkSecondsToTimestamp::default())), diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index ef3dd58efb4..052e2b57968 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -61,7 +61,7 @@ pub mod jvm_udf; mod conditional_funcs; mod conversion_funcs; mod map_funcs; -pub use map_funcs::{spark_map_extract, spark_map_sort, SparkMapExtract}; +pub use map_funcs::{spark_map_sort, SparkMapExtract}; mod math_funcs; mod nondetermenistic_funcs; pub mod url_funcs; diff --git a/native/spark-expr/src/map_funcs/map_extract.rs b/native/spark-expr/src/map_funcs/map_extract.rs index 1d95b22a154..a2a5df6ed61 100644 --- a/native/spark-expr/src/map_funcs/map_extract.rs +++ b/native/spark-expr/src/map_funcs/map_extract.rs @@ -51,9 +51,16 @@ use std::sync::Arc; /// Key types whose Spark equality this cannot reproduce (floating point, non-default collations, /// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala` declines them so the /// expression falls back to Spark. +/// +/// One capability of DataFusion's kernel is dropped deliberately: it special-cases a +/// `DataType::Null` first argument in `return_type`, `coerce_types` and its inner loop and answers +/// `NULL`, where this rejects it as not a map. Spark's `ExtractValue.apply` requires a `MapType` +/// child and `CometElementAt.getSupportLevel` declines anything that is neither an array nor a +/// map, so a `Null`-typed map argument cannot reach here from Comet. #[derive(Debug, Hash, Eq, PartialEq)] pub struct SparkMapExtract { signature: Signature, + aliases: Vec, } impl Default for SparkMapExtract { @@ -68,6 +75,7 @@ impl SparkMapExtract { // `user_defined` so `coerce_types` runs and casts the lookup key to the map's key // type; Comet's planner applies that coercion to the argument expression. signature: Signature::user_defined(Volatility::Immutable), + aliases: vec!["element_at".to_string()], } } } @@ -81,6 +89,15 @@ impl ScalarUDFImpl for SparkMapExtract { &self.signature } + /// The same alias DataFusion's `map_extract` declares. `register_udf` inserts one registry + /// entry per alias, so without this the override would be partial: `map_extract` would resolve + /// here while `element_at` still resolved to the list-returning DataFusion kernel. Comet emits + /// only the name `map_extract` today, so this keeps the registry consistent rather than fixing + /// a live bug. + fn aliases(&self) -> &[String] { + &self.aliases + } + fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult { let [map_type, _] = take_function_args(self.name(), arg_types)?; Ok(map_entry_fields(map_type)?.1.data_type().clone()) @@ -112,11 +129,17 @@ fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef, &FieldR } /// Look up `key_arg` in each row of `map_arg`, returning the matched value or `NULL`. -pub fn spark_map_extract( +fn spark_map_extract( map_arg: &ColumnarValue, key_arg: &ColumnarValue, number_rows: usize, ) -> DataFusionResult { + // Expanding a scalar map builds `number_rows` copies of its entries. No Comet path reaches + // that: the native `Literal` proto carries no map, so `CometLiteral` rebuilds a folded + // `MapType` literal as a `CreateMap` tree that `CometCreateMap` hands to the JVM codegen + // dispatcher, which yields an array; the one map-producing literal shape that stays native, + // `MapFromArrays` over two empty arrays, has no entries and stops at the fast path below. + // Other users of the crate get the general (if unoptimized) answer rather than an error. let map_ref: ArrayRef = match map_arg { ColumnarValue::Array(array) => Arc::clone(array), ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?, @@ -130,6 +153,7 @@ pub fn spark_map_extract( let num_rows = map_array.len(); let value_type = map_array.value_type(); + validate_lookup_key(key_arg, map_array.key_type(), num_rows)?; // Arrow keeps a sliced `MapArray`'s entries child intact and slices only the offsets, so the // offsets index the *unsliced* keys/values and the visible entries are the half-open range @@ -155,12 +179,6 @@ pub fn spark_map_extract( key_match_mask(&keys, &key, true)? } ColumnarValue::Array(key_array) => { - if key_array.len() != num_rows { - return exec_err!( - "map_extract: expected {num_rows} lookup keys, got {}", - key_array.len() - ); - } // One vectorized compare needs a lookup key per *entry*, not per row, so gather each // row's key across that row's entries. Entries in a gap between two rows (offsets are // only required to be monotonic) keep index 0; the per-row scan below never reads @@ -180,11 +198,14 @@ pub fn spark_map_extract( // fits in `u32`. // // A NULL map row reads NULL whatever its entries hold. Arrow does not require a null row's - // offset range to be empty, and neither Comet's struct-field helper (which adds a parent null - // mask while preserving the child buffers) nor the UDF execution layer clears those entries, so - // a null row can carry a live `a -> 7` that would otherwise match. Spark returns NULL for a - // NULL map under both ANSI modes, for `element_at` and for `GetMapValue` alike, and only - // `element_at` has a nullable-input guard upstream of this kernel. + // offset range to be empty, and nothing downstream of a producer re-establishes that: adding a + // parent null mask over intact child buffers (as Comet's struct-field helper does) leaves any + // entries the child already had in place, and the UDF execution layer does not clear them + // either. Every producer traced so far -- the Parquet readers, Spark's `ArrowWriter`, and + // Arrow's own `filter` / `take` / `concat` -- gives a null row an empty range, so this is a + // representation the format permits rather than one known to arrive here. Guard it anyway: + // Spark returns NULL for a NULL map under both ANSI modes, for `element_at` and `GetMapValue` + // alike, and only `element_at` has a nullable-input guard upstream of this kernel. let map_nulls = map_array.nulls(); let mut indices = vec![0u32; num_rows]; let mut nulls = NullBufferBuilder::new(num_rows); @@ -210,6 +231,39 @@ pub fn spark_map_extract( )?)) } +/// Reject a lookup key this kernel cannot compare against `key_type`. +/// +/// Both checks are properties of the call rather than of the data, so they run before the +/// data-dependent early returns in [`spark_map_extract`]. Validating them later would make the +/// same query succeed on a batch whose maps are all empty or all NULL and fail on the next batch. +fn validate_lookup_key( + key_arg: &ColumnarValue, + key_type: &DataType, + num_rows: usize, +) -> DataFusionResult<()> { + let lookup_type = match key_arg { + ColumnarValue::Array(key_array) => { + if key_array.len() != num_rows { + return exec_err!( + "map_extract: expected {num_rows} lookup keys, got {}", + key_array.len() + ); + } + key_array.data_type().clone() + } + ColumnarValue::Scalar(scalar) => scalar.data_type(), + }; + // The planner casts the lookup key to the map's declared key type, so a mismatch here means + // the runtime encoding is not the declared one (a dictionary-encoded key column, say). Reject + // it rather than comparing incomparable encodings and reporting every row as a miss. + if &lookup_type != key_type { + return exec_err!( + "map_extract: lookup key type {lookup_type} does not match the map key type {key_type}" + ); + } + Ok(()) +} + /// A bit per map entry: set where the stored key equals the lookup key. `lookup` is either a /// length-1 array broadcast over every entry (constant key) or one key per entry. fn key_match_mask( @@ -217,35 +271,26 @@ fn key_match_mask( lookup: &ArrayRef, lookup_is_scalar: bool, ) -> DataFusionResult { - // The planner casts the lookup key to the map's declared key type, so a mismatch here means - // the runtime encoding is not the declared one (a dictionary-encoded key column, say). Reject - // it rather than comparing incomparable encodings and reporting every row as a miss. - if keys.data_type() != lookup.data_type() { - return exec_err!( - "map_extract: lookup key type {} does not match the map key type {}", - lookup.data_type(), - keys.data_type() - ); + // `eq` rejects nested key types, so those take DataFusion's element-wise comparison instead. + // `MapKeySupport` declines every complex key type, so this backstop is unreachable from Comet + // and exists for other users of the crate. Decide from the type rather than from an `eq` + // failure: with the lookup's length and type already checked, nesting is the only thing left + // that `eq` rejects, so any other error is a real one and should surface rather than fall into + // a silent per-row throughput cliff. + if keys.data_type().is_nested() { + return Ok(elementwise_match_mask(keys, lookup, lookup_is_scalar)); } let compared = if lookup_is_scalar { - eq(keys, &Scalar::new(Arc::clone(lookup))) + eq(keys, &Scalar::new(Arc::clone(lookup)))? } else { - eq(keys, lookup) + eq(keys, lookup)? }; - match compared { - Ok(mask) => { - // A NULL on either side compares as NULL, which is not a match. - let (values, nulls) = mask.into_parts(); - Ok(match nulls { - Some(nulls) if nulls.null_count() > 0 => &values & nulls.inner(), - _ => values, - }) - } - // `eq` rejects nested key types. `MapKeySupport` declines those before they reach the - // native lookup, but keep DataFusion's element-wise comparison as a backstop so this - // kernel is never less capable than the one it replaces. - Err(_) => Ok(elementwise_match_mask(keys, lookup, lookup_is_scalar)), - } + // A NULL on either side compares as NULL, which is not a match. + let (values, nulls) = compared.into_parts(); + Ok(match nulls { + Some(nulls) if nulls.null_count() > 0 => &values & nulls.inner(), + _ => values, + }) } fn elementwise_match_mask( @@ -286,6 +331,30 @@ mod tests { ]) } + /// A `MapArray` over already-built children. The fixtures below vary only in their key/value + /// types, offsets and null mask, so they all build through here. + fn map_of( + keys: ArrayRef, + values: ArrayRef, + offsets: Vec, + nulls: Option, + ) -> MapArray { + let fields = Fields::from(vec![ + Arc::new(Field::new("key", keys.data_type().clone(), false)), + Arc::new(Field::new("value", values.data_type().clone(), true)), + ]); + let entries = StructArray::new(fields.clone(), vec![keys, values], None); + let entries_field = Arc::new(Field::new("entries", DataType::Struct(fields), false)); + MapArray::try_new( + entries_field, + OffsetBuffer::new(offsets.into()), + entries, + nulls, + false, + ) + .unwrap() + } + fn map_from(rows: Vec) -> MapArray { let mut keys = Vec::new(); let mut values = Vec::new(); @@ -305,29 +374,12 @@ mod tests { offsets.push(keys.len() as i32); } - let key_field = Arc::new(Field::new("key", DataType::Utf8, false)); - let value_field = Arc::new(Field::new("value", DataType::Int32, true)); - let entries = StructArray::new( - Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), - vec![ - Arc::new(StringArray::from(keys)) as ArrayRef, - Arc::new(Int32Array::from(values)) as ArrayRef, - ], - None, - ); - let entries_field = Arc::new(Field::new( - "entries", - DataType::Struct(Fields::from(vec![key_field, value_field])), - false, - )); - MapArray::try_new( - entries_field, - OffsetBuffer::new(offsets.into()), - entries, + map_of( + Arc::new(StringArray::from(keys)), + Arc::new(Int32Array::from(values)), + offsets, nulls.finish(), - false, ) - .unwrap() } fn extract(map: MapArray, key: ColumnarValue) -> Vec> { @@ -377,35 +429,20 @@ mod tests { ); } - /// A NULL row whose entries were retained rather than dropped. `map_from` gives a NULL row an - /// empty offset range, which is the shape Arrow's builders produce, but nothing in the format - /// requires it: adding a parent null mask over intact child buffers leaves the entries in - /// place. Such a row must still read NULL, not the value its live entry holds. + /// A NULL row whose entries were retained rather than dropped. Every builder and kernel traced + /// so far -- `map_from` here, Arrow's own builders, the Parquet readers, `filter` / `take` / + /// `concat` -- gives a NULL row an empty offset range, but nothing in the format requires it, + /// and adding a parent null mask over intact child buffers would not. Such a row must read + /// NULL, not the value its live entry holds, so pin the contract at the kernel boundary rather + /// than relying on every producer to normalize. fn null_row_with_retained_entries() -> MapArray { - let key_field = Arc::new(Field::new("key", DataType::Utf8, false)); - let value_field = Arc::new(Field::new("value", DataType::Int32, true)); - let entries = StructArray::new( - Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), - vec![ - Arc::new(StringArray::from(vec!["a", "a", "b"])) as ArrayRef, - Arc::new(Int32Array::from(vec![Some(7), Some(1), Some(2)])) as ArrayRef, - ], - None, - ); - let entries_field = Arc::new(Field::new( - "entries", - DataType::Struct(Fields::from(vec![key_field, value_field])), - false, - )); - MapArray::try_new( - entries_field, + map_of( + Arc::new(StringArray::from(vec!["a", "a", "b"])), + Arc::new(Int32Array::from(vec![Some(7), Some(1), Some(2)])), // Row 0 is NULL but still spans entry 0 (`a -> 7`); row 1 is a live `{a: 1, b: 2}`. - OffsetBuffer::new(vec![0i32, 1, 3].into()), - entries, + vec![0, 1, 3], Some(NullBuffer::from(vec![false, true])), - false, ) - .unwrap() } #[test] @@ -481,29 +518,12 @@ mod tests { fn non_string_keys() { // Integer keys go down the same vectorized compare; pin that the gathered value lines up // with the matching key rather than the key's position. - let key_field = Arc::new(Field::new("key", DataType::Int32, false)); - let value_field = Arc::new(Field::new("value", DataType::Utf8, true)); - let entries = StructArray::new( - Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), - vec![ - Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, - Arc::new(StringArray::from(vec!["x", "y", "z"])) as ArrayRef, - ], + let map = map_of( + Arc::new(Int32Array::from(vec![10, 20, 30])), + Arc::new(StringArray::from(vec!["x", "y", "z"])), + vec![0, 2, 3], None, ); - let entries_field = Arc::new(Field::new( - "entries", - DataType::Struct(Fields::from(vec![key_field, value_field])), - false, - )); - let map = MapArray::try_new( - entries_field, - OffsetBuffer::new(vec![0i32, 2, 3].into()), - entries, - None, - false, - ) - .unwrap(); let result = spark_map_extract( &ColumnarValue::Array(Arc::new(map)), @@ -577,6 +597,34 @@ mod tests { assert!(err.to_string().contains("does not match the map key type")); } + #[test] + fn argument_checks_do_not_depend_on_the_data() { + // The key type and the key count are properties of the call, so a batch that happens to + // hold no entries has to fail the same way as one that does. Checking them after the + // empty-window fast path would make the same query answer NULLs on one partition and + // error on the next. + let empty = || ColumnarValue::Array(Arc::new(map_from(vec![Some(vec![]), None]))); + let err = spark_map_extract( + &empty(), + &ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + 2, + ) + .unwrap_err(); + assert!(err.to_string().contains("does not match the map key type")); + + let short_keys: ArrayRef = Arc::new(StringArray::from(vec!["a"])); + let err = spark_map_extract(&empty(), &ColumnarValue::Array(short_keys), 2).unwrap_err(); + assert!(err.to_string().contains("expected 2 lookup keys, got 1")); + } + + #[test] + fn overrides_both_registry_entries_of_the_kernel_it_replaces() { + // `register_udf` inserts one entry per alias, and DataFusion's `map_extract` declares + // `element_at`. Without the same alias the override would leave `element_at` resolving to + // the list-returning kernel. + assert_eq!(SparkMapExtract::new().aliases(), ["element_at"]); + } + #[test] fn nested_key_falls_back_to_elementwise_comparison() { // `eq` refuses nested types. `MapKeySupport` keeps these on Spark, but the backstop has to @@ -588,29 +636,12 @@ mod tests { Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, None, ); - let key_field = Arc::new(Field::new("key", key_values.data_type().clone(), false)); - let value_field = Arc::new(Field::new("value", DataType::Int32, true)); - let entries = StructArray::new( - Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), - vec![ - Arc::new(key_values) as ArrayRef, - Arc::new(Int32Array::from(vec![11, 22])) as ArrayRef, - ], + let map = map_of( + Arc::new(key_values), + Arc::new(Int32Array::from(vec![11, 22])), + vec![0, 2], None, ); - let entries_field = Arc::new(Field::new( - "entries", - DataType::Struct(Fields::from(vec![key_field, value_field])), - false, - )); - let map = MapArray::try_new( - entries_field, - OffsetBuffer::new(vec![0i32, 2].into()), - entries, - None, - false, - ) - .unwrap(); let lookup: ArrayRef = Arc::new(arrow::array::ListArray::new( inner, diff --git a/native/spark-expr/src/map_funcs/mod.rs b/native/spark-expr/src/map_funcs/mod.rs index d9de257d1ad..644466d0320 100644 --- a/native/spark-expr/src/map_funcs/mod.rs +++ b/native/spark-expr/src/map_funcs/mod.rs @@ -17,5 +17,5 @@ mod map_extract; mod map_sort; -pub use map_extract::{spark_map_extract, SparkMapExtract}; +pub use map_extract::SparkMapExtract; pub use map_sort::spark_map_sort; diff --git a/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql b/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql index b9e505336e2..95450b81592 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/element_at_map.sql @@ -81,6 +81,60 @@ SELECT element_at(map(named_struct('a', 1), 7), named_struct('a', 1)) query SELECT element_at(map(CAST('a' AS BINARY), 1, CAST('b' AS BINARY), 2), CAST('b' AS BINARY)) +-- Every key type `MapKeySupport` admits reaches Arrow's `eq`, which is stricter about the exact +-- Arrow type than the `ArrayData` equality it replaced: it distinguishes decimal precision and +-- scale, timestamp time zone, and integer width. Comet's planner casts the lookup key to the type +-- `coerce_types` reports, and the native lookup errors rather than comparing across encodings, so +-- a disagreement between the two shows up as a query failure. Cover the admitted key types that +-- the string/int fixtures above do not. +statement +CREATE TABLE test_element_at_map_keys( + mb map, + mt map, + ms map, + ml map, + md map, + mdate map, + mts map, + mntz map) USING parquet + +statement +INSERT INTO test_element_at_map_keys VALUES ( + map(true, 1, false, 2), + map(CAST(1 AS TINYINT), 10, CAST(2 AS TINYINT), 20), + map(CAST(1 AS SMALLINT), 10, CAST(2 AS SMALLINT), 20), + map(CAST(1 AS BIGINT), 10, CAST(2 AS BIGINT), 20), + map(CAST(1.50 AS DECIMAL(10,2)), 10, CAST(2.25 AS DECIMAL(10,2)), 20), + map(DATE '2024-01-01', 10, DATE '2024-06-15', 20), + map(TIMESTAMP '2024-01-01 00:00:00', 10, TIMESTAMP '2024-06-15 12:30:45', 20), + map(CAST('2024-01-01 00:00:00' AS TIMESTAMP_NTZ), 10, + CAST('2024-06-15 12:30:45' AS TIMESTAMP_NTZ), 20)) + +query +SELECT element_at(mb, true), element_at(mb, false) FROM test_element_at_map_keys + +query +SELECT element_at(mt, CAST(2 AS TINYINT)), element_at(ms, CAST(2 AS SMALLINT)), + element_at(ml, CAST(2 AS BIGINT)), element_at(ml, CAST(9 AS BIGINT)) +FROM test_element_at_map_keys + +-- Spark requires a decimal lookup key to have the map's exact precision and scale (a `DECIMAL(5,2)` +-- key against a `MAP` fails analysis with MAP_FUNCTION_DIFF_TYPES), so the +-- planner hands the native lookup a `Decimal128(10, 2)` on both sides and Arrow's `eq` agrees. +query +SELECT element_at(md, CAST(2.25 AS DECIMAL(10,2))), element_at(md, CAST(9.99 AS DECIMAL(10,2))) +FROM test_element_at_map_keys + +query +SELECT element_at(mdate, DATE '2024-06-15'), element_at(mdate, DATE '2020-01-01') +FROM test_element_at_map_keys + +query +SELECT element_at(mts, TIMESTAMP '2024-06-15 12:30:45'), + element_at(mts, TIMESTAMP '2020-01-01 00:00:00'), + element_at(mntz, CAST('2024-06-15 12:30:45' AS TIMESTAMP_NTZ)) +FROM test_element_at_map_keys + -- Nested INT-keyed map: the inner `element_at` returns NULL for ids not in the outer map (2, 3), -- and the outer `element_at` looks it up with a per-row key `id % (id - 2)`. This harness runs with -- ANSI disabled, so the remainder-by-zero at id = 2 evaluates to NULL rather than throwing, and diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 58f17f2867d..ebbdce406a3 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -465,12 +465,21 @@ class CometMapExpressionSuite extends CometTestBase { // MapArray's original entry offsets, so the visible entries start part way into the keys child -- // the same trap `mapsort` hit below. Reading the mask from index 0 would answer every row with // some other row's entries. - test("element_at on a sliced map reads the visible entries") { - withParquetTable( - (0 until 20).map(i => (i, Map(s"a${i % 5}" -> i, "shared" -> (i * 10)))), - "t") { + // + // `_2[k]` is the load-bearing form here. `element_at` on a nullable operand is wrapped in + // `CASE WHEN _2 IS NOT NULL` under ANSI, and DataFusion's CaseExpr evaluates the THEN branch + // through `filter_record_batch`, which compacts the entries child and resets the first offset to + // 0 as soon as any row is NULL -- so an `element_at`-only test would quietly stop slicing. + // `GetMapValue` has no such guard, so the sliced map reaches the kernel on every Spark version + // and in both ANSI modes. The NULL rows are here to keep that distinction honest. + test("map lookup on a sliced map reads the visible entries") { + val rows = (0 until 20).map { i => + val map = if (i % 7 == 3) null else Map(s"a${i % 5}" -> i, "shared" -> (i * 10)) + (i, map) + } + withParquetTable(rows, "t") { checkSparkAnswerAndOperator( - "SELECT _1, element_at(_2, 'a3'), element_at(_2, 'shared') " + + "SELECT _1, _2['a3'], _2['shared'], element_at(_2, 'a3') " + "FROM (SELECT * FROM t ORDER BY _1 LIMIT 15 OFFSET 5)") } }