Skip to content

fix, perf: Use Arrow comparator for key comparison in map_extract #24999

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 1 commit into
base: main
Choose a base branch
from
+271 −47
Open

fix, perf: Use Arrow comparator for key comparison in map_extract #24999

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
114 changes: 109 additions & 5 deletions datafusion/functions-nested/benches/map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<i32>| -> 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::<Vec<_>>();
let return_type = udf
.return_type(
&args
.iter()
.map(ColumnarValue::data_type)
.collect::<Vec<_>>(),
)
.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);
139 changes: 113 additions & 26 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 @@ -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::{
Expand All @@ -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!(
Expand Down Expand Up @@ -149,43 +150,55 @@ fn general_map_extract_inner(
query_keys_array: &dyn Array,
) -> Result<ArrayRef> {
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::<i32>::new(offsets.into()),
Arc::new(make_array(data)),
None,
make_array(data),
map_array.nulls().cloned(),
)))
}

Expand All @@ -210,3 +223,77 @@ fn map_extract_inner(args: &[ArrayRef]) -> Result<ArrayRef> {

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<i32>, offsets: Vec<i32>) -> 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::<Int32Type, _, _>([
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::<Int32Type, _, _>([Some(
vec![Some(expected)],
)]);
assert_eq!(result.as_ref(), &expected);
}
Ok(())
}
}
Loading