This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 2af3d3a55b Implement user defined planner for position (#11243)
2af3d3a55b is described below
commit 2af3d3a55ba67726ff915743680380fb02a3e2c2
Author: Xin Li <[email protected]>
AuthorDate: Sat Jul 6 04:55:55 2024 -0700
Implement user defined planner for position (#11243)
* Implement user defined planner for position
* Fix format
* Move planner to session_state
* Extract function
---
datafusion/core/src/execution/session_state.rs | 2 ++
datafusion/expr/src/planner.rs | 6 ++++
datafusion/functions/src/unicode/mod.rs | 1 +
datafusion/functions/src/unicode/planner.rs | 36 +++++++++++++++++++++
datafusion/sql/src/expr/mod.rs | 44 +++++++++++++++-----------
5 files changed, 70 insertions(+), 19 deletions(-)
diff --git a/datafusion/core/src/execution/session_state.rs
b/datafusion/core/src/execution/session_state.rs
index a831f92def..ffaaa2df5e 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -240,6 +240,8 @@ impl SessionState {
Arc::new(functions_array::planner::FieldAccessPlanner),
#[cfg(feature = "datetime_expressions")]
Arc::new(functions::datetime::planner::ExtractPlanner),
+ #[cfg(feature = "unicode_expressions")]
+ Arc::new(functions::unicode::planner::PositionPlanner),
];
let mut new_self = SessionState {
diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs
index bba0228ae0..bcbf5eb203 100644
--- a/datafusion/expr/src/planner.rs
+++ b/datafusion/expr/src/planner.rs
@@ -116,6 +116,12 @@ pub trait UserDefinedSQLPlanner: Send + Sync {
Ok(PlannerResult::Original(exprs))
}
+ // Plan the POSITION expression, e.g., POSITION(<expr> in <expr>)
+ // returns origin expression arguments if not possible
+ fn plan_position(&self, args: Vec<Expr>) ->
Result<PlannerResult<Vec<Expr>>> {
+ Ok(PlannerResult::Original(args))
+ }
+
/// Plan the dictionary literal `{ key: value, ...}`
///
/// Returns origin expression arguments if not possible
diff --git a/datafusion/functions/src/unicode/mod.rs
b/datafusion/functions/src/unicode/mod.rs
index 9e8c07cd36..a391b8ba11 100644
--- a/datafusion/functions/src/unicode/mod.rs
+++ b/datafusion/functions/src/unicode/mod.rs
@@ -25,6 +25,7 @@ pub mod character_length;
pub mod find_in_set;
pub mod left;
pub mod lpad;
+pub mod planner;
pub mod reverse;
pub mod right;
pub mod rpad;
diff --git a/datafusion/functions/src/unicode/planner.rs
b/datafusion/functions/src/unicode/planner.rs
new file mode 100644
index 0000000000..4d6f73321b
--- /dev/null
+++ b/datafusion/functions/src/unicode/planner.rs
@@ -0,0 +1,36 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// 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
+//
+// 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
+// specific language governing permissions and limitations
+// under the License.
+
+//! SQL planning extensions like [`PositionPlanner`]
+
+use datafusion_common::Result;
+use datafusion_expr::{
+ expr::ScalarFunction,
+ planner::{PlannerResult, UserDefinedSQLPlanner},
+ Expr,
+};
+
+#[derive(Default)]
+pub struct PositionPlanner;
+
+impl UserDefinedSQLPlanner for PositionPlanner {
+ fn plan_position(&self, args: Vec<Expr>) ->
Result<PlannerResult<Vec<Expr>>> {
+ Ok(PlannerResult::Planned(Expr::ScalarFunction(
+ ScalarFunction::new_udf(crate::unicode::strpos(), args),
+ )))
+ }
+}
diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs
index 2ddd2d22c0..6295821fa9 100644
--- a/datafusion/sql/src/expr/mod.rs
+++ b/datafusion/sql/src/expr/mod.rs
@@ -629,6 +629,31 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}
}
+ fn sql_position_to_expr(
+ &self,
+ substr_expr: SQLExpr,
+ str_expr: SQLExpr,
+ schema: &DFSchema,
+ planner_context: &mut PlannerContext,
+ ) -> Result<Expr> {
+ 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];
+ for planner in self.planners.iter() {
+ match planner.plan_position(extract_args)? {
+ PlannerResult::Planned(expr) => return Ok(expr),
+ PlannerResult::Original(args) => {
+ extract_args = args;
+ }
+ }
+ }
+
+ not_impl_err!(
+ "Position not supported by UserDefinedExtensionPlanners:
{extract_args:?}"
+ )
+ }
+
fn try_plan_dictionary_literal(
&self,
fields: Vec<DictionaryField>,
@@ -924,25 +949,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
};
Ok(Expr::ScalarFunction(ScalarFunction::new_udf(fun, args)))
}
- fn sql_position_to_expr(
- &self,
- substr_expr: SQLExpr,
- str_expr: SQLExpr,
- schema: &DFSchema,
- planner_context: &mut PlannerContext,
- ) -> Result<Expr> {
- let fun = self
- .context_provider
- .get_function_meta("strpos")
- .ok_or_else(|| {
- internal_datafusion_err!("Unable to find expected 'strpos'
function")
- })?;
- 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 args = vec![fullstr, substr];
- Ok(Expr::ScalarFunction(ScalarFunction::new_udf(fun, args)))
- }
}
#[cfg(test)]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
