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