use optimizer rule instead

Signed-off-by: Ruihang Xia <[email protected]>

Author: waynexia

datafusion/optimizer/src/analyzer/mod.rs

@@ -29,15 +29,13 @@ use datafusion_expr::expr_rewriter::FunctionRewrite;
 use datafusion_expr::{InvariantLevel, LogicalPlan};
 
 use crate::analyzer::resolve_grouping_function::ResolveGroupingFunction;
-use crate::analyzer::set_comparison::RewriteSetComparison;
 use crate::analyzer::type_coercion::TypeCoercion;
 use crate::utils::log_plan;
 
 use self::function_rewrite::ApplyFunctionRewrites;
 
 pub mod function_rewrite;
 pub mod resolve_grouping_function;
-pub mod set_comparison;
 pub mod type_coercion;
 
 /// [`AnalyzerRule`]s transform [`LogicalPlan`]s in some way to make
@@ -89,7 +87,6 @@ impl Analyzer {
     pub fn new() -> Self {
         let rules: Vec<Arc<dyn AnalyzerRule + Send + Sync>> = vec![
             Arc::new(ResolveGroupingFunction::new()),
-            Arc::new(RewriteSetComparison::new()),
             Arc::new(TypeCoercion::new()),
         ];
         Self::with_rules(rules)

datafusion/optimizer/src/lib.rs

@@ -65,6 +65,7 @@ pub mod propagate_empty_relation;
 pub mod push_down_filter;
 pub mod push_down_limit;
 pub mod replace_distinct_aggregate;
+pub mod rewrite_set_comparison;
 pub mod scalar_subquery_to_join;
 pub mod simplify_expressions;
 pub mod single_distinct_to_groupby;

datafusion/optimizer/src/optimizer.rs

@@ -51,6 +51,7 @@ use crate::propagate_empty_relation::PropagateEmptyRelation;
 use crate::push_down_filter::PushDownFilter;
 use crate::push_down_limit::PushDownLimit;
 use crate::replace_distinct_aggregate::ReplaceDistinctWithAggregate;
+use crate::rewrite_set_comparison::RewriteSetComparison;
 use crate::scalar_subquery_to_join::ScalarSubqueryToJoin;
 use crate::simplify_expressions::SimplifyExpressions;
 use crate::single_distinct_to_groupby::SingleDistinctToGroupBy;
@@ -227,6 +228,7 @@ impl Optimizer {
     /// Create a new optimizer using the recommended list of rules
     pub fn new() -> Self {
         let rules: Vec<Arc<dyn OptimizerRule + Sync + Send>> = vec![
+            Arc::new(RewriteSetComparison::new()),
             Arc::new(OptimizeUnions::new()),
             Arc::new(SimplifyExpressions::new()),
             Arc::new(ReplaceDistinctWithAggregate::new()),

…optimizer/src/analyzer/set_comparison.rs …/optimizer/src/rewrite_set_comparison.rsdatafusion/optimizer/src/analyzer/set_comparison.rs renamed to datafusion/optimizer/src/rewrite_set_comparison.rs datafusion/optimizer/src/analyzer/set_comparison.rs renamed to datafusion/optimizer/src/rewrite_set_comparison.rs

@@ -15,14 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! Rewrite `SetComparison` subqueries (e.g. `= ANY`, `> ALL`) into
-//! boolean expressions built from `EXISTS` subqueries that capture SQL
-//! three-valued logic.
+//! Optimizer rule rewriting `SetComparison` subqueries (e.g. `= ANY`,
+//! `> ALL`) into boolean expressions built from `EXISTS` subqueries
+//! that capture SQL three-valued logic.
 
-use super::AnalyzerRule;
-use datafusion_common::config::ConfigOptions;
+use crate::{OptimizerConfig, OptimizerRule};
 use datafusion_common::tree_node::{Transformed, TreeNode};
-use datafusion_common::{plan_datafusion_err, DFSchema, ExprSchema, Result, ScalarValue};
+use datafusion_common::ExprSchema;
+use datafusion_common::{plan_datafusion_err, DFSchema, Result, ScalarValue};
 use datafusion_expr::expr::{self, Exists, SetComparison, SetQuantifier};
 use datafusion_expr::logical_plan::builder::LogicalPlanBuilder;
 use datafusion_expr::logical_plan::Subquery;
@@ -52,14 +52,17 @@ impl RewriteSetComparison {
     }
 }
 
-impl AnalyzerRule for RewriteSetComparison {
+impl OptimizerRule for RewriteSetComparison {
     fn name(&self) -> &str {
         "rewrite_set_comparison"
     }
 
-    fn analyze(&self, plan: LogicalPlan, _config: &ConfigOptions) -> Result<LogicalPlan> {
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        _config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
         plan.transform_up_with_subqueries(|plan| self.rewrite_plan(plan))
-            .map(|t| t.data)
     }
 }
 

datafusion/sql/Cargo.toml

@@ -61,6 +61,7 @@ log = { workspace = true }
 recursive = { workspace = true, optional = true }
 regex = { workspace = true }
 sqlparser = { workspace = true }
+datafusion-functions-nested = { workspace = true, features = ["sql"] }
 
 [dev-dependencies]
 ctor = { workspace = true }

datafusion/sql/src/expr/mod.rs

@@ -40,6 +40,7 @@ use datafusion_expr::{
 };
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
+use datafusion_functions_nested::expr_fn::array_has;
 
 mod binary_op;
 mod function;
@@ -605,7 +606,17 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                     planner_context,
                 ),
                 _ => {
-                    not_impl_err!("ANY/SOME only supports subquery comparison currently")
+                    if compare_op != BinaryOperator::Eq {
+                        plan_err!(
+                            "Unsupported AnyOp: '{compare_op}', only '=' is supported"
+                        )
+                    } else {
+                        let left_expr =
+                            self.sql_to_expr(*left, schema, planner_context)?;
+                        let right_expr =
+                            self.sql_to_expr(*right, schema, planner_context)?;
+                        Ok(array_has(right_expr, left_expr))
+                    }
                 }
             },
             SQLExpr::AllOp {