Backport #51610 to 23.4: Fix for moving 'IN' conditions to PREWHERE

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();

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())

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.
       */

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

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;