CNDB-12937: Update jvector to 4.0.0-beta.2; add new graph construction parameters to index config (#1676)

### What is the issue

Fixes: riptano/cndb#12937

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

This pull request upgrades jvector from **4.0.0-beta.1** to
**4.0.0-beta.2** and introduces three new configuration options to
influence graph construction:

1. `neighborhood_overflow`  
2. `alpha`  
3. `enable_hierarchy`

The defaults for these hyperparameters vary between in memory and on
disk, but when these are configured, they will be uniformly applied to
graphs built by a memtable and by compaction.

**Details**

- **jvector 4.0.0-beta.2**  
- Minor changes to the graph index architecture, including some code now
under `...disk.feature...` packages.
- Removed the old `CachingGraphIndex` code in Cassandra, which is no
longer used.
- New constructor arguments for controlling `neighborhood_overflow`,
`alpha`, and hierarchical levels in the HNSW graph.

- **New configuration options**  
1. **`neighborhood_overflow`**: A `float` >= 1.0 controlling how
aggressively the graph tries to insert extra neighbors on each HNSW
layer.
2. **`alpha`**: A `float` > 0 used in the neighbor selection phase for
HNSW.
3. **`enable_hierarchy`**: A `boolean` indicating whether HNSW should
allow multiple layers (true) or a single-layer approximate graph
(false). Defaults to false.

I manually verified that the jvector upgrade is backwards compatible,
meaning that when we build using the new jvector version, we can read
with an old jvector version, so I did not create a new SAI on disk file
format version.

Author: michaeljmarshall

build.xml

@@ -754,7 +754,7 @@
           <dependency groupId="org.apache.lucene" artifactId="lucene-core" version="9.8.0-5ea8bb4f21" />
           <dependency groupId="org.apache.lucene" artifactId="lucene-analysis-common" version="9.8.0-5ea8bb4f21" />
           <dependency groupId="org.apache.lucene" artifactId="lucene-backward-codecs" version="9.8.0-5ea8bb4f21" />
-          <dependency groupId="io.github.jbellis" artifactId="jvector" version="4.0.0-beta.1" />
+          <dependency groupId="io.github.jbellis" artifactId="jvector" version="4.0.0-beta.2" />
           <dependency groupId="com.bpodgursky" artifactId="jbool_expressions" version="1.14" scope="test"/>
 
           <dependency groupId="com.carrotsearch.randomizedtesting" artifactId="randomizedtesting-runner" version="2.1.2" scope="test">

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

@@ -362,9 +362,6 @@ public enum CassandraRelevantProperties
     /** Controls the hnsw vector cache size, in bytes, per index segment. 0 to disable */
     SAI_HNSW_VECTOR_CACHE_BYTES("cassandra.sai.vector_search.vector_cache_bytes", String.valueOf(4 * 1024 * 1024)),
 
-    /** Whether to allow the user to specify custom options to the hnsw index */
-    SAI_HNSW_ALLOW_CUSTOM_PARAMETERS("cassandra.sai.hnsw.allow_custom_parameters", "false"),
-
     /** Whether to validate terms that will be SAI indexed at the coordinator */
     SAI_VALIDATE_TERMS_AT_COORDINATOR("cassandra.sai.validate_terms_at_coordinator", "true"),
 

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

@@ -206,6 +206,9 @@ public List<SecondaryIndexBuilder> getParallelIndexBuildTasks(ColumnFamilyStore
                                                                      IndexWriterConfig.POSTING_LIST_LVL_SKIP_OPTION,
                                                                      IndexWriterConfig.MAXIMUM_NODE_CONNECTIONS,
                                                                      IndexWriterConfig.CONSTRUCTION_BEAM_WIDTH,
+                                                                     IndexWriterConfig.NEIGHBORHOOD_OVERFLOW,
+                                                                     IndexWriterConfig.ALPHA,
+                                                                     IndexWriterConfig.ENABLE_HIERARCHY,
                                                                      IndexWriterConfig.SIMILARITY_FUNCTION,
                                                                      IndexWriterConfig.SOURCE_MODEL,
                                                                      IndexWriterConfig.OPTIMIZE_FOR,

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

@@ -45,6 +45,9 @@ public class IndexWriterConfig
 
     public static final String MAXIMUM_NODE_CONNECTIONS = "maximum_node_connections";
     public static final String CONSTRUCTION_BEAM_WIDTH = "construction_beam_width";
+    public static final String NEIGHBORHOOD_OVERFLOW = "neighborhood_overflow";
+    public static final String ALPHA = "alpha";
+    public static final String ENABLE_HIERARCHY = "enable_hierarchy";
     public static final String SIMILARITY_FUNCTION = "similarity_function";
     public static final String SOURCE_MODEL = "source_model";
     public static final String OPTIMIZE_FOR = "optimize_for"; // unused, retained for compatibility w/ old schemas
@@ -54,6 +57,7 @@ public class IndexWriterConfig
 
     public static final int DEFAULT_MAXIMUM_NODE_CONNECTIONS = 16;
     public static final int DEFAULT_CONSTRUCTION_BEAM_WIDTH = 100;
+    public static final boolean DEFAULT_ENABLE_HIERARCHY = false;
 
     public static final int MAX_TOP_K = SAI_VECTOR_SEARCH_MAX_TOP_K.getInt();
 
@@ -87,6 +91,10 @@ public class IndexWriterConfig
     private final VectorSimilarityFunction similarityFunction;
     private final VectorSourceModel sourceModel;
 
+    private final Float neighborhoodOverflow; // default varies for in memory/compaction build
+    private final Float alpha; // default varies for in memory/compaction build
+    private final boolean enableHierarchy; // defaults to false
+
     public IndexWriterConfig(String indexName,
                              int bkdPostingsSkip,
                              int bkdPostingsMinLeaves)
@@ -108,6 +116,21 @@ public IndexWriterConfig(String indexName,
                              int constructionBeamWidth,
                              VectorSimilarityFunction similarityFunction,
                              VectorSourceModel sourceModel)
+    {
+        this(indexName, bkdPostingsSkip, bkdPostingsMinLeaves, maximumNodeConnections, constructionBeamWidth,
+             similarityFunction, sourceModel, null, null, false);
+    }
+
+    public IndexWriterConfig(String indexName,
+                             int bkdPostingsSkip,
+                             int bkdPostingsMinLeaves,
+                             int maximumNodeConnections,
+                             int constructionBeamWidth,
+                             VectorSimilarityFunction similarityFunction,
+                             VectorSourceModel sourceModel,
+                             Float neighborhoodOverflow,
+                             Float alpha,
+                             boolean enableHierarchy)
     {
         this.indexName = indexName;
         this.bkdPostingsSkip = bkdPostingsSkip;
@@ -116,6 +139,9 @@ public IndexWriterConfig(String indexName,
         this.constructionBeamWidth = constructionBeamWidth;
         this.similarityFunction = similarityFunction;
         this.sourceModel = sourceModel;
+        this.neighborhoodOverflow = neighborhoodOverflow;
+        this.alpha = alpha;
+        this.enableHierarchy = enableHierarchy;
     }
 
     public String getIndexName()
@@ -163,6 +189,21 @@ public VectorSourceModel getSourceModel()
         return sourceModel;
     }
 
+    public float getNeighborhoodOverflow(float defaultValue)
+    {
+        return neighborhoodOverflow == null ? defaultValue : neighborhoodOverflow;
+    }
+
+    public float getAlpha(float defaultValue)
+    {
+        return alpha == null ? defaultValue : alpha;
+    }
+
+    public boolean isHierarchyEnabled()
+    {
+        return enableHierarchy;
+    }
+
     public static IndexWriterConfig fromOptions(String indexName, AbstractType<?> type, Map<String, String> options)
     {
         int minLeaves = DEFAULT_POSTING_LIST_MIN_LEAVES;
@@ -172,6 +213,10 @@ public static IndexWriterConfig fromOptions(String indexName, AbstractType<?> ty
         VectorSourceModel sourceModel = DEFAULT_SOURCE_MODEL;
         VectorSimilarityFunction similarityFunction = sourceModel.defaultSimilarityFunction; // don't leave null in case no options at all are given
 
+        Float neighborhoodOverflow = null;
+        Float alpha = null;
+        boolean enableHierarchy = DEFAULT_ENABLE_HIERARCHY;
+
         if (options.get(POSTING_LIST_LVL_MIN_LEAVES) != null || options.get(POSTING_LIST_LVL_SKIP_OPTION) != null)
         {
             if (TypeUtil.isLiteral(type))
@@ -213,16 +258,16 @@ else if (options.get(MAXIMUM_NODE_CONNECTIONS) != null ||
                  options.get(CONSTRUCTION_BEAM_WIDTH) != null ||
                  options.get(OPTIMIZE_FOR) != null ||
                  options.get(SIMILARITY_FUNCTION) != null ||
-                 options.get(SOURCE_MODEL) != null)
+                 options.get(SOURCE_MODEL) != null ||
+                 options.get(NEIGHBORHOOD_OVERFLOW) != null ||
+                 options.get(ALPHA) != null ||
+                 options.get(ENABLE_HIERARCHY) != null)
         {
             if (!type.isVector())
                 throw new InvalidRequestException(String.format("CQL type %s cannot have vector options", type.asCQL3Type()));
 
             if (options.containsKey(MAXIMUM_NODE_CONNECTIONS))
             {
-                if (!CassandraRelevantProperties.SAI_HNSW_ALLOW_CUSTOM_PARAMETERS.getBoolean())
-                    throw new InvalidRequestException(String.format("Maximum node connections cannot be set without enabling %s", CassandraRelevantProperties.SAI_HNSW_ALLOW_CUSTOM_PARAMETERS.name()));
-
                 try
                 {
                     maximumNodeConnections = Integer.parseInt(options.get(MAXIMUM_NODE_CONNECTIONS));
@@ -237,9 +282,6 @@ else if (options.get(MAXIMUM_NODE_CONNECTIONS) != null ||
             }
             if (options.containsKey(CONSTRUCTION_BEAM_WIDTH))
             {
-                if (!CassandraRelevantProperties.SAI_HNSW_ALLOW_CUSTOM_PARAMETERS.getBoolean())
-                    throw new InvalidRequestException(String.format("Construction beam width cannot be set without enabling %s", CassandraRelevantProperties.SAI_HNSW_ALLOW_CUSTOM_PARAMETERS.name()));
-
                 try
                 {
                     queueSize = Integer.parseInt(options.get(CONSTRUCTION_BEAM_WIDTH));
@@ -285,9 +327,51 @@ else if (options.get(MAXIMUM_NODE_CONNECTIONS) != null ||
             {
                 similarityFunction = sourceModel.defaultSimilarityFunction;
             }
+
+            if (options.containsKey(NEIGHBORHOOD_OVERFLOW))
+            {
+                try
+                {
+                    neighborhoodOverflow = Float.parseFloat(options.get(NEIGHBORHOOD_OVERFLOW));
+                    if (neighborhoodOverflow < 1.0f)
+                        throw new InvalidRequestException(String.format("Neighborhood overflow for index %s must be >= 1.0, was %s",
+                                                                      indexName, neighborhoodOverflow));
+                }
+                catch (NumberFormatException e)
+                {
+                    throw new InvalidRequestException(String.format("Neighborhood overflow %s is not a valid float for index %s",
+                                                                    options.get(NEIGHBORHOOD_OVERFLOW), indexName));
+                }
+            }
+
+            if (options.containsKey(ALPHA))
+            {
+                try
+                {
+                    alpha = Float.parseFloat(options.get(ALPHA));
+                    if (alpha <= 0)
+                        throw new InvalidRequestException(String.format("Alpha for index %s must be > 0, was %s", 
+                                                                      indexName, alpha));
+                }
+                catch (NumberFormatException e)
+                {
+                    throw new InvalidRequestException(String.format("Alpha %s is not a valid float for index %s",
+                                                                  options.get(ALPHA), indexName));
+                }
+            }
+
+            if (options.containsKey(ENABLE_HIERARCHY))
+            {
+                String value = options.get(ENABLE_HIERARCHY).toLowerCase();
+                if (!value.equals("true") && !value.equals("false"))
+                    throw new InvalidRequestException(String.format("Enable hierarchy must be 'true' or 'false' for index %s, was '%s'",
+                                                                  indexName, value));
+                enableHierarchy = Boolean.parseBoolean(value);
+            }
         }
 
-        return new IndexWriterConfig(indexName, skip, minLeaves, maximumNodeConnections, queueSize, similarityFunction, sourceModel);
+        return new IndexWriterConfig(indexName, skip, minLeaves, maximumNodeConnections, queueSize,
+                                   similarityFunction, sourceModel, neighborhoodOverflow, alpha, enableHierarchy);
     }
 
     public static IndexWriterConfig defaultConfig(String indexName)
@@ -310,12 +394,15 @@ public static IndexWriterConfig emptyConfig()
     @Override
     public String toString()
     {
-        return String.format("IndexWriterConfig{%s=%d, %s=%d, %s=%d, %s=%d, %s=%s, %s=%s}",
+        return String.format("IndexWriterConfig{%s=%d, %s=%d, %s=%d, %s=%d, %s=%s, %s=%s, %s=%f, %s=%f, %s=%b}",
                              POSTING_LIST_LVL_SKIP_OPTION, bkdPostingsSkip,
                              POSTING_LIST_LVL_MIN_LEAVES, bkdPostingsMinLeaves,
                              MAXIMUM_NODE_CONNECTIONS, maximumNodeConnections,
                              CONSTRUCTION_BEAM_WIDTH, constructionBeamWidth,
                              SIMILARITY_FUNCTION, similarityFunction,
-                             SOURCE_MODEL, sourceModel);
+                             SOURCE_MODEL, sourceModel,
+                             NEIGHBORHOOD_OVERFLOW, neighborhoodOverflow,
+                             ALPHA, alpha,
+                             ENABLE_HIERARCHY, enableHierarchy);
     }
 }

src/java/org/apache/cassandra/index/sai/disk/vector/CassandraDiskAnn.java

@@ -29,8 +29,7 @@
 
 import io.github.jbellis.jvector.graph.GraphIndex;
 import io.github.jbellis.jvector.graph.GraphSearcher;
-import io.github.jbellis.jvector.graph.disk.CachingGraphIndex;
-import io.github.jbellis.jvector.graph.disk.FeatureId;
+import io.github.jbellis.jvector.graph.disk.feature.FeatureId;
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
 import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider;
 import io.github.jbellis.jvector.quantization.BQVectors;
@@ -57,7 +56,6 @@
 import org.apache.cassandra.tracing.Tracing;
 import org.apache.cassandra.utils.CloseableIterator;
 
-import static java.lang.Math.min;
 
 public class CassandraDiskAnn
 {
@@ -94,7 +92,7 @@ public CassandraDiskAnn(SSTableContext sstableContext, SegmentMetadata.Component
         graphHandle = indexFiles.termsData();
         var rawGraph = OnDiskGraphIndex.load(graphHandle::createReader, termsMetadata.offset);
         features = rawGraph.getFeatureSet();
-        graph = V3OnDiskFormat.ENABLE_EDGES_CACHE ? cachingGraphFor(rawGraph) : rawGraph;
+        graph = rawGraph;
 
         long pqSegmentOffset = this.componentMetadatas.get(IndexComponentType.PQ).offset;
         try (var pqFile = indexFiles.pq();
@@ -175,27 +173,6 @@ public ProductQuantization getPQ()
         return pq;
     }
 
-    private GraphIndex cachingGraphFor(OnDiskGraphIndex rawGraph)
-    {
-        // cache edges around the entry point
-        // we can easily hold 1% of the edges in memory for typical index sizes, but
-        // there is a lot of redundancy in the nodes we observe in practice around the entry point
-        // (only 10%-20% are unique), so use 5% as our target.
-        //
-        // 32**3 = 32k, which would be 4MB if all the nodes are unique, so 3 levels deep is a safe upper bound
-        int distance = min(logBaseX(0.05d * rawGraph.size(), rawGraph.maxDegree()), 3);
-        var result = new CachingGraphIndex(rawGraph, distance);
-        logger.debug("Cached {}@{} to distance {} in {}B",
-                     this, graphHandle.path(), distance, result.ramBytesUsed());
-        return result;
-    }
-
-    private static int logBaseX(double val, double base) {
-        if (base <= 1.0d || val <= 1.0d)
-            return 0;
-        return (int)Math.floor(Math.log(val) / Math.log(base));
-    }
-
     public long ramBytesUsed()
     {
         return graph.ramBytesUsed();

src/java/org/apache/cassandra/index/sai/disk/vector/CassandraOnHeapGraph.java

@@ -43,11 +43,11 @@
 import io.github.jbellis.jvector.graph.GraphSearcher;
 import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
 import io.github.jbellis.jvector.graph.SearchResult;
-import io.github.jbellis.jvector.graph.disk.Feature;
-import io.github.jbellis.jvector.graph.disk.FeatureId;
-import io.github.jbellis.jvector.graph.disk.InlineVectors;
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter;
 import io.github.jbellis.jvector.graph.disk.OrdinalMapper;
+import io.github.jbellis.jvector.graph.disk.feature.Feature;
+import io.github.jbellis.jvector.graph.disk.feature.FeatureId;
+import io.github.jbellis.jvector.graph.disk.feature.InlineVectors;
 import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider;
 import io.github.jbellis.jvector.quantization.BinaryQuantization;
 import io.github.jbellis.jvector.quantization.CompressedVectors;
@@ -157,8 +157,9 @@ public CassandraOnHeapGraph(IndexContext context, boolean forSearching, Memtable
                                         similarityFunction,
                                         indexConfig.getAnnMaxDegree(),
                                         indexConfig.getConstructionBeamWidth(),
-                                        1.0f, // no overflow means add will be a bit slower but flush will be faster
-                                        dimension > 3 ? 1.2f : 2.0f);
+                                        indexConfig.getNeighborhoodOverflow(1.0f), // no overflow means add will be a bit slower but flush will be faster
+                                        indexConfig.getAlpha(dimension > 3 ? 1.2f : 2.0f),
+                                        indexConfig.isHierarchyEnabled());
         searchers = ThreadLocal.withInitial(() -> new GraphSearcherAccessManager(new GraphSearcher(builder.getGraph())));
     }
 

src/java/org/apache/cassandra/index/sai/disk/vector/CompactionGraph.java

@@ -39,10 +39,10 @@
 
 import io.github.jbellis.jvector.graph.GraphIndexBuilder;
 import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
-import io.github.jbellis.jvector.graph.disk.Feature;
-import io.github.jbellis.jvector.graph.disk.FeatureId;
-import io.github.jbellis.jvector.graph.disk.FusedADC;
-import io.github.jbellis.jvector.graph.disk.InlineVectors;
+import io.github.jbellis.jvector.graph.disk.feature.Feature;
+import io.github.jbellis.jvector.graph.disk.feature.FeatureId;
+import io.github.jbellis.jvector.graph.disk.feature.FusedADC;
+import io.github.jbellis.jvector.graph.disk.feature.InlineVectors;
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter;
 import io.github.jbellis.jvector.graph.disk.OrdinalMapper;
@@ -204,8 +204,9 @@ else if (compressor instanceof BinaryQuantization)
                                         dimension,
                                         indexConfig.getAnnMaxDegree(),
                                         indexConfig.getConstructionBeamWidth(),
-                                        1.2f,
-                                        dimension > 3 ? 1.2f : 1.4f,
+                                        indexConfig.getNeighborhoodOverflow(1.2f),
+                                        indexConfig.getAlpha(dimension > 3 ? 1.2f : 1.4f),
+                                        indexConfig.isHierarchyEnabled(),
                                         compactionSimdPool, compactionFjp);
 
         termsFile = perIndexComponents.addOrGet(IndexComponentType.TERMS_DATA).file();
@@ -316,8 +317,7 @@ public InsertionResult maybeAddVector(ByteBuffer term, int segmentRowId) throws
                     compressedVectors = new MutablePQVectors((ProductQuantization) compressor);
                     compactionFjp.submit(() -> {
                         IntStream.range(0, encodedVectorCount)
-                                 // FIXME parallel is disabled until 4.0.0 beta2 (encodeAndSet is not threadsafe before then)
-//                                 .parallel()
+                                 .parallel()
                                  .forEach(i -> {
                                      var v = vectorsByOrdinal.get(i);
                                      if (v == null)

test/unit/org/apache/cassandra/index/sai/cql/VectorCompactionTest.java

@@ -77,7 +77,9 @@ public void testPQRefine()
             flush();
         }
 
-        CompactionGraph.PQ_TRAINING_SIZE = 2 * MIN_PQ_ROWS;
+//        TODO deterimine proper design for PQ training on vectors of dimension < 100.
+//        see https://github.com/riptano/cndb/issues/13630
+//        CompactionGraph.PQ_TRAINING_SIZE = 2 * MIN_PQ_ROWS;
         compact();
 
         // Confirm we can query the data with reasonable recall

test/unit/org/apache/cassandra/index/sai/cql/VectorTester.java

@@ -84,7 +84,8 @@ public static double rawIndexedRecall(Collection<float[]> rawVectors, float[] ra
                                                  16,
                                                  100,
                                                  1.2f,
-                                                 1.4f);
+                                                 1.4f,
+                                                 false);
 
         for (float[] raw : rawVectors)
         {