Allow not in be used in partition pruning

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

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

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;