Skip to content

Commit de0765a

Browse files
jayzhan211alamb
jayzhan211
and
alamb
authored
Get expr planners when creating new planner (#11485)
* get expr planners when creating new planner Signed-off-by: jayzhan211 <[email protected]> * get expr planner when creating planner Signed-off-by: jayzhan211 <[email protected]> * no planners in sqltorel Signed-off-by: jayzhan211 <[email protected]> * Add docs about SessionContextProvider * Use Slice rather than Vec to access expr planners * add test Signed-off-by: jayzhan211 <[email protected]> * clippy Signed-off-by: jayzhan211 <[email protected]> --------- Signed-off-by: jayzhan211 <[email protected]> Co-authored-by: Andrew Lamb <[email protected]>
1 parent d67b0fb commit de0765a

File tree

6 files changed

+68
-35
lines changed

6 files changed

+68
-35
lines changed

‎datafusion/core/src/execution/session_state.rs

Lines changed: 54 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -516,7 +516,7 @@ impl SessionState {
516516
}
517517
}
518518

519-
let query = self.build_sql_query_planner(&provider);
519+
let query = SqlToRel::new_with_options(&provider, self.get_parser_options());
520520
query.statement_to_plan(statement)
521521
}
522522

@@ -569,7 +569,7 @@ impl SessionState {
569569
tables: HashMap::new(),
570570
};
571571

572-
let query = self.build_sql_query_planner(&provider);
572+
let query = SqlToRel::new_with_options(&provider, self.get_parser_options());
573573
query.sql_to_expr(sql_expr, df_schema, &mut PlannerContext::new())
574574
}
575575

@@ -854,20 +854,6 @@ impl SessionState {
854854
let udtf = self.table_functions.remove(name);
855855
Ok(udtf.map(|x| x.function().clone()))
856856
}
857-
858-
fn build_sql_query_planner<'a, S>(&self, provider: &'a S) -> SqlToRel<'a, S>
859-
where
860-
S: ContextProvider,
861-
{
862-
let mut query = SqlToRel::new_with_options(provider, self.get_parser_options());
863-
864-
// custom planners are registered first, so they're run first and take precedence over built-in planners
865-
for planner in self.expr_planners.iter() {
866-
query = query.with_user_defined_planner(planner.clone());
867-
}
868-
869-
query
870-
}
871857
}
872858

873859
/// A builder to be used for building [`SessionState`]'s. Defaults will
@@ -1597,12 +1583,20 @@ impl SessionStateDefaults {
15971583
}
15981584
}
15991585

1586+
/// Adapter that implements the [`ContextProvider`] trait for a [`SessionState`]
1587+
///
1588+
/// This is used so the SQL planner can access the state of the session without
1589+
/// having a direct dependency on the [`SessionState`] struct (and core crate)
16001590
struct SessionContextProvider<'a> {
16011591
state: &'a SessionState,
16021592
tables: HashMap<String, Arc<dyn TableSource>>,
16031593
}
16041594

16051595
impl<'a> ContextProvider for SessionContextProvider<'a> {
1596+
fn get_expr_planners(&self) -> &[Arc<dyn ExprPlanner>] {
1597+
&self.state.expr_planners
1598+
}
1599+
16061600
fn get_table_source(
16071601
&self,
16081602
name: TableReference,
@@ -1898,3 +1892,47 @@ impl<'a> SimplifyInfo for SessionSimplifyProvider<'a> {
18981892
expr.get_type(self.df_schema)
18991893
}
19001894
}
1895+
1896+
#[cfg(test)]
1897+
mod tests {
1898+
use std::collections::HashMap;
1899+
1900+
use arrow_schema::{DataType, Field, Schema};
1901+
use datafusion_common::DFSchema;
1902+
use datafusion_common::Result;
1903+
use datafusion_expr::Expr;
1904+
use datafusion_sql::planner::{PlannerContext, SqlToRel};
1905+
1906+
use crate::execution::context::SessionState;
1907+
1908+
use super::{SessionContextProvider, SessionStateBuilder};
1909+
1910+
#[test]
1911+
fn test_session_state_with_default_features() {
1912+
// test array planners with and without builtin planners
1913+
fn sql_to_expr(state: &SessionState) -> Result<Expr> {
1914+
let provider = SessionContextProvider {
1915+
state,
1916+
tables: HashMap::new(),
1917+
};
1918+
1919+
let sql = "[1,2,3]";
1920+
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
1921+
let df_schema = DFSchema::try_from(schema)?;
1922+
let dialect = state.config.options().sql_parser.dialect.as_str();
1923+
let sql_expr = state.sql_to_expr(sql, dialect)?;
1924+
1925+
let query = SqlToRel::new_with_options(&provider, state.get_parser_options());
1926+
query.sql_to_expr(sql_expr, &df_schema, &mut PlannerContext::new())
1927+
}
1928+
1929+
let state = SessionStateBuilder::new().with_default_features().build();
1930+
1931+
assert!(sql_to_expr(&state).is_ok());
1932+
1933+
// if no builtin planners exist, you should register your own, otherwise returns error
1934+
let state = SessionStateBuilder::new().build();
1935+
1936+
assert!(sql_to_expr(&state).is_err())
1937+
}
1938+
}

‎datafusion/expr/src/planner.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,11 @@ pub trait ContextProvider {
6060
not_impl_err!("Recursive CTE is not implemented")
6161
}
6262

63+
/// Getter for expr planners
64+
fn get_expr_planners(&self) -> &[Arc<dyn ExprPlanner>] {
65+
&[]
66+
}
67+
6368
/// Getter for a UDF description
6469
fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>>;
6570
/// Getter for a UDAF description

‎datafusion/sql/src/expr/mod.rs

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
111111
) -> Result<Expr> {
112112
// try extension planers
113113
let mut binary_expr = datafusion_expr::planner::RawBinaryExpr { op, left, right };
114-
for planner in self.planners.iter() {
114+
for planner in self.context_provider.get_expr_planners() {
115115
match planner.plan_binary_op(binary_expr, schema)? {
116116
PlannerResult::Planned(expr) => {
117117
return Ok(expr);
@@ -184,7 +184,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
184184
self.sql_expr_to_logical_expr(*expr, schema, planner_context)?,
185185
];
186186

187-
for planner in self.planners.iter() {
187+
for planner in self.context_provider.get_expr_planners() {
188188
match planner.plan_extract(extract_args)? {
189189
PlannerResult::Planned(expr) => return Ok(expr),
190190
PlannerResult::Original(args) => {
@@ -283,7 +283,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
283283
};
284284

285285
let mut field_access_expr = RawFieldAccessExpr { expr, field_access };
286-
for planner in self.planners.iter() {
286+
for planner in self.context_provider.get_expr_planners() {
287287
match planner.plan_field_access(field_access_expr, schema)? {
288288
PlannerResult::Planned(expr) => return Ok(expr),
289289
PlannerResult::Original(expr) => {
@@ -653,7 +653,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
653653
self.create_struct_expr(values, schema, planner_context)?
654654
};
655655

656-
for planner in self.planners.iter() {
656+
for planner in self.context_provider.get_expr_planners() {
657657
match planner.plan_struct_literal(create_struct_args, is_named_struct)? {
658658
PlannerResult::Planned(expr) => return Ok(expr),
659659
PlannerResult::Original(args) => create_struct_args = args,
@@ -673,7 +673,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
673673
self.sql_expr_to_logical_expr(substr_expr, schema, planner_context)?;
674674
let fullstr = self.sql_expr_to_logical_expr(str_expr, schema, planner_context)?;
675675
let mut position_args = vec![fullstr, substr];
676-
for planner in self.planners.iter() {
676+
for planner in self.context_provider.get_expr_planners() {
677677
match planner.plan_position(position_args)? {
678678
PlannerResult::Planned(expr) => return Ok(expr),
679679
PlannerResult::Original(args) => {
@@ -703,7 +703,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
703703

704704
let mut raw_expr = RawDictionaryExpr { keys, values };
705705

706-
for planner in self.planners.iter() {
706+
for planner in self.context_provider.get_expr_planners() {
707707
match planner.plan_dictionary_literal(raw_expr, schema)? {
708708
PlannerResult::Planned(expr) => {
709709
return Ok(expr);
@@ -927,7 +927,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
927927
}
928928
None => vec![arg, what_arg, from_arg],
929929
};
930-
for planner in self.planners.iter() {
930+
for planner in self.context_provider.get_expr_planners() {
931931
match planner.plan_overlay(overlay_args)? {
932932
PlannerResult::Planned(expr) => return Ok(expr),
933933
PlannerResult::Original(args) => overlay_args = args,

‎datafusion/sql/src/expr/substring.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
6868
}
6969
};
7070

71-
for planner in self.planners.iter() {
71+
for planner in self.context_provider.get_expr_planners() {
7272
match planner.plan_substring(substring_args)? {
7373
PlannerResult::Planned(expr) => return Ok(expr),
7474
PlannerResult::Original(args) => {

‎datafusion/sql/src/expr/value.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
154154
schema: &DFSchema,
155155
) -> Result<Expr> {
156156
let mut exprs = values;
157-
for planner in self.planners.iter() {
157+
for planner in self.context_provider.get_expr_planners() {
158158
match planner.plan_array_literal(exprs, schema)? {
159159
PlannerResult::Planned(expr) => {
160160
return Ok(expr);

‎datafusion/sql/src/planner.rs

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ use arrow_schema::*;
2424
use datafusion_common::{
2525
field_not_found, internal_err, plan_datafusion_err, DFSchemaRef, SchemaError,
2626
};
27-
use datafusion_expr::planner::ExprPlanner;
2827
use sqlparser::ast::TimezoneInfo;
2928
use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo};
3029
use sqlparser::ast::{ColumnDef as SQLColumnDef, ColumnOption};
@@ -186,8 +185,6 @@ pub struct SqlToRel<'a, S: ContextProvider> {
186185
pub(crate) context_provider: &'a S,
187186
pub(crate) options: ParserOptions,
188187
pub(crate) normalizer: IdentNormalizer,
189-
/// user defined planner extensions
190-
pub(crate) planners: Vec<Arc<dyn ExprPlanner>>,
191188
}
192189

193190
impl<'a, S: ContextProvider> SqlToRel<'a, S> {
@@ -196,12 +193,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
196193
Self::new_with_options(context_provider, ParserOptions::default())
197194
}
198195

199-
/// add an user defined planner
200-
pub fn with_user_defined_planner(mut self, planner: Arc<dyn ExprPlanner>) -> Self {
201-
self.planners.push(planner);
202-
self
203-
}
204-
205196
/// Create a new query planner
206197
pub fn new_with_options(context_provider: &'a S, options: ParserOptions) -> Self {
207198
let normalize = options.enable_ident_normalization;
@@ -210,7 +201,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
210201
context_provider,
211202
options,
212203
normalizer: IdentNormalizer::new(normalize),
213-
planners: vec![],
214204
}
215205
}
216206

0 commit comments

Comments
 (0)