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-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/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); 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 ???