Add support for a vector search index in SAI

- Adds jbellis/jvector (1.0.2) library for DiskANN based indexes on floating point vectors
- Adds ORDER BY ANN OF capability to do ANN search and order the results by score

 patch by Mike Adamson; reviewed by Andrés de la Peña, Jonathon Ellis for CASSANDRA-18715

Co-authored-by Jonathon Ellis [email protected]
Co-authored-by Zhao Yang [email protected]

Author: mike-tr-adamson

.build/cassandra-deps-template.xml

@@ -372,5 +372,9 @@
       <groupId>org.apache.lucene</groupId>
       <artifactId>lucene-analysis-common</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.github.jbellis</groupId>
+      <artifactId>jvector</artifactId>
+    </dependency>
   </dependencies>
 </project>

.build/parent-pom-template.xml

@@ -1201,6 +1201,11 @@
         <artifactId>lucene-analysis-common</artifactId>
         <version>9.7.0</version>
       </dependency>
+      <dependency>
+        <groupId>io.github.jbellis</groupId>
+        <artifactId>jvector</artifactId>
+        <version>1.0.2</version>
+      </dependency>
       <dependency>
         <groupId>com.carrotsearch.randomizedtesting</groupId>
         <artifactId>randomizedtesting-runner</artifactId>

CHANGES.txt

@@ -1,4 +1,5 @@
 5.0-alpha2
+ * Add support for vector search in SAI (CASSANDRA-18715)
  * Remove crc_check_chance from CompressionParams (CASSANDRA-18872)
  * Fix schema loading of UDTs inside vectors inside UDTs (CASSANDRA-18964)
  * Add cqlsh autocompletion for the vector data type (CASSANDRA-18946)

NEWS.txt

@@ -86,6 +86,7 @@ New features
       src/java/org/apache/cassandra/db/compaction/UnifiedCompactionStrategy.md
     - New `VectorType` (cql `vector<element_type, dimension>`) which adds new fixed-length element arrays. See CASSANDRA-18504
     - Added new vector similarity functions `similarity_cosine`, `similarity_euclidean` and `similarity_dot_product`.
+    - Added ANN vector similarity search via ORDER BY ANN OF syntax on SAI indexes (using jvector library).
     - Removed UDT type migration logic for 3.6+ clusters upgrading to 4.0.  If migration has been disabled, it must be
       enabled before upgrading to 5.0 if the cluster used UDTs. See CASSANDRA-18504
     - Entended max expiration time from 2038-01-19T03:14:06+00:00 to 2106-02-07T06:28:13+00:00

src/antlr/Lexer.g

@@ -224,7 +224,8 @@ K_MASKED:      M A S K E D;
 K_UNMASK:      U N M A S K;
 K_SELECT_MASKED: S E L E C T '_' M A S K E D;
 
-K_VECTOR:       V E C T O R;
+K_VECTOR:      V E C T O R;
+K_ANN:         A N N;
 
 // Case-insensitive alpha characters
 fragment A: ('a'|'A');

src/antlr/Parser.g

@@ -268,7 +268,7 @@ selectStatement returns [SelectStatement.RawStatement expr]
     @init {
         Term.Raw limit = null;
         Term.Raw perPartitionLimit = null;
-        Map<ColumnIdentifier, Boolean> orderings = new LinkedHashMap<>();
+        List<Ordering.Raw> orderings = new ArrayList<>();
         List<Selectable.Raw> groups = new ArrayList<>();
         boolean allowFiltering = false;
         boolean isJson = false;
@@ -459,11 +459,17 @@ customIndexExpression [WhereClause.Builder clause]
     : 'expr(' idxName[name] ',' t=term ')' { clause.add(new CustomIndexExpression(name, t));}
     ;
 
-orderByClause[Map<ColumnIdentifier, Boolean> orderings]
+orderByClause[List<Ordering.Raw> orderings]
     @init{
-        boolean reversed = false;
+        Ordering.Direction direction = Ordering.Direction.ASC;
+    }
+    : c=cident (K_ANN K_OF t=term)? (K_ASC | K_DESC { direction = Ordering.Direction.DESC; })?
+    {
+        Ordering.Raw.Expression expr = (t == null)
+            ? new Ordering.Raw.SingleColumn(c)
+            : new Ordering.Raw.Ann(c, t);
+        orderings.add(new Ordering.Raw(expr, direction));
     }
-    : c=cident (K_ASC | K_DESC { reversed = true; })? { orderings.put(c, reversed); }
     ;
 
 groupByClause[List<Selectable.Raw> groups]
@@ -2014,5 +2020,6 @@ basic_unreserved_keyword returns [String str]
         | K_UNMASK
         | K_SELECT_MASKED
         | K_VECTOR
+        | K_ANN
         ) { $str = $k.text; }
     ;

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

@@ -429,8 +429,9 @@ public enum CassandraRelevantProperties
     SAI_INTERSECTION_CLAUSE_LIMIT("cassandra.sai.intersection_clause_limit", "2"),
     /** Latest version to be used for SAI index writing */
     SAI_LATEST_VERSION("cassandra.sai.latest_version", "aa"),
-    SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size_kb", "5"),
-    SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size_kb", "1"),
+    SAI_MAX_FROZEN_TERM_SIZE("cassandra.sai.max_frozen_term_size", "5KiB"),
+    SAI_MAX_STRING_TERM_SIZE("cassandra.sai.max_string_term_size", "1KiB"),
+    SAI_MAX_VECTOR_TERM_SIZE("cassandra.sai.max_vector_term_size", "32KiB"),
 
     /** Minimum number of reachable leaves for a given node to be eligible for an auxiliary posting list */
     SAI_MINIMUM_POSTINGS_LEAVES("cassandra.sai.minimum_postings_leaves", "64"),
@@ -455,6 +456,18 @@ public enum CassandraRelevantProperties
     SAI_TEST_BALANCED_TREE_DEBUG_ENABLED("cassandra.sai.test.balanced_tree_debug_enabled", "false"),
     SAI_TEST_DISABLE_TIMEOUT("cassandra.sai.test.disable.timeout", "false"),
 
+    /** Whether to allow the user to specify custom options to the hnsw index */
+    SAI_VECTOR_ALLOW_CUSTOM_PARAMETERS("cassandra.sai.vector.allow_custom_parameters", "false"),
+
+    /** Controls the maximum top-k limit for vector search */
+    SAI_VECTOR_SEARCH_MAX_TOP_K("cassandra.sai.vector_search.max_top_k", "1000"),
+
+    /**
+     * Controls the maximum number of PrimaryKeys that will be read into memory at one time when ordering/limiting
+     * the results of an ANN query constrained by non-ANN predicates.
+     */
+    SAI_VECTOR_SEARCH_ORDER_CHUNK_SIZE("cassandra.sai.vector_search.order_chunk_size", "100000"),
+
     SCHEMA_PULL_INTERVAL_MS("cassandra.schema_pull_interval_ms", "60000"),
     SCHEMA_UPDATE_HANDLER_FACTORY_CLASS("cassandra.schema.update_handler_factory.class"),
     SEARCH_CONCURRENCY_FACTOR("cassandra.search_concurrency_factor", "1"),

src/java/org/apache/cassandra/cql3/Operator.java

@@ -258,6 +258,23 @@ public boolean isSatisfiedBy(AbstractType<?> type, ByteBuffer leftOperand, ByteB
         {
             throw new UnsupportedOperationException();
         }
+    },
+    ANN(15)
+    {
+        @Override
+        public String toString()
+        {
+            return "ANN";
+        }
+
+        @Override
+        public boolean isSatisfiedBy(AbstractType<?> type, ByteBuffer leftOperand, ByteBuffer rightOperand)
+        {
+            // The ANN operator is only supported by the vector index so, normally, should never be called directly.
+            // In networked queries (non-local) the coordinator will end up calling the row filter directly. So, this
+            // needs to return true so that the returned values are allowed through to the VectorTopKProcessor
+            return true;
+        }
     };
 
     /**

src/java/org/apache/cassandra/cql3/Ordering.java

@@ -0,0 +1,180 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.cql3;
+
+import org.apache.cassandra.cql3.restrictions.SingleColumnRestriction;
+import org.apache.cassandra.cql3.restrictions.SingleRestriction;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.schema.TableMetadata;
+
+/**
+ * A single element of an ORDER BY clause.
+ * <code>ORDER BY ordering1 [, ordering2 [, ...]] </code>
+ * <p>
+ * An ordering comprises an expression that produces the values to compare against each other
+ * and a sorting direction (ASC, DESC).
+ */
+public class Ordering
+{
+    public final Expression expression;
+    public final Direction direction;
+
+    public Ordering(Expression expression, Direction direction)
+    {
+        this.expression = expression;
+        this.direction = direction;
+    }
+
+    public static abstract class Expression
+    {
+        protected final ColumnMetadata columnMetadata;
+
+        public Expression(ColumnMetadata columnMetadata)
+        {
+            this.columnMetadata = columnMetadata;
+        }
+
+        public boolean hasNonClusteredOrdering()
+        {
+            return false;
+        }
+
+        public SingleRestriction toRestriction()
+        {
+            throw new UnsupportedOperationException();
+        }
+
+        public ColumnMetadata getColumn()
+        {
+            return columnMetadata;
+        }
+    }
+
+    /**
+     * Represents a single column in <code>ORDER BY column</code>
+     */
+    public static class SingleColumn extends Expression
+    {
+        public SingleColumn(ColumnMetadata columnMetadata)
+        {
+            super(columnMetadata);
+        }
+    }
+
+    /**
+     * An expression used in Approximate Nearest Neighbor ordering. <code>ORDER BY column ANN OF value</code>
+     */
+    public static class Ann extends Expression
+    {
+        final Term vectorValue;
+
+        public Ann(ColumnMetadata columnMetadata, Term vectorValue)
+        {
+            super(columnMetadata);
+            this.vectorValue = vectorValue;
+        }
+
+        @Override
+        public boolean hasNonClusteredOrdering()
+        {
+            return true;
+        }
+
+        @Override
+        public SingleRestriction toRestriction()
+        {
+            return new SingleColumnRestriction.AnnRestriction(columnMetadata, vectorValue);
+        }
+    }
+
+    public enum Direction
+    {ASC, DESC}
+
+
+    /**
+     * Represents ANTLR's abstract syntax tree of a single element in the {@code ORDER BY} clause.
+     * This comes directly out of CQL parser.
+     */
+    public static class Raw
+    {
+
+        final Expression expression;
+        final Direction direction;
+
+        public Raw(Expression expression, Direction direction)
+        {
+            this.expression = expression;
+            this.direction = direction;
+        }
+
+        /**
+         * Resolves column identifiers against the table schema.
+         * Binds markers (?) to columns.
+         */
+        public Ordering bind(TableMetadata table, VariableSpecifications boundNames)
+        {
+            return new Ordering(expression.bind(table, boundNames), direction);
+        }
+
+        public interface Expression
+        {
+            Ordering.Expression bind(TableMetadata table, VariableSpecifications boundNames);
+        }
+
+        public static class SingleColumn implements Expression
+        {
+            final ColumnIdentifier column;
+
+            SingleColumn(ColumnIdentifier column)
+            {
+                this.column = column;
+            }
+
+            @Override
+            public Ordering.Expression bind(TableMetadata table, VariableSpecifications boundNames)
+            {
+                return new Ordering.SingleColumn(table.getExistingColumn(column));
+            }
+        }
+
+        public static class Ann implements Expression
+        {
+            final ColumnIdentifier columnId;
+            final Term.Raw vectorValue;
+
+            Ann(ColumnIdentifier column, Term.Raw vectorValue)
+            {
+                this.columnId = column;
+                this.vectorValue = vectorValue;
+            }
+
+            @Override
+            public Ordering.Expression bind(TableMetadata table, VariableSpecifications boundNames)
+            {
+                ColumnMetadata column = table.getExistingColumn(columnId);
+                Term value = vectorValue.prepare(table.keyspace, column);
+                value.collectMarkerSpecification(boundNames);
+                return new Ordering.Ann(column, value);
+            }
+        }
+    }
+}
+
+
+

src/java/org/apache/cassandra/cql3/QueryOptions.java

@@ -114,6 +114,16 @@ public static QueryOptions addColumnSpecifications(QueryOptions options, List<Co
         return new OptionsWithColumnSpecifications(options, columnSpecs);
     }
 
+    public static QueryOptions withConsistencyLevel(QueryOptions options, ConsistencyLevel consistencyLevel)
+    {
+        return new OptionsWithConsistencyLevel(options, consistencyLevel);
+    }
+
+    public static QueryOptions withPageSize(QueryOptions options, int pageSize)
+    {
+        return new OptionsWithPageSize(options, pageSize);
+    }
+
     public abstract ConsistencyLevel getConsistency();
     public abstract List<ByteBuffer> getValues();
     public abstract boolean skipMetadata();
@@ -424,6 +434,40 @@ public QueryOptions prepare(List<ColumnSpecification> specs)
         }
     }
 
+    static class OptionsWithConsistencyLevel extends QueryOptionsWrapper
+    {
+        private final ConsistencyLevel consistencyLevel;
+
+        OptionsWithConsistencyLevel(QueryOptions wrapped, ConsistencyLevel consistencyLevel)
+        {
+            super(wrapped);
+            this.consistencyLevel = consistencyLevel;
+        }
+
+        @Override
+        public ConsistencyLevel getConsistency()
+        {
+            return consistencyLevel;
+        }
+    }
+
+    static class OptionsWithPageSize extends QueryOptionsWrapper
+    {
+        private final int pageSize;
+
+        OptionsWithPageSize(QueryOptions wrapped, int pageSize)
+        {
+            super(wrapped);
+            this.pageSize = pageSize;
+        }
+
+        @Override
+        public int getPageSize()
+        {
+            return pageSize;
+        }
+    }
+
     /**
      * <code>QueryOptions</code> decorator that provides access to the column specifications.
      */

src/java/org/apache/cassandra/cql3/Relation.java

@@ -155,6 +155,8 @@ public final Restriction toRestriction(TableMetadata table, VariableSpecificatio
             case LIKE_MATCHES:
             case LIKE:
                 return newLikeRestriction(table, boundNames, relationType);
+            case ANN:
+                throw invalidRequest("ANN is only supported in ORDER BY");
             default: throw invalidRequest("Unsupported \"!=\" relation: %s", this);
         }
     }