use coercion in starts_with

Author: zjregee

datafusion/expr-common/src/type_coercion/binary.rs

@@ -1256,7 +1256,7 @@ fn list_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataType> {
 /// Coercion rules for binary (Binary/LargeBinary) to string (Utf8/LargeUtf8):
 /// If one argument is binary and the other is a string then coerce to string
 /// (e.g. for `like`)
-fn binary_to_string_coercion(
+pub fn binary_to_string_coercion(
     lhs_type: &DataType,
     rhs_type: &DataType,
 ) -> Option<DataType> {

datafusion/functions/src/string/starts_with.rs

@@ -20,23 +20,39 @@ use std::sync::Arc;
 
 use arrow::array::ArrayRef;
 use arrow::datatypes::DataType;
+use datafusion_expr::binary::{binary_to_string_coercion, string_coercion};
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
 
 use crate::utils::make_scalar_function;
 use datafusion_common::types::logical_string;
 use datafusion_common::{internal_err, Result, ScalarValue};
 use datafusion_expr::{
-    Coercion, ColumnarValue, Documentation, Expr, Like, ScalarFunctionArgs,
+    cast, Coercion, ColumnarValue, Documentation, Expr, Like, ScalarFunctionArgs,
     ScalarUDFImpl, Signature, TypeSignatureClass, Volatility,
 };
 use datafusion_macros::user_doc;
 
 /// Returns true if string starts with prefix.
 /// starts_with('alphabet', 'alph') = 't'
 pub fn starts_with(args: &[ArrayRef]) -> Result<ArrayRef> {
-    let arg0 = arrow::compute::kernels::cast::cast(&args[0], args[1].data_type())?;
-    let result = arrow::compute::kernels::comparison::starts_with(&arg0, &args[1])?;
-    Ok(Arc::new(result) as ArrayRef)
+    if args[0].data_type() == args[1].data_type() {
+        let result =
+            arrow::compute::kernels::comparison::starts_with(&args[0], &args[1])?;
+        return Ok(Arc::new(result) as ArrayRef);
+    }
+
+    if let Some(coercion_data_type) =
+        string_coercion(args[0].data_type(), args[1].data_type()).or_else(|| {
+            binary_to_string_coercion(args[0].data_type(), args[1].data_type())
+        })
+    {
+        let arg0 = arrow::compute::kernels::cast::cast(&args[0], &coercion_data_type)?;
+        let arg1 = arrow::compute::kernels::cast::cast(&args[1], &coercion_data_type)?;
+        let result = arrow::compute::kernels::comparison::starts_with(&arg0, &arg1)?;
+        return Ok(Arc::new(result) as ArrayRef);
+    } else {
+        return internal_err!("Unsupported data types for starts_with. Expected Utf8, LargeUtf8 or Utf8View");
+    }
 }
 
 #[user_doc(
@@ -108,7 +124,7 @@ impl ScalarUDFImpl for StartsWithFunc {
     fn simplify(
         &self,
         args: Vec<Expr>,
-        _info: &dyn SimplifyInfo,
+        info: &dyn SimplifyInfo,
     ) -> Result<ExprSimplifyResult> {
         if let Expr::Literal(scalar_value) = &args[1] {
             // Convert starts_with(col, 'prefix') to col LIKE 'prefix%' with proper escaping
@@ -117,31 +133,42 @@ impl ScalarUDFImpl for StartsWithFunc {
             //   2. 'ja\%'        (escape special char '%')
             //   3. 'ja\%%'       (add suffix for starts_with)
             let like_expr = match scalar_value {
-                ScalarValue::Utf8(Some(pattern)) => {
+                ScalarValue::Utf8(Some(pattern))
+                | ScalarValue::LargeUtf8(Some(pattern))
+                | ScalarValue::Utf8View(Some(pattern)) => {
                     let escaped_pattern = pattern.replace("%", "\\%");
                     let like_pattern = format!("{}%", escaped_pattern);
                     Expr::Literal(ScalarValue::Utf8(Some(like_pattern)))
                 }
-                ScalarValue::LargeUtf8(Some(pattern)) => {
-                    let escaped_pattern = pattern.replace("%", "\\%");
-                    let like_pattern = format!("{}%", escaped_pattern);
-                    Expr::Literal(ScalarValue::LargeUtf8(Some(like_pattern)))
-                }
-                ScalarValue::Utf8View(Some(pattern)) => {
-                    let escaped_pattern = pattern.replace("%", "\\%");
-                    let like_pattern = format!("{}%", escaped_pattern);
-                    Expr::Literal(ScalarValue::Utf8View(Some(like_pattern)))
-                }
                 _ => return Ok(ExprSimplifyResult::Original(args)),
             };
 
-            return Ok(ExprSimplifyResult::Simplified(Expr::Like(Like {
-                negated: false,
-                expr: Box::new(args[0].clone()),
-                pattern: Box::new(like_expr),
-                escape_char: None,
-                case_insensitive: false,
-            })));
+            let expr_data_type = info.get_data_type(&args[0])?;
+            let pattern_data_type = info.get_data_type(&like_expr)?;
+
+            if expr_data_type == pattern_data_type {
+                return Ok(ExprSimplifyResult::Simplified(Expr::Like(Like {
+                    negated: false,
+                    expr: Box::new(args[0].clone()),
+                    pattern: Box::new(like_expr),
+                    escape_char: None,
+                    case_insensitive: false,
+                })));
+            }
+
+            if let Some(coercion_data_type) =
+                string_coercion(&expr_data_type, &pattern_data_type).or_else(|| {
+                    binary_to_string_coercion(&expr_data_type, &pattern_data_type)
+                })
+            {
+                return Ok(ExprSimplifyResult::Simplified(Expr::Like(Like {
+                    negated: false,
+                    expr: Box::new(cast(args[0].clone(), coercion_data_type.clone())),
+                    pattern: Box::new(cast(like_expr, coercion_data_type)),
+                    escape_char: None,
+                    case_insensitive: false,
+                })));
+            }
         }
 
         Ok(ExprSimplifyResult::Original(args))

datafusion/physical-expr/src/expressions/like.rs

@@ -19,7 +19,6 @@ use std::hash::Hash;
 use std::{any::Any, sync::Arc};
 
 use crate::PhysicalExpr;
-use arrow::compute::can_cast_types;
 use arrow::datatypes::{DataType, Schema};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::{internal_err, Result};
@@ -122,10 +121,7 @@ impl PhysicalExpr for LikeExpr {
     fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
         use arrow::compute::*;
         let lhs = self.expr.evaluate(batch)?;
-        let rhs = self
-            .pattern
-            .evaluate(batch)?
-            .cast_to(&lhs.data_type(), None)?;
+        let rhs = self.pattern.evaluate(batch)?;
         match (self.negated, self.case_insensitive) {
             (false, false) => apply_cmp(&lhs, &rhs, like),
             (false, true) => apply_cmp(&lhs, &rhs, ilike),
@@ -169,10 +165,7 @@ pub fn like(
 ) -> Result<Arc<dyn PhysicalExpr>> {
     let expr_type = &expr.data_type(input_schema)?;
     let pattern_type = &pattern.data_type(input_schema)?;
-    if !expr_type.eq(pattern_type)
-        && !can_cast_types(expr_type, pattern_type)
-        && !can_like_type(expr_type)
-    {
+    if !expr_type.eq(pattern_type) && !can_like_type(expr_type) {
         return internal_err!(
             "The type of {expr_type} AND {pattern_type} of like physical should be same"
         );

datafusion/sqllogictest/test_files/parquet.slt

@@ -619,8 +619,8 @@ query TT
 explain select * from foo where starts_with(column1, 'f');
 ----
 logical_plan
-01)Filter: foo.column1 LIKE Utf8("f%")
-02)--TableScan: foo projection=[column1], partial_filters=[foo.column1 LIKE Utf8("f%")]
+01)Filter: CAST(foo.column1 AS Utf8View) LIKE Utf8View("f%")
+02)--TableScan: foo projection=[column1], partial_filters=[CAST(foo.column1 AS Utf8View) LIKE Utf8View("f%")]
 physical_plan
 01)CoalesceBatchesExec: target_batch_size=8192
 02)--FilterExec: column1@0 LIKE f%

datafusion/sqllogictest/test_files/string/string_view.slt

@@ -355,8 +355,9 @@ EXPLAIN SELECT
 FROM test;
 ----
 logical_plan
-01)Projection: test.column1_utf8view LIKE Utf8("äöüß%") AS c1, CASE test.column1_utf8view IS NOT NULL WHEN Boolean(true) THEN Boolean(true) END AS c2, starts_with(test.column1_utf8view, Utf8View(NULL)) AS c3, starts_with(Utf8View(NULL), test.column1_utf8view) AS c4
-02)--TableScan: test projection=[column1_utf8view]
+01)Projection: __common_expr_1 LIKE Utf8View("äöüß%") AS c1, CASE __common_expr_1 IS NOT NULL WHEN Boolean(true) THEN Boolean(true) END AS c2, starts_with(test.column1_utf8view, Utf8View(NULL)) AS c3, starts_with(Utf8View(NULL), test.column1_utf8view) AS c4
+02)--Projection: CAST(test.column1_utf8view AS Utf8View) AS __common_expr_1, test.column1_utf8view
+03)----TableScan: test projection=[column1_utf8view]
 
 ## Test STARTS_WITH is rewitten to LIKE when the pattern is a constant
 query TT
@@ -370,8 +371,9 @@ EXPLAIN SELECT
 FROM test;
 ----
 logical_plan
-01)Projection: test.column1_utf8 LIKE Utf8("foo\%%") AS c1, test.column1_large_utf8 LIKE Utf8("foo\%%") AS c2, test.column1_utf8view LIKE Utf8("foo\%%") AS c3, test.column1_utf8 LIKE Utf8("f_o%") AS c4, test.column1_large_utf8 LIKE Utf8("f_o%") AS c5, test.column1_utf8view LIKE Utf8("f_o%") AS c6
-02)--TableScan: test projection=[column1_utf8, column1_large_utf8, column1_utf8view]
+01)Projection: test.column1_utf8 LIKE Utf8("foo\%%") AS c1, __common_expr_1 LIKE LargeUtf8("foo\%%") AS c2, __common_expr_2 LIKE Utf8View("foo\%%") AS c3, test.column1_utf8 LIKE Utf8("f_o%") AS c4, __common_expr_1 LIKE LargeUtf8("f_o%") AS c5, __common_expr_2 LIKE Utf8View("f_o%") AS c6
+02)--Projection: CAST(test.column1_large_utf8 AS LargeUtf8) AS __common_expr_1, CAST(test.column1_utf8view AS Utf8View) AS __common_expr_2, test.column1_utf8
+03)----TableScan: test projection=[column1_utf8, column1_large_utf8, column1_utf8view]
 
 ## Test STARTS_WITH works with column arguments
 query TT
@@ -940,7 +942,7 @@ EXPLAIN SELECT
 FROM test;
 ----
 logical_plan
-01)Projection: test.column1_utf8view LIKE Utf8("foo%") AS c, starts_with(test.column1_utf8view, test.column2_utf8view) AS c2
+01)Projection: CAST(test.column1_utf8view AS Utf8View) LIKE Utf8View("foo%") AS c, starts_with(test.column1_utf8view, test.column2_utf8view) AS c2
 02)--TableScan: test projection=[column1_utf8view, column2_utf8view]
 
 ## Ensure no casts for TRANSLATE