diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d66638ba32d..aae364cf9f3 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -707,17 +707,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 987d2f59359..f00f1a7c495 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -219,6 +219,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..1914c016f89 --- /dev/null +++ b/native/spark-expr/benches/map_extract.rs @@ -0,0 +1,138 @@ +// 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. +//! +//! 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}; +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..35aa5c6ab37 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,10 @@ 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). 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())), 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..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_sort; +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 new file mode 100644 index 00000000000..a2a5df6ed61 --- /dev/null +++ b/native/spark-expr/src/map_funcs/map_extract.rs @@ -0,0 +1,663 @@ +// 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. +/// +/// 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 { + 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), + aliases: vec!["element_at".to_string()], + } + } +} + +impl ScalarUDFImpl for SparkMapExtract { + fn name(&self) -> &str { + "map_extract" + } + + fn signature(&self) -> &Signature { + &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()) + } + + 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`. +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)?, + }; + 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(); + 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 + // [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) => { + // 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`. + // + // A NULL map row reads NULL whatever its entries hold. Arrow does not require a null row's + // 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); + 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)); + 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, + )?)) +} + +/// 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( + keys: &ArrayRef, + lookup: &ArrayRef, + lookup_is_scalar: bool, +) -> DataFusionResult { + // `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)))? + } else { + eq(keys, lookup)? + }; + // 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( + 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::{NullBuffer, 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)]), + ]) + } + + /// 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(); + 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); + } + + map_of( + Arc::new(StringArray::from(keys)), + Arc::new(Int32Array::from(values)), + offsets, + nulls.finish(), + ) + } + + 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] + ); + } + + /// 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 { + 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}`. + vec![0, 1, 3], + Some(NullBuffer::from(vec![false, true])), + ) + } + + #[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. + 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 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 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 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 + // 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 map = map_of( + Arc::new(key_values), + Arc::new(Int32Array::from(vec![11, 22])), + vec![0, 2], + None, + ); + + 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..644466d0320 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::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 f4a559b872b..ebbdce406a3 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -460,6 +460,43 @@ 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. + // + // `_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, _2['a3'], _2['shared'], element_at(_2, 'a3') " + + "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