Skip to content

perf: Unify, optimize map[k] and map_extract(key) #25201

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

Open
neilconway wants to merge 4 commits into
base: main
Choose a base branch
from
Open

perf: Unify, optimize map[k] and map_extract(key) #25201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d029ea1
bench: measure get_field and map_extract map lookups on the same shapes
neilconway Sep 9, 2026
b33a1c0
perf, chore: Consolidate and optimize `get_field` and `map_extract`
neilconway Sep 9, 2026
d3e41be
Merge remote-tracking branch 'origin/main' into neilc/map-lookup-kernel
neilconway Sep 11, 2026
76d7c62
Merge branch 'main' into neilc/map-lookup-kernel
comphead Sep 12, 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
100 changes: 66 additions & 34 deletions datafusion/functions-nested/benches/map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<i32>| -> ArrayRef {
match key_type {
"int32" => Arc::new(Int32Array::from(keys)),
Expand All @@ -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);
Expand All @@ -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())
Expand Down
102 changes: 30 additions & 72 deletions datafusion/functions-nested/src/map_extract.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -119,7 +118,11 @@ impl ScalarUDFImpl for MapExtract {
}

fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
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] {
Expand Down Expand Up @@ -149,79 +152,34 @@ fn general_map_extract_inner(
map_array: &MapArray,
query_keys_array: &dyn Array,
) -> Result<ArrayRef> {
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::<i32>::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<ArrayRef> {
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)]
Expand Down
Loading