This is an automated email from the ASF dual-hosted git repository.
jayzhan 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 1e39a8575b Implement user defined planner for `create_struct` &
`create_named_struct` (#11273)
1e39a8575b is described below
commit 1e39a8575be54c6cfbd9af967705509299c80991
Author: Dharan Aditya <[email protected]>
AuthorDate: Mon Jul 8 10:39:39 2024 +0530
Implement user defined planner for `create_struct` & `create_named_struct`
(#11273)
* add UserDefinedSQLPlanner for create struct
* fix linting
* add create name struct user defined sql planner
* simplify
* refactor
* refactor
* remove named_struct from functions
* formatting
* revert 953ad31
* update docs
---
datafusion/expr/src/planner.rs | 15 +++++
datafusion/functions/src/core/mod.rs | 1 -
datafusion/functions/src/core/planner.rs | 19 +++++++
datafusion/sql/src/expr/mod.rs | 96 ++++++++++++++------------------
4 files changed, 75 insertions(+), 56 deletions(-)
diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs
index bcbf5eb203..4086964ab8 100644
--- a/datafusion/expr/src/planner.rs
+++ b/datafusion/expr/src/planner.rs
@@ -139,6 +139,21 @@ pub trait UserDefinedSQLPlanner: Send + Sync {
fn plan_extract(&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
+ /// the struct uses the optional name
+ ///
+ /// Returns a `PlannerResult` containing either the planned struct
expressions or the original
+ /// input expressions if planning is not possible.
+ fn plan_struct_literal(
+ &self,
+ args: Vec<Expr>,
+ _is_named_struct: bool,
+ ) -> Result<PlannerResult<Vec<Expr>>> {
+ Ok(PlannerResult::Original(args))
+ }
}
/// An operator with two arguments to plan
diff --git a/datafusion/functions/src/core/mod.rs
b/datafusion/functions/src/core/mod.rs
index 5761a4bd39..062a4a104d 100644
--- a/datafusion/functions/src/core/mod.rs
+++ b/datafusion/functions/src/core/mod.rs
@@ -93,7 +93,6 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
nvl(),
nvl2(),
arrow_typeof(),
- r#struct(),
named_struct(),
get_field(),
coalesce(),
diff --git a/datafusion/functions/src/core/planner.rs
b/datafusion/functions/src/core/planner.rs
index 71f6c7f8ac..e803c92dd0 100644
--- a/datafusion/functions/src/core/planner.rs
+++ b/datafusion/functions/src/core/planner.rs
@@ -17,7 +17,9 @@
use datafusion_common::DFSchema;
use datafusion_common::Result;
+use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::planner::{PlannerResult, RawDictionaryExpr,
UserDefinedSQLPlanner};
+use datafusion_expr::Expr;
use super::named_struct;
@@ -37,4 +39,21 @@ impl UserDefinedSQLPlanner for CoreFunctionPlanner {
}
Ok(PlannerResult::Planned(named_struct().call(args)))
}
+
+ fn plan_struct_literal(
+ &self,
+ args: Vec<Expr>,
+ is_named_struct: bool,
+ ) -> Result<PlannerResult<Vec<Expr>>> {
+ Ok(PlannerResult::Planned(Expr::ScalarFunction(
+ ScalarFunction::new_udf(
+ if is_named_struct {
+ crate::core::named_struct()
+ } else {
+ crate::core::r#struct()
+ },
+ args,
+ ),
+ )))
+ }
}
diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs
index 6295821fa9..eba93735b5 100644
--- a/datafusion/sql/src/expr/mod.rs
+++ b/datafusion/sql/src/expr/mod.rs
@@ -21,7 +21,8 @@ use datafusion_expr::planner::PlannerResult;
use datafusion_expr::planner::RawDictionaryExpr;
use datafusion_expr::planner::RawFieldAccessExpr;
use sqlparser::ast::{
- CastKind, DictionaryField, Expr as SQLExpr, Subscript, TrimWhereField,
Value,
+ CastKind, DictionaryField, Expr as SQLExpr, StructField, Subscript,
TrimWhereField,
+ Value,
};
use datafusion_common::{
@@ -597,7 +598,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}
SQLExpr::Struct { values, fields } => {
- self.parse_struct(values, fields, schema, planner_context)
+ self.parse_struct(schema, planner_context, values, fields)
}
SQLExpr::Position { expr, r#in } => {
self.sql_position_to_expr(*expr, *r#in, schema,
planner_context)
@@ -629,6 +630,36 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}
}
+ /// Parses a struct(..) expression and plans it creation
+ fn parse_struct(
+ &self,
+ schema: &DFSchema,
+ planner_context: &mut PlannerContext,
+ values: Vec<sqlparser::ast::Expr>,
+ fields: Vec<StructField>,
+ ) -> Result<Expr> {
+ if !fields.is_empty() {
+ return not_impl_err!("Struct fields are not supported yet");
+ }
+ let is_named_struct = values
+ .iter()
+ .any(|value| matches!(value, SQLExpr::Named { .. }));
+
+ let mut create_struct_args = if is_named_struct {
+ self.create_named_struct_expr(values, schema, planner_context)?
+ } else {
+ self.create_struct_expr(values, schema, planner_context)?
+ };
+
+ for planner in self.planners.iter() {
+ match planner.plan_struct_literal(create_struct_args,
is_named_struct)? {
+ PlannerResult::Planned(expr) => return Ok(expr),
+ PlannerResult::Original(args) => create_struct_args = args,
+ }
+ }
+ not_impl_err!("Struct not supported by UserDefinedExtensionPlanners:
{create_struct_args:?}")
+ }
+
fn sql_position_to_expr(
&self,
substr_expr: SQLExpr,
@@ -683,37 +714,15 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
not_impl_err!("Unsupported dictionary literal: {raw_expr:?}")
}
- /// Parses a struct(..) expression
- fn parse_struct(
- &self,
- values: Vec<SQLExpr>,
- fields: Vec<sqlparser::ast::StructField>,
- input_schema: &DFSchema,
- planner_context: &mut PlannerContext,
- ) -> Result<Expr> {
- if !fields.is_empty() {
- return not_impl_err!("Struct fields are not supported yet");
- }
-
- if values
- .iter()
- .any(|value| matches!(value, SQLExpr::Named { .. }))
- {
- self.create_named_struct(values, input_schema, planner_context)
- } else {
- self.create_struct(values, input_schema, planner_context)
- }
- }
-
// Handles a call to struct(...) where the arguments are named. For example
// `struct (v as foo, v2 as bar)` by creating a call to the `named_struct`
function
- fn create_named_struct(
+ fn create_named_struct_expr(
&self,
values: Vec<SQLExpr>,
input_schema: &DFSchema,
planner_context: &mut PlannerContext,
- ) -> Result<Expr> {
- let args = values
+ ) -> Result<Vec<Expr>> {
+ Ok(values
.into_iter()
.enumerate()
.map(|(i, value)| {
@@ -742,47 +751,24 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
.collect::<Result<Vec<_>>>()?
.into_iter()
.flatten()
- .collect();
-
- let named_struct_func = self
- .context_provider
- .get_function_meta("named_struct")
- .ok_or_else(|| {
- internal_datafusion_err!("Unable to find expected 'named_struct'
function")
- })?;
-
- Ok(Expr::ScalarFunction(ScalarFunction::new_udf(
- named_struct_func,
- args,
- )))
+ .collect())
}
// Handles a call to struct(...) where the arguments are not named. For
example
// `struct (v, v2)` by creating a call to the `struct` function
// which will create a struct with fields named `c0`, `c1`, etc.
- fn create_struct(
+ fn create_struct_expr(
&self,
values: Vec<SQLExpr>,
input_schema: &DFSchema,
planner_context: &mut PlannerContext,
- ) -> Result<Expr> {
- let args = values
+ ) -> Result<Vec<Expr>> {
+ values
.into_iter()
.map(|value| {
self.sql_expr_to_logical_expr(value, input_schema,
planner_context)
})
- .collect::<Result<Vec<_>>>()?;
- let struct_func = self
- .context_provider
- .get_function_meta("struct")
- .ok_or_else(|| {
- internal_datafusion_err!("Unable to find expected 'struct'
function")
- })?;
-
- Ok(Expr::ScalarFunction(ScalarFunction::new_udf(
- struct_func,
- args,
- )))
+ .collect::<Result<Vec<_>>>()
}
fn sql_in_list_to_expr(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
