From d029ea1df008bea1dcf9b0c7cadc8defe3eddfc1 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Wed, 9 Sep 2026 11:26:22 -0400 Subject: [PATCH 1/2] bench: measure get_field and map_extract map lookups on the same shapes Both suites cover Int32, Utf8View, and struct keys on maps of 4 and 32 entries, named `{key type}/{lookup}/{rows}x{entries}`. The lookups are the first and last entry of every row, a key that every row holds at a different position (`shuffled`), and a key present in no row, plus one key per row for `map_extract`. Cases that only rescaled another case, such as 8192 rows or one-entry maps, are dropped; `map_extract` keeps two single-row cases that measure per-batch fixed cost. --- datafusion/functions-nested/benches/map.rs | 100 ++++++---- .../functions/benches/misc/get_field.rs | 184 +++++++++++++----- 2 files changed, 197 insertions(+), 87 deletions(-) diff --git a/datafusion/functions-nested/benches/map.rs b/datafusion/functions-nested/benches/map.rs index c65696aefe404..0848803288a9d 100644 --- a/datafusion/functions-nested/benches/map.rs +++ b/datafusion/functions-nested/benches/map.rs @@ -214,8 +214,27 @@ fn bench_map_extract(c: &mut Criterion) { 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"] { + // Cases are named `{key type}/{lookup}/{rows}x{entries}`. The single-row + // shapes measure per-batch fixed cost. `shuffled` looks up a key that + // every row holds at a different position, and `varying` looks up a + // different key per row, mixing matches and misses. + let shapes: &[(usize, usize, &[&str])] = &[ + (1, 0, &["last"]), + (1, 1, &["last"]), + (1024, 4, &["last", "shuffled", "missing", "varying"]), + ( + 1024, + 32, + &["first", "last", "shuffled", "missing", "varying"], + ), + ]; + for &(rows, width, lookups) in shapes { + let key_types: &[&str] = if rows == 1 { + &["int32"] + } else { + &["int32", "utf8_view", "struct"] + }; + for &key_type in key_types { let make_keys = |keys: Vec| -> ArrayRef { match key_type { "int32" => Arc::new(Int32Array::from(keys)), @@ -229,39 +248,52 @@ fn bench_map_extract(c: &mut Criterion) { _ => 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"] + // Every row holds the keys `0..width`. With `shuffled`, each + // row's entries are rotated by the row number. + let make_map = |shuffled: bool| -> ArrayRef { + let keys = (0..rows) + .flat_map(|row| { + (0..width).map(move |position| { + if shuffled { + ((position + row) % width) as i32 + } else { + position as i32 + } + }) + }) + .collect(); + let keys = make_keys(keys); + 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, + ), + ]); + 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 map = make_map(false); + let shuffled_map = make_map(true); 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() - } + let (map, query_keys) = match lookup { + "first" => (&map, vec![0]), + "last" => (&map, vec![width.saturating_sub(1) as i32]), + "shuffled" => (&shuffled_map, vec![0]), + "missing" => (&map, vec![width as i32]), + "varying" => ( + &map, + (0..rows).map(|row| (row % (width + 1)) as i32).collect(), + ), _ => unreachable!(), }; let query_keys = make_keys(query_keys); @@ -272,7 +304,7 @@ fn bench_map_extract(c: &mut Criterion) { ScalarValue::try_from_array(&query_keys, 0).unwrap(), ) }; - let args = vec![ColumnarValue::Array(Arc::clone(&map)), query_keys]; + 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()) diff --git a/datafusion/functions/benches/misc/get_field.rs b/datafusion/functions/benches/misc/get_field.rs index b405e235420a5..6350f505d3992 100644 --- a/datafusion/functions/benches/misc/get_field.rs +++ b/datafusion/functions/benches/misc/get_field.rs @@ -15,9 +15,10 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, Int32Builder, MapBuilder, StringBuilder}; -use arrow::datatypes::{DataType, Field}; -use criterion::{Criterion, criterion_group}; +use arrow::array::{Array, ArrayRef, Int32Array, MapArray, StringViewArray, StructArray}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use criterion::{Bencher, BenchmarkId, Criterion, criterion_group}; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; @@ -25,68 +26,145 @@ use datafusion_functions::core::get_field; use std::hint::black_box; use std::sync::Arc; -/// A map array with `size` rows, each holding `entries` key/value pairs. -/// Every tenth row is null. -fn map_array(size: usize, entries: usize) -> ArrayRef { - let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); +const ROWS: usize = 1024; + +/// Map key types covered by the benchmarks. Struct keys are nested, so their +/// lookups can never switch from the comparator to the vectorized `eq`. +#[derive(Clone, Copy)] +enum KeyType { + Int32, + Utf8View, + Struct, +} + +impl KeyType { + fn name(self) -> &'static str { + match self { + KeyType::Int32 => "int32", + KeyType::Utf8View => "utf8_view", + KeyType::Struct => "struct", + } + } + + /// Builds a key array holding one key per element of `keys`. + fn make_keys(self, keys: &[i32]) -> ArrayRef { + match self { + KeyType::Int32 => Arc::new(Int32Array::from(keys.to_vec())), + KeyType::Utf8View => Arc::new(StringViewArray::from_iter_values( + keys.iter().map(|key| format!("key_{key:016}")), + )), + KeyType::Struct => Arc::new(StructArray::from(vec![( + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Int32Array::from(keys.to_vec())) as ArrayRef, + )])), + } + } +} + +/// A map array with `size` rows, each holding `entries` key/value pairs whose +/// keys are `0..entries`. Every tenth row is null and has no entries. With +/// `shuffled`, each row's entries are rotated by the row number, so a given +/// key sits at a different position in every row. +fn map_array(key_type: KeyType, size: usize, entries: usize, shuffled: bool) -> ArrayRef { + let mut keys = Vec::with_capacity(size * entries); + let mut values = Vec::with_capacity(size * entries); + let mut lengths = Vec::with_capacity(size); + let mut valid = Vec::with_capacity(size); for row in 0..size { - if row % 10 == 0 { - builder.append(false).unwrap(); + let is_null = row % 10 == 0; + valid.push(!is_null); + if is_null { + lengths.push(0); continue; } - for entry in 0..entries { - builder.keys().append_value(format!("key_{entry}")); - builder.values().append_value((row * entry) as i32); + lengths.push(entries); + for position in 0..entries { + let key = if shuffled { + (position + row) % entries + } else { + position + }; + keys.push(key as i32); + values.push((row * position) as i32); } - builder.append(true).unwrap(); } - Arc::new(builder.finish()) + let keys = key_type.make_keys(&keys); + let entries = StructArray::from(vec![ + ( + Arc::new(Field::new("keys", keys.data_type().clone(), false)), + keys, + ), + ( + Arc::new(Field::new("values", DataType::Int32, true)), + Arc::new(Int32Array::from(values)) as ArrayRef, + ), + ]); + Arc::new(MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::from_lengths(lengths), + entries, + Some(NullBuffer::from(valid)), + false, + )) } -fn bench_get_field( - c: &mut Criterion, - name: &str, - size: usize, - entries: usize, - key: &str, -) { - let udf = get_field(); - let args = vec![ - ColumnarValue::Array(map_array(size, entries)), - ColumnarValue::Scalar(ScalarValue::Utf8(Some(key.to_string()))), - ]; - let arg_fields = vec![ - Field::new("map", args[0].data_type(), true).into(), - Field::new("key", DataType::Utf8, false).into(), +/// Runs `get_field(map, key)` over a map built by [`map_array`]. `lookup` +/// picks the key: the first or last entry of every row, a key that every row +/// holds at a different position (`shuffled`), or one present in no row. +fn bench_get_field(b: &mut Bencher<'_>, key_type: KeyType, entries: usize, lookup: &str) { + let (key, shuffled) = match lookup { + "first" => (0, false), + "last" => (entries as i32 - 1, false), + "shuffled" => (0, true), + "missing" => (entries as i32, false), + _ => unreachable!(), + }; + let map = map_array(key_type, ROWS, entries, shuffled); + let lookup = ScalarValue::try_from_array(&key_type.make_keys(&[key]), 0) + .expect("lookup key should convert to a scalar"); + let arg_fields: Vec = vec![ + Field::new("map", map.data_type().clone(), true).into(), + Field::new("key", lookup.data_type(), false).into(), ]; + let args = vec![ColumnarValue::Array(map), ColumnarValue::Scalar(lookup)]; + let udf = get_field(); let config_options = Arc::new(ConfigOptions::default()); - - c.bench_function(name, |b| { - b.iter(|| { - black_box( - udf.invoke_with_args(ScalarFunctionArgs { - args: args.clone(), - arg_fields: arg_fields.clone(), - number_rows: size, - return_field: Field::new("f", DataType::Int32, true).into(), - config_options: Arc::clone(&config_options), - }) - .unwrap(), - ) - }) - }); + let return_field: FieldRef = Field::new("f", DataType::Int32, true).into(); + 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(), + ) + }) } fn criterion_benchmark(c: &mut Criterion) { - // First key: the match is found immediately, so the per-row overhead - // dominates. - bench_get_field(c, "get_field_map_1024_entries_4_first", 1024, 4, "key_0"); - // Last key: every entry of the row is compared before the match. - bench_get_field(c, "get_field_map_1024_entries_4_last", 1024, 4, "key_3"); - bench_get_field(c, "get_field_map_1024_entries_16_last", 1024, 16, "key_15"); - // Key that is not present in any row. - bench_get_field(c, "get_field_map_1024_entries_4_missing", 1024, 4, "key_9"); - bench_get_field(c, "get_field_map_8192_entries_4_last", 8192, 4, "key_3"); + // Cases are named `{key type}/{lookup}/{rows}x{entries}`. + let mut group = c.benchmark_group("get_field_map"); + let shapes: &[(usize, &[&str])] = &[ + (4, &["last", "shuffled", "missing"]), + (32, &["first", "last", "shuffled", "missing"]), + ]; + for key_type in [KeyType::Int32, KeyType::Utf8View, KeyType::Struct] { + for &(entries, lookups) in shapes { + for &lookup in lookups { + group.bench_function( + BenchmarkId::new( + format!("{}/{lookup}", key_type.name()), + format!("{ROWS}x{entries}"), + ), + |b| bench_get_field(b, key_type, entries, lookup), + ); + } + } + } + group.finish(); } criterion_group!(benches, criterion_benchmark); From b33a1c00c6645feefec1efd7b9d2600f93943a19 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Wed, 9 Sep 2026 12:17:11 -0400 Subject: [PATCH 2/2] perf, chore: Consolidate and optimize `get_field` and `map_extract` `map[key]` (aka `get_field`) and `map_extract(map, key)` collectively had three ways to lookup keys in a map: (1): `get_field` used a per-batch `eq` kernel for scalar keys. This is efficient for maps with many entries where `key` is not found, but slower for maps in which the key can be found quickly (because `eq` does not allow early-stopping). On my local machine, `eq` only beats a comparator-based approach if the latter required touching more than 60% of the keys in a row. (2): `get_field` used a comparator-based approach for nested keys. (3): `map_extract` used a comparator-based approach for all keys. Conceptually, these two functions only differ in how the result is represented, so it makes sense to consolidate them. We can also adopt a hybrid strategy that gets the best of the previous approaches for most inputs: * Start with a comparator-based approach. * After the first row, remember the index at which the matching key was found, and check that index first for subsequent rows. This takes advantage of the observation that most map rows have their keys in the same order. * After 32 rows, check whether the comparator looked at more than 75% of the entries in those rows. If it did, the early stopping that the comparator approach allows is not useful and we switch to an `eq` kernel for the remainder of the batch (as long as the map key is not a nested type). Add `datafusion_functions::utils::map_lookup`, which returns for each map row the index of the first matching entry or null, and build both functions on it with `take`. The lookup scans each row with a comparator, which stops at the first match, and tries the position where the previous row matched first, since rows in a batch usually share key order; a hit then usually costs one comparison wherever the key sits. A miss has no match to stop at, so every entry of the row is compared at the comparator's higher per-comparison cost. After sampling 32 rows, if most of them missed, the remaining entries are compared with one vectorized `eq` instead, which does the same full comparison more cheaply. A lookup key that differs from the map key type only in dictionary encoding is cast rather than rejected. `map_extract` now passes a scalar key through as a single row instead of expanding it to the batch size. `map[key]` reports a mismatched key type as an execution error instead of an Arrow comparison error, and a null map row that still carries entries yields NULL, as it already did from `map_extract`. Closes #25083 --- .../functions-nested/src/map_extract.rs | 102 ++--- datafusion/functions/src/core/getfield.rs | 146 +------ datafusion/functions/src/utils.rs | 403 +++++++++++++++++- datafusion/sqllogictest/test_files/map.slt | 30 +- 4 files changed, 469 insertions(+), 212 deletions(-) diff --git a/datafusion/functions-nested/src/map_extract.rs b/datafusion/functions-nested/src/map_extract.rs index a9ef716514d52..6064b66f47a1c 100644 --- a/datafusion/functions-nested/src/map_extract.rs +++ b/datafusion/functions-nested/src/map_extract.rs @@ -17,20 +17,19 @@ //! [`ScalarUDFImpl`] definitions for map_extract functions. -use crate::utils::{get_map_entry_field, make_scalar_function}; -use arrow::array::{ - Array, ArrayRef, ListArray, MapArray, MutableArrayData, make_array, new_empty_array, -}; +use crate::utils::get_map_entry_field; +use arrow::array::{Array, ArrayRef, ListArray, MapArray, UInt32Array}; use arrow::buffer::OffsetBuffer; -use arrow::compute::SortOptions; +use arrow::compute::take; 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::function::Hint; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; +use datafusion_functions::utils::{make_scalar_function, map_lookup}; use datafusion_macros::user_doc; use std::sync::Arc; @@ -119,7 +118,11 @@ impl ScalarUDFImpl for MapExtract { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(map_extract_inner)(&args.args) + // A scalar key is passed through as a single row rather than expanded + // to the batch size; the lookup applies it to every map row. + make_scalar_function(map_extract_inner, vec![Hint::Pad, Hint::AcceptsSingular])( + &args.args, + ) } fn aliases(&self) -> &[String] { @@ -149,79 +152,34 @@ fn general_map_extract_inner( map_array: &MapArray, query_keys_array: &dyn Array, ) -> Result { - let keys = map_array.keys(); - let values = map_array.values(); - 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(), - ))); - } - - // 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); - - 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 mut offset = offsets[row_index]; - - 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; - } - - // A missing key results in an empty list. - offsets.push(offset); - } - - let data = mutable.freeze(); - + let indices = map_lookup(map_array, query_keys_array)?; + // Each matched row contributes one list element. Every other row is an + // empty list, or NULL when the map itself is NULL. + let lengths = indices.iter().map(|index| usize::from(index.is_some())); + let mut matched = Vec::with_capacity(indices.len() - indices.null_count()); + matched.extend(indices.iter().flatten()); + let values = take( + map_array.values().as_ref(), + &UInt32Array::from(matched), + None, + )?; Ok(Arc::new(ListArray::new( - field, - OffsetBuffer::::new(offsets.into()), - make_array(data), + Arc::new(Field::new_list_field(map_array.value_type().clone(), true)), + OffsetBuffer::from_lengths(lengths), + values, map_array.nulls().cloned(), ))) } fn map_extract_inner(args: &[ArrayRef]) -> Result { let [map_arg, key_arg] = take_function_args("map_extract", args)?; - - let map_array = match map_arg.data_type() { - DataType::Map(_, _) => as_map_array(&map_arg)?, - DataType::Null => return Ok(Arc::clone(map_arg)), - _ => return exec_err!("The first argument in map_extract must be a map"), - }; - - let key_type = map_array.key_type(); - - if key_type != key_arg.data_type() { - return exec_err!( - "The key type {} does not match the map key type {}", - key_arg.data_type(), - key_type - ); + match map_arg.data_type() { + DataType::Map(_, _) => { + general_map_extract_inner(as_map_array(map_arg.as_ref())?, key_arg.as_ref()) + } + DataType::Null => Ok(Arc::clone(map_arg)), + _ => exec_err!("The first argument in map_extract must be a map"), } - - general_map_extract_inner(map_array, key_arg) } #[cfg(test)] diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index a0f024bbc7ea4..64606af803aae 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -17,11 +17,8 @@ use std::sync::{Arc, OnceLock}; -use arrow::array::{ - Array, Capacities, MutableArrayData, Scalar, cast::AsArray, make_array, - make_comparator, -}; -use arrow::compute::SortOptions; +use arrow::array::{Array, cast::AsArray}; +use arrow::compute::take; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::cast::{as_map_array, as_struct_array}; @@ -38,6 +35,7 @@ use datafusion_macros::user_doc; use super::named_struct::NamedStructFunc; use super::r#struct::StructFunc; +use crate::utils::map_lookup; #[user_doc( doc_section(label = "Other Functions"), @@ -99,89 +97,6 @@ impl Default for GetFieldFunc { } } -/// Process a map array with a non-nested key type by comparing the single -/// lookup key against every map key with the `eq` kernel, then scanning the -/// result for each row. -/// -/// `eq` does not support nested types, so list, struct, and map keys go -/// through [`process_map_with_nested_key`] instead. -fn process_map_array( - array: &dyn Array, - key_array: Arc, -) -> Result { - let map_array = as_map_array(array)?; - let be_compared = Scalar::new(key_array); - let keys = arrow::compute::kernels::cmp::eq(&be_compared, map_array.keys())?; - - let original_data = map_array.entries().column(1).to_data(); - let capacity = Capacities::Array(original_data.len()); - let mut mutable = - MutableArrayData::with_capacities(vec![&original_data], true, capacity); - - let offsets = map_array.value_offsets(); - // Scan the comparison result in place: slicing it per entry would allocate - // a new array for every row of the map. Map keys are non-null by - // definition, so the comparison result carries no nulls to check here. - let matches = keys.values(); - - for entry in 0..map_array.len() { - let start = offsets[entry] as usize; - let end = offsets[entry + 1] as usize; - - let matched = (start..end).find(|&i| matches.value(i)); - - match matched { - Some(i) => mutable.try_extend(0, i, i + 1)?, - None => mutable.try_extend_nulls(1)?, - } - } - - let data = mutable.freeze(); - let data = make_array(data); - Ok(ColumnarValue::Array(data)) -} - -/// Process a map array with a nested key type by iterating through entries -/// and using a comparator for key matching. -/// -/// This specialized version is used when the key type is nested (e.g., struct, list). -fn process_map_with_nested_key( - array: &dyn Array, - key_array: &dyn Array, -) -> Result { - let map_array = as_map_array(array)?; - - let comparator = - make_comparator(map_array.keys().as_ref(), key_array, SortOptions::default())?; - - let original_data = map_array.entries().column(1).to_data(); - let capacity = Capacities::Array(original_data.len()); - let mut mutable = - MutableArrayData::with_capacities(vec![&original_data], true, capacity); - - for entry in 0..map_array.len() { - let start = map_array.value_offsets()[entry] as usize; - let end = map_array.value_offsets()[entry + 1] as usize; - - let mut found_match = false; - for i in start..end { - if comparator(i, 0).is_eq() { - mutable.try_extend(0, i, i + 1)?; - found_match = true; - break; - } - } - - if !found_match { - mutable.try_extend_nulls(1)?; - } - } - - let data = mutable.freeze(); - let data = make_array(data); - Ok(ColumnarValue::Array(data)) -} - /// Extract a single field from a struct or map array fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result { let arrays = ColumnarValue::values_to_arrays(&[base])?; @@ -209,14 +124,11 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result { - // The lookup key is a single scalar. `eq` does not support nested - // key types, so those are matched with a comparator instead. - let key_array = key.to_array()?; - if key_array.data_type().is_nested() { - process_map_with_nested_key(&array, key_array.as_ref()) - } else { - process_map_array(&array, key_array) - } + // The lookup key is a single scalar + let map_array = as_map_array(array.as_ref())?; + let indices = map_lookup(map_array, key.to_array()?.as_ref())?; + let values = take(map_array.values().as_ref(), &indices, None)?; + Ok(ColumnarValue::Array(values)) } (DataType::Struct(_), _, Some(k)) => { let as_struct_array = as_struct_array(&array)?; @@ -656,11 +568,8 @@ impl ScalarUDFImpl for GetFieldFunc { #[cfg(test)] mod tests { use super::*; - use arrow::array::{ - ArrayRef, Int32Array, Int32Builder, ListArray, ListBuilder, MapBuilder, - StructArray, - }; - use arrow::datatypes::{Fields, Int32Type}; + use arrow::array::{ArrayRef, Int32Array, StructArray}; + use arrow::datatypes::Fields; #[test] fn test_get_field_utf8view_key() -> Result<()> { @@ -698,41 +607,6 @@ mod tests { Ok(()) } - #[test] - fn test_get_field_map_list_key() -> Result<()> { - // One map row with two list keys. The lookup key matches the second - // entry, so the match is not at the first entry of the row. - let mut builder = MapBuilder::new( - None, - ListBuilder::new(Int32Builder::new()), - Int32Builder::new(), - ); - builder.keys().append_value([Some(1), Some(2)]); - builder.values().append_value(1); - builder.keys().append_value([Some(3), Some(4)]); - builder.values().append_value(2); - builder.append(true)?; - let base = ColumnarValue::Array(Arc::new(builder.finish())); - - let list_key = |values: Vec| { - ScalarValue::List(Arc::new( - ListArray::from_iter_primitive::([Some( - values.into_iter().map(Some), - )]), - )) - }; - - let result = extract_single_field(base.clone(), list_key(vec![3, 4]))?; - let expected = Int32Array::from(vec![Some(2)]); - assert_eq!(result.into_array(1)?.as_ref(), &expected as &dyn Array); - - let result = extract_single_field(base, list_key(vec![9, 9]))?; - let expected = Int32Array::from(vec![None]); - assert_eq!(result.into_array(1)?.as_ref(), &expected as &dyn Array); - - Ok(()) - } - #[test] fn test_get_field_dict_encoded_struct() -> Result<()> { use arrow::array::{DictionaryArray, StringArray, UInt32Array}; diff --git a/datafusion/functions/src/utils.rs b/datafusion/functions/src/utils.rs index 8e1dcb3b91733..498a0f21c5f58 100644 --- a/datafusion/functions/src/utils.rs +++ b/datafusion/functions/src/utils.rs @@ -15,14 +15,20 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, ArrayRef, ArrowPrimitiveType, AsArray, PrimitiveArray}; -use arrow::compute::try_binary; +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, AsArray, MapArray, PrimitiveArray, Scalar, + UInt32Array, UInt32Builder, make_comparator, +}; +use arrow::buffer::NullBuffer; +use arrow::compute::kernels::cmp::eq; +use arrow::compute::{SortOptions, cast, try_binary}; use arrow::datatypes::{DataType, DecimalType}; use arrow::error::ArrowError; -use datafusion_common::{DataFusionError, Result, ScalarValue}; +use datafusion_common::{DataFusionError, Result, ScalarValue, exec_err, internal_err}; use datafusion_expr::ColumnarValue; use datafusion_expr::function::Hint; use std::cmp::Ordering; +use std::ops::Range; use std::sync::Arc; /// Creates a function to identify the optimal return type of a string function given @@ -347,6 +353,180 @@ pub fn decimal64_to_i64(value: i64, scale: i8) -> Result { } } +/// Finds, for each row of `map`, the first entry whose key equals that row's +/// lookup key. +/// +/// `keys` holds either a single key, which every row is looked up with, or +/// one key per map row. The result has one element per map row: the index of +/// the matching entry into `map.values()`, or null when the row is null, the +/// lookup key is null, or no entry matches. It can be passed directly to +/// [`arrow::compute::take`] on `map.values()`. +/// +/// Non-nested keys must have the map's key type, up to dictionary encoding. +/// Nested keys must have the same structure, and may differ in field names +/// and nullability. Keys are compared the way `ORDER BY` compares values: +/// floating point keys use total ordering, so `-0.0` and `0.0` are different +/// keys and NaN matches NaN. +pub fn map_lookup(map: &MapArray, keys: &dyn Array) -> Result { + let map_keys = map.keys(); + let single_key = match keys.len() { + 1 => true, + len if len == map.len() => false, + len => { + return internal_err!( + "map_lookup expects one lookup key or one per map row ({}), got {len}", + map.len() + ); + } + }; + let key_type = map_keys.data_type(); + // A nested lookup key only has to be nested here; `make_comparator` + // checks its structure. A non-nested lookup key must have the map's + // key type, ignoring dictionary encoding. + let compatible = if key_type.is_nested() { + keys.data_type().is_nested() + } else { + strip_dictionary(key_type).equals_datatype(strip_dictionary(keys.data_type())) + }; + if !compatible { + return exec_err!( + "The key type {} does not match the map key type {}", + keys.data_type(), + key_type + ); + } + // The comparison kernels need both sides to use the same encoding. + let cast_keys; + let keys: &dyn Array = if key_type.is_nested() || keys.data_type() == key_type { + keys + } else { + cast_keys = cast(keys, key_type)?; + cast_keys.as_ref() + }; + + let offsets = map.value_offsets(); + let (first, last) = (offsets[0] as usize, offsets[map.len()] as usize); + // No row has any entries, so nothing can match. Map keys are never + // null, so a null lookup key matches nothing either. + if first == last || (single_key && keys.logical_null_count() > 0) { + return Ok(UInt32Array::new_null(map.len())); + } + let key_nulls = if single_key { + None + } else { + keys.logical_nulls() + }; + let mut scanner = RowScanner::new(map, key_nulls.as_ref()); + + // Scan with a comparator, which stops at the first match in each row. + // Count the comparisons over a sample of rows to see whether stopping + // early pays off. + let cmp = make_comparator(map_keys.as_ref(), keys, SortOptions::default())?; + let compare = + |entry: usize, row: usize| cmp(entry, if single_key { 0 } else { row }).is_eq(); + let sample = map.len().min(SAMPLE_ROWS); + let mut comparisons = 0; + let sampled_entries = scanner.scan(0..sample, |entry, row| { + comparisons += 1; + compare(entry, row) + }); + + // If the sampled rows compared more than half of their entries, stopping + // early is not paying off, so the remaining rows are cheaper to compare all + // at once with the vectorized `eq`. We can only use `eq` when we have a + // single, non-nested key. The exact break-even point depends on the key + // type and the hardware; half keeps the cost of a wrong guess to about a + // third in either direction. + let rest = sample..map.len(); + if single_key + && !key_type.is_nested() + && !rest.is_empty() + && comparisons * 2 > sampled_entries + { + let range_start = offsets[sample] as usize; + let in_range = map_keys.slice(range_start, last - range_start); + let matches = eq(&Scalar::new(keys.slice(0, 1)), &in_range)?; + // Neither side has nulls, so the value bits alone are meaningful. + let bits = matches.values(); + scanner.scan(rest, |entry, _| bits.value(entry - range_start)); + } else { + scanner.scan(rest, compare); + } + Ok(scanner.finish()) +} + +/// Number of rows [`map_lookup`] scans with the comparator before deciding +/// whether the rest of the batch is better served by the vectorized `eq`. +const SAMPLE_ROWS: usize = 32; + +/// The value type of a dictionary-encoded type, or the type itself. +fn strip_dictionary(data_type: &DataType) -> &DataType { + match data_type { + DataType::Dictionary(_, value_type) => value_type, + other => other, + } +} + +/// Scans map rows for the first entry that satisfies a predicate. +struct RowScanner<'a> { + offsets: &'a [i32], + /// Rows to skip: null map rows and rows whose lookup key is null. + skip: Option, + found: UInt32Builder, + /// Position within its row of the most recent match. + hint: usize, +} + +impl<'a> RowScanner<'a> { + fn new(map: &'a MapArray, key_nulls: Option<&NullBuffer>) -> Self { + Self { + offsets: map.value_offsets(), + skip: NullBuffer::union(map.nulls(), key_nulls), + found: UInt32Builder::with_capacity(map.len()), + hint: 0, + } + } + + /// Scans `rows`, recording the first entry for which `is_match(entry, row)` + /// holds, or null for a skipped row or a row without a match. Returns the + /// number of entries in the rows that were scanned. + fn scan( + &mut self, + rows: Range, + mut is_match: impl FnMut(usize, usize) -> bool, + ) -> usize { + let mut entries = 0; + for row in rows { + if self.skip.as_ref().is_some_and(|skip| skip.is_null(row)) { + self.found.append_null(); + continue; + } + let start = self.offsets[row] as usize; + let end = self.offsets[row + 1] as usize; + entries += end - start; + + // Rows in a batch usually share the same key order, so try the + // position where the previous row matched first. When that guess + // is right, the lookup costs one comparison wherever the key sits. + let hinted = start + self.hint; + let found = if hinted < end && is_match(hinted, row) { + Some(hinted) + } else { + (start..end).find(|&entry| entry != hinted && is_match(entry, row)) + }; + if let Some(entry) = found { + self.hint = entry - start; + } + self.found.append_option(found.map(|entry| entry as u32)); + } + entries + } + + fn finish(mut self) -> UInt32Array { + self.found.finish() + } +} + #[cfg(test)] pub mod test { /// $FUNC ScalarUDFImpl to test @@ -626,3 +806,220 @@ pub mod test { } } } + +#[cfg(test)] +mod map_lookup_tests { + use super::*; + use arrow::array::{ + DictionaryArray, Int32Array, Int64Array, ListArray, StringArray, StructArray, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::{Field, Int32Type}; + + /// A map whose rows have the given lengths, drawing keys and values in + /// order from `keys` and `values`. Rows flagged false in `valid` are null. + fn make_map( + keys: ArrayRef, + values: ArrayRef, + lengths: &[usize], + valid: Option<&[bool]>, + ) -> MapArray { + let entries = StructArray::from(vec![ + ( + Arc::new(Field::new("key", keys.data_type().clone(), false)), + keys, + ), + ( + Arc::new(Field::new("value", values.data_type().clone(), true)), + values, + ), + ]); + MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::from_lengths(lengths.iter().copied()), + entries, + valid.map(|valid| NullBuffer::from(valid.to_vec())), + false, + ) + } + + /// Rows: `{1: 10, 2: 20}`, `{}`, `{2: 30}`, and a null row that still + /// carries the entry `{1: 40}`. + fn int_map() -> MapArray { + make_map( + Arc::new(Int32Array::from(vec![1, 2, 2, 1])), + Arc::new(Int32Array::from(vec![10, 20, 30, 40])), + &[2, 0, 1, 1], + Some(&[true, true, true, false]), + ) + } + + #[test] + fn single_key() -> Result<()> { + let map = int_map(); + let result = map_lookup(&map, &Int32Array::from(vec![2]))?; + assert_eq!( + result, + UInt32Array::from(vec![Some(1), None, Some(2), None]) + ); + + // The null row's entry is never matched. + let result = map_lookup(&map, &Int32Array::from(vec![1]))?; + assert_eq!(result, UInt32Array::from(vec![Some(0), None, None, None])); + + let result = map_lookup(&map, &Int32Array::from(vec![9]))?; + assert_eq!(result, UInt32Array::from(vec![None; 4])); + Ok(()) + } + + #[test] + fn single_null_key_matches_nothing() -> Result<()> { + let map = int_map(); + let result = map_lookup(&map, &Int32Array::from(vec![None]))?; + assert_eq!(result, UInt32Array::from(vec![None; 4])); + Ok(()) + } + + #[test] + fn one_key_per_row() -> Result<()> { + let map = int_map(); + let keys = Int32Array::from(vec![Some(2), Some(1), None, Some(1)]); + let result = map_lookup(&map, &keys)?; + assert_eq!(result, UInt32Array::from(vec![Some(1), None, None, None])); + Ok(()) + } + + #[test] + fn nested_single_key() -> Result<()> { + let keys = ListArray::from_iter_primitive::([ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(3), Some(4)]), + ]); + let map = make_map( + Arc::new(keys), + Arc::new(Int32Array::from(vec![10, 20])), + &[2], + None, + ); + let list_key = |values: Vec| { + ListArray::from_iter_primitive::([Some( + values.into_iter().map(Some), + )]) + }; + let result = map_lookup(&map, &list_key(vec![3, 4]))?; + assert_eq!(result, UInt32Array::from(vec![Some(1)])); + + let result = map_lookup(&map, &list_key(vec![9, 9]))?; + assert_eq!(result, UInt32Array::from(vec![None])); + Ok(()) + } + + #[test] + fn sliced_map_and_keys() -> Result<()> { + let map = int_map(); + let keys = Int32Array::from(vec![0, 0, 2, 0]); + + // Rows 1 and 2 of the map with keys 1 and 2 of the key array. Entry + // indices stay relative to the unsliced entries. + let result = map_lookup(&map.slice(1, 2), &keys.slice(1, 2))?; + assert_eq!(result, UInt32Array::from(vec![None, Some(2)])); + + // A single key against a slice whose entries start past offset 0. + let result = map_lookup(&map.slice(2, 1), &Int32Array::from(vec![2]))?; + assert_eq!(result, UInt32Array::from(vec![Some(2)])); + + // An empty slice still sees the unsliced entries buffer. + let result = map_lookup(&map.slice(1, 0), &keys.slice(1, 0))?; + assert_eq!(result.len(), 0); + Ok(()) + } + + #[test] + fn dictionary_keys_match_plain_keys() -> Result<()> { + // Rows: `{a: 10, b: 20}`, `{a: 30}`. + let keys: DictionaryArray = vec!["a", "b", "a"].into_iter().collect(); + let map = make_map( + Arc::new(keys), + Arc::new(Int32Array::from(vec![10, 20, 30])), + &[2, 1], + None, + ); + let result = map_lookup(&map, &StringArray::from(vec!["b"]))?; + assert_eq!(result, UInt32Array::from(vec![Some(1), None])); + + let result = map_lookup(&map, &StringArray::from(vec!["b", "a"]))?; + assert_eq!(result, UInt32Array::from(vec![Some(1), Some(2)])); + Ok(()) + } + + #[test] + fn vectorized_scan_after_missing_sample() -> Result<()> { + // The first 40 rows are `{1: 0, 2: 0, 3: 0}` and the rest are + // `{9: 0, 1: 0, 2: 0}`, so a lookup of 9 misses throughout the sampled + // rows and must still be found afterwards. + let rows = 100; + let switch = 40; + let keys: Vec = (0..rows) + .flat_map(|row| if row < switch { [1, 2, 3] } else { [9, 1, 2] }) + .collect(); + let map = make_map( + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(vec![0; rows * 3])), + &vec![3; rows], + None, + ); + let result = map_lookup(&map, &Int32Array::from(vec![9]))?; + let expected: UInt32Array = (0..rows) + .map(|row| (row >= switch).then_some(row as u32 * 3)) + .collect(); + assert_eq!(result, expected); + + // A key present in every row keeps the comparator scan. + let result = map_lookup(&map, &Int32Array::from(vec![2]))?; + let expected: UInt32Array = (0..rows) + .map(|row| Some(row as u32 * 3 + if row < switch { 1 } else { 2 })) + .collect(); + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn nested_keys_may_differ_in_field_nullability() -> Result<()> { + // A map read from a schema with a non-null struct field, looked up + // with a struct literal whose fields are nullable. + let map_keys = StructArray::from(vec![( + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + )]); + let map = make_map( + Arc::new(map_keys), + Arc::new(Int32Array::from(vec![10, 20])), + &[2], + None, + ); + let lookup = StructArray::from(vec![( + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![2])) as ArrayRef, + )]); + let result = map_lookup(&map, &lookup)?; + assert_eq!(result, UInt32Array::from(vec![Some(1)])); + Ok(()) + } + + #[test] + fn key_type_mismatch_is_an_error() { + let map = int_map(); + let err = map_lookup(&map, &Int64Array::from(vec![2])).unwrap_err(); + assert!( + err.to_string() + .contains("The key type Int64 does not match the map key type Int32"), + "{err}" + ); + } + + #[test] + fn wrong_key_count_is_an_error() { + let map = int_map(); + assert!(map_lookup(&map, &Int32Array::from(vec![1, 2])).is_err()); + } +} diff --git a/datafusion/sqllogictest/test_files/map.slt b/datafusion/sqllogictest/test_files/map.slt index debaf7f8b344f..f802d5a714387 100644 --- a/datafusion/sqllogictest/test_files/map.slt +++ b/datafusion/sqllogictest/test_files/map.slt @@ -627,6 +627,34 @@ NULL NULL statement ok DROP TABLE map_list_keys; +# a typed NULL key matches nothing +query T +SELECT MAP {1:'a', 2:'b'}[arrow_cast(NULL, 'Int64')]; +---- +NULL + +# dictionary-encoded map keys are looked up by value +query I +SELECT MAP {arrow_cast('a', 'Dictionary(Int32, Utf8)'):1, arrow_cast('b', 'Dictionary(Int32, Utf8)'):2}['b']; +---- +2 + +# NULLIF turns a map row NULL without clearing its entries; the lookups must +# report NULL rather than read those entries +statement ok +CREATE TABLE map_nullif AS +SELECT MAP {'a': 1, 'b': 2} AS m, MAP {'a': 1, 'b': 2} AS n +UNION ALL SELECT MAP {'a': 3, 'b': 4}, MAP {'x': 9}; + +query ?I? rowsort +SELECT NULLIF(m, n), NULLIF(m, n)['b'], map_extract(NULLIF(m, n), 'b') FROM map_nullif; +---- +NULL NULL NULL +{a: 3, b: 4} 4 [4] + +statement ok +DROP TABLE map_nullif; + # accessing map with non-string key query I SELECT MAKE_MAP(1, null, 2, 33, 3, null)[2]; @@ -754,7 +782,7 @@ NULL NULL -query error DataFusion error: Arrow error: Invalid argument error +query error DataFusion error: Execution error: The key type Null does not match the map key type Int64 SELECT column1[NULL] FROM map_array_table_1; query ???