This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25130-224cc563cc38617f85a5c22f63265961dde2b57e in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 301235e7a84780c88192ef041a2071c313657f84 Author: Neil Conway <[email protected]> AuthorDate: Fri Sep 11 14:53:56 2026 +0000 fix: `map[key]` panics for nested map keys (#25130) ## Which issue does this PR close? - Closes #25082. ## Rationale for this change `get_field` has two map lookup paths: * `process_map_with_nested_key` uses a comparator and works for all key types * `process_map_array` uses the `eq` kernel, and therefore does not support map keys that are nested types (e.g., lists). The previous implementation arranged to call `process_map_with_nested_key` for map-typed keys, but that was too narrow: this path should be taken for list- and struct-typed keys as well. ## What changes are included in this PR? * Route every nested key type through `process_map_with_nested_key` * Add unit test and SLTs ## What is the testing strategy for this PR? Existing tests pass; new test cases added. ## Are there any user-facing changes? No, aside from addressing the bug. --- datafusion/functions/src/core/getfield.rs | 88 +++++++++++++++++++----------- datafusion/sqllogictest/test_files/map.slt | 30 ++++++++++ 2 files changed, 85 insertions(+), 33 deletions(-) diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index 388daa32bf..a0f024bbc7 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -18,12 +18,11 @@ use std::sync::{Arc, OnceLock}; use arrow::array::{ - Array, BooleanArray, Capacities, MutableArrayData, Scalar, cast::AsArray, make_array, + Array, Capacities, MutableArrayData, Scalar, cast::AsArray, make_array, make_comparator, }; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, FieldRef}; -use arrow_buffer::NullBuffer; use datafusion_common::cast::{as_map_array, as_struct_array}; use datafusion_common::{ @@ -100,29 +99,19 @@ impl Default for GetFieldFunc { } } -/// Process a map array by finding matching keys and extracting corresponding values. +/// 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. /// -/// This function handles both simple (scalar) and nested key types by using -/// appropriate comparison strategies. +/// `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<dyn Array>, ) -> Result<ColumnarValue> { let map_array = as_map_array(array)?; - let keys = if key_array.data_type().is_nested() { - let comparator = make_comparator( - map_array.keys().as_ref(), - key_array.as_ref(), - SortOptions::default(), - )?; - let len = map_array.keys().len().min(key_array.len()); - let values = (0..len).map(|i| comparator(i, i).is_eq()).collect(); - let nulls = NullBuffer::union(map_array.keys().nulls(), key_array.nulls()); - BooleanArray::new(values, nulls) - } else { - let be_compared = Scalar::new(key_array); - arrow::compute::kernels::cmp::eq(&be_compared, map_array.keys())? - }; + 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()); @@ -219,19 +208,14 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result<Column dict.with_values(Arc::clone(field_col)), )) } - (DataType::Map(_, _), ScalarValue::List(arr), _) => { - let key_array: Arc<dyn Array> = arr; - process_map_array(&array, key_array) - } - (DataType::Map(_, _), ScalarValue::Struct(arr), _) => { - process_map_array(&array, arr as Arc<dyn Array>) - } - (DataType::Map(_, _), other, _) => { - let data_type = other.data_type(); - if data_type.is_nested() { - process_map_with_nested_key(&array, &other.to_array()?) + (DataType::Map(_, _), key, _) => { + // 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, other.to_array()?) + process_map_array(&array, key_array) } } (DataType::Struct(_), _, Some(k)) => { @@ -672,8 +656,11 @@ impl ScalarUDFImpl for GetFieldFunc { #[cfg(test)] mod tests { use super::*; - use arrow::array::{ArrayRef, Int32Array, StructArray}; - use arrow::datatypes::Fields; + use arrow::array::{ + ArrayRef, Int32Array, Int32Builder, ListArray, ListBuilder, MapBuilder, + StructArray, + }; + use arrow::datatypes::{Fields, Int32Type}; #[test] fn test_get_field_utf8view_key() -> Result<()> { @@ -711,6 +698,41 @@ 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<i32>| { + ScalarValue::List(Arc::new( + ListArray::from_iter_primitive::<Int32Type, _, _>([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/sqllogictest/test_files/map.slt b/datafusion/sqllogictest/test_files/map.slt index 32fdeaac1d..debaf7f8b3 100644 --- a/datafusion/sqllogictest/test_files/map.slt +++ b/datafusion/sqllogictest/test_files/map.slt @@ -597,6 +597,36 @@ NULL # TODO(https://github.com/apache/datafusion/pull/18394): Test accessing map with empty map as key # TODO(https://github.com/apache/datafusion/pull/18394): Test accessing map with null map as key +# accessing map with list key: match on a later entry, absent key, match on the first entry +query III +SELECT MAP {[1,2]:1, [3,4]:2}[[3,4]], MAP {[1,2]:1, [3,4]:2}[[9,9]], MAP {[1,2]:1, [3,4]:2}[[1,2]]; +---- +2 NULL 1 + +# accessing map with struct key +query II +SELECT MAP {named_struct('a',1):1, named_struct('a',2):2}[named_struct('a',2)], + MAP {named_struct('a',1):1, named_struct('a',2):2}[named_struct('a',3)]; +---- +2 NULL + +# accessing map with list key from a column, so the lookup runs on a batch rather than at planning time +statement ok +CREATE TABLE map_list_keys AS +SELECT MAP {[1,2]:1, [3,4]:2} AS m +UNION ALL SELECT MAP {[3,4]:5, [1,2]:6} +UNION ALL SELECT NULL; + +query II rowsort +SELECT m[[3,4]], m[[9,9]] FROM map_list_keys; +---- +2 NULL +5 NULL +NULL NULL + +statement ok +DROP TABLE map_list_keys; + # accessing map with non-string key query I SELECT MAKE_MAP(1, null, 2, 33, 3, null)[2]; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
