diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 94830ee360585..df64df652abcb 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -306,6 +306,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 0000000000000..013c7c2081668 --- /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> = (0..N_ROWS) + .map(|_| { + let r = rng.random::(); + 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 = (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::>(); + 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 0a83eb3ed61ef..fa23532406ce1 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::::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> = + (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 = (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::; 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 = (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::(&strings, &long_refs).unwrap(); + let long = long.as_any().downcast_ref::().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::(&strings, &short_refs).unwrap(); + let short = short.as_any().downcast_ref::().unwrap(); + assert_eq!(short.value(0), 0); + assert!(short.is_null(4)); + assert_eq!(short.value(5), 2); // "b" at position 2 + } }