From b518dfceadec0cbf8998e38ccff69e4b464f5446 Mon Sep 17 00:00:00 2001
From: Xin Li <bluer906@gmail.com>
Date: Mon, 8 Jul 2024 16:20:54 +0800
Subject: [PATCH] Implement user defined planner for sql_substring_to_expr

---
 datafusion/expr/src/planner.rs       |  7 +++++++
 datafusion/functions/src/planner.rs  |  7 +++++++
 datafusion/sql/src/expr/mod.rs       |  8 ++++----
 datafusion/sql/src/expr/substring.rs | 24 ++++++++++++++----------
 4 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs
index 4086964ab8ed..c255edbea5ae 100644
--- a/datafusion/expr/src/planner.rs
+++ b/datafusion/expr/src/planner.rs
@@ -140,6 +140,13 @@ pub trait UserDefinedSQLPlanner: Send + Sync {
         Ok(PlannerResult::Original(args))
     }
 
+    /// Plan an substring expression, e.g., `SUBSTRING(<expr> [FROM <expr>] [FOR <expr>])`
+    ///
+    /// Returns origin expression arguments if not possible
+    fn plan_substring(&self, args: Vec<Expr>) -> Result<PlannerResult<Vec<Expr>>> {
+        Ok(PlannerResult::Original(args))
+    }
+
     /// Plans a struct `struct(expression1[, ..., expression_n])`
     /// literal based on the given input expressions.
     /// This function takes a vector of expressions and a boolean flag indicating whether
diff --git a/datafusion/functions/src/planner.rs b/datafusion/functions/src/planner.rs
index b00d5cf60810..41ff92f26111 100644
--- a/datafusion/functions/src/planner.rs
+++ b/datafusion/functions/src/planner.rs
@@ -41,4 +41,11 @@ impl UserDefinedSQLPlanner for UserDefinedFunctionPlanner {
             ScalarFunction::new_udf(crate::unicode::strpos(), args),
         )))
     }
+
+    #[cfg(feature = "unicode_expressions")]
+    fn plan_substring(&self, args: Vec<Expr>) -> Result<PlannerResult<Vec<Expr>>> {
+        Ok(PlannerResult::Planned(Expr::ScalarFunction(
+            ScalarFunction::new_udf(crate::unicode::substr(), args),
+        )))
+    }
 }
diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs
index eba93735b552..0546a101fcb2 100644
--- a/datafusion/sql/src/expr/mod.rs
+++ b/datafusion/sql/src/expr/mod.rs
@@ -670,18 +670,18 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         let substr =
             self.sql_expr_to_logical_expr(substr_expr, schema, planner_context)?;
         let fullstr = self.sql_expr_to_logical_expr(str_expr, schema, planner_context)?;
-        let mut extract_args = vec![fullstr, substr];
+        let mut position_args = vec![fullstr, substr];
         for planner in self.planners.iter() {
-            match planner.plan_position(extract_args)? {
+            match planner.plan_position(position_args)? {
                 PlannerResult::Planned(expr) => return Ok(expr),
                 PlannerResult::Original(args) => {
-                    extract_args = args;
+                    position_args = args;
                 }
             }
         }
 
         not_impl_err!(
-            "Position not supported by UserDefinedExtensionPlanners: {extract_args:?}"
+            "Position not supported by UserDefinedExtensionPlanners: {position_args:?}"
         )
     }
 
diff --git a/datafusion/sql/src/expr/substring.rs b/datafusion/sql/src/expr/substring.rs
index f58c6f3b94d0..a0dfee1b9d90 100644
--- a/datafusion/sql/src/expr/substring.rs
+++ b/datafusion/sql/src/expr/substring.rs
@@ -16,9 +16,9 @@
 // under the License.
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
-use datafusion_common::{internal_datafusion_err, plan_err};
+use datafusion_common::{not_impl_err, plan_err};
 use datafusion_common::{DFSchema, Result, ScalarValue};
-use datafusion_expr::expr::ScalarFunction;
+use datafusion_expr::planner::PlannerResult;
 use datafusion_expr::Expr;
 use sqlparser::ast::Expr as SQLExpr;
 
@@ -31,7 +31,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         schema: &DFSchema,
         planner_context: &mut PlannerContext,
     ) -> Result<Expr> {
-        let args = match (substring_from, substring_for) {
+        let mut substring_args = match (substring_from, substring_for) {
             (Some(from_expr), Some(for_expr)) => {
                 let arg =
                     self.sql_expr_to_logical_expr(*expr, schema, planner_context)?;
@@ -68,13 +68,17 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             }
         };
 
-        let fun = self
-            .context_provider
-            .get_function_meta("substr")
-            .ok_or_else(|| {
-                internal_datafusion_err!("Unable to find expected 'substr' function")
-            })?;
+        for planner in self.planners.iter() {
+            match planner.plan_substring(substring_args)? {
+                PlannerResult::Planned(expr) => return Ok(expr),
+                PlannerResult::Original(args) => {
+                    substring_args = args;
+                }
+            }
+        }
 
-        Ok(Expr::ScalarFunction(ScalarFunction::new_udf(fun, args)))
+        not_impl_err!(
+            "Substring not supported by UserDefinedExtensionPlanners: {substring_args:?}"
+        )
     }
 }