From c2eb9e7e782ac9b4733f87334619f6f18055959b Mon Sep 17 00:00:00 2001
From: Anton Popov <pad11rus@gmail.com>
Date: Wed, 10 Nov 2021 00:44:51 +0300
Subject: [PATCH] optimize mapContains to subcolumn

---
 .../RewriteFunctionToSubcolumnVisitor.cpp     | 46 +++++++++++++++----
 .../0_stateless/02115_map_contains.reference  |  4 ++
 .../0_stateless/02115_map_contains.sql        | 12 +++++
 3 files changed, 52 insertions(+), 10 deletions(-)
 create mode 100644 tests/queries/0_stateless/02115_map_contains.reference
 create mode 100644 tests/queries/0_stateless/02115_map_contains.sql

diff --git a/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp b/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp
index b0647baaa7d..debbd9ae8fe 100644
--- a/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp
+++ b/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp
@@ -38,7 +38,13 @@ ASTPtr transformCountNullableToSubcolumn(const String & name_in_storage, const S
     return makeASTFunction("sum", makeASTFunction("not", ast));
 }
 
-const std::unordered_map<String, std::tuple<TypeIndex, String, decltype(&transformToSubcolumn)>> function_to_subcolumn =
+ASTPtr transformMapContainsToSubcolumn(const String & name_in_storage, const String & subcolumn_name, const ASTPtr & arg)
+{
+    auto ast = transformToSubcolumn(name_in_storage, subcolumn_name);
+    return makeASTFunction("has", ast, arg);
+}
+
+const std::unordered_map<String, std::tuple<TypeIndex, String, decltype(&transformToSubcolumn)>> unary_function_to_subcolumn =
 {
     {"length",    {TypeIndex::Array, "size0", transformToSubcolumn}},
     {"empty",     {TypeIndex::Array, "size0", transformEmptyToSubcolumn}},
@@ -50,30 +56,50 @@ const std::unordered_map<String, std::tuple<TypeIndex, String, decltype(&transfo
     {"mapValues", {TypeIndex::Map, "values", transformToSubcolumn}},
 };
 
+const std::unordered_map<String, std::tuple<TypeIndex, String, decltype(&transformMapContainsToSubcolumn)>> binary_function_to_subcolumn
+{
+    {"mapContains", {TypeIndex::Map, "keys", transformMapContainsToSubcolumn}},
+};
+
 }
 
 void RewriteFunctionToSubcolumnData::visit(ASTFunction & function, ASTPtr & ast) const
 {
     const auto & arguments = function.arguments->children;
-    if (arguments.size() != 1)
+    if (arguments.empty() || arguments.size() > 2)
         return;
 
     const auto * identifier = arguments[0]->as<ASTIdentifier>();
     if (!identifier)
         return;
 
-    auto it = function_to_subcolumn.find(function.name);
-    if (it == function_to_subcolumn.end())
-        return;
-
-    const auto & [type_id, subcolumn_name, transformer] = it->second;
     const auto & columns = metadata_snapshot->getColumns();
     const auto & name_in_storage = identifier->name();
 
-    if (columns.has(name_in_storage)
-        && columns.get(name_in_storage).type->getTypeId() == type_id)
+    if (!columns.has(name_in_storage))
+        return;
+
+    TypeIndex column_type_id = columns.get(name_in_storage).type->getTypeId();
+
+    if (arguments.size() == 1)
     {
-        ast = transformer(name_in_storage, subcolumn_name);
+        auto it = unary_function_to_subcolumn.find(function.name);
+        if (it != unary_function_to_subcolumn.end())
+        {
+            const auto & [type_id, subcolumn_name, transformer] = it->second;
+            if (column_type_id == type_id)
+                ast = transformer(name_in_storage, subcolumn_name);
+        }
+    }
+    else
+    {
+        auto it = binary_function_to_subcolumn.find(function.name);
+        if (it != binary_function_to_subcolumn.end())
+        {
+            const auto & [type_id, subcolumn_name, transformer] = it->second;
+            if (column_type_id == type_id)
+                ast = transformer(name_in_storage, subcolumn_name, arguments[1]);
+        }
     }
 }
 
diff --git a/tests/queries/0_stateless/02115_map_contains.reference b/tests/queries/0_stateless/02115_map_contains.reference
new file mode 100644
index 00000000000..975e9876237
--- /dev/null
+++ b/tests/queries/0_stateless/02115_map_contains.reference
@@ -0,0 +1,4 @@
+SELECT has(`m.keys`, \'a\')
+FROM t_map_contains
+1
+0
diff --git a/tests/queries/0_stateless/02115_map_contains.sql b/tests/queries/0_stateless/02115_map_contains.sql
new file mode 100644
index 00000000000..3c7f21cb4f1
--- /dev/null
+++ b/tests/queries/0_stateless/02115_map_contains.sql
@@ -0,0 +1,12 @@
+DROP TABLE IF EXISTS t_map_contains;
+
+CREATE TABLE t_map_contains (m Map(String, UInt32)) ENGINE = Memory;
+
+INSERT INTO t_map_contains VALUES (map('a', 1, 'b', 2)), (map('c', 3, 'd', 4));
+
+SET optimize_functions_to_subcolumns = 1;
+
+EXPLAIN SYNTAX SELECT mapContains(m, 'a') FROM t_map_contains;
+SELECT mapContains(m, 'a') FROM t_map_contains;
+
+DROP TABLE t_map_contains;