Skip to content

Commit 970c6ee

Browse files
jayzhan211alamb
authored and
xinlifoobar
committed
Get expr planners when creating new planner (apache#11485)
* get expr planners when creating new planner Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * get expr planner when creating planner Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * no planners in sqltorel Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * Add docs about SessionContextProvider * Use Slice rather than Vec to access expr planners * add test Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * clippy Signed-off-by: jayzhan211 <jayzhan211@gmail.com> --------- Signed-off-by: jayzhan211 <jayzhan211@gmail.com> Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent 9ca0d44 commit 970c6ee

6 files changed

Lines changed: 70 additions & 36 deletions

File tree

‎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

@@ -571,7 +571,7 @@ impl SessionState {
571571
tables: HashMap::new(),
572572
};
573573

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

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

875861
/// A builder to be used for building [`SessionState`]'s. Defaults will
@@ -1599,12 +1585,20 @@ impl SessionStateDefaults {
15991585
}
16001586
}
16011587

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

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

‎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: 2 additions & 11 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::{ArrayElemTypeDef, ExactNumberInfo};
2928
use sqlparser::ast::{ColumnDef as SQLColumnDef, ColumnOption};
3029
use sqlparser::ast::{DataType as SQLDataType, Ident, ObjectName, TableAlias};
@@ -216,8 +215,7 @@ impl PlannerContext {
216215
pub struct SqlToRel<'a, S: ContextProvider> {
217216
pub(crate) context_provider: &'a S,
218217
pub(crate) options: ParserOptions,
219-
pub(crate) ident_normalizer: IdentNormalizer,
220-
pub(crate) value_normalizer: ValueNormalizer,
218+
pub(crate) normalizer: IdentNormalizer,
221219
/// user defined planner extensions
222220
pub(crate) planners: Vec<Arc<dyn ExprPlanner>>,
223221
}
@@ -228,12 +226,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
228226
Self::new_with_options(context_provider, ParserOptions::default())
229227
}
230228

231-
/// add an user defined planner
232-
pub fn with_user_defined_planner(mut self, planner: Arc<dyn ExprPlanner>) -> Self {
233-
self.planners.push(planner);
234-
self
235-
}
236-
237229
/// Create a new query planner
238230
pub fn new_with_options(context_provider: &'a S, options: ParserOptions) -> Self {
239231
let ident_normalize = options.enable_ident_normalization;
@@ -242,8 +234,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
242234
SqlToRel {
243235
context_provider,
244236
options,
245-
ident_normalizer: IdentNormalizer::new(ident_normalize),
246-
value_normalizer: ValueNormalizer::new(options_value_normalize),
237+
normalizer: IdentNormalizer::new(normalize),
247238
planners: vec![],
248239
}
249240
}

0 commit comments

Comments
 (0)