Skip to content

Commit

Permalink
Migrate testing optimizer rules to use rewrite API (#10576)
Browse files Browse the repository at this point in the history
* Migrate testing optimizer rules

* Remove error

* fmt
  • Loading branch information
lewiszlw committed May 21, 2024
1 parent 21bce33 commit f807947
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 18 deletions.
58 changes: 42 additions & 16 deletions datafusion/optimizer/src/optimizer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -475,6 +475,7 @@ pub(crate) fn assert_schema_is_the_same(
mod tests {
use std::sync::{Arc, Mutex};

use datafusion_common::tree_node::Transformed;
use datafusion_common::{plan_err, DFSchema, DFSchemaRef, Result};
use datafusion_expr::logical_plan::EmptyRelation;
use datafusion_expr::{col, lit, LogicalPlan, LogicalPlanBuilder, Projection};
Expand Down Expand Up @@ -676,13 +677,27 @@ mod tests {
_: &LogicalPlan,
_: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
let table_scan = test_table_scan()?;
Ok(Some(LogicalPlanBuilder::from(table_scan).build()?))
unreachable!()
}

fn name(&self) -> &str {
"get table_scan rule"
}

fn supports_rewrite(&self) -> bool {
true
}

fn rewrite(
&self,
_plan: LogicalPlan,
_config: &dyn OptimizerConfig,
) -> Result<Transformed<LogicalPlan>> {
let table_scan = test_table_scan()?;
Ok(Transformed::yes(
LogicalPlanBuilder::from(table_scan).build()?,
))
}
}

/// A goofy rule doing rotation of columns in all projections.
Expand All @@ -704,12 +719,32 @@ mod tests {
impl OptimizerRule for RotateProjectionRule {
fn try_optimize(
&self,
plan: &LogicalPlan,
_plan: &LogicalPlan,
_: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
unreachable!()
}

fn name(&self) -> &str {
"rotate_projection"
}

fn apply_order(&self) -> Option<ApplyOrder> {
Some(ApplyOrder::TopDown)
}

fn supports_rewrite(&self) -> bool {
true
}

fn rewrite(
&self,
plan: LogicalPlan,
_config: &dyn OptimizerConfig,
) -> Result<Transformed<LogicalPlan>> {
let projection = match plan {
LogicalPlan::Projection(p) if p.expr.len() >= 2 => p,
_ => return Ok(None),
_ => return Ok(Transformed::no(plan)),
};

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

Ok(Some(LogicalPlan::Projection(Projection::try_new(
exprs,
projection.input.clone(),
)?)))
}

fn apply_order(&self) -> Option<ApplyOrder> {
Some(ApplyOrder::TopDown)
}

fn name(&self) -> &str {
"rotate_projection"
Ok(Transformed::yes(LogicalPlan::Projection(
Projection::try_new(exprs, projection.input.clone())?,
)))
}
}
}
2 changes: 0 additions & 2 deletions datafusion/proto/proto/datafusion.proto
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
Expand Down

0 comments on commit f807947

Please sign in to comment.