From df86f8ed369125802a3e4140542fe8610a7d8cb1 Mon Sep 17 00:00:00 2001
From: Amos Bird <amosbird@gmail.com>
Date: Thu, 3 Jun 2021 10:46:01 +0800
Subject: [PATCH] Allow not in be used in partition pruning

---
 src/Interpreters/Set.cpp                      | 31 +++++++++++++++++++
 .../01891_not_in_partition_prune.reference    |  6 ++++
 .../01891_not_in_partition_prune.sql          | 10 ++++++
 3 files changed, 47 insertions(+)
 create mode 100644 tests/queries/0_stateless/01891_not_in_partition_prune.reference
 create mode 100644 tests/queries/0_stateless/01891_not_in_partition_prune.sql

diff --git a/src/Interpreters/Set.cpp b/src/Interpreters/Set.cpp
index 65896edfa69..b9ccf2a8a5e 100644
--- a/src/Interpreters/Set.cpp
+++ b/src/Interpreters/Set.cpp
@@ -543,6 +543,37 @@ BoolMask MergeTreeSetIndex::checkInRange(const std::vector<Range> & key_ranges,
     auto left_lower = std::lower_bound(indices.begin(), indices.end(), left_point, less);
     auto right_lower = std::lower_bound(indices.begin(), indices.end(), right_point, less);
 
+    /// A special case of 1-element KeyRange. It's useful for partition pruning
+    bool one_element_range = true;
+    for (size_t i = 0; i < tuple_size; ++i)
+    {
+        auto & left = left_point[i];
+        auto & right = right_point[i];
+        if (left.getType() == right.getType())
+        {
+            if (left.getType() == ValueWithInfinity::NORMAL)
+            {
+                if (0 != left.getColumnIfFinite().compareAt(0, 0, right.getColumnIfFinite(), 1))
+                {
+                    one_element_range = false;
+                    break;
+                }
+            }
+        }
+        else
+        {
+            one_element_range = false;
+            break;
+        }
+    }
+    if (one_element_range)
+    {
+        if (left_lower != indices.end() && equals(*left_lower, left_point))
+            return {true, false};
+        else
+            return {false, true};
+    }
+
     return
     {
         left_lower != right_lower
diff --git a/tests/queries/0_stateless/01891_not_in_partition_prune.reference b/tests/queries/0_stateless/01891_not_in_partition_prune.reference
new file mode 100644
index 00000000000..628053cd4f8
--- /dev/null
+++ b/tests/queries/0_stateless/01891_not_in_partition_prune.reference
@@ -0,0 +1,6 @@
+7
+0	100
+6	106
+7	107
+8	108
+9	109
diff --git a/tests/queries/0_stateless/01891_not_in_partition_prune.sql b/tests/queries/0_stateless/01891_not_in_partition_prune.sql
new file mode 100644
index 00000000000..edbfad93e5d
--- /dev/null
+++ b/tests/queries/0_stateless/01891_not_in_partition_prune.sql
@@ -0,0 +1,10 @@
+drop table if exists test1;
+
+create table test1 (i int, j int) engine MergeTree partition by i order by tuple() settings index_granularity = 1;
+
+insert into test1 select number, number + 100 from numbers(10);
+select count() from test1 where i not in (1,2,3);
+set max_rows_to_read = 5;
+select * from test1 where i not in (1,2,3,4,5) order by i;
+
+drop table test1;