Skip to content

Commit feccbd2

Browse files
neilconwayneilconway
committed
fix: Use Arrow comparator for key comparison in map_extract
The previous implementation of `map_extract` did the following for row: 1. Create a one-element array slice containing the row's search key 2. Scan the map's entries. For each entry, create a one-element array slice and compare the two slices using Arrow's array equality 3. Stop at the first match; if no matches, append a NULL instead This had three shortcomings: 1. It was very inefficient, because a lot of allocations are done for every element of every map. 2. It got the equality semantics wrong for some corner-cases. In particular, maps with dictionary-valued keys might encode a logical NULL in two physically distinct ways (apache#24983). Arrow's array equality also considers sparse unions that have different values in unselected child fields to be distinct; this is arguably a bug in Arrow though. 3. It returned `[NULL]` for missing map keys instead of an empty list, which is the behavior implemented by DuckDB (apache#24981). Instead, we can implement `map_extract` with a single arrow-ord comparator. This enables comparing the search key with each map element directly by index, without allocating. It also avoids the differences in comparison semantics outlined above. Finally, this PR fixes the behavior for absent map keys to be consistent with DuckDB. Benchmark results (M4 Max): - int32/first/1024x32, 135.629 µs -> 6.234 µs, -95.40% - int32/last/1024x1, 133.110 µs -> 6.157 µs, -95.37% - int32/last/1024x32, 3477.393 µs -> 38.562 µs, -98.89% - int32/last/1x0, 0.498 µs -> 0.313 µs, -37.23% - int32/last/1x1, 0.593 µs -> 0.384 µs, -35.34% - int32/missing/1024x32, 3472.698 µs -> 34.889 µs, -99.00% - int32/varying/1024x32, 1844.686 µs -> 25.307 µs, -98.63% - struct/first/1024x32, 335.421 µs -> 8.769 µs, -97.39% - struct/last/1024x1, 335.097 µs -> 8.647 µs, -97.42% - struct/last/1024x32, 8325.830 µs -> 71.611 µs, -99.14% - struct/last/1x0, 0.466 µs -> 0.242 µs, -48.08% - struct/last/1x1, 0.747 µs -> 0.391 µs, -47.67% - struct/missing/1024x32, 8451.153 µs -> 61.480 µs, -99.27% - struct/varying/1024x32, 4498.016 µs -> 41.343 µs, -99.08% - utf8_view/first/1024x32, 218.809 µs -> 9.877 µs, -95.49% - utf8_view/last/1024x1, 190.468 µs -> 8.464 µs, -95.56% - utf8_view/last/1024x32, 6016.140 µs -> 124.405 µs, -97.93% - utf8_view/last/1x0, 0.526 µs -> 0.353 µs, -32.97% - utf8_view/last/1x1, 0.762 µs -> 0.523 µs, -31.42% - utf8_view/missing/1024x32, 5999.511 µs -> 114.082 µs, -98.10% - utf8_view/varying/1024x32, 3226.583 µs -> 71.732 µs, -97.78% ("1024x32" means 1024 rows and each row is a map with 32 entries.)
1 parent 262936e commit feccbd2

3 files changed

Lines changed: 271 additions & 47 deletions

File tree

‎datafusion/functions-nested/benches/map.rs‎

Lines changed: 109 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,17 +16,18 @@
1616
// under the License.
1717

1818
use arrow::array::{
19-
ArrayRef, BinaryArray, BinaryViewArray, Int32Array, ListArray, StringArray,
20-
StringViewArray,
19+
Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, ListArray, MapArray,
20+
StringArray, StringViewArray, StructArray,
2121
};
2222
use arrow::buffer::{OffsetBuffer, ScalarBuffer};
23-
use arrow::datatypes::Field;
24-
use criterion::{Criterion, criterion_group, criterion_main};
23+
use arrow::datatypes::{DataType, Field};
24+
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
2525
use datafusion_common::ScalarValue;
2626
use datafusion_common::config::ConfigOptions;
2727
use datafusion_expr::planner::ExprPlanner;
2828
use datafusion_expr::{ColumnarValue, Expr, ScalarFunctionArgs};
2929
use datafusion_functions_nested::map::map_udf;
30+
use datafusion_functions_nested::map_extract::map_extract_udf;
3031
use datafusion_functions_nested::planner::NestedFunctionPlanner;
3132
use rand::prelude::*;
3233
use std::collections::HashSet;
@@ -208,5 +209,108 @@ fn criterion_benchmark(c: &mut Criterion) {
208209
}
209210
}
210211

211-
criterion_group!(benches, criterion_benchmark);
212+
fn bench_map_extract(c: &mut Criterion) {
213+
let udf = map_extract_udf();
214+
let config_options = Arc::new(ConfigOptions::default());
215+
let mut group = c.benchmark_group("map_extract");
216+
217+
for (rows, width) in [(1, 0), (1, 1), (1024, 1), (1024, 32)] {
218+
for key_type in ["int32", "utf8_view", "struct"] {
219+
let make_keys = |keys: Vec<i32>| -> ArrayRef {
220+
match key_type {
221+
"int32" => Arc::new(Int32Array::from(keys)),
222+
"utf8_view" => Arc::new(StringViewArray::from_iter_values(
223+
keys.iter().map(|key| format!("key_{key:016}")),
224+
)),
225+
"struct" => Arc::new(StructArray::from(vec![(
226+
Arc::new(Field::new("key", DataType::Int32, false)),
227+
Arc::new(Int32Array::from(keys)) as ArrayRef,
228+
)])),
229+
_ => unreachable!(),
230+
}
231+
};
232+
let keys = make_keys((0..rows).flat_map(|_| 0..width as i32).collect());
233+
let entries = StructArray::from(vec![
234+
(
235+
Arc::new(Field::new("key", keys.data_type().clone(), false)),
236+
keys,
237+
),
238+
(
239+
Arc::new(Field::new("value", DataType::Int32, false)),
240+
Arc::new(Int32Array::from_iter_values(0..(rows * width) as i32))
241+
as ArrayRef,
242+
),
243+
]);
244+
let map: ArrayRef = Arc::new(MapArray::new(
245+
Arc::new(Field::new("entries", entries.data_type().clone(), false)),
246+
OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows)),
247+
entries,
248+
None,
249+
false,
250+
));
251+
let lookups: &[&str] = if width <= 1 {
252+
&["last"]
253+
} else {
254+
&["first", "last", "missing", "varying"]
255+
};
256+
for &lookup in lookups {
257+
let query_keys = match lookup {
258+
"first" => vec![0],
259+
"last" => vec![width.saturating_sub(1) as i32],
260+
"missing" => vec![width as i32],
261+
// Mix matches and misses with a different lookup key per row.
262+
"varying" => {
263+
(0..rows).map(|row| (row % (width + 1)) as i32).collect()
264+
}
265+
_ => unreachable!(),
266+
};
267+
let query_keys = make_keys(query_keys);
268+
let query_keys = if lookup == "varying" {
269+
ColumnarValue::Array(query_keys)
270+
} else {
271+
ColumnarValue::Scalar(
272+
ScalarValue::try_from_array(&query_keys, 0).unwrap(),
273+
)
274+
};
275+
let args = vec![ColumnarValue::Array(Arc::clone(&map)), query_keys];
276+
let arg_fields = args
277+
.iter()
278+
.map(|arg| Field::new("arg", arg.data_type(), true).into())
279+
.collect::<Vec<_>>();
280+
let return_type = udf
281+
.return_type(
282+
&args
283+
.iter()
284+
.map(ColumnarValue::data_type)
285+
.collect::<Vec<_>>(),
286+
)
287+
.unwrap();
288+
let return_field = Arc::new(Field::new("result", return_type, true));
289+
group.bench_function(
290+
BenchmarkId::new(
291+
format!("{key_type}/{lookup}"),
292+
format!("{rows}x{width}"),
293+
),
294+
|b| {
295+
b.iter(|| {
296+
black_box(
297+
udf.invoke_with_args(ScalarFunctionArgs {
298+
args: args.clone(),
299+
arg_fields: arg_fields.clone(),
300+
number_rows: rows,
301+
return_field: Arc::clone(&return_field),
302+
config_options: Arc::clone(&config_options),
303+
})
304+
.unwrap(),
305+
)
306+
});
307+
},
308+
);
309+
}
310+
}
311+
}
312+
group.finish();
313+
}
314+
315+
criterion_group!(benches, criterion_benchmark, bench_map_extract);
212316
criterion_main!(benches);

‎datafusion/functions-nested/src/map_extract.rs‎

Lines changed: 113 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,12 @@
1919
2020
use crate::utils::{get_map_entry_field, make_scalar_function};
2121
use arrow::array::{
22-
Array, ArrayRef, Capacities, ListArray, MapArray, MutableArrayData, make_array,
22+
Array, ArrayRef, ListArray, MapArray, MutableArrayData, make_array, new_empty_array,
2323
};
2424
use arrow::buffer::OffsetBuffer;
25+
use arrow::compute::SortOptions;
2526
use arrow::datatypes::{DataType, Field};
27+
use arrow_ord::ord::make_comparator;
2628
use datafusion_common::utils::take_function_args;
2729
use datafusion_common::{Result, cast::as_map_array, exec_err};
2830
use datafusion_expr::{
@@ -31,7 +33,6 @@ use datafusion_expr::{
3133
};
3234
use datafusion_macros::user_doc;
3335
use std::sync::Arc;
34-
use std::vec;
3536

3637
// Create static instances of ScalarUDFs for each function
3738
make_udf_expr_and_func!(
@@ -149,43 +150,55 @@ fn general_map_extract_inner(
149150
query_keys_array: &dyn Array,
150151
) -> Result<ArrayRef> {
151152
let keys = map_array.keys();
152-
let mut offsets = vec![0_i32];
153-
154153
let values = map_array.values();
155-
let original_data = values.to_data();
156-
let capacity = Capacities::Array(original_data.len());
154+
let field = Arc::new(Field::new_list_field(map_array.value_type().clone(), true));
155+
let map_offsets = map_array.value_offsets();
156+
if map_offsets.first() == map_offsets.last() {
157+
return Ok(Arc::new(ListArray::new(
158+
field,
159+
OffsetBuffer::new_zeroed(map_array.len()),
160+
new_empty_array(values.data_type()),
161+
map_array.nulls().cloned(),
162+
)));
163+
}
157164

158-
let mut mutable =
159-
MutableArrayData::with_capacities(vec![&original_data], true, capacity);
165+
// Compare keys by index using a single comparator for the batch.
166+
let compare =
167+
make_comparator(keys.as_ref(), query_keys_array, SortOptions::default())?;
168+
let mut offsets = Vec::with_capacity(map_array.len() + 1);
169+
offsets.push(0_i32);
160170

161-
for (row_index, offset_window) in map_array.value_offsets().windows(2).enumerate() {
171+
let original_data = values.to_data();
172+
// There is at most one output value per map row.
173+
let mut mutable = MutableArrayData::new(
174+
vec![&original_data],
175+
false,
176+
map_array.len().min(values.len()),
177+
);
178+
179+
for (row_index, offset_window) in map_offsets.windows(2).enumerate() {
162180
let start = offset_window[0] as usize;
163181
let end = offset_window[1] as usize;
164-
let len = end - start;
165-
166-
let query_key = query_keys_array.slice(row_index, 1);
182+
let mut offset = offsets[row_index];
167183

168-
let value_index =
169-
(0..len).find(|&i| keys.slice(start + i, 1).as_ref() == query_key.as_ref());
170-
171-
match value_index {
172-
Some(index) => {
173-
mutable.try_extend(0, start + index, start + index + 1)?;
174-
}
175-
None => {
176-
mutable.try_extend_nulls(1)?;
177-
}
184+
if map_array.is_valid(row_index)
185+
&& let Some(index) = (start..end).find(|&i| compare(i, row_index).is_eq())
186+
{
187+
mutable.try_extend(0, index, index + 1)?;
188+
offset += 1;
178189
}
179-
offsets.push(offsets[row_index] + 1);
190+
191+
// A missing key results in an empty list.
192+
offsets.push(offset);
180193
}
181194

182195
let data = mutable.freeze();
183196

184197
Ok(Arc::new(ListArray::new(
185-
Arc::new(Field::new_list_field(map_array.value_type().clone(), true)),
198+
field,
186199
OffsetBuffer::<i32>::new(offsets.into()),
187-
Arc::new(make_array(data)),
188-
None,
200+
make_array(data),
201+
map_array.nulls().cloned(),
189202
)))
190203
}
191204

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

211224
general_map_extract_inner(map_array, key_arg)
212225
}
226+
227+
#[cfg(test)]
228+
mod tests {
229+
use super::*;
230+
use arrow::array::{Float64Array, Int32Array, StructArray};
231+
use arrow::datatypes::Int32Type;
232+
233+
fn make_map(keys: ArrayRef, values: Vec<i32>, offsets: Vec<i32>) -> MapArray {
234+
let entries = StructArray::from(vec![
235+
(
236+
Arc::new(Field::new("key", keys.data_type().clone(), false)),
237+
keys,
238+
),
239+
(
240+
Arc::new(Field::new("value", DataType::Int32, true)),
241+
Arc::new(Int32Array::from(values)) as ArrayRef,
242+
),
243+
]);
244+
MapArray::new(
245+
Arc::new(Field::new("entries", entries.data_type().clone(), false)),
246+
OffsetBuffer::new(offsets.into()),
247+
entries,
248+
None,
249+
false,
250+
)
251+
}
252+
253+
#[test]
254+
fn map_extract_sliced_maps() -> Result<()> {
255+
let map = make_map(
256+
Arc::new(Int32Array::from(vec![0, 1, 2, 3])),
257+
vec![0, 10, 20, 30],
258+
vec![0, 1, 3, 4],
259+
);
260+
let query_keys = Int32Array::from(vec![0, 2, 9]);
261+
262+
// Map offsets address the original entries; query indices address the slice.
263+
let result =
264+
general_map_extract_inner(&map.slice(1, 2), &query_keys.slice(1, 2))?;
265+
let expected = ListArray::from_iter_primitive::<Int32Type, _, _>([
266+
Some(vec![Some(20)]),
267+
Some(vec![]),
268+
]);
269+
assert_eq!(result.as_ref(), &expected);
270+
271+
// Empty slices may retain the original nonempty keys and values buffers.
272+
let result =
273+
general_map_extract_inner(&map.slice(1, 0), &query_keys.slice(1, 0))?;
274+
assert_eq!(result.len(), 0);
275+
Ok(())
276+
}
277+
278+
#[test]
279+
fn map_extract_float_keys() -> Result<()> {
280+
let nan = f64::NAN;
281+
let other_nan = f64::from_bits(nan.to_bits() + 1);
282+
let map = make_map(
283+
Arc::new(Float64Array::from(vec![-0.0, 0.0, nan, other_nan])),
284+
vec![1, 2, 3, 4],
285+
vec![0, 4],
286+
);
287+
288+
// Signed zeros and distinct NaN payloads identify different keys.
289+
for (query, expected) in [(-0.0, 1), (0.0, 2), (nan, 3), (other_nan, 4)] {
290+
let result =
291+
general_map_extract_inner(&map, &Float64Array::from(vec![query]))?;
292+
let expected = ListArray::from_iter_primitive::<Int32Type, _, _>([Some(
293+
vec![Some(expected)],
294+
)]);
295+
assert_eq!(result.as_ref(), &expected);
296+
}
297+
Ok(())
298+
}
299+
}

0 commit comments

Comments
 (0)