From 426f5179851497a0525dd1f3ec658f9f7d09782d Mon Sep 17 00:00:00 2001
From: robot-clickhouse <robot-clickhouse@users.noreply.github.com>
Date: Mon, 3 Jul 2023 16:03:11 +0000
Subject: [PATCH] Backport #51610 to 23.4: Fix for moving 'IN' conditions to
 PREWHERE

---
 .../MergeTree/MergeTreeWhereOptimizer.cpp     |  3 +
 src/Storages/MergeTree/RPNBuilder.cpp         | 15 +++++
 src/Storages/MergeTree/RPNBuilder.h           |  2 +
 .../02809_prewhere_and_in.reference           |  8 +++
 .../0_stateless/02809_prewhere_and_in.sql     | 56 +++++++++++++++++++
 5 files changed, 84 insertions(+)
 create mode 100644 tests/queries/0_stateless/02809_prewhere_and_in.reference
 create mode 100644 tests/queries/0_stateless/02809_prewhere_and_in.sql

diff --git a/src/Storages/MergeTree/MergeTreeWhereOptimizer.cpp b/src/Storages/MergeTree/MergeTreeWhereOptimizer.cpp
index 25a4579c73ea..1620ba98d589 100644
--- a/src/Storages/MergeTree/MergeTreeWhereOptimizer.cpp
+++ b/src/Storages/MergeTree/MergeTreeWhereOptimizer.cpp
@@ -110,6 +110,9 @@ static void collectColumns(const RPNBuilderTreeNode & node, const NameSet & colu
     if (node.isConstant())
         return;
 
+    if (node.isSubqueryOrSet())
+        return;
+
     if (!node.isFunction())
     {
         auto column_name = node.getColumnName();
diff --git a/src/Storages/MergeTree/RPNBuilder.cpp b/src/Storages/MergeTree/RPNBuilder.cpp
index e49459d3d17d..847695aec7b9 100644
--- a/src/Storages/MergeTree/RPNBuilder.cpp
+++ b/src/Storages/MergeTree/RPNBuilder.cpp
@@ -180,6 +180,21 @@ bool RPNBuilderTreeNode::isConstant() const
     }
 }
 
+bool RPNBuilderTreeNode::isSubqueryOrSet() const
+{
+    if (ast_node)
+    {
+        return
+            typeid_cast<const ASTSubquery *>(ast_node) ||
+            typeid_cast<const ASTTableIdentifier *>(ast_node);
+    }
+    else
+    {
+        const auto * node_without_alias = getNodeWithoutAlias(dag_node);
+        return node_without_alias->result_type->getTypeId() == TypeIndex::Set;
+    }
+}
+
 ColumnWithTypeAndName RPNBuilderTreeNode::getConstantColumn() const
 {
     if (!isConstant())
diff --git a/src/Storages/MergeTree/RPNBuilder.h b/src/Storages/MergeTree/RPNBuilder.h
index 626eb2884934..63155e1125d0 100644
--- a/src/Storages/MergeTree/RPNBuilder.h
+++ b/src/Storages/MergeTree/RPNBuilder.h
@@ -98,6 +98,8 @@ class RPNBuilderTreeNode
     /// Is node constant
     bool isConstant() const;
 
+    bool isSubqueryOrSet() const;
+
     /** Get constant as constant column.
       * Node must be constant before calling these method, otherwise logical exception is thrown.
       */
diff --git a/tests/queries/0_stateless/02809_prewhere_and_in.reference b/tests/queries/0_stateless/02809_prewhere_and_in.reference
new file mode 100644
index 000000000000..3080ae862bb6
--- /dev/null
+++ b/tests/queries/0_stateless/02809_prewhere_and_in.reference
@@ -0,0 +1,8 @@
+PREWHERE a IN
+PREWHERE a IN
+PREWHERE a IN
+PREWHERE a IN
+PREWHERE b NOT IN
+PREWHERE b NOT IN
+PREWHERE b NOT IN
+PREWHERE b NOT IN
diff --git a/tests/queries/0_stateless/02809_prewhere_and_in.sql b/tests/queries/0_stateless/02809_prewhere_and_in.sql
new file mode 100644
index 000000000000..345577d6c7c6
--- /dev/null
+++ b/tests/queries/0_stateless/02809_prewhere_and_in.sql
@@ -0,0 +1,56 @@
+DROP TABLE IF EXISTS t_02809;
+
+CREATE TABLE t_02809(a Int64, b Int64, s String)
+ENGINE=MergeTree order by tuple()
+AS SELECT number, number%10, toString(arrayMap(i-> cityHash64(i*number), range(50))) FROM numbers(10000);
+
+CREATE TABLE t_02809_set(c Int64)
+ENGINE=Set()
+AS SELECT * FROM numbers(10);
+
+CREATE TABLE t_02809_aux(c Int64)
+ENGINE=Memory()
+AS SELECT * FROM numbers(10);
+
+
+SET optimize_move_to_prewhere=1;
+
+-- Queries with 'IN'
+SELECT substring(explain, 1, 13) FROM (EXPLAIN SYNTAX
+     SELECT * FROM t_02809 WHERE a IN (SELECT * FROM system.one)
+) WHERE explain LIKE '%WHERE%';
+
+SELECT substring(explain, 1, 13) FROM (EXPLAIN SYNTAX
+     SELECT * FROM t_02809 WHERE a IN (1,2,3)
+) WHERE explain LIKE '%WHERE%';
+
+SELECT substring(explain, 1, 13) FROM (EXPLAIN SYNTAX
+     SELECT * FROM t_02809 WHERE a IN t_02809_set
+) WHERE explain LIKE '%WHERE%';
+
+SELECT substring(explain, 1, 13) FROM (EXPLAIN SYNTAX
+     SELECT * FROM t_02809 WHERE a IN t_02809_aux
+) WHERE explain LIKE '%WHERE%';
+
+
+-- Queries with 'NOT IN'
+SELECT substring(explain, 1, 17) FROM (EXPLAIN SYNTAX
+     SELECT * FROM t_02809 WHERE b NOT IN (SELECT * FROM system.one)
+) WHERE explain LIKE '%WHERE%';
+
+SELECT substring(explain, 1, 17) FROM (EXPLAIN SYNTAX
+     SELECT * FROM t_02809 WHERE b NOT IN (1,2,3)
+) WHERE explain LIKE '%WHERE%';
+
+SELECT substring(explain, 1, 17) FROM (EXPLAIN SYNTAX
+     SELECT * FROM t_02809 WHERE b NOT IN t_02809_set
+) WHERE explain LIKE '%WHERE%';
+
+SELECT substring(explain, 1, 17) FROM (EXPLAIN SYNTAX
+     SELECT * FROM t_02809 WHERE b NOT IN t_02809_aux
+) WHERE explain LIKE '%WHERE%';
+
+
+DROP TABLE t_02809;
+DROP TABLE t_02809_set;
+DROP TABLE t_02809_aux;