From ffa56bde24cf21234db2bda2023103ab6eb59580 Mon Sep 17 00:00:00 2001
From: Anton Popov <pad11rus@gmail.com>
Date: Mon, 21 Jun 2021 15:34:05 +0300
Subject: [PATCH] fix usage of index with array columns and ARRAY JOIN

---
 src/Storages/MergeTree/KeyCondition.cpp             | 13 +++++++++++++
 src/Storages/MergeTree/KeyCondition.h               |  2 ++
 .../01922_array_join_with_index.reference           |  1 +
 .../0_stateless/01922_array_join_with_index.sql     | 10 ++++++++++
 4 files changed, 26 insertions(+)
 create mode 100644 tests/queries/0_stateless/01922_array_join_with_index.reference
 create mode 100644 tests/queries/0_stateless/01922_array_join_with_index.sql

diff --git a/src/Storages/MergeTree/KeyCondition.cpp b/src/Storages/MergeTree/KeyCondition.cpp
index 246d00b831b..476032e66aa 100644
--- a/src/Storages/MergeTree/KeyCondition.cpp
+++ b/src/Storages/MergeTree/KeyCondition.cpp
@@ -423,6 +423,9 @@ KeyCondition::KeyCondition(
       */
     Block block_with_constants = getBlockWithConstants(query_info.query, query_info.syntax_analyzer_result, context);
 
+    for (const auto & [name, _] : query_info.syntax_analyzer_result->array_join_result_to_source)
+        array_joined_columns.insert(name);
+
     const ASTSelectQuery & select = query_info.query->as<ASTSelectQuery &>();
     if (select.where() || select.prewhere())
     {
@@ -610,6 +613,10 @@ bool KeyCondition::canConstantBeWrappedByMonotonicFunctions(
     DataTypePtr & out_type)
 {
     String expr_name = node->getColumnNameWithoutAlias();
+
+    if (array_joined_columns.count(expr_name))
+        return false;
+
     if (key_subexpr_names.count(expr_name) == 0)
         return false;
 
@@ -714,6 +721,9 @@ bool KeyCondition::canConstantBeWrappedByFunctions(
 {
     String expr_name = ast->getColumnNameWithoutAlias();
 
+    if (array_joined_columns.count(expr_name))
+        return false;
+
     if (key_subexpr_names.count(expr_name) == 0)
     {
         /// Let's check another one case.
@@ -1075,6 +1085,9 @@ bool KeyCondition::isKeyPossiblyWrappedByMonotonicFunctionsImpl(
     // Key columns should use canonical names for index analysis
     String name = node->getColumnNameWithoutAlias();
 
+    if (array_joined_columns.count(name))
+        return false;
+
     auto it = key_columns.find(name);
     if (key_columns.end() != it)
     {
diff --git a/src/Storages/MergeTree/KeyCondition.h b/src/Storages/MergeTree/KeyCondition.h
index 7e7b767b53b..c957c65fc40 100644
--- a/src/Storages/MergeTree/KeyCondition.h
+++ b/src/Storages/MergeTree/KeyCondition.h
@@ -459,6 +459,8 @@ private:
     const ExpressionActionsPtr key_expr;
     /// All intermediate columns are used to calculate key_expr.
     const NameSet key_subexpr_names;
+
+    NameSet array_joined_columns;
     PreparedSets prepared_sets;
 
     // If true, always allow key_expr to be wrapped by function
diff --git a/tests/queries/0_stateless/01922_array_join_with_index.reference b/tests/queries/0_stateless/01922_array_join_with_index.reference
new file mode 100644
index 00000000000..b7539310ef3
--- /dev/null
+++ b/tests/queries/0_stateless/01922_array_join_with_index.reference
@@ -0,0 +1 @@
+a	c
diff --git a/tests/queries/0_stateless/01922_array_join_with_index.sql b/tests/queries/0_stateless/01922_array_join_with_index.sql
new file mode 100644
index 00000000000..1444c639692
--- /dev/null
+++ b/tests/queries/0_stateless/01922_array_join_with_index.sql
@@ -0,0 +1,10 @@
+DROP TABLE IF EXISTS t_array_index;
+
+CREATE TABLE t_array_index (n Nested(key String, value String))
+ENGINE = MergeTree ORDER BY n.key;
+
+INSERT INTO t_array_index VALUES (['a', 'b'], ['c', 'd']);
+
+SELECT * FROM t_array_index ARRAY JOIN n WHERE n.key = 'a';
+
+DROP TABLE IF EXISTS t_array_index;