Skip to content

Commit f62b46b

Browse files
waralexromwaralexrom
committed
refactor
1 parent b566c6e commit f62b46b

17 files changed

+417
-153
lines changed

‎rust/cubesqlplanner/cubesqlplanner/src/plan/select.rs

-2
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
use itertools::Itertools;
2-
31
use super::{Cte, Expr, Filter, From, OrderBy, Schema, SchemaColumn};
42
use crate::planner::sql_templates::PlanSqlTemplates;
53
use crate::planner::sql_templates::{

‎rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/mod.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ pub mod collectors;
22
pub mod compiler;
33
mod dependecy;
44
pub mod evaluation_node;
5+
pub mod sql_node_transformers;
56
pub mod sql_nodes;
67
pub mod sql_visitor;
78
pub mod symbols;
@@ -16,4 +17,4 @@ pub use symbols::{
1617
DimensionSymbol, DimensionSymbolFactory, MeasureSymbol, MeasureSymbolFactory, MemberSymbol,
1718
MemberSymbolType, SimpleSqlSymbol, SimpleSqlSymbolFactory, SymbolFactory,
1819
};
19-
pub use visitor::{EvaluatorVisitor, TraversalVisitor};
20+
pub use visitor::TraversalVisitor;

‎rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/sql_node_transformers/mod.rs

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
pub mod set_schema;
2+
3+
pub use set_schema::set_schema;

‎rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/sql_node_transformers/set_schema.rs

+86
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
use crate::plan::schema::Schema;
2+
use crate::planner::sql_evaluator::sql_nodes::final_measure::FinalMeasureSqlNode;
3+
use crate::planner::sql_evaluator::sql_nodes::{
4+
AutoPrefixSqlNode, EvaluateSqlNode, MeasureFilterSqlNode, MultiStageRankNode,
5+
MultiStageWindowNode, RenderReferencesSqlNode, RootSqlNode, SqlNode, TimeShiftSqlNode,
6+
};
7+
use std::rc::Rc;
8+
9+
pub fn set_schema(node_processors: Rc<dyn SqlNode>, schema: Rc<Schema>) -> Rc<dyn SqlNode> {
10+
set_schema_impl(node_processors, schema)
11+
}
12+
13+
pub fn set_schema_impl(sql_node: Rc<dyn SqlNode>, schema: Rc<Schema>) -> Rc<dyn SqlNode> {
14+
if let Some(auto_prefix) = sql_node
15+
.clone()
16+
.as_any()
17+
.downcast_ref::<AutoPrefixSqlNode>()
18+
{
19+
let input = set_schema_impl(auto_prefix.input().clone(), schema.clone());
20+
AutoPrefixSqlNode::new_with_schema(input, schema)
21+
} else if let Some(_) = sql_node.clone().as_any().downcast_ref::<EvaluateSqlNode>() {
22+
sql_node
23+
} else if let Some(final_measure) = sql_node
24+
.clone()
25+
.as_any()
26+
.downcast_ref::<FinalMeasureSqlNode>()
27+
{
28+
let input = set_schema_impl(final_measure.input().clone(), schema.clone());
29+
FinalMeasureSqlNode::new(input)
30+
} else if let Some(measure_filter) = sql_node
31+
.clone()
32+
.as_any()
33+
.downcast_ref::<MeasureFilterSqlNode>()
34+
{
35+
let input = set_schema_impl(measure_filter.input().clone(), schema.clone());
36+
MeasureFilterSqlNode::new(input)
37+
} else if let Some(multi_stage_rank) = sql_node
38+
.clone()
39+
.as_any()
40+
.downcast_ref::<MultiStageRankNode>()
41+
{
42+
let else_processor =
43+
set_schema_impl(multi_stage_rank.else_processor().clone(), schema.clone());
44+
MultiStageRankNode::new(else_processor, multi_stage_rank.partition().clone())
45+
} else if let Some(multi_stage_window) = sql_node
46+
.clone()
47+
.as_any()
48+
.downcast_ref::<MultiStageWindowNode>()
49+
{
50+
let input = set_schema_impl(multi_stage_window.input().clone(), schema.clone());
51+
let else_processor =
52+
set_schema_impl(multi_stage_window.else_processor().clone(), schema.clone());
53+
MultiStageWindowNode::new(
54+
input,
55+
else_processor,
56+
multi_stage_window.partition().clone(),
57+
)
58+
} else if let Some(render_references) = sql_node
59+
.clone()
60+
.as_any()
61+
.downcast_ref::<RenderReferencesSqlNode>()
62+
{
63+
let input = set_schema_impl(render_references.input().clone(), schema.clone());
64+
RenderReferencesSqlNode::new_with_schema(input, schema)
65+
} else if let Some(root_node) = sql_node.clone().as_any().downcast_ref::<RootSqlNode>() {
66+
let dimension_processor =
67+
set_schema_impl(root_node.dimension_processor().clone(), schema.clone());
68+
let measure_processor =
69+
set_schema_impl(root_node.measure_processor().clone(), schema.clone());
70+
let cube_name_processor =
71+
set_schema_impl(root_node.cube_name_processor().clone(), schema.clone());
72+
let default_processor =
73+
set_schema_impl(root_node.default_processor().clone(), schema.clone());
74+
RootSqlNode::new(
75+
dimension_processor,
76+
measure_processor,
77+
cube_name_processor,
78+
default_processor,
79+
)
80+
} else if let Some(time_shift) = sql_node.clone().as_any().downcast_ref::<TimeShiftSqlNode>() {
81+
let input = set_schema_impl(time_shift.input().clone(), schema.clone());
82+
TimeShiftSqlNode::new(time_shift.shifts().clone(), input)
83+
} else {
84+
unreachable!("Not all nodes are implemented");
85+
}
86+
}

‎rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/sql_nodes/auto_prefix.rs

+33-6
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,35 @@
11
use super::SqlNode;
2+
use crate::plan::Schema;
23
use crate::planner::query_tools::QueryTools;
34
use crate::planner::sql_evaluator::SqlEvaluatorVisitor;
45
use crate::planner::sql_evaluator::{EvaluationNode, MemberSymbol, MemberSymbolType};
56
use cubenativeutils::CubeError;
7+
use std::any::Any;
68
use std::rc::Rc;
79

810
pub struct AutoPrefixSqlNode {
911
input: Rc<dyn SqlNode>,
12+
schema: Rc<Schema>,
1013
}
1114

1215
impl AutoPrefixSqlNode {
1316
pub fn new(input: Rc<dyn SqlNode>) -> Rc<Self> {
14-
Rc::new(Self { input })
17+
Rc::new(Self {
18+
input,
19+
schema: Rc::new(Schema::empty()),
20+
})
21+
}
22+
23+
pub fn new_with_schema(input: Rc<dyn SqlNode>, schema: Rc<Schema>) -> Rc<Self> {
24+
Rc::new(Self { input, schema })
25+
}
26+
27+
pub fn input(&self) -> &Rc<dyn SqlNode> {
28+
&self.input
29+
}
30+
31+
pub fn schema(&self) -> &Rc<Schema> {
32+
&self.schema
1533
}
1634
}
1735

@@ -21,24 +39,33 @@ impl SqlNode for AutoPrefixSqlNode {
2139
visitor: &mut SqlEvaluatorVisitor,
2240
node: &Rc<EvaluationNode>,
2341
query_tools: Rc<QueryTools>,
42+
node_processor: Rc<dyn SqlNode>,
2443
) -> Result<String, CubeError> {
25-
let source_schema = visitor.source_schema().clone();
26-
let input = self.input.to_sql(visitor, node, query_tools.clone())?;
44+
let input =
45+
self.input
46+
.to_sql(visitor, node, query_tools.clone(), node_processor.clone())?;
2747
let res = match node.symbol() {
2848
MemberSymbolType::Dimension(ev) => {
29-
let cube_alias = source_schema.resolve_cube_alias(&ev.cube_name());
49+
let cube_alias = self.schema.resolve_cube_alias(&ev.cube_name());
3050
query_tools.auto_prefix_with_cube_name(&cube_alias, &input)
3151
}
3252
MemberSymbolType::Measure(ev) => {
33-
let cube_alias = source_schema.resolve_cube_alias(&ev.cube_name());
53+
let cube_alias = self.schema.resolve_cube_alias(&ev.cube_name());
3454
query_tools.auto_prefix_with_cube_name(&cube_alias, &input)
3555
}
3656
MemberSymbolType::CubeName(_) => {
37-
let cube_alias = source_schema.resolve_cube_alias(&input);
57+
let cube_alias = self.schema.resolve_cube_alias(&input);
3858
query_tools.escape_column_name(&cube_alias)
3959
}
4060
_ => input,
4161
};
4262
Ok(res)
4363
}
64+
fn as_any(self: Rc<Self>) -> Rc<dyn Any> {
65+
self.clone()
66+
}
67+
68+
fn childs(&self) -> Vec<Rc<dyn SqlNode>> {
69+
vec![self.input.clone()]
70+
}
4471
}

‎rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/sql_nodes/evaluate_sql.rs

+11-2
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
use super::SqlNode;
22
use crate::planner::query_tools::QueryTools;
3-
use crate::planner::sql_evaluator::visitor::EvaluatorVisitor;
43
use crate::planner::sql_evaluator::SqlEvaluatorVisitor;
54
use crate::planner::sql_evaluator::{EvaluationNode, MemberSymbolType};
65
use cubenativeutils::CubeError;
6+
use std::any::Any;
77
use std::rc::Rc;
88

99
pub struct EvaluateSqlNode {}
@@ -20,8 +20,9 @@ impl SqlNode for EvaluateSqlNode {
2020
visitor: &mut SqlEvaluatorVisitor,
2121
node: &Rc<EvaluationNode>,
2222
_query_tools: Rc<QueryTools>,
23+
node_processor: Rc<dyn SqlNode>,
2324
) -> Result<String, CubeError> {
24-
let args = visitor.evaluate_deps(node)?;
25+
let args = visitor.evaluate_deps(node, node_processor.clone())?;
2526
match node.symbol() {
2627
MemberSymbolType::Dimension(ev) => ev.evaluate_sql(args),
2728
MemberSymbolType::Measure(ev) => ev.evaluate_sql(args),
@@ -30,4 +31,12 @@ impl SqlNode for EvaluateSqlNode {
3031
MemberSymbolType::SimpleSql(ev) => ev.evaluate_sql(args),
3132
}
3233
}
34+
35+
fn as_any(self: Rc<Self>) -> Rc<dyn Any> {
36+
self.clone()
37+
}
38+
39+
fn childs(&self) -> Vec<Rc<dyn SqlNode>> {
40+
vec![]
41+
}
3342
}

‎rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/sql_nodes/final_measure.rs

+20-1
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use crate::planner::query_tools::QueryTools;
33
use crate::planner::sql_evaluator::SqlEvaluatorVisitor;
44
use crate::planner::sql_evaluator::{EvaluationNode, MemberSymbolType};
55
use cubenativeutils::CubeError;
6+
use std::any::Any;
67
use std::rc::Rc;
78

89
pub struct FinalMeasureSqlNode {
@@ -13,6 +14,10 @@ impl FinalMeasureSqlNode {
1314
pub fn new(input: Rc<dyn SqlNode>) -> Rc<Self> {
1415
Rc::new(Self { input })
1516
}
17+
18+
pub fn input(&self) -> &Rc<dyn SqlNode> {
19+
&self.input
20+
}
1621
}
1722

1823
impl SqlNode for FinalMeasureSqlNode {
@@ -21,10 +26,16 @@ impl SqlNode for FinalMeasureSqlNode {
2126
visitor: &mut SqlEvaluatorVisitor,
2227
node: &Rc<EvaluationNode>,
2328
query_tools: Rc<QueryTools>,
29+
node_processor: Rc<dyn SqlNode>,
2430
) -> Result<String, CubeError> {
2531
let res = match node.symbol() {
2632
MemberSymbolType::Measure(ev) => {
27-
let input = self.input.to_sql(visitor, node, query_tools.clone())?;
33+
let input = self.input.to_sql(
34+
visitor,
35+
node,
36+
query_tools.clone(),
37+
node_processor.clone(),
38+
)?;
2839

2940
if ev.is_calculated() {
3041
input
@@ -41,4 +52,12 @@ impl SqlNode for FinalMeasureSqlNode {
4152
};
4253
Ok(res)
4354
}
55+
56+
fn as_any(self: Rc<Self>) -> Rc<dyn Any> {
57+
self.clone()
58+
}
59+
60+
fn childs(&self) -> Vec<Rc<dyn SqlNode>> {
61+
vec![self.input.clone()]
62+
}
4463
}

‎rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/sql_nodes/measure_filter.rs

+21-3
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
use super::SqlNode;
22
use crate::planner::query_tools::QueryTools;
3-
use crate::planner::sql_evaluator::visitor::EvaluatorVisitor;
43
use crate::planner::sql_evaluator::SqlEvaluatorVisitor;
54
use crate::planner::sql_evaluator::{EvaluationNode, MemberSymbolType};
65
use cubenativeutils::CubeError;
6+
use std::any::Any;
77
use std::rc::Rc;
88

99
pub struct MeasureFilterSqlNode {
@@ -14,6 +14,10 @@ impl MeasureFilterSqlNode {
1414
pub fn new(input: Rc<dyn SqlNode>) -> Rc<Self> {
1515
Rc::new(Self { input })
1616
}
17+
18+
pub fn input(&self) -> &Rc<dyn SqlNode> {
19+
&self.input
20+
}
1721
}
1822

1923
impl SqlNode for MeasureFilterSqlNode {
@@ -22,16 +26,22 @@ impl SqlNode for MeasureFilterSqlNode {
2226
visitor: &mut SqlEvaluatorVisitor,
2327
node: &Rc<EvaluationNode>,
2428
query_tools: Rc<QueryTools>,
29+
node_processor: Rc<dyn SqlNode>,
2530
) -> Result<String, CubeError> {
26-
let input = self.input.to_sql(visitor, node, query_tools.clone())?;
31+
let input =
32+
self.input
33+
.to_sql(visitor, node, query_tools.clone(), node_processor.clone())?;
2734
let res = match node.symbol() {
2835
MemberSymbolType::Measure(ev) => {
2936
let measure_filters = ev.measure_filters();
3037
if !measure_filters.is_empty() {
3138
let filters = measure_filters
3239
.iter()
3340
.map(|filter| -> Result<String, CubeError> {
34-
Ok(format!("({})", visitor.apply(filter)?))
41+
Ok(format!(
42+
"({})",
43+
visitor.apply(filter, node_processor.clone())?
44+
))
3545
})
3646
.collect::<Result<Vec<_>, _>>()?
3747
.join(" AND ");
@@ -53,4 +63,12 @@ impl SqlNode for MeasureFilterSqlNode {
5363
};
5464
Ok(res)
5565
}
66+
67+
fn as_any(self: Rc<Self>) -> Rc<dyn Any> {
68+
self.clone()
69+
}
70+
71+
fn childs(&self) -> Vec<Rc<dyn SqlNode>> {
72+
vec![self.input.clone()]
73+
}
5674
}

0 commit comments

Comments
 (0)