From feccbd271d15b8d1504086ea801a9593ab212da7 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Sun, 6 Sep 2026 16:17:43 -0400 Subject: [PATCH] fix: Use Arrow comparator for key comparison in `map_extract` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous implementation of `map_extract` did the following for row: 1. Create a one-element array slice containing the row's search key 2. Scan the map's entries. For each entry, create a one-element array slice and compare the two slices using Arrow's array equality 3. Stop at the first match; if no matches, append a NULL instead This had three shortcomings: 1. It was very inefficient, because a lot of allocations are done for every element of every map. 2. It got the equality semantics wrong for some corner-cases. In particular, maps with dictionary-valued keys might encode a logical NULL in two physically distinct ways (#24983). Arrow's array equality also considers sparse unions that have different values in unselected child fields to be distinct; this is arguably a bug in Arrow though. 3. It returned `[NULL]` for missing map keys instead of an empty list, which is the behavior implemented by DuckDB (#24981). Instead, we can implement `map_extract` with a single arrow-ord comparator. This enables comparing the search key with each map element directly by index, without allocating. It also avoids the differences in comparison semantics outlined above. Finally, this PR fixes the behavior for absent map keys to be consistent with DuckDB. Benchmark results (M4 Max): - int32/first/1024x32, 135.629 µs -> 6.234 µs, -95.40% - int32/last/1024x1, 133.110 µs -> 6.157 µs, -95.37% - int32/last/1024x32, 3477.393 µs -> 38.562 µs, -98.89% - int32/last/1x0, 0.498 µs -> 0.313 µs, -37.23% - int32/last/1x1, 0.593 µs -> 0.384 µs, -35.34% - int32/missing/1024x32, 3472.698 µs -> 34.889 µs, -99.00% - int32/varying/1024x32, 1844.686 µs -> 25.307 µs, -98.63% - struct/first/1024x32, 335.421 µs -> 8.769 µs, -97.39% - struct/last/1024x1, 335.097 µs -> 8.647 µs, -97.42% - struct/last/1024x32, 8325.830 µs -> 71.611 µs, -99.14% - struct/last/1x0, 0.466 µs -> 0.242 µs, -48.08% - struct/last/1x1, 0.747 µs -> 0.391 µs, -47.67% - struct/missing/1024x32, 8451.153 µs -> 61.480 µs, -99.27% - struct/varying/1024x32, 4498.016 µs -> 41.343 µs, -99.08% - utf8_view/first/1024x32, 218.809 µs -> 9.877 µs, -95.49% - utf8_view/last/1024x1, 190.468 µs -> 8.464 µs, -95.56% - utf8_view/last/1024x32, 6016.140 µs -> 124.405 µs, -97.93% - utf8_view/last/1x0, 0.526 µs -> 0.353 µs, -32.97% - utf8_view/last/1x1, 0.762 µs -> 0.523 µs, -31.42% - utf8_view/missing/1024x32, 5999.511 µs -> 114.082 µs, -98.10% - utf8_view/varying/1024x32, 3226.583 µs -> 71.732 µs, -97.78% ("1024x32" means 1024 rows and each row is a map with 32 entries.) --- datafusion/functions-nested/benches/map.rs | 114 +++++++++++++- .../functions-nested/src/map_extract.rs | 139 ++++++++++++++---- datafusion/sqllogictest/test_files/map.slt | 65 ++++++-- 3 files changed, 271 insertions(+), 47 deletions(-) diff --git a/datafusion/functions-nested/benches/map.rs b/datafusion/functions-nested/benches/map.rs index 9cc4289ca1f1c..c65696aefe404 100644 --- a/datafusion/functions-nested/benches/map.rs +++ b/datafusion/functions-nested/benches/map.rs @@ -16,17 +16,18 @@ // under the License. use arrow::array::{ - ArrayRef, BinaryArray, BinaryViewArray, Int32Array, ListArray, StringArray, - StringViewArray, + Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, ListArray, MapArray, + StringArray, StringViewArray, StructArray, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; -use arrow::datatypes::Field; -use criterion::{Criterion, criterion_group, criterion_main}; +use arrow::datatypes::{DataType, Field}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::planner::ExprPlanner; use datafusion_expr::{ColumnarValue, Expr, ScalarFunctionArgs}; use datafusion_functions_nested::map::map_udf; +use datafusion_functions_nested::map_extract::map_extract_udf; use datafusion_functions_nested::planner::NestedFunctionPlanner; use rand::prelude::*; use std::collections::HashSet; @@ -208,5 +209,108 @@ fn criterion_benchmark(c: &mut Criterion) { } } -criterion_group!(benches, criterion_benchmark); +fn bench_map_extract(c: &mut Criterion) { + let udf = map_extract_udf(); + let config_options = Arc::new(ConfigOptions::default()); + let mut group = c.benchmark_group("map_extract"); + + for (rows, width) in [(1, 0), (1, 1), (1024, 1), (1024, 32)] { + for key_type in ["int32", "utf8_view", "struct"] { + let make_keys = |keys: Vec| -> ArrayRef { + match key_type { + "int32" => Arc::new(Int32Array::from(keys)), + "utf8_view" => Arc::new(StringViewArray::from_iter_values( + keys.iter().map(|key| format!("key_{key:016}")), + )), + "struct" => Arc::new(StructArray::from(vec![( + Arc::new(Field::new("key", DataType::Int32, false)), + Arc::new(Int32Array::from(keys)) as ArrayRef, + )])), + _ => unreachable!(), + } + }; + let keys = make_keys((0..rows).flat_map(|_| 0..width as i32).collect()); + let entries = StructArray::from(vec![ + ( + Arc::new(Field::new("key", keys.data_type().clone(), false)), + keys, + ), + ( + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Int32Array::from_iter_values(0..(rows * width) as i32)) + as ArrayRef, + ), + ]); + let map: ArrayRef = Arc::new(MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows)), + entries, + None, + false, + )); + let lookups: &[&str] = if width <= 1 { + &["last"] + } else { + &["first", "last", "missing", "varying"] + }; + for &lookup in lookups { + let query_keys = match lookup { + "first" => vec![0], + "last" => vec![width.saturating_sub(1) as i32], + "missing" => vec![width as i32], + // Mix matches and misses with a different lookup key per row. + "varying" => { + (0..rows).map(|row| (row % (width + 1)) as i32).collect() + } + _ => unreachable!(), + }; + let query_keys = make_keys(query_keys); + let query_keys = if lookup == "varying" { + ColumnarValue::Array(query_keys) + } else { + ColumnarValue::Scalar( + ScalarValue::try_from_array(&query_keys, 0).unwrap(), + ) + }; + let args = vec![ColumnarValue::Array(Arc::clone(&map)), query_keys]; + let arg_fields = args + .iter() + .map(|arg| Field::new("arg", arg.data_type(), true).into()) + .collect::>(); + let return_type = udf + .return_type( + &args + .iter() + .map(ColumnarValue::data_type) + .collect::>(), + ) + .unwrap(); + let return_field = Arc::new(Field::new("result", return_type, true)); + group.bench_function( + BenchmarkId::new( + format!("{key_type}/{lookup}"), + format!("{rows}x{width}"), + ), + |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: rows, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }); + }, + ); + } + } + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark, bench_map_extract); criterion_main!(benches); diff --git a/datafusion/functions-nested/src/map_extract.rs b/datafusion/functions-nested/src/map_extract.rs index 40340ec2cf635..d313043b922a4 100644 --- a/datafusion/functions-nested/src/map_extract.rs +++ b/datafusion/functions-nested/src/map_extract.rs @@ -19,10 +19,12 @@ use crate::utils::{get_map_entry_field, make_scalar_function}; use arrow::array::{ - Array, ArrayRef, Capacities, ListArray, MapArray, MutableArrayData, make_array, + Array, ArrayRef, ListArray, MapArray, MutableArrayData, make_array, new_empty_array, }; use arrow::buffer::OffsetBuffer; +use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field}; +use arrow_ord::ord::make_comparator; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, cast::as_map_array, exec_err}; use datafusion_expr::{ @@ -31,7 +33,6 @@ use datafusion_expr::{ }; use datafusion_macros::user_doc; use std::sync::Arc; -use std::vec; // Create static instances of ScalarUDFs for each function make_udf_expr_and_func!( @@ -149,43 +150,55 @@ fn general_map_extract_inner( query_keys_array: &dyn Array, ) -> Result { let keys = map_array.keys(); - let mut offsets = vec![0_i32]; - let values = map_array.values(); - let original_data = values.to_data(); - let capacity = Capacities::Array(original_data.len()); + let field = Arc::new(Field::new_list_field(map_array.value_type().clone(), true)); + let map_offsets = map_array.value_offsets(); + if map_offsets.first() == map_offsets.last() { + return Ok(Arc::new(ListArray::new( + field, + OffsetBuffer::new_zeroed(map_array.len()), + new_empty_array(values.data_type()), + map_array.nulls().cloned(), + ))); + } - let mut mutable = - MutableArrayData::with_capacities(vec![&original_data], true, capacity); + // Compare keys by index using a single comparator for the batch. + let compare = + make_comparator(keys.as_ref(), query_keys_array, SortOptions::default())?; + let mut offsets = Vec::with_capacity(map_array.len() + 1); + offsets.push(0_i32); - for (row_index, offset_window) in map_array.value_offsets().windows(2).enumerate() { + let original_data = values.to_data(); + // There is at most one output value per map row. + let mut mutable = MutableArrayData::new( + vec![&original_data], + false, + map_array.len().min(values.len()), + ); + + for (row_index, offset_window) in map_offsets.windows(2).enumerate() { let start = offset_window[0] as usize; let end = offset_window[1] as usize; - let len = end - start; - - let query_key = query_keys_array.slice(row_index, 1); + let mut offset = offsets[row_index]; - let value_index = - (0..len).find(|&i| keys.slice(start + i, 1).as_ref() == query_key.as_ref()); - - match value_index { - Some(index) => { - mutable.try_extend(0, start + index, start + index + 1)?; - } - None => { - mutable.try_extend_nulls(1)?; - } + if map_array.is_valid(row_index) + && let Some(index) = (start..end).find(|&i| compare(i, row_index).is_eq()) + { + mutable.try_extend(0, index, index + 1)?; + offset += 1; } - offsets.push(offsets[row_index] + 1); + + // A missing key results in an empty list. + offsets.push(offset); } let data = mutable.freeze(); Ok(Arc::new(ListArray::new( - Arc::new(Field::new_list_field(map_array.value_type().clone(), true)), + field, OffsetBuffer::::new(offsets.into()), - Arc::new(make_array(data)), - None, + make_array(data), + map_array.nulls().cloned(), ))) } @@ -210,3 +223,77 @@ fn map_extract_inner(args: &[ArrayRef]) -> Result { general_map_extract_inner(map_array, key_arg) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Float64Array, Int32Array, StructArray}; + use arrow::datatypes::Int32Type; + + fn make_map(keys: ArrayRef, values: Vec, offsets: Vec) -> MapArray { + let entries = StructArray::from(vec![ + ( + Arc::new(Field::new("key", keys.data_type().clone(), false)), + keys, + ), + ( + Arc::new(Field::new("value", DataType::Int32, true)), + Arc::new(Int32Array::from(values)) as ArrayRef, + ), + ]); + MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::new(offsets.into()), + entries, + None, + false, + ) + } + + #[test] + fn map_extract_sliced_maps() -> Result<()> { + let map = make_map( + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + vec![0, 10, 20, 30], + vec![0, 1, 3, 4], + ); + let query_keys = Int32Array::from(vec![0, 2, 9]); + + // Map offsets address the original entries; query indices address the slice. + let result = + general_map_extract_inner(&map.slice(1, 2), &query_keys.slice(1, 2))?; + let expected = ListArray::from_iter_primitive::([ + Some(vec![Some(20)]), + Some(vec![]), + ]); + assert_eq!(result.as_ref(), &expected); + + // Empty slices may retain the original nonempty keys and values buffers. + let result = + general_map_extract_inner(&map.slice(1, 0), &query_keys.slice(1, 0))?; + assert_eq!(result.len(), 0); + Ok(()) + } + + #[test] + fn map_extract_float_keys() -> Result<()> { + let nan = f64::NAN; + let other_nan = f64::from_bits(nan.to_bits() + 1); + let map = make_map( + Arc::new(Float64Array::from(vec![-0.0, 0.0, nan, other_nan])), + vec![1, 2, 3, 4], + vec![0, 4], + ); + + // Signed zeros and distinct NaN payloads identify different keys. + for (query, expected) in [(-0.0, 1), (0.0, 2), (nan, 3), (other_nan, 4)] { + let result = + general_map_extract_inner(&map, &Float64Array::from(vec![query]))?; + let expected = ListArray::from_iter_primitive::([Some( + vec![Some(expected)], + )]); + assert_eq!(result.as_ref(), &expected); + } + Ok(()) + } +} diff --git a/datafusion/sqllogictest/test_files/map.slt b/datafusion/sqllogictest/test_files/map.slt index 59f340a083c80..32fdeaac1d936 100644 --- a/datafusion/sqllogictest/test_files/map.slt +++ b/datafusion/sqllogictest/test_files/map.slt @@ -618,21 +618,21 @@ query ???? select map_extract(MAP {'a': 1, 'b': NULL, 'c': 3}, 'a'), map_extract(MAP {'a': 1, 'b': NULL, 'c': 3}, 'b'), map_extract(MAP {'a': 1, 'b': NULL, 'c': 3}, 'c'), map_extract(MAP {'a': 1, 'b': NULL, 'c': 3}, 'd'); ---- -[1] [NULL] [3] [NULL] +[1] [NULL] [3] [] # key is integer query ???? select map_extract(MAP {1: 1, 2: NULL, 3:3}, 1), map_extract(MAP {1: 1, 2: NULL, 3:3}, 2), map_extract(MAP {1: 1, 2: NULL, 3:3}, 3), map_extract(MAP {1: 1, 2: NULL, 3:3}, 4); ---- -[1] [NULL] [3] [NULL] +[1] [NULL] [3] [] # value is list query ???? select map_extract(MAP {1: [1, 2], 2: NULL, 3:[3]}, 1), map_extract(MAP {1: [1, 2], 2: NULL, 3:[3]}, 2), map_extract(MAP {1: [1, 2], 2: NULL, 3:[3]}, 3), map_extract(MAP {1: [1, 2], 2: NULL, 3:[3]}, 4); ---- -[[1, 2]] [NULL] [[3]] [NULL] +[[1, 2]] [NULL] [[3]] [] # key in map and query key are different types query ????? @@ -640,7 +640,7 @@ select map_extract(MAP {1: 1, 2: 2, 3:3}, '1'), map_extract(MAP {1: 1, 2: 2, 3:3 map_extract(MAP {1.0: 1, 2: 2, 3:3}, '1'), map_extract(MAP {'1': 1, '2': 2, '3':3}, 1.0), map_extract(MAP {arrow_cast('1', 'Utf8View'): 1, arrow_cast('2', 'Utf8View'): 2, arrow_cast('3', 'Utf8View'):3}, '1'); ---- -[1] [1] [1] [NULL] [1] +[1] [1] [1] [] [1] # null arg query ? @@ -648,14 +648,47 @@ select map_extract(NULL, 'a'); ---- NULL +# Empty maps, null lookup keys, and the element_at alias +query ???? +select map_extract(map(CAST([] AS VARCHAR[]), CAST([] AS INT[])), 'a'), map_extract(MAP {'a': 1}, NULL), + element_at(MAP {'a': 1, 'b': NULL}, 'missing'), element_at(MAP {'a': 1, 'b': NULL}, 'b'); +---- +[] [] [] [NULL] + +# Materialize a dictionary of structs so extracting d preserves a valid dictionary +# index pointing to a null value. A typed null instead has a null dictionary index. +statement ok +CREATE TABLE map_extract_dictionary_structs AS +SELECT arrow_cast( + named_struct('d', CAST(NULL AS VARCHAR)), + 'Dictionary(Int8, Struct("d": Utf8))' +) AS s; + +# These non-null struct keys compare equal despite their different null encodings. +query B?? +WITH keys AS ( + SELECT named_struct('d', s['d']) AS stored_key, + named_struct('d', arrow_cast(NULL, 'Dictionary(Int8, Utf8)')) AS query_key + FROM map_extract_dictionary_structs +) +SELECT stored_key = query_key, + map_extract(MAP {stored_key: 42}, query_key), + map_extract(MAP {query_key: 42}, stored_key) +FROM keys; +---- +true [42] [42] + +statement ok +DROP TABLE map_extract_dictionary_structs; + # map_extract with columns query ??? select map_extract(column1, 1), map_extract(column1, 5), map_extract(column1, 7) from map_array_table_1; ---- -[[1, NULL, 3]] [NULL] [NULL] -[NULL] [[4, NULL, 6]] [NULL] -[NULL] [NULL] [[1, NULL, 3]] -[NULL] [NULL] [NULL] +[[1, NULL, 3]] [] [] +[] [[4, NULL, 6]] [] +[] [] [[1, NULL, 3]] +NULL NULL NULL query ? select column1[1] from map_array_table_1; @@ -699,22 +732,22 @@ select map_extract(column1, column2), map_extract(column1, column3), map_extract ---- [[1, NULL, 3]] [[1, NULL, 3]] [[1, NULL, 3]] [[4, NULL, 6]] [[4, NULL, 6]] [[4, NULL, 6]] -[NULL] [NULL] [NULL] -[NULL] [NULL] [NULL] +[] [] [] +NULL NULL NULL query ??? select map_extract(column1, column2), map_extract(column1, column3), map_extract(column1, column4) from map_array_table_2; ---- -[[1, NULL, 3]] [NULL] [[1, NULL, 3]] -[[4, NULL, 6]] [NULL] [[4, NULL, 6]] -[NULL] [NULL] [NULL] +[[1, NULL, 3]] [] [[1, NULL, 3]] +[[4, NULL, 6]] [] [[4, NULL, 6]] +[] [] [] query ??? select map_extract(column1, 1), map_extract(column1, 5), map_extract(column1, 7) from map_array_table_2; ---- -[[1, NULL, 3]] [NULL] [NULL] -[NULL] [[4, NULL, 6]] [NULL] -[NULL] [NULL] [[1, NULL, 3]] +[[1, NULL, 3]] [] [] +[] [[4, NULL, 6]] [] +[] [] [[1, NULL, 3]] # Tests for map_entries