Skip to content

Commit f807947

Browse files
authored
Migrate testing optimizer rules to use rewrite API (#10576)
* Migrate testing optimizer rules * Remove error * fmt
1 parent 21bce33 commit f807947

File tree

2 files changed

+42
-18
lines changed

2 files changed

+42
-18
lines changed

datafusion/optimizer/src/optimizer.rs

Lines changed: 42 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -475,6 +475,7 @@ pub(crate) fn assert_schema_is_the_same(
475475
mod tests {
476476
use std::sync::{Arc, Mutex};
477477

478+
use datafusion_common::tree_node::Transformed;
478479
use datafusion_common::{plan_err, DFSchema, DFSchemaRef, Result};
479480
use datafusion_expr::logical_plan::EmptyRelation;
480481
use datafusion_expr::{col, lit, LogicalPlan, LogicalPlanBuilder, Projection};
@@ -676,13 +677,27 @@ mod tests {
676677
_: &LogicalPlan,
677678
_: &dyn OptimizerConfig,
678679
) -> Result<Option<LogicalPlan>> {
679-
let table_scan = test_table_scan()?;
680-
Ok(Some(LogicalPlanBuilder::from(table_scan).build()?))
680+
unreachable!()
681681
}
682682

683683
fn name(&self) -> &str {
684684
"get table_scan rule"
685685
}
686+
687+
fn supports_rewrite(&self) -> bool {
688+
true
689+
}
690+
691+
fn rewrite(
692+
&self,
693+
_plan: LogicalPlan,
694+
_config: &dyn OptimizerConfig,
695+
) -> Result<Transformed<LogicalPlan>> {
696+
let table_scan = test_table_scan()?;
697+
Ok(Transformed::yes(
698+
LogicalPlanBuilder::from(table_scan).build()?,
699+
))
700+
}
686701
}
687702

688703
/// A goofy rule doing rotation of columns in all projections.
@@ -704,12 +719,32 @@ mod tests {
704719
impl OptimizerRule for RotateProjectionRule {
705720
fn try_optimize(
706721
&self,
707-
plan: &LogicalPlan,
722+
_plan: &LogicalPlan,
708723
_: &dyn OptimizerConfig,
709724
) -> Result<Option<LogicalPlan>> {
725+
unreachable!()
726+
}
727+
728+
fn name(&self) -> &str {
729+
"rotate_projection"
730+
}
731+
732+
fn apply_order(&self) -> Option<ApplyOrder> {
733+
Some(ApplyOrder::TopDown)
734+
}
735+
736+
fn supports_rewrite(&self) -> bool {
737+
true
738+
}
739+
740+
fn rewrite(
741+
&self,
742+
plan: LogicalPlan,
743+
_config: &dyn OptimizerConfig,
744+
) -> Result<Transformed<LogicalPlan>> {
710745
let projection = match plan {
711746
LogicalPlan::Projection(p) if p.expr.len() >= 2 => p,
712-
_ => return Ok(None),
747+
_ => return Ok(Transformed::no(plan)),
713748
};
714749

715750
let mut exprs = projection.expr.clone();
@@ -722,18 +757,9 @@ mod tests {
722757
exprs.rotate_left(1);
723758
}
724759

725-
Ok(Some(LogicalPlan::Projection(Projection::try_new(
726-
exprs,
727-
projection.input.clone(),
728-
)?)))
729-
}
730-
731-
fn apply_order(&self) -> Option<ApplyOrder> {
732-
Some(ApplyOrder::TopDown)
733-
}
734-
735-
fn name(&self) -> &str {
736-
"rotate_projection"
760+
Ok(Transformed::yes(LogicalPlan::Projection(
761+
Projection::try_new(exprs, projection.input.clone())?,
762+
)))
737763
}
738764
}
739765
}

datafusion/proto/proto/datafusion.proto

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
11
/*
2-
3-
42
* Licensed to the Apache Software Foundation (ASF) under one
53
* or more contributor license agreements. See the NOTICE file
64
* distributed with this work for additional information

0 commit comments

Comments
 (0)