Skip to content

fix: map[key] panics for nested map keys #25130

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
neilconway merged 2 commits into from
Sep 11, 2026
Merged

fix: map[key] panics for nested map keys #25130

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b9ea60b
fix: `map[key]` panics for list and struct keys when the match is not…
neilconway Sep 9, 2026
b346690
Merge remote-tracking branch 'origin/main' into neilc/fix-map-getfield
neilconway Sep 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 55 additions & 33 deletions datafusion/functions/src/core/getfield.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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)) => {
Expand Down Expand Up @@ -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<()> {
Expand Down Expand Up @@ -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};
Expand Down
30 changes: 30 additions & 0 deletions datafusion/sqllogictest/test_files/map.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down