Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: Expr::Exists to use a struct. #6292

Merged
merged 1 commit into from
May 9, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 23 additions & 12 deletions datafusion/expr/src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -170,12 +170,7 @@ pub enum Expr {
negated: bool,
},
/// EXISTS subquery
Exists {
/// subquery that will produce a single column of data
subquery: Subquery,
/// Whether the expression is negated
negated: bool,
},
Exists(Exists),
/// IN subquery
InSubquery {
/// The expression to compare
Expand Down Expand Up @@ -500,6 +495,22 @@ impl WindowFunction {
}
}

// Exists expression.
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct Exists {
/// subquery that will produce a single column of data
pub subquery: Subquery,
/// Whether the expression is negated
pub negated: bool,
}

impl Exists {
// Create a new Exists expression.
pub fn new(subquery: Subquery, negated: bool) -> Self {
Self { subquery, negated }
}
}

#[derive(Clone, PartialEq, Eq, Hash)]
pub struct AggregateUDF {
/// The function
Expand Down Expand Up @@ -926,14 +937,14 @@ impl fmt::Debug for Expr {
Expr::IsNotTrue(expr) => write!(f, "{expr:?} IS NOT TRUE"),
Expr::IsNotFalse(expr) => write!(f, "{expr:?} IS NOT FALSE"),
Expr::IsNotUnknown(expr) => write!(f, "{expr:?} IS NOT UNKNOWN"),
Expr::Exists {
Expr::Exists(Exists {
subquery,
negated: true,
} => write!(f, "NOT EXISTS ({subquery:?})"),
Expr::Exists {
}) => write!(f, "NOT EXISTS ({subquery:?})"),
Expr::Exists(Exists {
subquery,
negated: false,
} => write!(f, "EXISTS ({subquery:?})"),
}) => write!(f, "EXISTS ({subquery:?})"),
Expr::InSubquery {
expr,
subquery,
Expand Down Expand Up @@ -1310,8 +1321,8 @@ fn create_name(e: &Expr) -> Result<String> {
let expr = create_name(expr)?;
Ok(format!("{expr} IS NOT UNKNOWN"))
}
Expr::Exists { negated: true, .. } => Ok("NOT EXISTS".to_string()),
Expr::Exists { negated: false, .. } => Ok("EXISTS".to_string()),
Expr::Exists(Exists { negated: true, .. }) => Ok("NOT EXISTS".to_string()),
Expr::Exists(Exists { negated: false, .. }) => Ok("EXISTS".to_string()),
Expr::InSubquery { negated: true, .. } => Ok("NOT IN".to_string()),
Expr::InSubquery { negated: false, .. } => Ok("IN".to_string()),
Expr::ScalarSubquery(subquery) => {
Expand Down
10 changes: 5 additions & 5 deletions datafusion/expr/src/expr_fn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
//! Functions for creating logical expressions

use crate::expr::{
AggregateFunction, BinaryExpr, Cast, GroupingSet, ScalarFunction, TryCast,
AggregateFunction, BinaryExpr, Cast, Exists, GroupingSet, ScalarFunction, TryCast,
};
use crate::{
aggregate_function, built_in_function, conditional_expressions::CaseBuilder,
Expand Down Expand Up @@ -307,25 +307,25 @@ pub fn approx_percentile_cont_with_weight(
/// Create an EXISTS subquery expression
pub fn exists(subquery: Arc<LogicalPlan>) -> Expr {
let outer_ref_columns = subquery.all_out_ref_exprs();
Expr::Exists {
Expr::Exists(Exists {
subquery: Subquery {
subquery,
outer_ref_columns,
},
negated: false,
}
})
}

/// Create a NOT EXISTS subquery expression
pub fn not_exists(subquery: Arc<LogicalPlan>) -> Expr {
let outer_ref_columns = subquery.all_out_ref_exprs();
Expr::Exists {
Expr::Exists(Exists {
subquery: Subquery {
subquery,
outer_ref_columns,
},
negated: true,
}
})
}

/// Create an IN subquery expression
Expand Down
5 changes: 3 additions & 2 deletions datafusion/expr/src/logical_plan/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.

use crate::expr::Exists;
///! Logical plan types
use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor};
use crate::logical_plan::extension::UserDefinedLogicalNode;
Expand Down Expand Up @@ -544,7 +545,7 @@ impl LogicalPlan {
// recursively look for subqueries
inspect_expr_pre(expr, |expr| {
match expr {
Expr::Exists { subquery, .. }
Expr::Exists(Exists { subquery, .. })
| Expr::InSubquery { subquery, .. }
| Expr::ScalarSubquery(subquery) => {
// use a synthetic plan so the collector sees a
Expand All @@ -570,7 +571,7 @@ impl LogicalPlan {
// recursively look for subqueries
inspect_expr_pre(expr, |expr| {
match expr {
Expr::Exists { subquery, .. }
Expr::Exists(Exists { subquery, .. })
| Expr::InSubquery { subquery, .. }
| Expr::ScalarSubquery(subquery) => {
// use a synthetic plan so the visitor sees a
Expand Down
8 changes: 4 additions & 4 deletions datafusion/optimizer/src/analyzer/count_wildcard_rule.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
use datafusion_common::{Column, DFField, DFSchema, DFSchemaRef, Result};
use datafusion_expr::expr::AggregateFunction;
use datafusion_expr::utils::COUNT_STAR_EXPANSION;
use datafusion_expr::Expr::{Exists, InSubquery, ScalarSubquery};
use datafusion_expr::Expr::{InSubquery, ScalarSubquery};
use datafusion_expr::{
aggregate_function, count, expr, lit, window_function, Aggregate, Expr, Filter,
LogicalPlan, Projection, Sort, Subquery, Window,
Expand Down Expand Up @@ -211,20 +211,20 @@ impl TreeNodeRewriter for CountWildcardRewriter {
negated,
}
}
Exists { subquery, negated } => {
Expr::Exists(expr::Exists { subquery, negated }) => {
let new_plan = subquery
.subquery
.as_ref()
.clone()
.transform_down(&analyze_internal)?;

Exists {
Expr::Exists(expr::Exists {
subquery: Subquery {
subquery: Arc::new(new_plan),
outer_ref_columns: subquery.outer_ref_columns,
},
negated,
}
})
}
_ => old_expr,
};
Expand Down
5 changes: 3 additions & 2 deletions datafusion/optimizer/src/analyzer/inline_table_scan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use crate::analyzer::AnalyzerRule;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::Result;
use datafusion_expr::expr::Exists;
use datafusion_expr::{
logical_plan::LogicalPlan, Expr, Filter, LogicalPlanBuilder, TableScan,
};
Expand Down Expand Up @@ -84,11 +85,11 @@ fn analyze_internal(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {

fn rewrite_subquery(expr: Expr) -> Result<Transformed<Expr>> {
match expr {
Expr::Exists { subquery, negated } => {
Expr::Exists(Exists { subquery, negated }) => {
let plan = subquery.subquery.as_ref().clone();
let new_plan = plan.transform_up(&analyze_internal)?;
let subquery = subquery.with_plan(Arc::new(new_plan));
Ok(Transformed::Yes(Expr::Exists { subquery, negated }))
Ok(Transformed::Yes(Expr::Exists(Exists { subquery, negated })))
}
Expr::InSubquery {
expr,
Expand Down
3 changes: 2 additions & 1 deletion datafusion/optimizer/src/analyzer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ use crate::utils::log_plan;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{TreeNode, VisitRecursion};
use datafusion_common::{DataFusionError, Result};
use datafusion_expr::expr::Exists;
use datafusion_expr::utils::inspect_expr_pre;
use datafusion_expr::{Expr, LogicalPlan};
use log::debug;
Expand Down Expand Up @@ -119,7 +120,7 @@ fn check_plan(plan: &LogicalPlan) -> Result<()> {
for expr in plan.expressions().iter() {
// recursively look for subqueries
inspect_expr_pre(expr, |expr| match expr {
Expr::Exists { subquery, .. }
Expr::Exists(Exists { subquery, .. })
| Expr::InSubquery { subquery, .. }
| Expr::ScalarSubquery(subquery) => {
check_subquery_expr(plan, &subquery.subquery, expr)
Expand Down
9 changes: 5 additions & 4 deletions datafusion/optimizer/src/analyzer/type_coercion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,8 @@ use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{RewriteRecursion, TreeNodeRewriter};
use datafusion_common::{DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue};
use datafusion_expr::expr::{
self, Between, BinaryExpr, Case, Like, ScalarFunction, ScalarUDF, WindowFunction,
self, Between, BinaryExpr, Case, Exists, Like, ScalarFunction, ScalarUDF,
WindowFunction,
};
use datafusion_expr::expr_schema::cast_subquery;
use datafusion_expr::logical_plan::Subquery;
Expand Down Expand Up @@ -133,15 +134,15 @@ impl TreeNodeRewriter for TypeCoercionRewriter {
outer_ref_columns,
}))
}
Expr::Exists { subquery, negated } => {
Expr::Exists(Exists { subquery, negated }) => {
let new_plan = analyze_internal(&self.schema, &subquery.subquery)?;
Ok(Expr::Exists {
Ok(Expr::Exists(Exists {
subquery: Subquery {
subquery: Arc::new(new_plan),
outer_ref_columns: subquery.outer_ref_columns,
},
negated,
})
}))
}
Expr::InSubquery {
expr,
Expand Down
3 changes: 2 additions & 1 deletion datafusion/optimizer/src/decorrelate_where_exists.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ use crate::utils::{
};
use crate::{OptimizerConfig, OptimizerRule};
use datafusion_common::{Column, DataFusionError, Result};
use datafusion_expr::expr::Exists;
use datafusion_expr::{
logical_plan::{Distinct, Filter, JoinType, Subquery},
Expr, LogicalPlan, LogicalPlanBuilder,
Expand Down Expand Up @@ -57,7 +58,7 @@ impl DecorrelateWhereExists {
let mut others = vec![];
for it in filters.iter() {
match it {
Expr::Exists { subquery, negated } => {
Expr::Exists(Exists { subquery, negated }) => {
let subquery_plan = self
.try_optimize(&subquery.subquery, config)?
.map(Arc::new)
Expand Down
5 changes: 3 additions & 2 deletions datafusion/sql/src/expr/subquery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
use datafusion_common::{DFSchema, Result};
use datafusion_expr::expr::Exists;
use datafusion_expr::{Expr, Subquery};
use sqlparser::ast::Expr as SQLExpr;
use sqlparser::ast::Query;
Expand All @@ -35,13 +36,13 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
let sub_plan = self.query_to_plan(subquery, planner_context)?;
let outer_ref_columns = sub_plan.all_out_ref_exprs();
planner_context.set_outer_query_schema(old_outer_query_schema);
Ok(Expr::Exists {
Ok(Expr::Exists(Exists {
subquery: Subquery {
subquery: Arc::new(sub_plan),
outer_ref_columns,
},
negated,
})
}))
}

pub(super) fn parse_in_subquery(
Expand Down