Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: Rename UserDefinedSQLPlanner to ExprPlanner #11338

Merged
merged 4 commits into from
Jul 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 5 additions & 7 deletions datafusion/core/src/execution/context/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ use datafusion_execution::registry::SerializerRegistry;
use datafusion_expr::{
expr_rewriter::FunctionRewrite,
logical_plan::{DdlStatement, Statement},
planner::UserDefinedSQLPlanner,
planner::ExprPlanner,
Expr, UserDefinedLogicalNode, WindowUDF,
};

Expand Down Expand Up @@ -1392,17 +1392,15 @@ impl FunctionRegistry for SessionContext {
self.state.write().register_function_rewrite(rewrite)
}

fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
self.state.read().expr_planners()
}

fn register_user_defined_sql_planner(
fn register_expr_planner(
&mut self,
user_defined_sql_planner: Arc<dyn UserDefinedSQLPlanner>,
expr_planner: Arc<dyn ExprPlanner>,
) -> Result<()> {
self.state
.write()
.register_user_defined_sql_planner(user_defined_sql_planner)
self.state.write().register_expr_planner(expr_planner)
}
}

Expand Down
21 changes: 10 additions & 11 deletions datafusion/core/src/execution/session_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_execution::TaskContext;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::expr_rewriter::FunctionRewrite;
use datafusion_expr::planner::UserDefinedSQLPlanner;
use datafusion_expr::planner::ExprPlanner;
use datafusion_expr::registry::{FunctionRegistry, SerializerRegistry};
use datafusion_expr::simplify::SimplifyInfo;
use datafusion_expr::var_provider::{is_system_variables, VarType};
Expand Down Expand Up @@ -101,7 +101,7 @@ pub struct SessionState {
/// Responsible for analyzing and rewrite a logical plan before optimization
analyzer: Analyzer,
/// Provides support for customising the SQL planner, e.g. to add support for custom operators like `->>` or `?`
user_defined_sql_planners: Vec<Arc<dyn UserDefinedSQLPlanner>>,
expr_planners: Vec<Arc<dyn ExprPlanner>>,
/// Responsible for optimizing a logical plan
optimizer: Optimizer,
/// Responsible for optimizing a physical execution plan
Expand Down Expand Up @@ -231,7 +231,7 @@ impl SessionState {
);
}

let user_defined_sql_planners: Vec<Arc<dyn UserDefinedSQLPlanner>> = vec![
let expr_planners: Vec<Arc<dyn ExprPlanner>> = vec![
Arc::new(functions::core::planner::CoreFunctionPlanner::default()),
// register crate of array expressions (if enabled)
#[cfg(feature = "array_expressions")]
Expand All @@ -248,7 +248,7 @@ impl SessionState {
let mut new_self = SessionState {
session_id,
analyzer: Analyzer::new(),
user_defined_sql_planners,
expr_planners,
optimizer: Optimizer::new(),
physical_optimizers: PhysicalOptimizer::new(),
query_planner: Arc::new(DefaultQueryPlanner {}),
Expand Down Expand Up @@ -968,7 +968,7 @@ impl SessionState {
let mut query = SqlToRel::new_with_options(provider, self.get_parser_options());

// custom planners are registered first, so they're run first and take precedence over built-in planners
for planner in self.user_defined_sql_planners.iter() {
for planner in self.expr_planners.iter() {
query = query.with_user_defined_planner(planner.clone());
}

Expand Down Expand Up @@ -1184,16 +1184,15 @@ impl FunctionRegistry for SessionState {
Ok(())
}

fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
self.user_defined_sql_planners.clone()
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
self.expr_planners.clone()
}

fn register_user_defined_sql_planner(
fn register_expr_planner(
&mut self,
user_defined_sql_planner: Arc<dyn UserDefinedSQLPlanner>,
expr_planner: Arc<dyn ExprPlanner>,
) -> datafusion_common::Result<()> {
self.user_defined_sql_planners
.push(user_defined_sql_planner);
self.expr_planners.push(expr_planner);
Ok(())
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,12 @@ use datafusion::prelude::*;
use datafusion::sql::sqlparser::ast::BinaryOperator;
use datafusion_common::ScalarValue;
use datafusion_expr::expr::Alias;
use datafusion_expr::planner::{PlannerResult, RawBinaryExpr, UserDefinedSQLPlanner};
use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr};
use datafusion_expr::BinaryExpr;

struct MyCustomPlanner;

impl UserDefinedSQLPlanner for MyCustomPlanner {
impl ExprPlanner for MyCustomPlanner {
fn plan_binary_op(
&self,
expr: RawBinaryExpr,
Expand Down Expand Up @@ -68,7 +68,7 @@ async fn plan_and_collect(sql: &str) -> Result<Vec<RecordBatch>> {
let config =
SessionConfig::new().set_str("datafusion.sql_parser.dialect", "postgres");
let mut ctx = SessionContext::new_with_config(config);
ctx.register_user_defined_sql_planner(Arc::new(MyCustomPlanner))?;
ctx.register_expr_planner(Arc::new(MyCustomPlanner))?;
ctx.sql(sql).await?.collect().await
}

Expand Down
4 changes: 2 additions & 2 deletions datafusion/core/tests/user_defined/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,5 +30,5 @@ mod user_defined_window_functions;
/// Tests for User Defined Table Functions
mod user_defined_table_functions;

/// Tests for User Defined SQL Planner
mod user_defined_sql_planner;
/// Tests for Expression Planner
mod expr_planner;
Original file line number Diff line number Diff line change
Expand Up @@ -575,9 +575,9 @@ async fn test_user_defined_functions_cast_to_i64() -> Result<()> {
async fn test_user_defined_sql_functions() -> Result<()> {
let ctx = SessionContext::new();

let sql_planners = ctx.expr_planners();
let expr_planners = ctx.expr_planners();

assert!(!sql_planners.is_empty());
assert!(!expr_planners.is_empty());

Ok(())
}
Expand Down
4 changes: 2 additions & 2 deletions datafusion/execution/src/task.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use crate::{
runtime_env::{RuntimeConfig, RuntimeEnv},
};
use datafusion_common::{plan_datafusion_err, DataFusionError, Result};
use datafusion_expr::planner::UserDefinedSQLPlanner;
use datafusion_expr::planner::ExprPlanner;
use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};

/// Task Execution Context
Expand Down Expand Up @@ -192,7 +192,7 @@ impl FunctionRegistry for TaskContext {
Ok(self.scalar_functions.insert(udf.name().into(), udf))
}

fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}
}
Expand Down
12 changes: 6 additions & 6 deletions datafusion/expr/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.

//! [`ContextProvider`] and [`UserDefinedSQLPlanner`] APIs to customize SQL query planning
//! [`ContextProvider`] and [`ExprPlanner`] APIs to customize SQL query planning

use std::sync::Arc;

Expand Down Expand Up @@ -83,7 +83,7 @@ pub trait ContextProvider {
}

/// This trait allows users to customize the behavior of the SQL planner
pub trait UserDefinedSQLPlanner: Send + Sync {
pub trait ExprPlanner: Send + Sync {
/// Plan the binary operation between two expressions, returns original
/// BinaryExpr if not possible
fn plan_binary_op(
Expand Down Expand Up @@ -168,7 +168,7 @@ pub trait UserDefinedSQLPlanner: Send + Sync {
/// Note `left` and `right` are DataFusion [`Expr`]s but the `op` is the SQL AST
/// operator.
///
/// This structure is used by [`UserDefinedSQLPlanner`] to plan operators with
/// This structure is used by [`ExprPlanner`] to plan operators with
/// custom expressions.
#[derive(Debug, Clone)]
pub struct RawBinaryExpr {
Expand All @@ -179,7 +179,7 @@ pub struct RawBinaryExpr {

/// An expression with GetFieldAccess to plan
///
/// This structure is used by [`UserDefinedSQLPlanner`] to plan operators with
/// This structure is used by [`ExprPlanner`] to plan operators with
/// custom expressions.
#[derive(Debug, Clone)]
pub struct RawFieldAccessExpr {
Expand All @@ -189,15 +189,15 @@ pub struct RawFieldAccessExpr {

/// A Dictionary literal expression `{ key: value, ...}`
///
/// This structure is used by [`UserDefinedSQLPlanner`] to plan operators with
/// This structure is used by [`ExprPlanner`] to plan operators with
/// custom expressions.
#[derive(Debug, Clone)]
pub struct RawDictionaryExpr {
pub keys: Vec<Expr>,
pub values: Vec<Expr>,
}

/// Result of planning a raw expr with [`UserDefinedSQLPlanner`]
/// Result of planning a raw expr with [`ExprPlanner`]
#[derive(Debug, Clone)]
pub enum PlannerResult<T> {
/// The raw expression was successfully planned as a new [`Expr`]
Expand Down
16 changes: 8 additions & 8 deletions datafusion/expr/src/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
//! FunctionRegistry trait

use crate::expr_rewriter::FunctionRewrite;
use crate::planner::UserDefinedSQLPlanner;
use crate::planner::ExprPlanner;
use crate::{AggregateUDF, ScalarUDF, UserDefinedLogicalNode, WindowUDF};
use datafusion_common::{not_impl_err, plan_datafusion_err, Result};
use std::collections::HashMap;
Expand Down Expand Up @@ -110,15 +110,15 @@ pub trait FunctionRegistry {
not_impl_err!("Registering FunctionRewrite")
}

/// Set of all registered [`UserDefinedSQLPlanner`]s
fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>>;
/// Set of all registered [`ExprPlanner`]s
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>>;

/// Registers a new [`UserDefinedSQLPlanner`] with the registry.
fn register_user_defined_sql_planner(
/// Registers a new [`ExprPlanner`] with the registry.
fn register_expr_planner(
&mut self,
_user_defined_sql_planner: Arc<dyn UserDefinedSQLPlanner>,
_expr_planner: Arc<dyn ExprPlanner>,
) -> Result<()> {
not_impl_err!("Registering UserDefinedSQLPlanner")
not_impl_err!("Registering ExprPlanner")
}
}

Expand Down Expand Up @@ -196,7 +196,7 @@ impl FunctionRegistry for MemoryFunctionRegistry {
Ok(self.udwfs.insert(udaf.name().into(), udaf))
}

fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}
}
6 changes: 3 additions & 3 deletions datafusion/functions-array/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

use datafusion_common::{utils::list_ndims, DFSchema, Result};
use datafusion_expr::{
planner::{PlannerResult, RawBinaryExpr, RawFieldAccessExpr, UserDefinedSQLPlanner},
planner::{ExprPlanner, PlannerResult, RawBinaryExpr, RawFieldAccessExpr},
sqlparser, AggregateFunction, Expr, ExprSchemable, GetFieldAccess,
};
use datafusion_functions::expr_fn::get_field;
Expand All @@ -34,7 +34,7 @@ use crate::{

pub struct ArrayFunctionPlanner;

impl UserDefinedSQLPlanner for ArrayFunctionPlanner {
impl ExprPlanner for ArrayFunctionPlanner {
fn plan_binary_op(
&self,
expr: RawBinaryExpr,
Expand Down Expand Up @@ -101,7 +101,7 @@ impl UserDefinedSQLPlanner for ArrayFunctionPlanner {

pub struct FieldAccessPlanner;

impl UserDefinedSQLPlanner for FieldAccessPlanner {
impl ExprPlanner for FieldAccessPlanner {
fn plan_field_access(
&self,
expr: RawFieldAccessExpr,
Expand Down
4 changes: 2 additions & 2 deletions datafusion/functions/src/core/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,15 @@
use datafusion_common::DFSchema;
use datafusion_common::Result;
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::planner::{PlannerResult, RawDictionaryExpr, UserDefinedSQLPlanner};
use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawDictionaryExpr};
use datafusion_expr::Expr;

use super::named_struct;

#[derive(Default)]
pub struct CoreFunctionPlanner {}

impl UserDefinedSQLPlanner for CoreFunctionPlanner {
impl ExprPlanner for CoreFunctionPlanner {
fn plan_dictionary_literal(
&self,
expr: RawDictionaryExpr,
Expand Down
4 changes: 2 additions & 2 deletions datafusion/functions/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,14 @@
use datafusion_common::Result;
use datafusion_expr::{
expr::ScalarFunction,
planner::{PlannerResult, UserDefinedSQLPlanner},
planner::{ExprPlanner, PlannerResult},
Expr,
};

#[derive(Default)]
pub struct UserDefinedFunctionPlanner;

impl UserDefinedSQLPlanner for UserDefinedFunctionPlanner {
impl ExprPlanner for UserDefinedFunctionPlanner {
#[cfg(feature = "datetime_expressions")]
fn plan_extract(&self, args: Vec<Expr>) -> Result<PlannerResult<Vec<Expr>>> {
Ok(PlannerResult::Planned(Expr::ScalarFunction(
Expand Down
4 changes: 2 additions & 2 deletions datafusion/proto/src/bytes/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ use std::sync::Arc;
use datafusion::execution::registry::FunctionRegistry;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::SessionContext;
use datafusion_expr::planner::UserDefinedSQLPlanner;
use datafusion_expr::planner::ExprPlanner;

mod registry;

Expand Down Expand Up @@ -167,7 +167,7 @@ impl Serializeable for Expr {
)
}

fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}
}
Expand Down
4 changes: 2 additions & 2 deletions datafusion/proto/src/bytes/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ use std::{collections::HashSet, sync::Arc};
use datafusion::execution::registry::FunctionRegistry;
use datafusion_common::plan_err;
use datafusion_common::Result;
use datafusion_expr::planner::UserDefinedSQLPlanner;
use datafusion_expr::planner::ExprPlanner;
use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};

/// A default [`FunctionRegistry`] registry that does not resolve any
Expand Down Expand Up @@ -56,7 +56,7 @@ impl FunctionRegistry for NoRegistry {
plan_err!("No function registry provided to deserialize, so can not deserialize User Defined Window Function '{}'", udwf.inner().name())
}

fn expr_planners(&self) -> Vec<Arc<dyn UserDefinedSQLPlanner>> {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}
}
9 changes: 3 additions & 6 deletions datafusion/sql/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use arrow_schema::*;
use datafusion_common::{
field_not_found, internal_err, plan_datafusion_err, DFSchemaRef, SchemaError,
};
use datafusion_expr::planner::UserDefinedSQLPlanner;
use datafusion_expr::planner::ExprPlanner;
use sqlparser::ast::TimezoneInfo;
use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo};
use sqlparser::ast::{ColumnDef as SQLColumnDef, ColumnOption};
Expand Down Expand Up @@ -187,7 +187,7 @@ pub struct SqlToRel<'a, S: ContextProvider> {
pub(crate) options: ParserOptions,
pub(crate) normalizer: IdentNormalizer,
/// user defined planner extensions
pub(crate) planners: Vec<Arc<dyn UserDefinedSQLPlanner>>,
pub(crate) planners: Vec<Arc<dyn ExprPlanner>>,
}

impl<'a, S: ContextProvider> SqlToRel<'a, S> {
Expand All @@ -197,10 +197,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}

/// add an user defined planner
pub fn with_user_defined_planner(
mut self,
planner: Arc<dyn UserDefinedSQLPlanner>,
) -> Self {
pub fn with_user_defined_planner(mut self, planner: Arc<dyn ExprPlanner>) -> Self {
self.planners.push(planner);
self
}
Expand Down