CNDB-12922: Implement rerank_k in SAI ANN queries (#1562)

### What is the issue
Fixes riptano/cndb#12922

### What does this PR fix and why was it fixed

This PR follows up on #1525 by integrating the CQL configured `rerank_k`
in to SAI so that the ANN query can take the user input and let it
influence query execution.

Key changes:
* Bump `MessagingService` to `VERSION_DS_11`
* Fixes implementation of
`MessagingService.instance().endpointsWithVersionBelow` by using the
`keyspace` variant, which is necessary for CNDB.
* When configured, the `rerank_k` is used as the graph's `rerankK`
parameter. The only detail worth mentioning here is that we ignore the
segment proportionality computation if `rerank_k` is provided. This
diverges from the smart default design, but not too greatly.
* Added a guardrail for `rerank_k` to prevent it from exceeding 4 times
the `cassandra.sai.vector_search.max_top_k`. This is debatable and
easily changed, so please let me know if we want something different.
* Made the in memory brute force cost comparison logic use the `rerankK`
value instead of the `limit`. I am pretty sure this is correct, but
please review closely.

Instead of validating `rerank_k`'s value within the query, I added a
test that ensures that as we increase the rerank k, we increase the
recall.

CNDB pr: riptano/cndb#13095

Author: michaeljmarshall

src/java/org/apache/cassandra/config/CassandraRelevantProperties.java

@@ -588,7 +588,7 @@ public enum CassandraRelevantProperties
      * The current messaging version. This is used when we add new messaging versions without adopting them immediately,
      * or to force the node to use a specific version for testing purposes.
      */
-    DS_CURRENT_MESSAGING_VERSION("ds.current_messaging_version", Integer.toString(MessagingService.VERSION_DS_10)),
+    DS_CURRENT_MESSAGING_VERSION("ds.current_messaging_version", Integer.toString(MessagingService.VERSION_DS_11)),
 
     /**
      * Which compression algorithm to use for SSTable compression when not specified explicitly in the sstable options.
@@ -602,7 +602,6 @@ public enum CassandraRelevantProperties
      */
     SKIP_OPTIMAL_STREAMING_CANDIDATES_CALCULATION("cassandra.skip_optimal_streaming_candidates_calculation", "false");
 
-
     CassandraRelevantProperties(String key, String defaultVal)
     {
         this.key = key;

src/java/org/apache/cassandra/cql3/statements/SelectOptions.java

@@ -22,6 +22,7 @@
 import org.apache.cassandra.db.filter.ANNOptions;
 import org.apache.cassandra.exceptions.InvalidRequestException;
 import org.apache.cassandra.exceptions.RequestValidationException;
+import org.apache.cassandra.service.QueryState;
 
 /**
  * {@code WITH option1=... AND option2=...} options for SELECT statements.
@@ -36,16 +37,19 @@ public class SelectOptions extends PropertyDefinitions
     /**
      * Validates all the {@code SELECT} options.
      *
+     * @param state the query state
      * @param limit the {@code SELECT} query user-provided limit
      * @throws InvalidRequestException if any of the options are invalid
      */
-    public void validate(int limit) throws RequestValidationException
+    public void validate(QueryState state, String keyspace, int limit) throws RequestValidationException
     {
         validate(keywords, Collections.emptySet());
-        parseANNOptions().validate(limit);
+        parseANNOptions().validate(state, keyspace, limit);
     }
 
     /**
+     * Parse the ANN Options. Does not validate values of the options or whether peers will be able to process them.
+     *
      * @return the ANN options within these options, or {@link ANNOptions#NONE} if no options are present
      * @throws InvalidRequestException if the ANN options are invalid
      */

src/java/org/apache/cassandra/cql3/statements/SelectStatement.java

@@ -433,7 +433,7 @@ public ReadQuery getQuery(QueryState queryState,
             checkFalse(userOffset != NO_OFFSET, String.format(TOPK_OFFSET_ERROR, userOffset));
         }
 
-        selectOptions.validate(userLimit);
+        selectOptions.validate(queryState, table.keyspace, userLimit);
 
         return query;
     }

src/java/org/apache/cassandra/db/filter/ANNOptions.java

@@ -24,10 +24,12 @@
 
 import org.apache.cassandra.db.TypeSizes;
 import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.guardrails.Guardrails;
 import org.apache.cassandra.io.util.DataInputPlus;
 import org.apache.cassandra.io.util.DataOutputPlus;
 import org.apache.cassandra.locator.InetAddressAndPort;
 import org.apache.cassandra.net.MessagingService;
+import org.apache.cassandra.service.QueryState;
 import org.apache.cassandra.utils.FBUtilities;
 
 /**
@@ -60,10 +62,26 @@ public static ANNOptions create(@Nullable Integer rerankK)
         return rerankK == null ? NONE : new ANNOptions(rerankK);
     }
 
-    public void validate(int limit)
+    /**
+     * Validates the ANN options by checking that they are within the guardrails and that peers support the options.
+     */
+    public void validate(QueryState state, String keyspace, int limit)
     {
-        if (rerankK != null && rerankK > 0 && rerankK < limit)
+        if (rerankK == null)
+            return;
+
+        if (rerankK < limit)
             throw new InvalidRequestException(String.format("Invalid rerank_k value %d lesser than limit %d", rerankK, limit));
+
+        Guardrails.annRerankKMaxValue.guard(rerankK, "ANN options", false, state);
+
+        // Ensure that all nodes in the cluster are in a version that supports ANN options, including this one
+        assert keyspace != null;
+        Set<InetAddressAndPort> badNodes = MessagingService.instance().endpointsWithVersionBelow(keyspace, MessagingService.VERSION_DS_11);
+        if (MessagingService.current_version < MessagingService.VERSION_DS_11)
+            badNodes.add(FBUtilities.getBroadcastAddressAndPort());
+        if (!badNodes.isEmpty())
+            throw new InvalidRequestException("ANN options are not supported in clusters below DS 11.");
     }
 
     /**
@@ -74,13 +92,6 @@ public void validate(int limit)
      */
     public static ANNOptions fromMap(Map<String, String> map)
     {
-        // ensure that all nodes in the cluster are in a version that supports ANN options, including this one
-        Set<InetAddressAndPort> badNodes = MessagingService.instance().endpointsWithVersionBelow(MessagingService.VERSION_DS_11);
-        if (MessagingService.current_version < MessagingService.VERSION_DS_11)
-            badNodes.add(FBUtilities.getBroadcastAddressAndPort());
-        if (!badNodes.isEmpty())
-            throw new InvalidRequestException("ANN options are not supported in clusters below DS 11.");
-
         Integer rerankK = null;
 
         for (Map.Entry<String, String> entry : map.entrySet())

src/java/org/apache/cassandra/guardrails/Guardrails.java

@@ -120,6 +120,17 @@ what, formatSize(v), formatSize(t)))
                     format("%s has a vector of %s dimensions, this exceeds the %s threshold of %s.",
                           what, value, isWarning ? "warning" : "failure", threshold));
 
+    /**
+     * Guardrail on the maximum value for the rerank_k parameter, an ANN query option.
+     */
+    public static final Threshold annRerankKMaxValue =
+            factory.threshold("sai_ann_rerank_k_max_value",
+                    () -> config.sai_ann_rerank_k_warn_threshold,
+                    () -> config.sai_ann_rerank_k_fail_threshold,
+                    (isWarning, what, value, threshold) ->
+                    format("%s specifies rerank_k=%s, this exceeds the %s threshold of %s.",
+                          what, value, isWarning ? "warning" : "failure", threshold));
+
     public static final DisableFlag readBeforeWriteListOperationsEnabled =
             factory.disableFlag("read_before_write_list_operations",
                             () -> !config.read_before_write_list_operations_enabled,

src/java/org/apache/cassandra/guardrails/GuardrailsConfig.java

@@ -29,6 +29,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 
+import org.apache.cassandra.config.CassandraRelevantProperties;
 import org.apache.cassandra.config.Config;
 import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.cql3.statements.schema.TableAttributes;
@@ -81,6 +82,8 @@ public class GuardrailsConfig
     public volatile Boolean read_before_write_list_operations_enabled;
     public volatile Integer vector_dimensions_warn_threshold;
     public volatile Integer vector_dimensions_failure_threshold;
+    public volatile Integer sai_ann_rerank_k_warn_threshold;
+    public volatile Integer sai_ann_rerank_k_fail_threshold;
 
     // Legacy 2i guardrail
     public volatile Integer secondary_index_per_table_failure_threshold;
@@ -165,6 +168,11 @@ public void applyConfig()
         enforceDefault(tombstone_warn_threshold, v -> tombstone_warn_threshold = v, 1000, 1000);
         enforceDefault(tombstone_failure_threshold, v -> tombstone_failure_threshold = v, 100000, 100000);
 
+        // Default to no warning and failure at 4 times the maxTopK value
+        int maxTopK = CassandraRelevantProperties.SAI_VECTOR_SEARCH_MAX_TOP_K.getInt();
+        enforceDefault(sai_ann_rerank_k_warn_threshold, v -> sai_ann_rerank_k_warn_threshold = v, -1, -1);
+        enforceDefault(sai_ann_rerank_k_fail_threshold, v -> sai_ann_rerank_k_fail_threshold = v, 4 * maxTopK, 4 * maxTopK);
+
         // for write requests
         enforceDefault(logged_batch_enabled, v -> logged_batch_enabled = v, true, true);
         enforceDefault(batch_size_warn_threshold_in_kb, v -> batch_size_warn_threshold_in_kb = v, 64, 64);
@@ -269,6 +277,10 @@ public void validate()
         validateStrictlyPositiveInteger(vector_dimensions_failure_threshold, "vector_dimensions_failure_threshold");
         validateWarnLowerThanFail(vector_dimensions_warn_threshold, vector_dimensions_failure_threshold, "vector_dimensions");
 
+        validateStrictlyPositiveInteger(sai_ann_rerank_k_warn_threshold, "sai_ann_rerank_k_warn_threshold");
+        validateStrictlyPositiveInteger(sai_ann_rerank_k_fail_threshold, "sai_ann_rerank_k_fail_threshold");
+        validateWarnLowerThanFail(sai_ann_rerank_k_warn_threshold, sai_ann_rerank_k_fail_threshold, "sai_ann_rerank_k");
+
         validateStrictlyPositiveInteger(tables_warn_threshold, "tables_warn_threshold");
         validateStrictlyPositiveInteger(tables_failure_threshold, "tables_failure_threshold");
         validateWarnLowerThanFail(tables_warn_threshold, tables_failure_threshold, "tables");

src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java

@@ -717,10 +717,6 @@ public void validate(ReadCommand command) throws InvalidRequestException
             throw new InvalidRequestException(String.format("SAI based ORDER BY clause requires a LIMIT that is not greater than %s. LIMIT was %s",
                                                             MAX_TOP_K, command.limits().isUnlimited() ? "NO LIMIT" : command.limits().count()));
 
-        ANNOptions annOptions = command.rowFilter().annOptions();
-        if (annOptions != ANNOptions.NONE)
-            throw new InvalidRequestException("SAI doesn't support ANN options yet.");
-
         indexContext.validate(command.rowFilter());
     }
 

src/java/org/apache/cassandra/index/sai/disk/v1/Segment.java

@@ -205,10 +205,11 @@ public String toString()
      * the number of candidates, the more nodes we expect to visit just to find
      * results that are in that set.)
      */
-    public double estimateAnnSearchCost(int limit, int candidates)
+    public double estimateAnnSearchCost(Orderer orderer, int limit, int candidates)
     {
-        IndexSearcher searcher = getIndexSearcher();
-        return ((V2VectorIndexSearcher) searcher).estimateAnnSearchCost(limit, candidates);
+        V2VectorIndexSearcher searcher = (V2VectorIndexSearcher) getIndexSearcher();
+        int rerankK = orderer.rerankKFor(limit, searcher.getCompression());
+        return searcher.estimateAnnSearchCost(rerankK, candidates);
     }
 
     /**

src/java/org/apache/cassandra/index/sai/disk/v1/V1SearchableIndex.java

@@ -207,6 +207,8 @@ public List<CloseableIterator<PrimaryKeyWithSortKey>> orderBy(Orderer orderer, E
             {
                 if (segment.intersects(keyRange))
                 {
+                    // Note that the proportionality is not used when the user supplies a rerank_k value in the
+                    // ANN_OPTIONS map.
                     var segmentLimit = segment.proportionalAnnLimit(limit, totalRows);
                     iterators.add(segment.orderBy(orderer, slice, keyRange, context, segmentLimit));
                 }

src/java/org/apache/cassandra/index/sai/disk/v2/V2VectorIndexSearcher.java

@@ -163,7 +163,7 @@ public CloseableIterator<PrimaryKeyWithSortKey> orderBy(Orderer orderer, Express
         if (!orderer.isANN())
             throw new IllegalArgumentException(indexContext.logMessage("Unsupported expression during ANN index query: " + orderer));
 
-        int rerankK = indexContext.getIndexWriterConfig().getSourceModel().rerankKFor(limit, graph.getCompression());
+        int rerankK = orderer.rerankKFor(limit, graph.getCompression());
         var queryVector = vts.createFloatVector(orderer.getVectorTerm());
 
         var result = searchInternal(keyRange, context, queryVector, limit, rerankK, 0);
@@ -428,9 +428,8 @@ public double cost()
         }
     }
 
-    public double estimateAnnSearchCost(int limit, int candidates)
+    public double estimateAnnSearchCost(int rerankK, int candidates)
     {
-        int rerankK = indexContext.getIndexWriterConfig().getSourceModel().rerankKFor(limit, graph.getCompression());
         var estimate = estimateCost(rerankK, candidates);
         return estimate.cost();
     }
@@ -472,7 +471,7 @@ public CloseableIterator<PrimaryKeyWithSortKey> orderResultsBy(SSTableReader rea
         if (keys.isEmpty())
             return CloseableIterator.emptyIterator();
 
-        int rerankK = indexContext.getIndexWriterConfig().getSourceModel().rerankKFor(limit, graph.getCompression());
+        int rerankK = orderer.rerankKFor(limit, graph.getCompression());
         // Convert PKs to segment row ids and map to ordinals, skipping any that don't exist in this segment
         var segmentOrdinalPairs = flatmapPrimaryKeysToBitsAndRows(keys);
         var numRows = segmentOrdinalPairs.size();
@@ -611,9 +610,9 @@ public static double logBase2(double number) {
         return Math.log(number) / Math.log(2);
     }
 
-    private int getRawExpectedNodes(int limit, int nPermittedOrdinals)
+    private int getRawExpectedNodes(int rerankK, int nPermittedOrdinals)
     {
-        return VectorMemtableIndex.expectedNodesVisited(limit, nPermittedOrdinals, graph.size());
+        return VectorMemtableIndex.expectedNodesVisited(rerankK, nPermittedOrdinals, graph.size());
     }
 
     @Override