diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 032fe2524096e..a68d84f2f524e 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -43,7 +43,7 @@ use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema; use datafusion_expr::expr_schema::cast_subquery; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::type_coercion::binary::{ - comparison_coercion, like_coercion, type_union_coercion, + comparison_coercion, like_coercion, regex_coercion, type_union_coercion, }; use datafusion_expr::type_coercion::functions::{ UDFCoercionExt, fields_with_udf, value_fields_with_higher_order_udf_and_lambdas, @@ -613,6 +613,33 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { case_insensitive, )))) } + Expr::SimilarTo(Like { + negated, + expr, + pattern, + escape_char, + case_insensitive, + }) => { + let left_type = expr.get_type(self.schema)?; + let right_type = pattern.get_type(self.schema)?; + let coerced_type = regex_coercion(&left_type, &right_type).ok_or_else(|| { + plan_datafusion_err!( + "There isn't a common type to coerce {left_type} and {right_type} in SIMILAR TO expression" + ) + })?; + let expr = match left_type { + DataType::Dictionary(_, inner) if *inner == DataType::Utf8 => expr, + _ => Box::new(expr.cast_to(&coerced_type, self.schema)?), + }; + let pattern = Box::new(pattern.cast_to(&coerced_type, self.schema)?); + Ok(Transformed::yes(Expr::SimilarTo(Like::new( + negated, + expr, + pattern, + escape_char, + case_insensitive, + )))) + } Expr::BinaryExpr(BinaryExpr { left, op, right }) => { let (left, right) = self.coerce_binary_op(*left, self.schema, op, *right, self.schema)?; @@ -812,7 +839,6 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { | Expr::Column(_) | Expr::ScalarVariable(_, _) | Expr::Literal(_, _) - | Expr::SimilarTo(_) | Expr::IsNotNull(_) | Expr::IsNull(_) | Expr::Cast(_) @@ -2239,6 +2265,75 @@ mod test { Ok(()) } + #[test] + fn similar_to_for_type_coercion() -> Result<()> { + // utf8 SIMILAR TO utf8 needs no coercion + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::new_utf8("abc"))); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(Utf8); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r#" + Projection: a SIMILAR TO Utf8("abc") + EmptyRelation: rows=0 + "# + )?; + + // utf8 SIMILAR TO NULL casts the pattern to utf8 + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::Null)); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(Utf8); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r" + Projection: a SIMILAR TO CAST(NULL AS Utf8) + EmptyRelation: rows=0 + " + )?; + + // utf8view SIMILAR TO utf8 casts the pattern to utf8view + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::new_utf8("abc"))); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(DataType::Utf8View); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r#" + Projection: a SIMILAR TO CAST(Utf8("abc") AS Utf8View) + EmptyRelation: rows=0 + "# + )?; + + // non-string operands are rejected + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::new_utf8("abc"))); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(DataType::Int64); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + assert_type_coercion_error( + plan, + "There isn't a common type to coerce Int64 and Utf8 in SIMILAR TO expression", + )?; + + Ok(()) + } + #[test] fn unknown_for_type_coercion() -> Result<()> { // unknown diff --git a/datafusion/physical-expr/src/expressions/binary/kernels.rs b/datafusion/physical-expr/src/expressions/binary/kernels.rs index e573d7ece2afa..3302073687aad 100644 --- a/datafusion/physical-expr/src/expressions/binary/kernels.rs +++ b/datafusion/physical-expr/src/expressions/binary/kernels.rs @@ -25,11 +25,12 @@ use arrow::compute::kernels::bitwise::{ bitwise_xor, bitwise_xor_scalar, }; use arrow::compute::kernels::boolean::not; +use arrow::compute::kernels::cast::cast; use arrow::compute::kernels::comparison::{regexp_is_match, regexp_is_match_scalar}; use arrow::datatypes::DataType; use arrow::error::ArrowError; use datafusion_common::{Result, ScalarValue}; -use datafusion_common::{internal_err, plan_err}; +use datafusion_common::{downcast_value, internal_err, plan_err}; use std::sync::Arc; @@ -251,14 +252,8 @@ pub fn concat_elements_binary_view_array( /// Invoke a compute kernel on a pair of binary data arrays with flags macro_rules! regexp_is_match_flag { ($LEFT:expr, $RIGHT:expr, $ARRAYTYPE:ident, $NOT:expr, $FLAG:expr) => {{ - let ll = $LEFT - .as_any() - .downcast_ref::<$ARRAYTYPE>() - .expect("failed to downcast array"); - let rr = $RIGHT - .as_any() - .downcast_ref::<$ARRAYTYPE>() - .expect("failed to downcast array"); + let ll = downcast_value!($LEFT, $ARRAYTYPE); + let rr = downcast_value!($RIGHT, $ARRAYTYPE); let flag = if $FLAG { Some($ARRAYTYPE::from(vec!["i"; ll.len()])) @@ -273,12 +268,37 @@ macro_rules! regexp_is_match_flag { }}; } +/// Unpack a dictionary array to its value type, leaving other arrays as-is +fn flatten_dictionary(array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Dictionary(_, value_type) => Ok(cast(array, value_type)?), + _ => Ok(Arc::clone(array)), + } +} + pub(crate) fn regex_match_dyn( left: &ArrayRef, right: &ArrayRef, not_match: bool, flag: bool, ) -> Result { + // Type coercion deliberately preserves `Dictionary(_, Utf8)` inputs so + // that literal patterns can use the dictionary fast path in + // `regex_match_dyn_scalar`. Here every row has its own pattern, so the + // dictionary encoding offers no shortcut: flatten the inputs, and if the + // dictionary's value type differs from the coerced pattern type (e.g. + // `Dictionary(Int32, Utf8)` matched against `Utf8View` patterns), cast to + // a common string type. + if matches!(left.data_type(), DataType::Dictionary(_, _)) + || matches!(right.data_type(), DataType::Dictionary(_, _)) + { + let left = flatten_dictionary(left)?; + let mut right = flatten_dictionary(right)?; + if right.data_type() != left.data_type() { + right = cast(&right, left.data_type())?; + } + return regex_match_dyn(&left, &right, not_match, flag); + } match left.data_type() { DataType::Utf8 => { regexp_is_match_flag!(left, right, StringArray, not_match, flag) @@ -299,27 +319,32 @@ pub(crate) fn regex_match_dyn( /// Invoke a compute kernel on a data array and a scalar value with flag macro_rules! regexp_is_match_flag_scalar { ($LEFT:expr, $RIGHT:expr, $ARRAYTYPE:ident, $NOT:expr, $FLAG:expr) => {{ - let ll = $LEFT - .as_any() - .downcast_ref::<$ARRAYTYPE>() - .expect("failed to downcast array"); - - if let Some(Some(string_value)) = $RIGHT.try_as_str() { - let flag = $FLAG.then_some("i"); - match regexp_is_match_scalar(ll, &string_value, flag) { - Ok(mut array) => { - if $NOT { - array = not(&array).unwrap(); + match $LEFT.as_any().downcast_ref::<$ARRAYTYPE>() { + None => internal_err!( + "failed to downcast array to {} for operation 'regex_match_dyn_scalar'", + stringify!($ARRAYTYPE) + ), + Some(ll) => { + if let Some(Some(string_value)) = $RIGHT.try_as_str() { + let flag = $FLAG.then_some("i"); + match regexp_is_match_scalar(ll, &string_value, flag) { + Ok(mut array) => { + if $NOT { + array = not(&array).unwrap(); + } + Ok(Arc::new(array)) + } + Err(e) => { + internal_err!("failed to call 'regex_match_dyn_scalar' {}", e) + } } - Ok(Arc::new(array)) + } else { + internal_err!( + "failed to cast literal value {} for operation 'regex_match_dyn_scalar'", + $RIGHT + ) } - Err(e) => internal_err!("failed to call 'regex_match_dyn_scalar' {}", e), } - } else { - internal_err!( - "failed to cast literal value {} for operation 'regex_match_dyn_scalar'", - $RIGHT - ) } }}; } diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index 01e5ec4f149a6..280f507567707 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -927,7 +927,10 @@ impl SqlToRel<'_, S> { ) -> Result { let pattern = self.sql_expr_to_logical_expr(pattern, schema, planner_context)?; let pattern_type = pattern.get_type(schema)?; - if pattern_type != DataType::Utf8 && pattern_type != DataType::Null { + if !matches!( + pattern_type, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View | DataType::Null + ) { return plan_err!("Invalid pattern in SIMILAR TO expression"); } let escape_char = match escape_char.map(|v| v.value) { diff --git a/datafusion/sqllogictest/test_files/strings.slt b/datafusion/sqllogictest/test_files/strings.slt index 9fa453fa02523..9cba7253cdd35 100644 --- a/datafusion/sqllogictest/test_files/strings.slt +++ b/datafusion/sqllogictest/test_files/strings.slt @@ -91,6 +91,111 @@ P1e1 P1m1e1 e1 +# SIMILAR TO with a NULL pattern or NULL input is NULL, not an error +# (regression tests for https://github.com/apache/datafusion/issues/22886) +query B +SELECT 'a' SIMILAR TO NULL; +---- +NULL + +query B +SELECT 'a' NOT SIMILAR TO NULL; +---- +NULL + +query B +SELECT NULL SIMILAR TO 'a'; +---- +NULL + +query B +SELECT s SIMILAR TO NULL FROM test WHERE s = 'p1'; +---- +NULL + +# SIMILAR TO coerces mismatched string types instead of panicking +# (regression tests for https://github.com/apache/datafusion/issues/22886) +statement ok +CREATE TABLE patterns(pat TEXT) AS VALUES ('p[12].*'); + +# Utf8View input matched against a non-literal Utf8 pattern +query T rowsort +SELECT s FROM test, patterns WHERE arrow_cast(s, 'Utf8View') SIMILAR TO pat; +---- +p1 +p1e1 +p1m1e1 +p2 +p2e1 +p2m1e1 + +# LargeUtf8 input matched against a non-literal Utf8 pattern +query T rowsort +SELECT s FROM test, patterns WHERE arrow_cast(s, 'LargeUtf8') SIMILAR TO pat; +---- +p1 +p1e1 +p1m1e1 +p2 +p2e1 +p2m1e1 + +# Utf8 input matched against a non-literal Utf8View pattern +query T rowsort +SELECT s FROM test, patterns WHERE s SIMILAR TO arrow_cast(pat, 'Utf8View'); +---- +p1 +p1e1 +p1m1e1 +p2 +p2e1 +p2m1e1 + +# Dictionary input matched against a literal pattern (scalar fast path: the +# regex runs once per distinct dictionary value, so the dictionary input is +# deliberately not cast away by type coercion) +query T rowsort +SELECT s FROM test WHERE arrow_cast(s, 'Dictionary(Int32, Utf8)') SIMILAR TO 'p[12].*'; +---- +p1 +p1e1 +p1m1e1 +p2 +p2e1 +p2m1e1 + +# Dictionary input matched against a NULL pattern is NULL, not an error +query B +SELECT arrow_cast(s, 'Dictionary(Int32, Utf8)') SIMILAR TO NULL FROM test WHERE s = 'p1'; +---- +NULL + +# Dictionary input matched against a non-literal pattern +query T rowsort +SELECT s FROM test, patterns WHERE arrow_cast(s, 'Dictionary(Int32, Utf8)') SIMILAR TO pat; +---- +p1 +p1e1 +p1m1e1 +p2 +p2e1 +p2m1e1 + +# Dictionary input matched against a non-literal pattern of a different +# string type than the dictionary's value type +query T rowsort +SELECT s FROM test, patterns WHERE arrow_cast(s, 'Dictionary(Int32, Utf8)') SIMILAR TO arrow_cast(pat, 'Utf8View'); +---- +p1 +p1e1 +p1m1e1 +p2 +p2e1 +p2m1e1 + +statement ok +DROP TABLE patterns; + # NOT LIKE query T rowsort SELECT s FROM test WHERE s NOT LIKE 'p1%';