Skip to content

Commit fd350fa

Browse files
alambalamb
authored
minor(sqlparser): encapsulate PlanerContext, reduce some clones (#5814)
* minor(sqlparser): encapsulate PlanerContext, reduce some clones * clippy
1 parent 2191a69 commit fd350fa

File tree

9 files changed

+76
-43
lines changed

9 files changed

+76
-43
lines changed

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

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -58,9 +58,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
5858
}
5959
Err(_) => {
6060
// check the outer_query_schema and try to find a match
61-
let outer_query_schema_opt =
62-
planner_context.outer_query_schema.as_ref();
63-
if let Some(outer) = outer_query_schema_opt {
61+
if let Some(outer) = planner_context.outer_query_schema() {
6462
match outer.field_with_unqualified_name(normalize_ident.as_str())
6563
{
6664
Ok(field) => {
@@ -159,10 +157,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
159157
"Unsupported compound identifier: {ids:?}"
160158
)))
161159
} else {
162-
let outer_query_schema_opt =
163-
planner_context.outer_query_schema.as_ref();
164160
// check the outer_query_schema and try to find a match
165-
if let Some(outer) = outer_query_schema_opt {
161+
if let Some(outer) = planner_context.outer_query_schema() {
166162
let search_result = search_dfschema(&ids, outer);
167163
match search_result {
168164
// found matching field with spare identifier(s) for nested field(s) in structure

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
120120
) -> Result<Expr> {
121121
match sql {
122122
SQLExpr::Value(value) => {
123-
self.parse_value(value, &planner_context.prepare_param_data_types)
123+
self.parse_value(value, planner_context.prepare_param_data_types())
124124
}
125125
SQLExpr::Extract { field, expr } => Ok(Expr::ScalarFunction {
126126
fun: BuiltinScalarFunction::DatePart,

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,11 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
3030
input_schema: &DFSchema,
3131
planner_context: &mut PlannerContext,
3232
) -> Result<Expr> {
33-
let old_outer_query_schema = planner_context.outer_query_schema.clone();
34-
planner_context.outer_query_schema = Some(input_schema.clone());
33+
let old_outer_query_schema =
34+
planner_context.set_outer_query_schema(Some(input_schema.clone()));
3535
let sub_plan = self.query_to_plan(subquery, planner_context)?;
3636
let outer_ref_columns = sub_plan.all_out_ref_exprs();
37-
planner_context.outer_query_schema = old_outer_query_schema;
37+
planner_context.set_outer_query_schema(old_outer_query_schema);
3838
Ok(Expr::Exists {
3939
subquery: Subquery {
4040
subquery: Arc::new(sub_plan),
@@ -52,11 +52,11 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
5252
input_schema: &DFSchema,
5353
planner_context: &mut PlannerContext,
5454
) -> Result<Expr> {
55-
let old_outer_query_schema = planner_context.outer_query_schema.clone();
56-
planner_context.outer_query_schema = Some(input_schema.clone());
55+
let old_outer_query_schema =
56+
planner_context.set_outer_query_schema(Some(input_schema.clone()));
5757
let sub_plan = self.query_to_plan(subquery, planner_context)?;
5858
let outer_ref_columns = sub_plan.all_out_ref_exprs();
59-
planner_context.outer_query_schema = old_outer_query_schema;
59+
planner_context.set_outer_query_schema(old_outer_query_schema);
6060
let expr = Box::new(self.sql_to_expr(expr, input_schema, planner_context)?);
6161
Ok(Expr::InSubquery {
6262
expr,
@@ -74,11 +74,11 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
7474
input_schema: &DFSchema,
7575
planner_context: &mut PlannerContext,
7676
) -> Result<Expr> {
77-
let old_outer_query_schema = planner_context.outer_query_schema.clone();
78-
planner_context.outer_query_schema = Some(input_schema.clone());
77+
let old_outer_query_schema =
78+
planner_context.set_outer_query_schema(Some(input_schema.clone()));
7979
let sub_plan = self.query_to_plan(subquery, planner_context)?;
8080
let outer_ref_columns = sub_plan.all_out_ref_exprs();
81-
planner_context.outer_query_schema = old_outer_query_schema;
81+
planner_context.set_outer_query_schema(old_outer_query_schema);
8282
Ok(Expr::ScalarSubquery(Subquery {
8383
subquery: Arc::new(sub_plan),
8484
outer_ref_columns,

‎datafusion/sql/src/planner.rs

Lines changed: 52 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -69,19 +69,21 @@ impl Default for ParserOptions {
6969
}
7070
}
7171

72-
#[derive(Debug, Clone)]
7372
/// Struct to store the states used by the Planner. The Planner will leverage the states to resolve
7473
/// CTEs, Views, subqueries and PREPARE statements. The states include
7574
/// Common Table Expression (CTE) provided with WITH clause and
7675
/// Parameter Data Types provided with PREPARE statement and the query schema of the
7776
/// outer query plan
77+
#[derive(Debug, Clone)]
7878
pub struct PlannerContext {
79-
/// Data type provided with prepare statement
80-
pub prepare_param_data_types: Vec<DataType>,
81-
/// Map of CTE name to logical plan of the WITH clause
82-
pub ctes: HashMap<String, LogicalPlan>,
79+
/// Data types for numbered parameters ($1, $2, etc), if supplied
80+
/// in `PREPARE` statement
81+
prepare_param_data_types: Vec<DataType>,
82+
/// Map of CTE name to logical plan of the WITH clause.
83+
/// Use Arc<LogicalPlan> to allow cheap cloning
84+
ctes: HashMap<String, Arc<LogicalPlan>>,
8385
/// The query schema of the outer query plan, used to resolve the columns in subquery
84-
pub outer_query_schema: Option<DFSchema>,
86+
outer_query_schema: Option<DFSchema>,
8587
}
8688

8789
impl Default for PlannerContext {
@@ -100,15 +102,52 @@ impl PlannerContext {
100102
}
101103
}
102104

103-
/// Create a new PlannerContext with provided prepare_param_data_types
104-
pub fn new_with_prepare_param_data_types(
105+
/// Update the PlannerContext with provided prepare_param_data_types
106+
pub fn with_prepare_param_data_types(
107+
mut self,
105108
prepare_param_data_types: Vec<DataType>,
106109
) -> Self {
107-
Self {
108-
prepare_param_data_types,
109-
ctes: HashMap::new(),
110-
outer_query_schema: None,
111-
}
110+
self.prepare_param_data_types = prepare_param_data_types;
111+
self
112+
}
113+
114+
// return a reference to the outer queries schema
115+
pub fn outer_query_schema(&self) -> Option<&DFSchema> {
116+
self.outer_query_schema.as_ref()
117+
}
118+
119+
/// sets the outer query schema, returning the existing one, if
120+
/// any
121+
pub fn set_outer_query_schema(
122+
&mut self,
123+
mut schema: Option<DFSchema>,
124+
) -> Option<DFSchema> {
125+
std::mem::swap(&mut self.outer_query_schema, &mut schema);
126+
schema
127+
}
128+
129+
/// Return the types of parameters (`$1`, `$2`, etc) if known
130+
pub fn prepare_param_data_types(&self) -> &[DataType] {
131+
&self.prepare_param_data_types
132+
}
133+
134+
/// returns true if there is a Common Table Expression (CTE) /
135+
/// Subquery for the specified name
136+
pub fn contains_cte(&self, cte_name: &str) -> bool {
137+
self.ctes.contains_key(cte_name)
138+
}
139+
140+
/// Inserts a LogicalPlan for the Common Table Expression (CTE) /
141+
/// Subquery for the specified name
142+
pub fn insert_cte(&mut self, cte_name: impl Into<String>, plan: LogicalPlan) {
143+
let cte_name = cte_name.into();
144+
self.ctes.insert(cte_name, Arc::new(plan));
145+
}
146+
147+
/// Return a plan for the Common Table Expression (CTE) / Subquery for the
148+
/// specified name
149+
pub fn get_cte(&self, cte_name: &str) -> Option<&LogicalPlan> {
150+
self.ctes.get(cte_name).map(|cte| cte.as_ref())
112151
}
113152
}
114153

‎datafusion/sql/src/query.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
5454
for cte in with.cte_tables {
5555
// A `WITH` block can't use the same name more than once
5656
let cte_name = normalize_ident(cte.alias.name.clone());
57-
if planner_context.ctes.contains_key(&cte_name) {
57+
if planner_context.contains_cte(&cte_name) {
5858
return Err(DataFusionError::SQL(ParserError(format!(
5959
"WITH query name {cte_name:?} specified more than once"
6060
))));
@@ -68,7 +68,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
6868
// projection (e.g. "WITH table(t1, t2) AS SELECT 1, 2").
6969
let logical_plan = self.apply_table_alias(logical_plan, cte.alias)?;
7070

71-
planner_context.ctes.insert(cte_name, logical_plan);
71+
planner_context.insert_cte(cte_name, logical_plan);
7272
}
7373
}
7474
let plan = self.set_expr_to_plan(*set_expr, planner_context)?;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
3333
// normalize name and alias
3434
let table_ref = self.object_name_to_table_reference(name)?;
3535
let table_name = table_ref.to_string();
36-
let cte = planner_context.ctes.get(&table_name);
36+
let cte = planner_context.get_cte(&table_name);
3737
(
3838
match (
3939
cte,

‎datafusion/sql/src/select.rs

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -233,13 +233,11 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
233233
match selection {
234234
Some(predicate_expr) => {
235235
let fallback_schemas = plan.fallback_normalize_schemas();
236-
let outer_query_schema = planner_context.outer_query_schema.clone();
237-
let outer_query_schema_vec =
238-
if let Some(outer) = outer_query_schema.as_ref() {
239-
vec![outer]
240-
} else {
241-
vec![]
242-
};
236+
let outer_query_schema = planner_context.outer_query_schema().cloned();
237+
let outer_query_schema_vec = outer_query_schema
238+
.as_ref()
239+
.map(|schema| vec![schema])
240+
.unwrap_or_else(Vec::new);
243241

244242
let filter_expr =
245243
self.sql_to_expr(predicate_expr, plan.schema(), planner_context)?;

‎datafusion/sql/src/set_expr.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
2929
match set_expr {
3030
SetExpr::Select(s) => self.select_to_plan(*s, planner_context),
3131
SetExpr::Values(v) => {
32-
self.sql_values_to_plan(v, &planner_context.prepare_param_data_types)
32+
self.sql_values_to_plan(v, planner_context.prepare_param_data_types())
3333
}
3434
SetExpr::SetOperation {
3535
op,

‎datafusion/sql/src/statement.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -279,8 +279,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
279279
.collect::<Result<_>>()?;
280280

281281
// Create planner context with parameters
282-
let mut planner_context =
283-
PlannerContext::new_with_prepare_param_data_types(data_types.clone());
282+
let mut planner_context = PlannerContext::new()
283+
.with_prepare_param_data_types(data_types.clone());
284284

285285
// Build logical plan for inner statement of the prepare statement
286286
let plan = self.sql_statement_to_plan_with_context(
@@ -961,7 +961,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
961961

962962
// Projection
963963
let mut planner_context =
964-
PlannerContext::new_with_prepare_param_data_types(prepare_param_data_types);
964+
PlannerContext::new().with_prepare_param_data_types(prepare_param_data_types);
965965
let source = self.query_to_plan(*source, &mut planner_context)?;
966966
if fields.len() != source.schema().fields().len() {
967967
Err(DataFusionError::Plan(

0 commit comments

Comments
 (0)