From 9c4fdd8aca426a76c820c49d0e9cf218929a0443 Mon Sep 17 00:00:00 2001 From: Xin Li Date: Tue, 2 Jul 2024 21:24:24 +0800 Subject: [PATCH] Implement user defined planner for extract --- .../core/src/execution/session_state.rs | 15 +++++--- datafusion/expr/src/planner.rs | 6 ++++ datafusion/functions/src/datetime/mod.rs | 1 + datafusion/functions/src/datetime/planner.rs | 36 +++++++++++++++++++ datafusion/sql/src/expr/mod.rs | 24 ++++++------- 5 files changed, 65 insertions(+), 17 deletions(-) create mode 100644 datafusion/functions/src/datetime/planner.rs diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index ac94ee61fcb2..17e001d25274 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -947,7 +947,7 @@ impl SessionState { where S: ContextProvider, { - let query = SqlToRel::new_with_options(provider, self.get_parser_options()); + let mut query = SqlToRel::new_with_options(provider, self.get_parser_options()); // register crate of array expressions (if enabled) #[cfg(feature = "array_expressions")] @@ -958,14 +958,19 @@ impl SessionState { let field_access_planner = Arc::new(functions_array::planner::FieldAccessPlanner::default()) as _; - query + query = query .with_user_defined_planner(array_planner) - .with_user_defined_planner(field_access_planner) + .with_user_defined_planner(field_access_planner); } - #[cfg(not(feature = "array_expressions"))] + #[cfg(feature = "datetime_expressions")] { - query + let extract_planner = + Arc::new(functions::datetime::planner::ExtractPlanner::default()) as _; + + query = query.with_user_defined_planner(extract_planner); } + + query } } diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs index 1febfbec7ef0..d1a940f1bb75 100644 --- a/datafusion/expr/src/planner.rs +++ b/datafusion/expr/src/planner.rs @@ -110,6 +110,12 @@ pub trait UserDefinedSQLPlanner { ) -> Result>> { Ok(PlannerResult::Original(exprs)) } + + // Plan the Extract expression, e.g., EXTRACT(month FROM foo) + // returns origin expression arguments if not possible + fn plan_extract(&self, args: Vec) -> Result>> { + Ok(PlannerResult::Original(args)) + } } /// An operator with two arguments to plan diff --git a/datafusion/functions/src/datetime/mod.rs b/datafusion/functions/src/datetime/mod.rs index 9c2f80856bf8..8365a38f41f2 100644 --- a/datafusion/functions/src/datetime/mod.rs +++ b/datafusion/functions/src/datetime/mod.rs @@ -30,6 +30,7 @@ pub mod date_trunc; pub mod from_unixtime; pub mod make_date; pub mod now; +pub mod planner; pub mod to_char; pub mod to_date; pub mod to_timestamp; diff --git a/datafusion/functions/src/datetime/planner.rs b/datafusion/functions/src/datetime/planner.rs new file mode 100644 index 000000000000..5f2a5d71d76b --- /dev/null +++ b/datafusion/functions/src/datetime/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 [`ExtractPlanner`] + +use datafusion_common::Result; +use datafusion_expr::{ + expr::ScalarFunction, + planner::{PlannerResult, UserDefinedSQLPlanner}, + Expr, +}; + +#[derive(Default)] +pub struct ExtractPlanner {} + +impl UserDefinedSQLPlanner for ExtractPlanner { + fn plan_extract(&self, args: Vec) -> Result>> { + Ok(PlannerResult::Planned(Expr::ScalarFunction( + ScalarFunction::new_udf(crate::datetime::date_part(), args), + ))) + } +} diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index 786ea288fa0e..24c62a2ea6b6 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -175,21 +175,21 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { self.parse_value(value, planner_context.prepare_param_data_types()) } SQLExpr::Extract { field, expr } => { - let date_part = self - .context_provider - .get_function_meta("date_part") - .ok_or_else(|| { - internal_datafusion_err!( - "Unable to find expected 'date_part' function" - ) - })?; - let args = vec![ + let mut extract_args = vec![ Expr::Literal(ScalarValue::from(format!("{field}"))), self.sql_expr_to_logical_expr(*expr, schema, planner_context)?, ]; - Ok(Expr::ScalarFunction(ScalarFunction::new_udf( - date_part, args, - ))) + + for planner in self.planners.iter() { + match planner.plan_extract(extract_args)? { + PlannerResult::Planned(expr) => return Ok(expr), + PlannerResult::Original(args) => { + extract_args = args; + } + } + } + + not_impl_err!("Extract not supported by UserDefinedExtensionPlanners: {extract_args:?}") } SQLExpr::Array(arr) => self.sql_array_literal(arr.elem, schema),