Fix exact count optimization for non-perfect string prefix

Author: amosbird

src/Storages/MergeTree/KeyCondition.cpp

@@ -66,7 +66,7 @@ extern const int BAD_TYPE_OF_FIELD;
 /// - (1) the pattern has a wildcard
 /// - (2) the first wildcard is '%' and is only followed by nothing or other '%'
 /// e.g. 'test%' or 'test%% has perfect prefix 'test', 'test%x', 'test%_' or 'test_' has no perfect prefix.
-String extractFixedPrefixFromLikePattern(std::string_view like_pattern, bool requires_perfect_prefix)
+std::tuple<String, bool> extractFixedPrefixFromLikePattern(std::string_view like_pattern, bool requires_perfect_prefix)
 {
     String fixed_prefix;
     fixed_prefix.reserve(like_pattern.size());
@@ -79,27 +79,39 @@ String extractFixedPrefixFromLikePattern(std::string_view like_pattern, bool req
         {
             case '%':
             case '_':
+            {
+                bool is_perfect_prefix = std::all_of(pos, end, [](auto c) { return c == '%'; });
                 if (requires_perfect_prefix)
                 {
-                    bool is_prefect_prefix = std::all_of(pos, end, [](auto c) { return c == '%'; });
-                    return is_prefect_prefix ? fixed_prefix : "";
+                    if (is_perfect_prefix)
+                        return {fixed_prefix, true};
+                    else
+                        return {"", false};
+                }
+                else
+                {
+                    return {fixed_prefix, is_perfect_prefix};
                 }
-            return fixed_prefix;
+            }
             case '\\':
+            {
                 ++pos;
-            if (pos == end)
-                break;
-            [[fallthrough]];
+                if (pos == end)
+                    break;
+                [[fallthrough]];
+            }
             default:
+            {
                 fixed_prefix += *pos;
+            }
         }
 
         ++pos;
     }
     /// If we can reach this code, it means there was no wildcard found in the pattern, so it is not a perfect prefix
     if (requires_perfect_prefix)
-        return "";
-    return fixed_prefix;
+        return {"", false};
+    return {fixed_prefix, false};
 }
 
 /// for "^prefix..." string it returns "prefix"
@@ -361,10 +373,13 @@ const KeyCondition::AtomMap KeyCondition::atom_map
                 if (value.getType() != Field::Types::String)
                     return false;
 
-                String prefix = extractFixedPrefixFromLikePattern(value.safeGet<String>(), /*requires_perfect_prefix*/ false);
+                auto [prefix, is_perfect] = extractFixedPrefixFromLikePattern(value.safeGet<String>(), /*requires_perfect_prefix*/ false);
                 if (prefix.empty())
                     return false;
 
+                if (!is_perfect)
+                    out.relaxed = true;
+
                 String right_bound = firstStringThatIsGreaterThanAllStringsWithPrefix(prefix);
 
                 out.function = RPNElement::FUNCTION_IN_RANGE;
@@ -382,10 +397,12 @@ const KeyCondition::AtomMap KeyCondition::atom_map
                 if (value.getType() != Field::Types::String)
                     return false;
 
-                String prefix = extractFixedPrefixFromLikePattern(value.safeGet<String>(), /*requires_perfect_prefix*/ true);
+                auto [prefix, is_perfect] = extractFixedPrefixFromLikePattern(value.safeGet<String>(), /*requires_perfect_prefix*/ true);
                 if (prefix.empty())
                     return false;
 
+                chassert(is_perfect);
+
                 String right_bound = firstStringThatIsGreaterThanAllStringsWithPrefix(prefix);
 
                 out.function = RPNElement::FUNCTION_NOT_IN_RANGE;
@@ -441,6 +458,7 @@ const KeyCondition::AtomMap KeyCondition::atom_map
                 out.range = !right_bound.empty()
                     ? Range(prefix, true, right_bound, false)
                     : Range::createLeftBounded(prefix, true);
+                out.relaxed = true;
 
                 return true;
             }
@@ -477,7 +495,6 @@ const KeyCondition::AtomMap KeyCondition::atom_map
         }
 };
 
-static const std::set<std::string_view> always_relaxed_atom_functions = {"match"};
 static const std::set<KeyCondition::RPNElement::Function> always_relaxed_atom_elements
     = {KeyCondition::RPNElement::FUNCTION_UNKNOWN, KeyCondition::RPNElement::FUNCTION_ARGS_IN_HYPERRECTANGLE, KeyCondition::RPNElement::FUNCTION_POINT_IN_POLYGON};
 
@@ -2073,9 +2090,6 @@ bool KeyCondition::extractAtomFromTree(const RPNBuilderTreeNode & node, RPNEleme
             return atom_it->second(out, const_value);
         };
 
-        if (always_relaxed_atom_functions.contains(func_name))
-            relaxed = true;
-
         bool allow_constant_transformation = !no_relaxed_atom_functions.contains(func_name);
         if (num_args == 1)
         {
@@ -2291,7 +2305,10 @@ bool KeyCondition::extractAtomFromTree(const RPNBuilderTreeNode & node, RPNEleme
         out.monotonic_functions_chain = std::move(chain);
         out.argument_num_of_space_filling_curve = argument_num_of_space_filling_curve;
 
-        return atom_it->second(out, const_value);
+        bool valid_atom = atom_it->second(out, const_value);
+        if (valid_atom && out.relaxed)
+            relaxed = true;
+        return valid_atom;
     }
     if (node.tryGetConstant(const_value, const_type))
     {

src/Storages/MergeTree/KeyCondition.h

@@ -218,6 +218,9 @@ class KeyCondition
 
         Function function = FUNCTION_UNKNOWN;
 
+        /// Whether to relax the key condition (e.g., for LIKE queries without a perfect prefix).
+        bool relaxed = false;
+
         /// For FUNCTION_IN_RANGE and FUNCTION_NOT_IN_RANGE.
         Range range = Range::createWholeUniverse();
         size_t key_column = 0;
@@ -481,6 +484,6 @@ class KeyCondition
     bool relaxed = false;
 };
 
-String extractFixedPrefixFromLikePattern(std::string_view like_pattern, bool requires_perfect_prefix);
+std::tuple<String, bool> extractFixedPrefixFromLikePattern(std::string_view like_pattern, bool requires_perfect_prefix);
 
 }

tests/queries/0_stateless/03461_pk_prefix_trivial_count.reference

@@ -0,0 +1,7 @@
+-- { echo ON }
+
+DROP TABLE IF EXISTS t;
+CREATE TABLE t(k String) ORDER BY k as select 'dst_'||number from numbers(1e6);
+SELECT count(*) FROM t WHERE k LIKE 'dst_kkkk_1111%';
+0
+DROP TABLE t;

tests/queries/0_stateless/03461_pk_prefix_trivial_count.sql

@@ -0,0 +1,9 @@
+-- { echo ON }
+
+DROP TABLE IF EXISTS t;
+
+CREATE TABLE t(k String) ORDER BY k as select 'dst_'||number from numbers(1e6);
+
+SELECT count(*) FROM t WHERE k LIKE 'dst_kkkk_1111%';
+
+DROP TABLE t;