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 d4da80bfe2 Stop copying LogicalPlan and Exprs in EliminateNestedUnion
(#10319)
d4da80bfe2 is described below
commit d4da80bfe2e53a98e49e58f0337f2e8cae010d4b
Author: Matt Green <[email protected]>
AuthorDate: Thu May 2 06:53:37 2024 -0700
Stop copying LogicalPlan and Exprs in EliminateNestedUnion (#10319)
* First pass of https://github.com/apache/datafusion/issues/10296
* Update datafusion/optimizer/src/eliminate_nested_union.rs
Co-authored-by: Andrew Lamb <[email protected]>
* fix formatting
* return original expr instead of clone()
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/expr/src/expr_rewriter/mod.rs | 2 +-
datafusion/optimizer/src/eliminate_nested_union.rs | 73 +++++++++++++---------
2 files changed, 45 insertions(+), 30 deletions(-)
diff --git a/datafusion/expr/src/expr_rewriter/mod.rs
b/datafusion/expr/src/expr_rewriter/mod.rs
index 14154189a1..700dd560ec 100644
--- a/datafusion/expr/src/expr_rewriter/mod.rs
+++ b/datafusion/expr/src/expr_rewriter/mod.rs
@@ -250,7 +250,7 @@ fn coerce_exprs_for_schema(
_ => expr.cast_to(new_type, src_schema),
}
} else {
- Ok(expr.clone())
+ Ok(expr)
}
})
.collect::<Result<_>>()
diff --git a/datafusion/optimizer/src/eliminate_nested_union.rs
b/datafusion/optimizer/src/eliminate_nested_union.rs
index da2a6a1721..aa6f2b4975 100644
--- a/datafusion/optimizer/src/eliminate_nested_union.rs
+++ b/datafusion/optimizer/src/eliminate_nested_union.rs
@@ -18,7 +18,8 @@
//! [`EliminateNestedUnion`]: flattens nested `Union` to a single `Union`
use crate::optimizer::ApplyOrder;
use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::Result;
+use datafusion_common::tree_node::Transformed;
+use datafusion_common::{internal_err, Result};
use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema;
use datafusion_expr::{Distinct, LogicalPlan, Union};
use std::sync::Arc;
@@ -37,9 +38,29 @@ impl EliminateNestedUnion {
impl OptimizerRule for EliminateNestedUnion {
fn try_optimize(
&self,
- plan: &LogicalPlan,
+ _plan: &LogicalPlan,
_config: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
+ internal_err!("Should have called EliminateNestedUnion::rewrite")
+ }
+
+ fn name(&self) -> &str {
+ "eliminate_nested_union"
+ }
+
+ fn apply_order(&self) -> Option<ApplyOrder> {
+ Some(ApplyOrder::BottomUp)
+ }
+
+ fn supports_rewrite(&self) -> bool {
+ true
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>> {
match plan {
LogicalPlan::Union(Union { inputs, schema }) => {
let inputs = inputs
@@ -47,39 +68,33 @@ impl OptimizerRule for EliminateNestedUnion {
.flat_map(extract_plans_from_union)
.collect::<Vec<_>>();
- Ok(Some(LogicalPlan::Union(Union {
+ Ok(Transformed::yes(LogicalPlan::Union(Union {
inputs,
- schema: schema.clone(),
+ schema,
})))
}
- LogicalPlan::Distinct(Distinct::All(plan)) => match plan.as_ref() {
- LogicalPlan::Union(Union { inputs, schema }) => {
- let inputs = inputs
- .iter()
- .map(extract_plan_from_distinct)
- .flat_map(extract_plans_from_union)
- .collect::<Vec<_>>();
-
- Ok(Some(LogicalPlan::Distinct(Distinct::All(Arc::new(
- LogicalPlan::Union(Union {
- inputs,
- schema: schema.clone(),
- }),
- )))))
+ LogicalPlan::Distinct(Distinct::All(ref nested_plan)) => {
+ match nested_plan.as_ref() {
+ LogicalPlan::Union(Union { inputs, schema }) => {
+ let inputs = inputs
+ .iter()
+ .map(extract_plan_from_distinct)
+ .flat_map(extract_plans_from_union)
+ .collect::<Vec<_>>();
+
+
Ok(Transformed::yes(LogicalPlan::Distinct(Distinct::All(
+ Arc::new(LogicalPlan::Union(Union {
+ inputs,
+ schema: schema.clone(),
+ })),
+ ))))
+ }
+ _ => Ok(Transformed::no(plan)),
}
- _ => Ok(None),
- },
- _ => Ok(None),
+ }
+ _ => Ok(Transformed::no(plan)),
}
}
-
- fn name(&self) -> &str {
- "eliminate_nested_union"
- }
-
- fn apply_order(&self) -> Option<ApplyOrder> {
- Some(ApplyOrder::BottomUp)
- }
}
fn extract_plans_from_union(plan: &Arc<LogicalPlan>) -> Vec<Arc<LogicalPlan>> {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
