diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index aa81d77cf682..084557d99e6e 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -967,14 +967,19 @@ impl SessionState { let field_access_planner = Arc::new(functions_array::planner::FieldAccessPlanner) 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 c928ab39194d..873085c585c2 100644 --- a/datafusion/expr/src/planner.rs +++ b/datafusion/expr/src/planner.rs @@ -110,6 +110,12 @@ pub trait UserDefinedSQLPlanner: Send + Sync { ) -> 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),