From fd5599939874cf512f2463110ff1b1ca87955ea6 Mon Sep 17 00:00:00 2001
From: Ivan <5627721+abyss7@users.noreply.github.com>
Date: Fri, 10 Jan 2020 12:24:05 +0300
Subject: [PATCH] Initialize set for index when table filter is presented
 (#8357)

---
 dbms/src/IO/S3Common.cpp                      |  2 +-
 .../Interpreters/InterpreterSelectQuery.cpp   | 33 ++++++++++---------
 .../tests/integration/test_row_policy/test.py |  3 +-
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/dbms/src/IO/S3Common.cpp b/dbms/src/IO/S3Common.cpp
index a9015ca5982..b981c34c2d2 100644
--- a/dbms/src/IO/S3Common.cpp
+++ b/dbms/src/IO/S3Common.cpp
@@ -29,7 +29,7 @@ const std::pair<LogsLevel, Message::Priority> & convertLogLevel(Aws::Utils::Logg
     return mapping.at(log_level);
 }
 
-class AWSLogger : public Aws::Utils::Logging::LogSystemInterface
+class AWSLogger final : public Aws::Utils::Logging::LogSystemInterface
 {
 public:
     ~AWSLogger() final = default;
diff --git a/dbms/src/Interpreters/InterpreterSelectQuery.cpp b/dbms/src/Interpreters/InterpreterSelectQuery.cpp
index 1d54900ae93..743e79f7373 100644
--- a/dbms/src/Interpreters/InterpreterSelectQuery.cpp
+++ b/dbms/src/Interpreters/InterpreterSelectQuery.cpp
@@ -503,28 +503,31 @@ Block InterpreterSelectQuery::getSampleBlockImpl()
 
     /// Do all AST changes here, because actions from analysis_result will be used later in readImpl.
 
-    /// PREWHERE optimization.
-    /// Turn off, if the table filter (row-level security) is applied.
-    if (storage && !context->getRowPolicy()->getCondition(storage->getDatabaseName(), storage->getTableName(), RowPolicy::SELECT_FILTER))
+    if (storage)
     {
         query_analyzer->makeSetsForIndex(query.where());
         query_analyzer->makeSetsForIndex(query.prewhere());
 
-        auto optimize_prewhere = [&](auto & merge_tree)
+        /// PREWHERE optimization.
+        /// Turn off, if the table filter (row-level security) is applied.
+        if (!context->getRowPolicy()->getCondition(storage->getDatabaseName(), storage->getTableName(), RowPolicy::SELECT_FILTER))
         {
-            SelectQueryInfo current_info;
-            current_info.query = query_ptr;
-            current_info.syntax_analyzer_result = syntax_analyzer_result;
-            current_info.sets = query_analyzer->getPreparedSets();
+            auto optimize_prewhere = [&](auto & merge_tree)
+            {
+                SelectQueryInfo current_info;
+                current_info.query = query_ptr;
+                current_info.syntax_analyzer_result = syntax_analyzer_result;
+                current_info.sets = query_analyzer->getPreparedSets();
 
-            /// Try transferring some condition from WHERE to PREWHERE if enabled and viable
-            if (settings.optimize_move_to_prewhere && query.where() && !query.prewhere() && !query.final())
-                MergeTreeWhereOptimizer{current_info, *context, merge_tree,
-                                        syntax_analyzer_result->requiredSourceColumns(), log};
-        };
+                /// Try transferring some condition from WHERE to PREWHERE if enabled and viable
+                if (settings.optimize_move_to_prewhere && query.where() && !query.prewhere() && !query.final())
+                    MergeTreeWhereOptimizer{current_info, *context, merge_tree,
+                                            syntax_analyzer_result->requiredSourceColumns(), log};
+            };
 
-        if (const auto * merge_tree_data = dynamic_cast<const MergeTreeData *>(storage.get()))
-            optimize_prewhere(*merge_tree_data);
+            if (const auto * merge_tree_data = dynamic_cast<const MergeTreeData *>(storage.get()))
+                optimize_prewhere(*merge_tree_data);
+        }
     }
 
     if (storage && !options.only_analyze)
diff --git a/dbms/tests/integration/test_row_policy/test.py b/dbms/tests/integration/test_row_policy/test.py
index 421a4b0510c..f5136faff54 100644
--- a/dbms/tests/integration/test_row_policy/test.py
+++ b/dbms/tests/integration/test_row_policy/test.py
@@ -34,7 +34,7 @@ def started_cluster():
             CREATE TABLE mydb.filtered_table3 (a UInt8, b UInt8, c UInt16 ALIAS a + b) ENGINE MergeTree ORDER BY a;
             INSERT INTO mydb.filtered_table3 values (0, 0), (0, 1), (1, 0), (1, 1);
         ''')
-        
+
         yield cluster
 
     finally:
@@ -58,6 +58,7 @@ def test_smoke():
     assert instance.query("SELECT a FROM mydb.filtered_table1") == "1\n1\n"
     assert instance.query("SELECT b FROM mydb.filtered_table1") == "0\n1\n"
     assert instance.query("SELECT a FROM mydb.filtered_table1 WHERE a = 1") == "1\n1\n"
+    assert instance.query("SELECT a FROM mydb.filtered_table1 WHERE a IN (1)") == "1\n1\n"
     assert instance.query("SELECT a = 1 FROM mydb.filtered_table1") == "1\n1\n"
 
     assert instance.query("SELECT a FROM mydb.filtered_table3") == "0\n1\n"