Skip to content

Commit 83ea3b7

Browse files
ovrovr
committed
chore: optimize
1 parent 341185e commit 83ea3b7

File tree

1 file changed

+40
-48
lines changed

1 file changed

+40
-48
lines changed

‎rust/cubesql/cubesql/src/compile/engine/df/optimizers/plan_normalize.rs‎

Lines changed: 40 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ use datafusion::{
2525
use crate::compile::{engine::CubeContext, rewrite::rules::utils::DatePartToken};
2626

2727
struct PlanNormalizeContext<'a> {
28-
optimizer: &'a PlanNormalize<'a>,
28+
cube_ctx: &'a CubeContext,
2929
schema: &'a DFSchema,
3030
remapped_columns: &'a HashMap<Column, Column>,
3131
#[allow(unused)]
@@ -60,7 +60,7 @@ impl OptimizerRule for PlanNormalize<'_> {
6060
plan: &LogicalPlan,
6161
optimizer_config: &OptimizerConfig,
6262
) -> Result<LogicalPlan> {
63-
plan_normalize(self, plan, &mut HashMap::new(), optimizer_config)
63+
plan_normalize(self.cube_ctx, plan, &mut HashMap::new(), optimizer_config)
6464
}
6565

6666
fn name(&self) -> &str {
@@ -73,7 +73,7 @@ impl OptimizerRule for PlanNormalize<'_> {
7373
///
7474
/// `remapped_columns` passed to the function is assumed to be empty unless stated otherwise.
7575
fn plan_normalize(
76-
optimizer: &PlanNormalize,
76+
cube_ctx: &CubeContext,
7777
plan: &LogicalPlan,
7878
remapped_columns: &mut HashMap<Column, Column>,
7979
optimizer_config: &OptimizerConfig,
@@ -85,10 +85,10 @@ fn plan_normalize(
8585
schema: _,
8686
alias,
8787
}) => {
88-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
88+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
8989
let schema = input.schema();
9090
let ctx = Box::new(PlanNormalizeContext {
91-
optimizer,
91+
cube_ctx,
9292
schema,
9393
remapped_columns,
9494
optimizer_config,
@@ -122,10 +122,10 @@ fn plan_normalize(
122122
}
123123

124124
LogicalPlan::Filter(Filter { predicate, input }) => {
125-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
125+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
126126
let schema = input.schema();
127127
let ctx = Box::new(PlanNormalizeContext {
128-
optimizer,
128+
cube_ctx,
129129
schema,
130130
remapped_columns,
131131
optimizer_config,
@@ -140,10 +140,10 @@ fn plan_normalize(
140140
window_expr,
141141
schema: _,
142142
}) => {
143-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
143+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
144144
let schema = input.schema();
145145
let ctx = Box::new(PlanNormalizeContext {
146-
optimizer,
146+
cube_ctx,
147147
schema,
148148
remapped_columns,
149149
optimizer_config,
@@ -174,10 +174,10 @@ fn plan_normalize(
174174
aggr_expr,
175175
schema: _,
176176
}) => {
177-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
177+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
178178
let schema = input.schema();
179179
let ctx = Box::new(PlanNormalizeContext {
180-
optimizer,
180+
cube_ctx,
181181
schema,
182182
remapped_columns,
183183
optimizer_config,
@@ -217,10 +217,10 @@ fn plan_normalize(
217217
}
218218

219219
LogicalPlan::Sort(Sort { expr, input }) => {
220-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
220+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
221221
let schema = input.schema();
222222
let ctx = Box::new(PlanNormalizeContext {
223-
optimizer,
223+
cube_ctx,
224224
schema,
225225
remapped_columns,
226226
optimizer_config,
@@ -244,26 +244,26 @@ fn plan_normalize(
244244
}) => {
245245
let mut right_remapped_columns = HashMap::new();
246246
let left = Arc::new(plan_normalize(
247-
optimizer,
247+
cube_ctx,
248248
left,
249249
remapped_columns,
250250
optimizer_config,
251251
)?);
252252
let right = Arc::new(plan_normalize(
253-
optimizer,
253+
cube_ctx,
254254
right,
255255
&mut right_remapped_columns,
256256
optimizer_config,
257257
)?);
258258
let empty_schema = DFSchema::empty();
259259
let left_ctx = Box::new(PlanNormalizeContext {
260-
optimizer,
260+
cube_ctx,
261261
schema: &empty_schema,
262262
remapped_columns,
263263
optimizer_config,
264264
});
265265
let right_ctx = Box::new(PlanNormalizeContext {
266-
optimizer,
266+
cube_ctx,
267267
schema: &empty_schema,
268268
remapped_columns: &right_remapped_columns,
269269
optimizer_config,
@@ -304,9 +304,9 @@ fn plan_normalize(
304304
schema: _,
305305
}) => {
306306
let mut right_remapped_columns = HashMap::new();
307-
let left = plan_normalize(optimizer, left, remapped_columns, optimizer_config)?;
307+
let left = plan_normalize(cube_ctx, left, remapped_columns, optimizer_config)?;
308308
let right = plan_normalize(
309-
optimizer,
309+
cube_ctx,
310310
right,
311311
&mut right_remapped_columns,
312312
optimizer_config,
@@ -321,10 +321,10 @@ fn plan_normalize(
321321
input,
322322
partitioning_scheme,
323323
}) => {
324-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
324+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
325325
let schema = input.schema();
326326
let ctx = Box::new(PlanNormalizeContext {
327-
optimizer,
327+
cube_ctx,
328328
schema,
329329
remapped_columns,
330330
optimizer_config,
@@ -353,12 +353,8 @@ fn plan_normalize(
353353
let mut plan = None;
354354
for input in inputs {
355355
let mut new_remapped_columns = HashMap::new();
356-
let input = plan_normalize(
357-
optimizer,
358-
input,
359-
&mut new_remapped_columns,
360-
optimizer_config,
361-
)?;
356+
let input =
357+
plan_normalize(cube_ctx, input, &mut new_remapped_columns, optimizer_config)?;
362358
if let Some(last_plan) = plan.take() {
363359
plan = Some(union_with_alias(last_plan, input, alias.clone())?);
364360
} else {
@@ -385,7 +381,7 @@ fn plan_normalize(
385381
let projection = projection.clone();
386382
let projected_schema = Arc::clone(projected_schema);
387383
let ctx = Box::new(PlanNormalizeContext {
388-
optimizer,
384+
cube_ctx,
389385
schema: &projected_schema,
390386
remapped_columns,
391387
optimizer_config,
@@ -411,7 +407,7 @@ fn plan_normalize(
411407
LogicalPlan::Limit(Limit { skip, fetch, input }) => {
412408
let skip = *skip;
413409
let fetch = *fetch;
414-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
410+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
415411

416412
LogicalPlanBuilder::from(input).limit(skip, fetch)?.build()
417413
}
@@ -422,12 +418,12 @@ fn plan_normalize(
422418
types,
423419
schema: _,
424420
}) => {
425-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
421+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
426422
let mut new_subqueries = Vec::with_capacity(subqueries.len());
427423
for subquery in subqueries {
428424
let mut subquery_remapped_columns = HashMap::new();
429425
let new_subquery = plan_normalize(
430-
optimizer,
426+
cube_ctx,
431427
subquery,
432428
&mut subquery_remapped_columns,
433429
optimizer_config,
@@ -447,7 +443,7 @@ fn plan_normalize(
447443
LogicalPlan::CreateMemoryTable(CreateMemoryTable { name, input }) => {
448444
let name = name.clone();
449445
let input = Arc::new(plan_normalize(
450-
optimizer,
446+
cube_ctx,
451447
input,
452448
remapped_columns,
453449
optimizer_config,
@@ -464,7 +460,7 @@ fn plan_normalize(
464460

465461
LogicalPlan::Values(Values { schema, values }) => {
466462
let ctx = Box::new(PlanNormalizeContext {
467-
optimizer,
463+
cube_ctx,
468464
schema,
469465
remapped_columns,
470466
optimizer_config,
@@ -488,7 +484,7 @@ fn plan_normalize(
488484
schema: _,
489485
}) => {
490486
let verbose = *verbose;
491-
let plan = plan_normalize(optimizer, plan, remapped_columns, optimizer_config)?;
487+
let plan = plan_normalize(cube_ctx, plan, remapped_columns, optimizer_config)?;
492488

493489
*remapped_columns = HashMap::new();
494490

@@ -503,7 +499,7 @@ fn plan_normalize(
503499
schema: _,
504500
}) => {
505501
let verbose = *verbose;
506-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
502+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
507503

508504
*remapped_columns = HashMap::new();
509505

@@ -518,14 +514,14 @@ fn plan_normalize(
518514
schema: _,
519515
}) => {
520516
let input = Arc::new(plan_normalize(
521-
optimizer,
517+
cube_ctx,
522518
input,
523519
remapped_columns,
524520
optimizer_config,
525521
)?);
526522
let schema = input.schema();
527523
let ctx = Box::new(PlanNormalizeContext {
528-
optimizer,
524+
cube_ctx,
529525
schema,
530526
remapped_columns,
531527
optimizer_config,
@@ -563,7 +559,7 @@ fn plan_normalize(
563559
}
564560

565561
LogicalPlan::Distinct(Distinct { input }) => {
566-
let input = plan_normalize(optimizer, input, remapped_columns, optimizer_config)?;
562+
let input = plan_normalize(cube_ctx, input, remapped_columns, optimizer_config)?;
567563

568564
LogicalPlanBuilder::from(input).distinct()?.build()
569565
}
@@ -992,15 +988,11 @@ fn binary_expr_normalize(
992988
let right_type = right.get_type(ctx.schema)?;
993989

994990
if left_type == DataType::Date32 && op == Operator::Minus && right_type == DataType::Date32 {
995-
let fun = ctx
996-
.optimizer
997-
.cube_ctx
998-
.get_function_meta("datediff")
999-
.ok_or_else(|| {
1000-
DataFusionError::Internal(
1001-
"Unable to find 'datediff' function in cube context".to_string(),
1002-
)
1003-
})?;
991+
let fun = ctx.cube_ctx.get_function_meta("datediff").ok_or_else(|| {
992+
DataFusionError::Internal(
993+
"Unable to find 'datediff' function in cube context".to_string(),
994+
)
995+
})?;
1004996
let args = vec![
1005997
Expr::Literal(ScalarValue::Utf8(Some("day".to_string()))),
1006998
*right,
@@ -1146,7 +1138,7 @@ fn in_list_expr_normalize(
11461138
}
11471139

11481140
fn evaluate_expr_stacked(ctx: &Box<PlanNormalizeContext>, expr: Expr) -> Result<Expr> {
1149-
let execution_props = &ctx.optimizer.cube_ctx.state.execution_props;
1141+
let execution_props = &ctx.cube_ctx.state.execution_props;
11501142
let mut const_evaluator = ConstEvaluator::new(execution_props);
11511143
expr.rewrite(&mut const_evaluator)
11521144
}

0 commit comments

Comments
 (0)