From c7495fa59ed7bef021b268bd8d4c248eba97ad77 Mon Sep 17 00:00:00 2001
From: Jan Ciolek <jan.ciolek@scylladb.com>
Date: Thu, 14 Jul 2022 16:52:10 +0200
Subject: [PATCH] cql3: Add expr::index_supports_some_column

Add a function that checks if there is an index
which supports one of the columns present in
the given expression.

This functionality will soon be needed for
clustering and nonprimary columns so it's
good to separate into a reusable function.

Signed-off-by: Jan Ciolek <jan.ciolek@scylladb.com>
---
 cql3/expr/expression.cc                     | 15 +++++++++++++++
 cql3/expr/expression.hh                     |  7 +++++++
 cql3/restrictions/statement_restrictions.cc | 15 +--------------
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/cql3/expr/expression.cc b/cql3/expr/expression.cc
index 64f877fac5..001912ab2a 100644
--- a/cql3/expr/expression.cc
+++ b/cql3/expr/expression.cc
@@ -1028,6 +1028,21 @@ bool has_supporting_index(
                     support);
 }
 
+bool index_supports_some_column(
+        const expression& e,
+        const secondary_index::secondary_index_manager& index_manager,
+        allow_local_index allow_local) {
+    single_column_restrictions_map single_col_restrictions = get_single_column_restrictions_map(e);
+
+    for (auto&& [col, col_restrictions] : single_col_restrictions) {
+        if (has_supporting_index(col_restrictions, index_manager, allow_local)) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 std::ostream& operator<<(std::ostream& os, const column_value& cv) {
     os << cv.col->name_as_text();
     return os;
diff --git a/cql3/expr/expression.hh b/cql3/expr/expression.hh
index d1d5a1ea9a..1796f5ef9d 100644
--- a/cql3/expr/expression.hh
+++ b/cql3/expr/expression.hh
@@ -499,6 +499,13 @@ extern bool is_supported_by(const expression&, const secondary_index::index&);
 extern bool has_supporting_index(
         const expression&, const secondary_index::secondary_index_manager&, allow_local_index allow_local);
 
+// Looks at each column indivudually and checks whether some index can support restrictions on this single column.
+// Expression has to consist only of single column restrictions.
+extern bool index_supports_some_column(
+    const expression&,
+    const secondary_index::secondary_index_manager&,
+    allow_local_index allow_local);
+
 extern sstring to_string(const expression&);
 
 extern std::ostream& operator<<(std::ostream&, const column_value&);
diff --git a/cql3/restrictions/statement_restrictions.cc b/cql3/restrictions/statement_restrictions.cc
index 723580d986..a40c6de12f 100644
--- a/cql3/restrictions/statement_restrictions.cc
+++ b/cql3/restrictions/statement_restrictions.cc
@@ -881,20 +881,7 @@ bool statement_restrictions::parition_key_restrictions_have_supporting_index(con
         return false;
     }
 
-    // Otherwise those are single column restrictions
-    std::vector<const column_definition*> restricted_pk_columns = expr::get_sorted_column_defs(_partition_key_restrictions);
-
-    for (const column_definition* cdef : restricted_pk_columns) {
-        expr::expression col_restrictions = expr::conjunction {
-            .children = expr::extract_single_column_restrictions_for_column(_partition_key_restrictions, *cdef)
-        };
-
-        if (expr::has_supporting_index(col_restrictions, index_manager, allow_local)) {
-            return true;
-        }
-    }
-
-    return false;
+    return expr::index_supports_some_column(_partition_key_restrictions, index_manager, allow_local);
 }
 
 void statement_restrictions::process_clustering_columns_restrictions(bool for_view, bool allow_filtering) {