datafusion_optimizer/
rewrite_set_comparison.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Optimizer rule rewriting `SetComparison` subqueries (e.g. `= ANY`,
19//! `> ALL`) into boolean expressions built from `EXISTS` subqueries
20//! that capture SQL three-valued logic.
21
22use crate::{OptimizerConfig, OptimizerRule};
23use datafusion_common::tree_node::{Transformed, TreeNode};
24use datafusion_common::{Column, DFSchema, ExprSchema, Result, ScalarValue, plan_err};
25use datafusion_expr::expr::{self, Exists, SetComparison, SetQuantifier};
26use datafusion_expr::logical_plan::Subquery;
27use datafusion_expr::logical_plan::builder::LogicalPlanBuilder;
28use datafusion_expr::{DmlStatement, Expr, LogicalPlan, WriteOp, lit};
29use std::sync::Arc;
30
31use datafusion_expr::utils::merge_schema;
32
33/// Rewrite `SetComparison` expressions to scalar subqueries that return the
34/// correct boolean value (including SQL NULL semantics). After this rule
35/// runs, later rules such as `ScalarSubqueryToJoin` can decorrelate and
36/// remove the remaining subquery.
37#[derive(Debug, Default)]
38pub struct RewriteSetComparison;
39
40impl RewriteSetComparison {
41    /// Create a new `RewriteSetComparison` optimizer rule.
42    pub fn new() -> Self {
43        Self
44    }
45
46    fn rewrite_plan(&self, plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
47        let mut schema = merge_schema(&plan.inputs());
48        if let LogicalPlan::Dml(DmlStatement {
49            op: WriteOp::MergeInto(_),
50            table_name,
51            target,
52            ..
53        }) = &plan
54        {
55            schema.merge(&DFSchema::try_from_qualified_schema(
56                table_name.clone(),
57                &target.schema(),
58            )?);
59        }
60        plan.map_expressions(|expr| {
61            expr.transform_up(|expr| rewrite_set_comparison(expr, &schema))
62        })
63    }
64}
65
66impl OptimizerRule for RewriteSetComparison {
67    fn name(&self) -> &str {
68        "rewrite_set_comparison"
69    }
70
71    fn rewrite(
72        &self,
73        plan: LogicalPlan,
74        _config: &dyn OptimizerConfig,
75    ) -> Result<Transformed<LogicalPlan>> {
76        plan.transform_up_with_subqueries(|plan| self.rewrite_plan(plan))
77    }
78}
79
80fn rewrite_set_comparison(
81    expr: Expr,
82    outer_schema: &DFSchema,
83) -> Result<Transformed<Expr>> {
84    match expr {
85        Expr::SetComparison(set_comparison) => {
86            let rewritten = build_set_comparison_subquery(set_comparison, outer_schema)?;
87            Ok(Transformed::yes(rewritten))
88        }
89        _ => Ok(Transformed::no(expr)),
90    }
91}
92
93fn build_set_comparison_subquery(
94    set_comparison: SetComparison,
95    outer_schema: &DFSchema,
96) -> Result<Expr> {
97    let SetComparison {
98        expr,
99        subquery,
100        op,
101        quantifier,
102    } = set_comparison;
103
104    let left_expr = to_outer_reference(*expr, outer_schema)?;
105    let subquery_schema = subquery.subquery.schema();
106    if subquery_schema.fields().is_empty() {
107        return plan_err!("single expression required.");
108    }
109    // avoid `head_output_expr` for aggr/window plan, it will gives group-by expr if exists
110    let right_expr = Expr::Column(Column::from(subquery_schema.qualified_field(0)));
111
112    let comparison = Expr::BinaryExpr(expr::BinaryExpr::new(
113        Box::new(left_expr),
114        op,
115        Box::new(right_expr),
116    ));
117
118    let true_exists =
119        exists_subquery(&subquery, Expr::IsTrue(Box::new(comparison.clone())))?;
120    let null_exists =
121        exists_subquery(&subquery, Expr::IsNull(Box::new(comparison.clone())))?;
122
123    let result_expr = match quantifier {
124        SetQuantifier::Any => Expr::Case(expr::Case {
125            expr: None,
126            when_then_expr: vec![
127                (Box::new(true_exists), Box::new(lit(true))),
128                (
129                    Box::new(null_exists),
130                    Box::new(Expr::Literal(ScalarValue::Boolean(None), None)),
131                ),
132            ],
133            else_expr: Some(Box::new(lit(false))),
134        }),
135        SetQuantifier::All => {
136            let false_exists =
137                exists_subquery(&subquery, Expr::IsFalse(Box::new(comparison.clone())))?;
138            Expr::Case(expr::Case {
139                expr: None,
140                when_then_expr: vec![
141                    (Box::new(false_exists), Box::new(lit(false))),
142                    (
143                        Box::new(null_exists),
144                        Box::new(Expr::Literal(ScalarValue::Boolean(None), None)),
145                    ),
146                ],
147                else_expr: Some(Box::new(lit(true))),
148            })
149        }
150    };
151
152    Ok(result_expr)
153}
154
155fn exists_subquery(subquery: &Subquery, filter: Expr) -> Result<Expr> {
156    let plan = LogicalPlanBuilder::from(subquery.subquery.as_ref().clone())
157        .filter(filter)?
158        .build()?;
159    let outer_ref_columns = plan.all_out_ref_exprs();
160    Ok(Expr::Exists(Exists {
161        subquery: Subquery {
162            subquery: Arc::new(plan),
163            outer_ref_columns,
164            spans: subquery.spans.clone(),
165        },
166        negated: false,
167    }))
168}
169
170fn to_outer_reference(expr: Expr, outer_schema: &DFSchema) -> Result<Expr> {
171    expr.transform_up(|expr| match expr {
172        Expr::Column(col) => {
173            let field = outer_schema.field_from_column(&col)?;
174            Ok(Transformed::yes(Expr::OuterReferenceColumn(
175                Arc::clone(field),
176                col,
177            )))
178        }
179        Expr::OuterReferenceColumn(_, _) => Ok(Transformed::no(expr)),
180        _ => Ok(Transformed::no(expr)),
181    })
182    .map(|t| t.data)
183}
datafusion_optimizer/rewrite_set_comparison.rs

datafusion_optimizer/
rewrite_set_comparison.rs
