diff --git a/native/spark-expr/benches/contains.rs b/native/spark-expr/benches/contains.rs index 457ce79db35..22fa8eec23a 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 537227efdfc..545d1475600 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 { - 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(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 { + 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 { @@ -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 { + // 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 { - // 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::().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::().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::().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::().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::().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::().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::().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::::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::().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::().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::().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::>::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}" + ); + } }