You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/26 17:28:59 UTC
(datafusion-comet) 01/04: perf: optimize string contains function (#5322)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 37b92fe56f4dba2d7bedbb0d9a80fd0f2452a1ed
Author: Kazantsev Maksim <ka...@gmail.com>
AuthorDate: Sat Sep 26 15:59:57 2026 +0000
perf: optimize string contains function (#5322)
* impl map_from_entries
* Revert "impl map_from_entries"
This reverts commit 768b3e90f261c7aea58bdb98dc698b90deeeae34.
* work
* work
* work
* fmt
* clippy
* clippy
* address comments
* address comments
* address comments
* address comments
---------
Co-authored-by: Kazantsev Maksim <mn...@gmail.com>
---
native/spark-expr/benches/contains.rs | 119 ++++++++--
native/spark-expr/src/string_funcs/contains.rs | 286 +++++++++++++++++++------
2 files changed, 326 insertions(+), 79 deletions(-)
diff --git a/native/spark-expr/benches/contains.rs b/native/spark-expr/benches/contains.rs
index 457ce79db3..22fa8eec23 100644
--- a/native/spark-expr/benches/contains.rs
+++ b/native/spark-expr/benches/contains.rs
@@ -28,9 +28,36 @@ use std::sync::Arc;
mod common;
use common::{string_array, NULL_RATIOS, ROW_COUNTS};
+/// Scalar used as the haystack in the scalar/array shape. The matching needle
+/// values below are chosen so the scalar/array shape performs real work.
+const HAYSTACK_SCALAR: &str = "datafusion-comet";
+
+/// Scalar used as the needle in the array/scalar shape.
+const NEEDLE_SCALAR: &str = "comet";
+
+fn build_args(
+ haystack: ColumnarValue,
+ needle: ColumnarValue,
+ number_rows: usize,
+) -> ScalarFunctionArgs {
+ ScalarFunctionArgs {
+ args: vec![haystack, needle],
+ arg_fields: vec![],
+ number_rows,
+ return_field: Arc::new(Field::new("result", DataType::Boolean, true)),
+ config_options: Arc::new(ConfigOptions::default()),
+ }
+}
+
fn criterion_benchmark(c: &mut Criterion) {
let udf = SparkContains::new();
- let mut group = c.benchmark_group("spark_contains");
+
+ // ------------------------------------------------------------------
+ // Shape 1: array haystack vs scalar needle (`contains_array_scalar`).
+ // This path already used a scalar representation on `main`; included as a
+ // regression control since this PR touches it incidentally.
+ // ------------------------------------------------------------------
+ let mut group = c.benchmark_group("spark_contains/array_scalar");
for rows in ROW_COUNTS {
for (null_ratio, tag) in NULL_RATIOS {
let haystack = string_array(rows, null_ratio, |_| "datafusion-comet".to_string());
@@ -40,22 +67,80 @@ fn criterion_benchmark(c: &mut Criterion) {
|b, haystack| {
b.iter(|| {
black_box(
- udf.invoke_with_args(ScalarFunctionArgs {
- args: vec![
- ColumnarValue::Array(Arc::clone(haystack)),
- ColumnarValue::Scalar(ScalarValue::Utf8(Some(
- "comet".to_string(),
- ))),
- ],
- arg_fields: vec![],
- number_rows: haystack.len(),
- return_field: Arc::new(Field::new(
- "result",
- DataType::Boolean,
- true,
- )),
- config_options: Arc::new(ConfigOptions::default()),
- })
+ udf.invoke_with_args(build_args(
+ ColumnarValue::Array(Arc::clone(haystack)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+ NEEDLE_SCALAR.to_string(),
+ ))),
+ haystack.len(),
+ ))
+ .unwrap(),
+ )
+ })
+ },
+ );
+ }
+ }
+ group.finish();
+
+ // ------------------------------------------------------------------
+ // Shape 2: scalar haystack vs array needle (`contains_scalar_array`).
+ // This is the path the PR actually optimizes (it replaced
+ // `to_array_of_size(N)` with an O(1) broadcast), so it must be measured.
+ // The needle array is varied per row so the kernel does non-trivial work.
+ // ------------------------------------------------------------------
+ let mut group = c.benchmark_group("spark_contains/scalar_array");
+ for rows in ROW_COUNTS {
+ for (null_ratio, tag) in NULL_RATIOS {
+ let needle = string_array(rows, null_ratio, |i| {
+ if i % 2 == 0 {
+ "comet".to_string()
+ } else {
+ format!("comet-{i}")
+ }
+ });
+ group.bench_with_input(
+ BenchmarkId::from_parameter(format!("{rows}/{tag}")),
+ &needle,
+ |b, needle| {
+ b.iter(|| {
+ black_box(
+ udf.invoke_with_args(build_args(
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+ HAYSTACK_SCALAR.to_string(),
+ ))),
+ ColumnarValue::Array(Arc::clone(needle)),
+ needle.len(),
+ ))
+ .unwrap(),
+ )
+ })
+ },
+ );
+ }
+ }
+ group.finish();
+
+ // ------------------------------------------------------------------
+ // Shape 3: array haystack vs array needle (`arrow_contains` directly).
+ // Regression control for the straight-through path the PR does not touch.
+ // ------------------------------------------------------------------
+ let mut group = c.benchmark_group("spark_contains/array_array");
+ for rows in ROW_COUNTS {
+ for (null_ratio, tag) in NULL_RATIOS {
+ let haystack = string_array(rows, null_ratio, |_| "datafusion-comet".to_string());
+ let needle = string_array(rows, null_ratio, |_| "comet".to_string());
+ group.bench_with_input(
+ BenchmarkId::from_parameter(format!("{rows}/{tag}")),
+ &(haystack, needle),
+ |b, (haystack, needle)| {
+ b.iter(|| {
+ black_box(
+ udf.invoke_with_args(build_args(
+ ColumnarValue::Array(Arc::clone(haystack)),
+ ColumnarValue::Array(Arc::clone(needle)),
+ haystack.len(),
+ ))
.unwrap(),
)
})
diff --git a/native/spark-expr/src/string_funcs/contains.rs b/native/spark-expr/src/string_funcs/contains.rs
index 537227efdf..545d147560 100644
--- a/native/spark-expr/src/string_funcs/contains.rs
+++ b/native/spark-expr/src/string_funcs/contains.rs
@@ -1,29 +1,25 @@
// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
+// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
+// 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
+// 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
+// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
-//! Optimized `contains` string function for Spark compatibility.
-//!
-//! Optimized for scalar pattern case by passing scalar directly to arrow_contains
-//! instead of expanding to arrays like DataFusion's built-in contains.
-
-use arrow::array::{Array, ArrayRef, BooleanArray, StringArray};
+use arrow::array::{Array, ArrayRef, BooleanArray, Scalar};
+use arrow::compute::kernels::cast::cast;
use arrow::compute::kernels::comparison::contains as arrow_contains;
use arrow::datatypes::DataType;
-use datafusion::common::{exec_err, Result, ScalarValue};
+use datafusion::common::{exec_err, DataFusionError, Result, ScalarValue};
use datafusion::logical_expr::{
ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
};
@@ -83,15 +79,14 @@ fn spark_contains(haystack: &ColumnarValue, needle: &ColumnarValue) -> Result<Co
// Array haystack, scalar needle - OPTIMIZED PATH
(ColumnarValue::Array(haystack_array), ColumnarValue::Scalar(needle_scalar)) => {
- let result = contains_with_arrow_scalar(haystack_array, needle_scalar)?;
+ let result = contains_array_scalar(haystack_array, needle_scalar)?;
Ok(ColumnarValue::Array(result))
}
// Scalar haystack, array needle - less common
(ColumnarValue::Scalar(haystack_scalar), ColumnarValue::Array(needle_array)) => {
- let haystack_array = haystack_scalar.to_array_of_size(needle_array.len())?;
- let result = arrow_contains(&haystack_array, needle_array)?;
- Ok(ColumnarValue::Array(Arc::new(result)))
+ let result = contains_scalar_array(haystack_scalar, needle_array)?;
+ Ok(ColumnarValue::Array(result))
}
// Both scalars - compute single result
@@ -102,9 +97,43 @@ fn spark_contains(haystack: &ColumnarValue, needle: &ColumnarValue) -> Result<Co
}
}
+/// Helper to safely extract string reference from `ScalarValue`.
+/// Unwraps dictionary scalars recursively.
+#[inline]
+fn get_string_scalar_value<'a>(scalar: &'a ScalarValue, arg_name: &str) -> Result<&'a str> {
+ match scalar {
+ ScalarValue::Utf8(Some(s))
+ | ScalarValue::LargeUtf8(Some(s))
+ | ScalarValue::Utf8View(Some(s)) => Ok(s.as_str()),
+ ScalarValue::Dictionary(_, inner) => get_string_scalar_value(inner, arg_name),
+ _ => exec_err!(
+ "contains function requires string type for {}, got {:?}",
+ arg_name,
+ scalar.data_type()
+ ),
+ }
+}
+
+/// Materialize a scalar into a length-1 array whose type matches `target_type`,
+/// so Arrow's CONTAINS kernel accepts the (scalar, array) pair.
+/// Cost is O(1): the cast touches a single element.
+fn scalar_to_aligned_array(
+ scalar: &ScalarValue,
+ target_type: &DataType,
+ arg_name: &str,
+) -> Result<ArrayRef> {
+ let _ = get_string_scalar_value(scalar, arg_name)?;
+ let array = scalar.to_array()?;
+ if array.data_type() == target_type {
+ Ok(array)
+ } else {
+ cast(&array, target_type).map_err(DataFusionError::from)
+ }
+}
+
/// Optimized contains for array haystack with scalar needle.
/// Uses Arrow's native scalar handling for better performance.
-fn contains_with_arrow_scalar(
+fn contains_array_scalar(
haystack_array: &ArrayRef,
needle_scalar: &ScalarValue,
) -> Result<ArrayRef> {
@@ -112,25 +141,24 @@ fn contains_with_arrow_scalar(
if needle_scalar.is_null() {
return Ok(Arc::new(BooleanArray::new_null(haystack_array.len())));
}
+ let needle_array =
+ scalar_to_aligned_array(needle_scalar, haystack_array.data_type(), "needle")?;
+ let result = arrow_contains(haystack_array, &Scalar::new(needle_array))?;
+ Ok(Arc::new(result))
+}
- // Extract the needle string
- let needle_str = match needle_scalar {
- ScalarValue::Utf8(Some(s))
- | ScalarValue::LargeUtf8(Some(s))
- | ScalarValue::Utf8View(Some(s)) => s.clone(),
- _ => {
- return exec_err!(
- "contains function requires string type for needle, got {:?}",
- needle_scalar.data_type()
- )
- }
- };
-
- // Create scalar array for needle - tells Arrow to use optimized paths
- let needle_scalar_array = StringArray::new_scalar(needle_str);
-
- // Use Arrow's contains which detects scalar case and uses optimized paths
- let result = arrow_contains(haystack_array, &needle_scalar_array)?;
+/// Contains for scalar haystack with array needle - less common path.
+fn contains_scalar_array(
+ haystack_scalar: &ScalarValue,
+ needle_array: &ArrayRef,
+) -> Result<ArrayRef> {
+ // Handle null haystack
+ if haystack_scalar.is_null() {
+ return Ok(Arc::new(BooleanArray::new_null(needle_array.len())));
+ }
+ let haystack_array =
+ scalar_to_aligned_array(haystack_scalar, needle_array.data_type(), "haystack")?;
+ let result = arrow_contains(&Scalar::new(haystack_array), needle_array)?;
Ok(Arc::new(result))
}
@@ -139,34 +167,12 @@ fn contains_scalar_scalar(
haystack_scalar: &ScalarValue,
needle_scalar: &ScalarValue,
) -> Result<ScalarValue> {
- // Handle nulls
if haystack_scalar.is_null() || needle_scalar.is_null() {
return Ok(ScalarValue::Boolean(None));
}
- let haystack_str = match haystack_scalar {
- ScalarValue::Utf8(Some(s))
- | ScalarValue::LargeUtf8(Some(s))
- | ScalarValue::Utf8View(Some(s)) => s.as_str(),
- _ => {
- return exec_err!(
- "contains function requires string type for haystack, got {:?}",
- haystack_scalar.data_type()
- )
- }
- };
-
- let needle_str = match needle_scalar {
- ScalarValue::Utf8(Some(s))
- | ScalarValue::LargeUtf8(Some(s))
- | ScalarValue::Utf8View(Some(s)) => s.as_str(),
- _ => {
- return exec_err!(
- "contains function requires string type for needle, got {:?}",
- needle_scalar.data_type()
- )
- }
- };
+ let haystack_str = get_string_scalar_value(haystack_scalar, "haystack")?;
+ let needle_str = get_string_scalar_value(needle_scalar, "needle")?;
Ok(ScalarValue::Boolean(Some(
haystack_str.contains(needle_str),
@@ -176,7 +182,8 @@ fn contains_scalar_scalar(
#[cfg(test)]
mod tests {
use super::*;
- use arrow::array::StringArray;
+ use arrow::array::{DictionaryArray, LargeStringArray, StringArray, StringViewArray};
+ use arrow::datatypes::Int32Type;
#[test]
fn test_contains_array_scalar() {
@@ -188,7 +195,7 @@ mod tests {
])) as ArrayRef;
let needle = ScalarValue::Utf8(Some("world".to_string()));
- let result = contains_with_arrow_scalar(&haystack, &needle).unwrap();
+ let result = contains_array_scalar(&haystack, &needle).unwrap();
let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
assert!(bool_array.value(0)); // "hello world" contains "world"
@@ -218,7 +225,7 @@ mod tests {
])) as ArrayRef;
let needle = ScalarValue::Utf8(None);
- let result = contains_with_arrow_scalar(&haystack, &needle).unwrap();
+ let result = contains_array_scalar(&haystack, &needle).unwrap();
let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
// Null needle should produce null results
@@ -231,11 +238,166 @@ mod tests {
let haystack = Arc::new(StringArray::from(vec![Some("hello world"), Some("")])) as ArrayRef;
let needle = ScalarValue::Utf8(Some("".to_string()));
- let result = contains_with_arrow_scalar(&haystack, &needle).unwrap();
+ let result = contains_array_scalar(&haystack, &needle).unwrap();
let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
// Empty string is contained in any string
assert!(bool_array.value(0));
assert!(bool_array.value(1));
}
+
+ #[test]
+ fn test_contains_scalar_array_null_haystack() {
+ let haystack = ScalarValue::Utf8(None);
+ let needle = Arc::new(StringArray::from(vec![
+ Some("hello world"),
+ Some("foo bar"),
+ ])) as ArrayRef;
+
+ let result = contains_scalar_array(&haystack, &needle).unwrap();
+ let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ // Null haystack should produce null results for all array elements
+ assert!(bool_array.is_null(0));
+ assert!(bool_array.is_null(1));
+ }
+
+ #[test]
+ fn test_spark_contains_dispatcher_scalar_array() {
+ let haystack = ColumnarValue::Scalar(ScalarValue::Utf8(Some("abc".to_string())));
+ let needle =
+ ColumnarValue::Array(
+ Arc::new(StringArray::from(vec![Some("a"), Some("bc"), Some("d")])) as ArrayRef,
+ );
+
+ let result = spark_contains(&haystack, &needle).unwrap();
+ let array = match result {
+ ColumnarValue::Array(arr) => arr,
+ _ => panic!("Expected ColumnarValue::Array"),
+ };
+ let bool_array = array.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ assert!(bool_array.value(0));
+ assert!(bool_array.value(1));
+ assert!(!bool_array.value(2));
+ }
+
+ #[test]
+ fn test_contains_scalar_large_utf8() {
+ let haystack = ScalarValue::LargeUtf8(Some("abc".to_string()));
+ let needle = Arc::new(LargeStringArray::from(vec![
+ Some("a"),
+ Some("bc"),
+ None,
+ Some(""),
+ Some("d"),
+ ])) as ArrayRef;
+
+ let res = contains_scalar_array(&haystack, &needle).unwrap();
+ let res = res.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let expected =
+ BooleanArray::from(vec![Some(true), Some(true), None, Some(true), Some(false)]);
+
+ assert_eq!(res, &expected);
+ }
+
+ #[test]
+ fn test_contains_scalar_utf8_view() {
+ let haystack = ScalarValue::Utf8View(Some("abc".to_string()));
+ let needle = Arc::new(StringViewArray::from(vec![
+ Some("a"),
+ Some("bc"),
+ None,
+ Some(""),
+ Some("d"),
+ ])) as ArrayRef;
+
+ let res = contains_scalar_array(&haystack, &needle).unwrap();
+ let res = res.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let expected =
+ BooleanArray::from(vec![Some(true), Some(true), None, Some(true), Some(false)]);
+
+ assert_eq!(res, &expected);
+ }
+
+ #[test]
+ fn test_contains_scalar_dictionary() {
+ // Regression: a non-null dictionary-string scalar previously worked before
+ // the optimization, then started failing at `get_string_scalar_value`.
+ let haystack = ScalarValue::Dictionary(
+ Box::new(DataType::Int32),
+ Box::new(ScalarValue::Utf8(Some("abc".to_string()))),
+ );
+ let needle = Arc::new(DictionaryArray::<Int32Type>::from_iter(vec![
+ Some("a"),
+ Some("bc"),
+ None,
+ Some(""),
+ Some("d"),
+ ])) as ArrayRef;
+
+ let res = contains_scalar_array(&haystack, &needle).unwrap();
+ let res = res.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let expected =
+ BooleanArray::from(vec![Some(true), Some(true), None, Some(true), Some(false)]);
+
+ assert_eq!(res, &expected);
+ }
+
+ #[test]
+ fn test_contains_array_scalar_large_utf8_haystack() {
+ // Symmetric case: scalar needle must be aligned to the array's type,
+ // so a Utf8 needle works against a LargeUtf8 haystack.
+ let haystack = Arc::new(LargeStringArray::from(vec![Some("abc"), Some("xyz")])) as ArrayRef;
+ let needle = ScalarValue::Utf8(Some("bc".to_string()));
+
+ let res = contains_array_scalar(&haystack, &needle).unwrap();
+ let res = res.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ assert_eq!(res, &BooleanArray::from(vec![Some(true), Some(false)]));
+ }
+
+ #[test]
+ fn test_contains_scalar_array_all_cases() {
+ let haystack = ScalarValue::Utf8(Some("hello world".to_string()));
+ let needle = Arc::new(StringArray::from(vec![
+ Some("hello"),
+ Some("world"),
+ Some("foo"),
+ None,
+ ])) as ArrayRef;
+
+ let res = contains_scalar_array(&haystack, &needle).unwrap();
+ let bool_arr = res.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ assert_eq!(
+ bool_arr,
+ &BooleanArray::from(vec![Some(true), Some(true), Some(false), None])
+ );
+ }
+
+ #[test]
+ fn test_contains_scalar_array_empty_needle() {
+ let haystack = ScalarValue::Utf8(Some("hello world".to_string()));
+ let needle = Arc::new(StringArray::from(Vec::<Option<&str>>::new())) as ArrayRef;
+
+ let res = contains_scalar_array(&haystack, &needle).unwrap();
+ assert_eq!(res.len(), 0);
+ }
+
+ #[test]
+ fn test_contains_scalar_array_invalid_type_error() {
+ let haystack = ScalarValue::Int32(Some(123));
+ let needle = Arc::new(StringArray::from(vec![Some("1")])) as ArrayRef;
+
+ let err = contains_scalar_array(&haystack, &needle).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("contains function requires string type for haystack"),
+ "unexpected error: {err}"
+ );
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org