This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 08030f37fa Implement rewrite for EliminateOneUnion and EliminateJoin
(#10184)
08030f37fa is described below
commit 08030f37fa34995af2ea43143922476ace82e39d
Author: Lordworms <[email protected]>
AuthorDate: Tue Apr 23 12:38:37 2024 -0500
Implement rewrite for EliminateOneUnion and EliminateJoin (#10184)
* implement rewrite for EliminateJoin
* implement rewrite for eliminate_on_union
---
datafusion/expr/src/logical_plan/tree_node.rs | 2 +-
datafusion/optimizer/src/eliminate_join.rs | 51 ++++++++++++++++---------
datafusion/optimizer/src/eliminate_one_union.rs | 30 ++++++++++-----
3 files changed, 54 insertions(+), 29 deletions(-)
diff --git a/datafusion/expr/src/logical_plan/tree_node.rs
b/datafusion/expr/src/logical_plan/tree_node.rs
index f5db5a2704..37a36c36ca 100644
--- a/datafusion/expr/src/logical_plan/tree_node.rs
+++ b/datafusion/expr/src/logical_plan/tree_node.rs
@@ -363,7 +363,7 @@ impl TreeNode for LogicalPlan {
/// Converts a `Arc<LogicalPlan>` without copying, if possible. Copies the plan
/// if there is a shared reference
-fn unwrap_arc(plan: Arc<LogicalPlan>) -> LogicalPlan {
+pub fn unwrap_arc(plan: Arc<LogicalPlan>) -> LogicalPlan {
Arc::try_unwrap(plan)
// if None is returned, there is another reference to this
// LogicalPlan, so we can not own it, and must clone instead
diff --git a/datafusion/optimizer/src/eliminate_join.rs
b/datafusion/optimizer/src/eliminate_join.rs
index caf45dda98..fea87e7587 100644
--- a/datafusion/optimizer/src/eliminate_join.rs
+++ b/datafusion/optimizer/src/eliminate_join.rs
@@ -18,7 +18,8 @@
//! [`EliminateJoin`] rewrites `INNER JOIN` with `true`/`null`
use crate::optimizer::ApplyOrder;
use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::{Result, ScalarValue};
+use datafusion_common::tree_node::Transformed;
+use datafusion_common::{internal_err, Result, ScalarValue};
use datafusion_expr::JoinType::Inner;
use datafusion_expr::{
logical_plan::{EmptyRelation, LogicalPlan},
@@ -39,38 +40,50 @@ impl EliminateJoin {
impl OptimizerRule for EliminateJoin {
fn try_optimize(
&self,
- plan: &LogicalPlan,
+ _plan: &LogicalPlan,
_config: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
+ internal_err!("Should have called EliminateJoin::rewrite")
+ }
+
+ fn name(&self) -> &str {
+ "eliminate_join"
+ }
+
+ fn apply_order(&self) -> Option<ApplyOrder> {
+ Some(ApplyOrder::TopDown)
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>> {
match plan {
LogicalPlan::Join(join) if join.join_type == Inner &&
join.on.is_empty() => {
match join.filter {
Some(Expr::Literal(ScalarValue::Boolean(Some(true)))) => {
- Ok(Some(LogicalPlan::CrossJoin(CrossJoin {
- left: join.left.clone(),
- right: join.right.clone(),
- schema: join.schema.clone(),
+ Ok(Transformed::yes(LogicalPlan::CrossJoin(CrossJoin {
+ left: join.left,
+ right: join.right,
+ schema: join.schema,
})))
}
- Some(Expr::Literal(ScalarValue::Boolean(Some(false)))) => {
- Ok(Some(LogicalPlan::EmptyRelation(EmptyRelation {
+ Some(Expr::Literal(ScalarValue::Boolean(Some(false)))) =>
Ok(
+
Transformed::yes(LogicalPlan::EmptyRelation(EmptyRelation {
produce_one_row: false,
- schema: join.schema.clone(),
- })))
- }
- _ => Ok(None),
+ schema: join.schema,
+ })),
+ ),
+ _ => Ok(Transformed::no(LogicalPlan::Join(join))),
}
}
- _ => Ok(None),
+ _ => Ok(Transformed::no(plan)),
}
}
- fn name(&self) -> &str {
- "eliminate_join"
- }
-
- fn apply_order(&self) -> Option<ApplyOrder> {
- Some(ApplyOrder::TopDown)
+ fn supports_rewrite(&self) -> bool {
+ true
}
}
diff --git a/datafusion/optimizer/src/eliminate_one_union.rs
b/datafusion/optimizer/src/eliminate_one_union.rs
index 95a3370ab1..11a9009cd9 100644
--- a/datafusion/optimizer/src/eliminate_one_union.rs
+++ b/datafusion/optimizer/src/eliminate_one_union.rs
@@ -17,8 +17,8 @@
//! [`EliminateOneUnion`] eliminates single element `Union`
use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::Result;
-use datafusion_expr::logical_plan::{LogicalPlan, Union};
+use datafusion_common::{internal_err, tree_node::Transformed, Result};
+use datafusion_expr::logical_plan::{tree_node::unwrap_arc, LogicalPlan, Union};
use crate::optimizer::ApplyOrder;
@@ -36,21 +36,33 @@ impl EliminateOneUnion {
impl OptimizerRule for EliminateOneUnion {
fn try_optimize(
&self,
- plan: &LogicalPlan,
+ _plan: &LogicalPlan,
_config: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
- match plan {
- LogicalPlan::Union(Union { inputs, .. }) if inputs.len() == 1 => {
- Ok(inputs.first().map(|input| input.as_ref().clone()))
- }
- _ => Ok(None),
- }
+ internal_err!("Should have called EliminateOneUnion::rewrite")
}
fn name(&self) -> &str {
"eliminate_one_union"
}
+ fn supports_rewrite(&self) -> bool {
+ true
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>> {
+ match plan {
+ LogicalPlan::Union(Union { mut inputs, .. }) if inputs.len() == 1
=> {
+ Ok(Transformed::yes(unwrap_arc(inputs.pop().unwrap())))
+ }
+ _ => Ok(Transformed::no(plan)),
+ }
+ }
+
fn apply_order(&self) -> Option<ApplyOrder> {
Some(ApplyOrder::TopDown)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
