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-5215-aae5650cdf5b0ec98c605a3a77d2e9ac66f484fe in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit e6e92f1a1a10e1bca2bcf32858113ab61b2ba685 Author: sam-1112 <[email protected]> AuthorDate: Wed Sep 23 21:19:57 2026 +0000 fix: support Utf8/LargeUtf8/Utf8View in native RLike without panicking (#5215) * fix: support Utf8/LargeUtf8/Utf8View in native RLike without panicking (#5102) * style: cargo fmt * fix: handle any dictionary key type in RLike and tighten review follow-ups Use as_any_dictionary() so non-Int32 keys work, unify non-string errors on dictionary layouts and string array coverage. * test: cover nulls in RLike dictionary values and drop stray doc comment * test: assert RLike all-valid output has no null buffer * test: cover long Utf8View, UInt64 keys, and sliced dictionaries in RLike --------- Co-authored-by: Hung <[email protected]> --- native/spark-expr/src/predicate_funcs/rlike.rs | 227 +++++++++++++++++++++---- 1 file changed, 192 insertions(+), 35 deletions(-) diff --git a/native/spark-expr/src/predicate_funcs/rlike.rs b/native/spark-expr/src/predicate_funcs/rlike.rs index ee005dd1ac..56a3307cae 100644 --- a/native/spark-expr/src/predicate_funcs/rlike.rs +++ b/native/spark-expr/src/predicate_funcs/rlike.rs @@ -16,11 +16,10 @@ // under the License. use crate::SparkError; -use arrow::array::builder::BooleanBuilder; -use arrow::array::types::Int32Type; -use arrow::array::{Array, BooleanArray, DictionaryArray, RecordBatch, StringArray}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, RecordBatch, StringArrayType}; use arrow::compute::take; use arrow::datatypes::{DataType, Schema}; +use datafusion::common::cast::{as_large_string_array, as_string_array, as_string_view_array}; use datafusion::common::{internal_err, Result, ScalarValue}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ColumnarValue; @@ -71,22 +70,29 @@ impl RLike { }) } - fn is_match(&self, inputs: &StringArray) -> BooleanArray { - let mut builder = BooleanBuilder::with_capacity(inputs.len()); - if inputs.is_nullable() { - for i in 0..inputs.len() { - if inputs.is_null(i) { - builder.append_null(); - } else { - builder.append_value(self.pattern.is_match(inputs.value(i))); - } - } - } else { - for i in 0..inputs.len() { - builder.append_value(self.pattern.is_match(inputs.value(i))); + /// Match the pre-compiled pattern against a string array of any Arrow string layout. + /// + /// Keeps the plan-time compiled [`Regex`] rather than calling Arrow's + /// `regexp_is_match(_scalar)`, which recompiles the pattern on every batch. + fn is_match<'a, S>(&self, inputs: &'a S) -> BooleanArray + where + &'a S: StringArrayType<'a>, + { + inputs + .iter() + .map(|v| v.map(|s| self.pattern.is_match(s))) + .collect() + } + + fn is_match_array(&self, array: &ArrayRef) -> Result<BooleanArray> { + match array.data_type() { + DataType::Utf8 => Ok(self.is_match(as_string_array(array)?)), + DataType::LargeUtf8 => Ok(self.is_match(as_large_string_array(array)?)), + DataType::Utf8View => Ok(self.is_match(as_string_view_array(array)?)), + other => { + internal_err!("RLike requires string type for input, got {other:?}") } } - builder.finish() } } @@ -111,29 +117,19 @@ impl PhysicalExpr for RLike { fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { match self.child.evaluate(batch)? { - ColumnarValue::Array(array) if array.as_any().is::<DictionaryArray<Int32Type>>() => { - let dict_array = array - .as_any() - .downcast_ref::<DictionaryArray<Int32Type>>() - .expect("dict array"); - let dict_values = dict_array - .values() - .as_any() - .downcast_ref::<StringArray>() - .expect("strings"); + ColumnarValue::Array(array) + if matches!(array.data_type(), DataType::Dictionary(_, _)) => + { + let dict_array = array.as_any_dictionary(); // evaluate the regexp pattern against the dictionary values - let new_values = self.is_match(dict_values); + let new_values = self.is_match_array(dict_array.values())?; // convert to conventional (not dictionary-encoded) array let result = take(&new_values, dict_array.keys(), None)?; Ok(ColumnarValue::Array(result)) } ColumnarValue::Array(array) => { - let inputs = array - .as_any() - .downcast_ref::<StringArray>() - .expect("string array"); - let array = self.is_match(inputs); - Ok(ColumnarValue::Array(Arc::new(array))) + let result = self.is_match_array(&array)?; + Ok(ColumnarValue::Array(Arc::new(result))) } ColumnarValue::Scalar(scalar) => { if scalar.is_null() { @@ -180,7 +176,32 @@ impl PhysicalExpr for RLike { #[cfg(test)] mod tests { use super::*; - use datafusion::physical_expr::expressions::Literal; + use arrow::array::{ + DictionaryArray, Int32Array, Int8Array, LargeStringArray, StringArray, StringViewArray, + UInt64Array, + }; + use arrow::datatypes::{Field, Int32Type, Int8Type, UInt64Type}; + use datafusion::physical_expr::expressions::{Column, Literal}; + + fn assert_bool_results(result: ColumnarValue, expected: &[Option<bool>]) { + let ColumnarValue::Array(arr) = result else { + panic!("expected array result"); + }; + let bools = arr + .as_any() + .downcast_ref::<BooleanArray>() + .expect("boolean array"); + assert_eq!(bools.len(), expected.len()); + for (i, exp) in expected.iter().enumerate() { + match exp { + Some(v) => { + assert!(!bools.is_null(i), "row {i} should not be null"); + assert_eq!(bools.value(i), *v, "row {i}"); + } + None => assert!(bools.is_null(i), "row {i} should be null"), + } + } + } #[test] fn test_rlike_scalar_string_variants() { @@ -225,4 +246,140 @@ mod tests { let result = expr.evaluate(&RecordBatch::new_empty(Arc::new(Schema::empty()))); assert!(result.is_err()); } + + #[test] + fn test_rlike_string_array_layouts() { + let pattern = "R[a-z]+"; + let cases: Vec<(DataType, ArrayRef)> = vec![ + ( + DataType::Utf8, + Arc::new(StringArray::from(vec![Some("Rose"), None, Some("Daisy")])), + ), + ( + DataType::LargeUtf8, + Arc::new(LargeStringArray::from(vec![ + Some("Rose"), + None, + Some("Daisy"), + ])), + ), + ( + DataType::Utf8View, + Arc::new(StringViewArray::from(vec![ + Some("Rhododendrons"), + None, + Some("Daisy"), + ])), + ), + ]; + + for (data_type, array) in cases { + let schema = Arc::new(Schema::new(vec![Field::new("s", data_type, true)])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array]).unwrap(); + let expr = RLike::try_new(Arc::new(Column::new("s", 0)), pattern).unwrap(); + assert_bool_results( + expr.evaluate(&batch).unwrap(), + &[Some(true), None, Some(false)], + ); + } + } + + #[test] + fn test_rlike_string_array_no_nulls() { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(vec!["Rose", "Daisy"]))], + ) + .unwrap(); + + let expr = RLike::try_new(Arc::new(Column::new("s", 0)), "R[a-z]+").unwrap(); + let ColumnarValue::Array(arr) = expr.evaluate(&batch).unwrap() else { + panic!("expected array result"); + }; + // Preserve a null-buffer-free output shape for all-valid input, avoiding an + // unnecessary len / 8-byte validity allocation. + assert!(arr.nulls().is_none()); + assert_bool_results(ColumnarValue::Array(arr), &[Some(true), Some(false)]); + } + + #[test] + fn test_rlike_dictionary_arrays() { + let pattern = "R[a-z]+"; + let expected = [Some(true), None, Some(false)]; + + let utf8_values: ArrayRef = Arc::new(StringArray::from(vec!["Rose", "Daisy"])); + let utf8_view_values: ArrayRef = + Arc::new(StringViewArray::from(vec!["Rhododendrons", "Daisy"])); + // Null in dictionary values (keys all valid): is_match emits null, take carries it. + let utf8_values_with_null: ArrayRef = + Arc::new(StringArray::from(vec![Some("Rose"), None, Some("Daisy")])); + let sliced_dictionary = DictionaryArray::<Int32Type>::new( + Int32Array::from(vec![Some(1), Some(0), None, Some(1), Some(0)]), + Arc::clone(&utf8_values), + ); + let sliced_dictionary: ArrayRef = Arc::new(sliced_dictionary.slice(1, 3)); + + let cases: Vec<(DataType, ArrayRef)> = vec![ + ( + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + Arc::new(DictionaryArray::<Int32Type>::new( + Int32Array::from(vec![Some(0), None, Some(1)]), + Arc::clone(&utf8_values), + )), + ), + ( + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8View)), + Arc::new(DictionaryArray::<Int32Type>::new( + Int32Array::from(vec![Some(0), None, Some(1)]), + Arc::clone(&utf8_view_values), + )), + ), + ( + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + Arc::new(DictionaryArray::<Int8Type>::new( + Int8Array::from(vec![Some(0), None, Some(1)]), + Arc::clone(&utf8_values), + )), + ), + ( + DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)), + Arc::new(DictionaryArray::<UInt64Type>::new( + UInt64Array::from(vec![Some(0), None, Some(1)]), + Arc::clone(&utf8_values), + )), + ), + ( + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + Arc::new(DictionaryArray::<Int32Type>::new( + Int32Array::from(vec![Some(0), Some(1), Some(2)]), + utf8_values_with_null, + )), + ), + ( + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + sliced_dictionary, + ), + ]; + + for (data_type, array) in cases { + let schema = Arc::new(Schema::new(vec![Field::new("s", data_type, true)])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array]).unwrap(); + let expr = RLike::try_new(Arc::new(Column::new("s", 0)), pattern).unwrap(); + assert_bool_results(expr.evaluate(&batch).unwrap(), &expected); + } + } + + #[test] + fn test_rlike_array_non_string_error() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Boolean, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(BooleanArray::from(vec![Some(true), None]))], + ) + .unwrap(); + + let expr = RLike::try_new(Arc::new(Column::new("b", 0)), "R[a-z]+").unwrap(); + assert!(expr.evaluate(&batch).is_err()); + } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
