Rewrite array exists to has

src/Analyzer/Passes/ArrayExistsToHasPass.cpp

@@ -0,0 +1,85 @@
+#include <Functions/FunctionFactory.h>
+
+#include <Interpreters/Context.h>
+
+#include <Analyzer/ColumnNode.h>
+#include <Analyzer/ConstantNode.h>
+#include <Analyzer/FunctionNode.h>
+#include <Analyzer/InDepthQueryTreeVisitor.h>
+#include <Analyzer/LambdaNode.h>
+
+#include "ArrayExistsToHasPass.h"
+
+namespace DB
+{
+namespace
+{
+    class RewriteArrayExistsToHasVisitor : public InDepthQueryTreeVisitorWithContext<RewriteArrayExistsToHasVisitor>
+    {
+    public:
+        using Base = InDepthQueryTreeVisitorWithContext<RewriteArrayExistsToHasVisitor>;
+        using Base::Base;
+
+        void visitImpl(QueryTreeNodePtr & node)
+        {
+            if (!getSettings().optimize_rewrite_array_exists_to_has)
+                return;
+
+            auto * function_node = node->as<FunctionNode>();
+            if (!function_node || function_node->getFunctionName() != "arrayExists")
+                return;
+
+            auto & function_arguments_nodes = function_node->getArguments().getNodes();
+            if (function_arguments_nodes.size() != 2)
+                return;
+
+            /// lambda function must be like: x -> x = elem
+            auto * lambda_node = function_arguments_nodes[0]->as<LambdaNode>();
+            if (!lambda_node)
+                return;
+
+            auto & lambda_arguments_nodes = lambda_node->getArguments().getNodes();
+            if (lambda_arguments_nodes.size() != 1)
+                return;
+            auto * column_node = lambda_arguments_nodes[0]->as<ColumnNode>();
+
+            auto * filter_node = lambda_node->getExpression()->as<FunctionNode>();
+            if (!filter_node || filter_node->getFunctionName() != "equals")
+                return;
+
+            auto filter_arguments_nodes = filter_node->getArguments().getNodes();
+            if (filter_arguments_nodes.size() != 2)
+                return;
+
+            ColumnNode * filter_column_node = nullptr;
+            if (filter_arguments_nodes[1]->as<ConstantNode>() && (filter_column_node = filter_arguments_nodes[0]->as<ColumnNode>())
+                && filter_column_node->getColumnName() == column_node->getColumnName())
+            {
+                /// Rewrite arrayExists(x -> x = elem, arr) -> has(arr, elem)
+                function_arguments_nodes[0] = std::move(function_arguments_nodes[1]);
+                function_arguments_nodes[1] = std::move(filter_arguments_nodes[1]);
+                function_node->resolveAsFunction(
+                    FunctionFactory::instance().get("has", getContext())->build(function_node->getArgumentColumns()));
+            }
+            else if (
+                filter_arguments_nodes[0]->as<ConstantNode>() && (filter_column_node = filter_arguments_nodes[1]->as<ColumnNode>())
+                && filter_column_node->getColumnName() == column_node->getColumnName())
+            {
+                /// Rewrite arrayExists(x -> elem = x, arr) -> has(arr, elem)
+                function_arguments_nodes[0] = std::move(function_arguments_nodes[1]);
+                function_arguments_nodes[1] = std::move(filter_arguments_nodes[0]);
+                function_node->resolveAsFunction(
+                    FunctionFactory::instance().get("has", getContext())->build(function_node->getArgumentColumns()));
+            }
+        }
+    };
+
+}
+
+void RewriteArrayExistsToHasPass::run(QueryTreeNodePtr query_tree_node, ContextPtr context)
+{
+    RewriteArrayExistsToHasVisitor visitor(context);
+    visitor.visit(query_tree_node);
+}
+
+}

src/Analyzer/Passes/ArrayExistsToHasPass.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <Analyzer/IQueryTreePass.h>
+
+namespace DB
+{
+/// Rewrite possible 'arrayExists(func, arr)' to 'has(arr, elem)' to improve performance
+/// arrayExists(x -> x = 1, arr) -> has(arr, 1)
+class RewriteArrayExistsToHasPass final : public IQueryTreePass
+{
+public:
+    String getName() override { return "RewriteArrayExistsToHas"; }
+
+    String getDescription() override { return "Rewrite arrayExists(func, arr) functions to has(arr, elem) when logically equivalent"; }
+
+    void run(QueryTreeNodePtr query_tree_node, ContextPtr context) override;
+};
+}

src/Analyzer/QueryTreePassManager.cpp

@@ -35,6 +35,7 @@
 #include <Analyzer/Passes/ConvertOrLikeChainPass.h>
 #include <Analyzer/Passes/OptimizeRedundantFunctionsInOrderByPass.h>
 #include <Analyzer/Passes/GroupingFunctionsResolvePass.h>
+#include <Analyzer/Passes/ArrayExistsToHasPass.h>
 
 namespace DB
 {
@@ -217,6 +218,7 @@ void addQueryTreePasses(QueryTreePassManager & manager)
     manager.addPass(std::make_unique<CountDistinctPass>());
     manager.addPass(std::make_unique<RewriteAggregateFunctionWithIfPass>());
     manager.addPass(std::make_unique<SumIfToCountIfPass>());
+    manager.addPass(std::make_unique<RewriteArrayExistsToHasPass>());
     manager.addPass(std::make_unique<NormalizeCountVariantsPass>());
 
     manager.addPass(std::make_unique<CustomizeFunctionsPass>());

src/Core/Settings.h

@@ -550,6 +550,7 @@ class IColumn;
     \
     M(Bool, optimize_rewrite_sum_if_to_count_if, false, "Rewrite sumIf() and sum(if()) function countIf() function when logically equivalent", 0) \
     M(Bool, optimize_rewrite_aggregate_function_with_if, true, "Rewrite aggregate functions with if expression as argument when logically equivalent. For example, avg(if(cond, col, null)) can be rewritten to avgIf(cond, col)", 0) \
+    M(Bool, optimize_rewrite_array_exists_to_has, true, "Rewrite arrayExists() functions to has() when logically equivalent. For example, arrayExists(x -> x = 1, arr) can be rewritten to has(arr, 1)", 0) \
     M(UInt64, insert_shard_id, 0, "If non zero, when insert into a distributed table, the data will be inserted into the shard `insert_shard_id` synchronously. Possible values range from 1 to `shards_number` of corresponding distributed table", 0) \
     \
     M(Bool, collect_hash_table_stats_during_aggregation, true, "Enable collecting hash table statistics to optimize memory allocation", 0) \

src/Interpreters/RewriteArrayExistsFunctionVisitor.cpp

@@ -0,0 +1,79 @@
+#include <Interpreters/RewriteArrayExistsFunctionVisitor.h>
+#include <Parsers/ASTFunction.h>
+#include <Parsers/ASTIdentifier.h>
+#include <Parsers/ASTLiteral.h>
+
+namespace DB
+{
+void RewriteArrayExistsFunctionMatcher::visit(ASTPtr & ast, Data & data)
+{
+    if (auto * func = ast->as<ASTFunction>())
+    {
+        if (func->is_window_function)
+            return;
+
+        visit(*func, ast, data);
+    }
+}
+
+void RewriteArrayExistsFunctionMatcher::visit(const ASTFunction & func, ASTPtr & ast, Data &)
+{
+    if (func.name != "arrayExists" || !func.arguments)
+        return;
+
+    auto & array_exists_arguments = func.arguments->children;
+    if (array_exists_arguments.size() != 2)
+        return;
+
+    /// lambda function must be like: x -> x = elem
+    const auto * lambda_func = array_exists_arguments[0]->as<ASTFunction>();
+    if (!lambda_func || !lambda_func->is_lambda_function)
+        return;
+
+    const auto & lambda_func_arguments = lambda_func->arguments->children;
+    if (lambda_func_arguments.size() != 2)
+        return;
+
+    const auto * tuple_func = lambda_func_arguments[0]->as<ASTFunction>();
+    if (!tuple_func || tuple_func->name != "tuple")
+        return;
+
+    const auto & tuple_arguments = tuple_func->arguments->children;
+    if (tuple_arguments.size() != 1)
+        return;
+
+    const auto * id = tuple_arguments[0]->as<ASTIdentifier>();
+    if (!id)
+        return;
+
+    const auto * filter_func = lambda_func_arguments[1]->as<ASTFunction>();
+    if (!filter_func || filter_func->name != "equals")
+        return;
+
+    auto & filter_arguments = filter_func->arguments->children;
+    if (filter_arguments.size() != 2)
+        return;
+
+    const ASTIdentifier * filter_id = nullptr;
+    if ((filter_id = filter_arguments[0]->as<ASTIdentifier>()) && filter_arguments[1]->as<ASTLiteral>()
+        && filter_id->full_name == id->full_name)
+    {
+        /// arrayExists(x -> x = elem, arr) -> has(arr, elem)
+        auto new_func = makeASTFunction("has", std::move(array_exists_arguments[1]), std::move(filter_arguments[1]));
+        new_func->setAlias(func.alias);
+        ast = std::move(new_func);
+        return;
+    }
+    else if (
+        (filter_id = filter_arguments[1]->as<ASTIdentifier>()) && filter_arguments[0]->as<ASTLiteral>()
+        && filter_id->full_name == id->full_name)
+    {
+        /// arrayExists(x -> elem = x, arr) -> has(arr, elem)
+        auto new_func = makeASTFunction("has", std::move(array_exists_arguments[1]), std::move(filter_arguments[0]));
+        new_func->setAlias(func.alias);
+        ast = std::move(new_func);
+        return;
+    }
+}
+
+}

src/Interpreters/RewriteArrayExistsFunctionVisitor.h

@@ -0,0 +1,25 @@
+#pragma once
+
+#include <Interpreters/InDepthNodeVisitor.h>
+#include <Parsers/IAST.h>
+
+namespace DB
+{
+class ASTFunction;
+
+/// Rewrite possible 'arrayExists(func, arr)' to 'has(arr, elem)' to improve performance
+/// arrayExists(x -> x = 1, arr) -> has(arr, 1)
+class RewriteArrayExistsFunctionMatcher
+{
+public:
+    struct Data
+    {
+    };
+
+    static void visit(ASTPtr & ast, Data &);
+    static void visit(const ASTFunction &, ASTPtr & ast, Data &);
+    static bool needChildVisit(const ASTPtr &, const ASTPtr &) { return true; }
+};
+
+using RewriteArrayExistsFunctionVisitor = InDepthNodeVisitor<RewriteArrayExistsFunctionMatcher, false>;
+}

src/Interpreters/TreeOptimizer.cpp

@@ -23,6 +23,8 @@
 #include <Interpreters/Context.h>
 #include <Interpreters/ExternalDictionariesLoader.h>
 #include <Interpreters/GatherFunctionQuantileVisitor.h>
+#include <Interpreters/RewriteSumIfFunctionVisitor.h>
+#include <Interpreters/RewriteArrayExistsFunctionVisitor.h>
 
 #include <Parsers/ASTExpressionList.h>
 #include <Parsers/ASTFunction.h>
@@ -37,7 +39,6 @@
 #include <Functions/UserDefined/UserDefinedExecutableFunctionFactory.h>
 #include <Storages/IStorage.h>
 
-#include <Interpreters/RewriteSumIfFunctionVisitor.h>
 
 namespace DB
 {
@@ -658,6 +659,12 @@ void optimizeSumIfFunctions(ASTPtr & query)
     RewriteSumIfFunctionVisitor(data).visit(query);
 }
 
+void optimizeArrayExistsFunctions(ASTPtr & query)
+{
+    RewriteArrayExistsFunctionVisitor::Data data = {};
+    RewriteArrayExistsFunctionVisitor(data).visit(query);
+}
+
 void optimizeMultiIfToIf(ASTPtr & query)
 {
     OptimizeMultiIfToIfVisitor::Data data;
@@ -790,6 +797,9 @@ void TreeOptimizer::apply(ASTPtr & query, TreeRewriterResult & result,
     if (settings.optimize_rewrite_sum_if_to_count_if)
         optimizeSumIfFunctions(query);
 
+    if (settings.optimize_rewrite_array_exists_to_has)
+        optimizeArrayExistsFunctions(query);
+
     /// Remove injective functions inside uniq
     if (settings.optimize_injective_functions_inside_uniq)
         optimizeInjectiveFunctionsInsideUniq(query, context);
@@ -799,9 +809,7 @@ void TreeOptimizer::apply(ASTPtr & query, TreeRewriterResult & result,
         && !select_query->group_by_with_totals
         && !select_query->group_by_with_rollup
         && !select_query->group_by_with_cube)
-    {
         optimizeAggregateFunctionsOfGroupByKeys(select_query, query);
-    }
 
     /// Remove duplicate ORDER BY and DISTINCT from subqueries.
     if (settings.optimize_duplicate_order_by_and_distinct)

tests/performance/rewrite_array_exists.xml

@@ -0,0 +1,4 @@
+<test>
+    <query>select arrayExists(x -> x = 5, materialize(range(10))) from numbers(10000000) format Null</query>
+    <query>select has(materialize(range(10)), 5) from numbers(10000000) format Null</query>
+</test>