This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23460-e9a75bf47a962cd7ee11e563bc948f3f0d0a51c8 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit b8d3b0b525afaf7049320c426a2be3884600dd77 Author: Andy Grove <[email protected]> AuthorDate: Fri Jul 24 06:21:28 2026 -0600 perf: optimize `find_in_set` (up to 24x faster) (#23460) ## Which issue does this PR close? N/A ## Rationale for this change Improve performance of existing expression. ## What changes are included in this PR? Replace per-row O(set_len) linear scan in find_in_set's constant-list path with a one-time HashMap lookup (threshold-guarded so short lists keep the linear scan), giving O(1) per-row probing for large sets. ## Are these changes tested? Existing tests + new tests Benchmark (criterion): - long_list_256: 95.827% faster (base 1246412ns -> cand 52009ns) - ~24x faster - short_list_4: 2.13% faster (base 64343ns -> cand 62972ns) - long_list_64: 88.562% faster (base 422083ns -> cand 48276ns) ## Are there any user-facing changes? No <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> Co-authored-by: Jeffrey Vo <[email protected]> --- datafusion/functions/Cargo.toml | 5 ++ .../functions/benches/find_in_set_literal.rs | 98 ++++++++++++++++++++++ datafusion/functions/src/unicode/find_in_set.rs | 75 +++++++++++++++-- 3 files changed, 173 insertions(+), 5 deletions(-) diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index d0ce0d0be3..83f8c0f2a3 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -330,6 +330,11 @@ harness = false name = "find_in_set" required-features = ["unicode_expressions"] +[[bench]] +harness = false +name = "find_in_set_literal" +required-features = ["unicode_expressions"] + [[bench]] harness = false name = "contains" diff --git a/datafusion/functions/benches/find_in_set_literal.rs b/datafusion/functions/benches/find_in_set_literal.rs new file mode 100644 index 0000000000..013c7c2081 --- /dev/null +++ b/datafusion/functions/benches/find_in_set_literal.rs @@ -0,0 +1,98 @@ +// 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. + +//! Benchmarks the `find_in_set(column, constant_list)` path where the set is a +//! scalar literal. A long list exercises the pre-built lookup; a short list +//! stays on the per-row linear scan. + +use arrow::array::StringArray; +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::{ColumnarValue, ScalarFunctionArgs}; +use rand::prelude::StdRng; +use rand::{Rng, SeedableRng}; +use std::hint::black_box; +use std::sync::Arc; + +const N_ROWS: usize = 8192; + +/// Builds a string column whose values are drawn from `entries` plus a small +/// fraction of misses, so both hits and misses are exercised. +fn build_column(entries: &[String]) -> StringArray { + let mut rng = StdRng::seed_from_u64(42); + let values: Vec<Option<String>> = (0..N_ROWS) + .map(|_| { + let r = rng.random::<f32>(); + if r < 0.1 { + None + } else if r < 0.4 { + Some("__miss__".to_string()) + } else { + let idx = rng.random_range(0..entries.len()); + Some(entries[idx].clone()) + } + }) + .collect(); + StringArray::from(values) +} + +fn bench_case(c: &mut Criterion, label: &str, num_entries: usize) { + let find_in_set = datafusion_functions::unicode::find_in_set(); + let entries: Vec<String> = (0..num_entries).map(|i| format!("item{i}")).collect(); + let list = entries.join(","); + + let column = build_column(&entries); + let args = vec![ + ColumnarValue::Array(Arc::new(column)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(list))), + ]; + let arg_fields = args + .iter() + .map(|arg| Field::new("a", arg.data_type().clone(), true).into()) + .collect::<Vec<_>>(); + let return_field = Arc::new(Field::new("f", DataType::Int32, true)); + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_with_input( + BenchmarkId::new("find_in_set_literal", label), + &num_entries, + |b, _| { + b.iter(|| { + black_box(find_in_set.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: N_ROWS, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + // Short list stays on the linear scan (below the lookup threshold). + bench_case(c, "short_list_4", 4); + // Long lists exercise the pre-built lookup. + bench_case(c, "long_list_64", 64); + bench_case(c, "long_list_256", 256); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/src/unicode/find_in_set.rs b/datafusion/functions/src/unicode/find_in_set.rs index 0a83eb3ed6..fa23532406 100644 --- a/datafusion/functions/src/unicode/find_in_set.rs +++ b/datafusion/functions/src/unicode/find_in_set.rs @@ -25,7 +25,7 @@ use arrow_buffer::NullBuffer; use crate::utils::utf8_to_int_type; use datafusion_common::{ - Result, ScalarValue, exec_err, internal_err, utils::take_function_args, + HashMap, Result, ScalarValue, exec_err, internal_err, utils::take_function_args, }; use datafusion_expr::TypeSignature::Exact; use datafusion_expr::{ @@ -316,6 +316,11 @@ where Ok(Arc::new(PrimitiveArray::<T>::new(values.into(), nulls)) as ArrayRef) } +/// Minimum set length at which a pre-built lookup beats a per-row linear scan. +/// Below this, the linear scan's small constant factor wins, so short sets are +/// left untouched to avoid regressing them. +const FIND_IN_SET_LOOKUP_THRESHOLD: usize = 16; + fn find_in_set_right_literal<'a, T, V>( string_array: V, str_list: &[&str], @@ -329,16 +334,34 @@ where let nulls = string_array.nulls().cloned(); let zero = T::Native::from_usize(0).unwrap(); + // The set (`str_list`) is constant across all rows. For a large set, the + // per-row `position` linear scan is O(set_len). Building a lookup from each + // distinct entry to its 1-based position once turns each row into an O(1) + // probe (first occurrence wins, exactly matching `position`). Below the + // threshold the linear scan's small constant factor is faster, so the map is + // built at most once here rather than per row. + let map: Option<HashMap<&str, usize>> = + (str_list.len() >= FIND_IN_SET_LOOKUP_THRESHOLD).then(|| { + let mut map = HashMap::with_capacity(str_list.len()); + for (idx, entry) in str_list.iter().enumerate() { + map.entry(*entry).or_insert(idx + 1); + } + map + }); + let values: Vec<T::Native> = (0..len) .map(|i| { if nulls.as_ref().is_some_and(|n| n.is_null(i)) { return zero; } let string = string_array.value(i); - let position = str_list - .iter() - .position(|s| *s == string) - .map_or(0, |idx| idx + 1); + let position = match &map { + Some(map) => map.get(string).copied().unwrap_or(0), + None => str_list + .iter() + .position(|s| *s == string) + .map_or(0, |idx| idx + 1), + }; T::Native::from_usize(position).unwrap() }) .collect(); @@ -545,4 +568,46 @@ mod tests { ], Int32Array::from(vec![None::<i32>; 3]) ); + + // Exercises both the lookup-map path (list length >= threshold) and the + // linear-scan path (short list), including a duplicate entry to confirm the + // first occurrence wins in both. + #[test] + fn test_right_literal_lookup_matches_linear() { + use super::find_in_set_right_literal; + use arrow::datatypes::Int32Type; + + // 40 unique entries plus a duplicate of "item5" appended at index 40, so + // the length is well over FIND_IN_SET_LOOKUP_THRESHOLD. + let mut long_list: Vec<String> = (0..40).map(|i| format!("item{i}")).collect(); + long_list.push("item5".to_string()); + let long_refs: Vec<&str> = long_list.iter().map(|s| s.as_str()).collect(); + let short_refs = ["a", "b", "c"]; + + let strings = StringArray::from(vec![ + Some("item0"), + Some("item39"), + Some("item5"), + Some("missing"), + None, + Some("b"), + ]); + + let long = + find_in_set_right_literal::<Int32Type, _>(&strings, &long_refs).unwrap(); + let long = long.as_any().downcast_ref::<Int32Array>().unwrap(); + assert_eq!(long.value(0), 1); + assert_eq!(long.value(1), 40); + assert_eq!(long.value(2), 6); // first occurrence of "item5" + assert_eq!(long.value(3), 0); + assert!(long.is_null(4)); + assert_eq!(long.value(5), 0); + + let short = + find_in_set_right_literal::<Int32Type, _>(&strings, &short_refs).unwrap(); + let short = short.as_any().downcast_ref::<Int32Array>().unwrap(); + assert_eq!(short.value(0), 0); + assert!(short.is_null(4)); + assert_eq!(short.value(5), 2); // "b" at position 2 + } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
