From 9b40f71af9e2b569a782d4fdbade304f3b343a4e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 13 Jul 2026 15:06:14 -0600 Subject: [PATCH] perf: optimize get_field in datafusion-functions --- datafusion/functions/Cargo.toml | 4 + datafusion/functions/benches/get_field.rs | 95 +++++++++++++++++++++++ datafusion/functions/src/core/getfield.rs | 26 ++++--- 3 files changed, 113 insertions(+), 12 deletions(-) create mode 100644 datafusion/functions/benches/get_field.rs diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 94830ee360585..90a6d2fc4f7b5 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -326,6 +326,10 @@ harness = false name = "regexp_count" required-features = ["regex_expressions"] +[[bench]] +harness = false +name = "get_field" + [[bench]] harness = false name = "crypto" diff --git a/datafusion/functions/benches/get_field.rs b/datafusion/functions/benches/get_field.rs new file mode 100644 index 0000000000000..8a5fd0a1e2fa9 --- /dev/null +++ b/datafusion/functions/benches/get_field.rs @@ -0,0 +1,95 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +extern crate criterion; + +use arrow::array::{ArrayRef, Int32Builder, MapBuilder, StringBuilder}; +use arrow::datatypes::{DataType, Field}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_functions::core::get_field; +use std::hint::black_box; +use std::sync::Arc; + +/// A map array with `size` rows, each holding `entries` key/value pairs. +/// Every tenth row is null. +fn map_array(size: usize, entries: usize) -> ArrayRef { + let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + for row in 0..size { + if row % 10 == 0 { + builder.append(false).unwrap(); + continue; + } + for entry in 0..entries { + builder.keys().append_value(format!("key_{entry}")); + builder.values().append_value((row * entry) as i32); + } + builder.append(true).unwrap(); + } + Arc::new(builder.finish()) +} + +fn bench_get_field( + c: &mut Criterion, + name: &str, + size: usize, + entries: usize, + key: &str, +) { + let udf = get_field(); + let args = vec![ + ColumnarValue::Array(map_array(size, entries)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(key.to_string()))), + ]; + let arg_fields = vec![ + Field::new("map", args[0].data_type(), true).into(), + Field::new("key", DataType::Utf8, false).into(), + ]; + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_function(name, |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Int32, true).into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + // First key: the match is found immediately, so the per-row overhead + // dominates. + bench_get_field(c, "get_field_map_1024_entries_4_first", 1024, 4, "key_0"); + // Last key: every entry of the row is compared before the match. + bench_get_field(c, "get_field_map_1024_entries_4_last", 1024, 4, "key_3"); + bench_get_field(c, "get_field_map_1024_entries_16_last", 1024, 16, "key_15"); + // Key that is not present in any row. + bench_get_field(c, "get_field_map_1024_entries_4_missing", 1024, 4, "key_9"); + bench_get_field(c, "get_field_map_8192_entries_4_last", 8192, 4, "key_3"); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index 70fc8bb0ea129..931111f9cfd8f 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -129,22 +129,24 @@ fn process_map_array( let mut mutable = MutableArrayData::with_capacities(vec![&original_data], true, capacity); + let offsets = map_array.value_offsets(); + // Scan the comparison result in place: slicing it per entry would allocate + // a new array for every row of the map. + let matches = keys.values(); + let match_nulls = keys.nulls(); + for entry in 0..map_array.len() { - let start = map_array.value_offsets()[entry] as usize; - let end = map_array.value_offsets()[entry + 1] as usize; + let start = offsets[entry] as usize; + let end = offsets[entry + 1] as usize; - let maybe_matched = keys - .slice(start, end - start) - .iter() - .enumerate() - .find(|(_, t)| t.unwrap()); + let matched = (start..end).find(|&i| { + matches.value(i) && match_nulls.is_none_or(|nulls| nulls.is_valid(i)) + }); - if maybe_matched.is_none() { - mutable.try_extend_nulls(1)?; - continue; + match matched { + Some(i) => mutable.try_extend(0, i, i + 1)?, + None => mutable.try_extend_nulls(1)?, } - let (match_offset, _) = maybe_matched.unwrap(); - mutable.try_extend(0, start + match_offset, start + match_offset + 1)?; } let data = mutable.freeze();