Dynamically skip sharding L0 when SAI Vector index present

This is a partial solution to IllegalStateException thrown by VectorPostings. It works by using a single shard at L0 when a vector index is present. As noted in the jira ticket, there are edge cases that may still produce errors, notably the case where there are multiple data directories.
The key trade offs here are related to the time complexity for search. Since graph search is log(n), and searching m graphs is m * log(n), we see better search performance by building bigger graphs which is essentially log(m * n). We could pre-shard, which comes at a cost of increased search time complexity.

patch by Michael Marshall,Dmitry Konstantinov; reviewed by Caleb Rackliffe,Dmitry Konstantinov,Michael Semb Wever for CASSANDRA-19661

Co-authored-by: Michael Marshall <mmarshall@apache.org>
Co-authored-by: Dmitry Konstantinov <netudima@gmail.com>

Author: michaeljmarshall

CHANGES.txt

@@ -1,4 +1,5 @@
 5.0.7
+ * Dynamically skip sharding L0 when SAI Vector index present (CASSANDRA-19661)
  * Optionally force IndexStatusManager to use the optimized index status format (CASSANDRA-21132) 
  * No need to evict already prepared statements, as it creates a race condition between multiple threads (CASSANDRA-17401)
  * Upgrade logback version to 1.5.18 and slf4j dependencies to 2.0.17 (CASSANDRA-21137)

src/java/org/apache/cassandra/db/compaction/UnifiedCompactionStrategy.java

@@ -270,7 +270,9 @@ public SSTableMultiWriter createSSTableMultiWriter(Descriptor descriptor,
     {
         ShardManager shardManager = getShardManager();
         double flushDensity = cfs.metric.flushSizeOnDisk.get() * shardManager.shardSetCoverage() / shardManager.localSpaceCoverage();
-        ShardTracker boundaries = shardManager.boundaries(controller.getNumShards(flushDensity));
+        boolean supportsSharding = sstableLevel > 0 || indexGroups.stream().allMatch(Index.Group::supportsL0Shards);
+        int numShards = supportsSharding ? controller.getNumShards(flushDensity) : 1;
+        ShardTracker boundaries = shardManager.boundaries(numShards);
         return new ShardedMultiWriter(cfs,
                                       descriptor,
                                       keyCount,

src/java/org/apache/cassandra/index/Index.java

@@ -890,6 +890,15 @@ default boolean validateSSTableAttachedIndexes(Collection<SSTableReader> sstable
         {
             return true;
         }
+
+        /**
+         * Whether this index group supports sharding when flushing memtables, e.g. level 0 of UCS.
+         * @return true iff all indexes in the group support L0 sharding.
+         */
+        default boolean supportsL0Shards()
+        {
+            return true;
+        }
     }
 
     /**

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

@@ -703,6 +703,15 @@ public BooleanSupplier isIndexValid()
         return () -> valid;
     }
 
+    /**
+     * Vector indexes do not supporrt L0 shards due to the cost associated with resharding at flush time.
+     * @return true iff the index supports sharding at L0.
+     */
+    public boolean supportsL0Shards()
+    {
+        return !indexTermType.isVector();
+    }
+
     public boolean hasClustering()
     {
         return baseCfs.getComparator().size() > 0;

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

@@ -382,6 +382,17 @@ else if (throwOnIncomplete)
         return complete;
     }
 
+    @Override
+    public boolean supportsL0Shards()
+    {
+        for (StorageAttachedIndex index : indexes)
+            if (!index.supportsL0Shards())
+                return false;
+
+        // All indexes must support L0 sharding for the flush to shard at L0
+        return true;
+    }
+
     /**
      * open index files by checking number of {@link SSTableContext} and {@link SSTableIndex},
      * so transient open files during validation and files that are still open for in-flight requests will not be tracked.

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

@@ -511,6 +511,29 @@ public void multipleNonAnnSegmentsTest() throws Throwable
         }
     }
 
+    @Test
+    public void flushSuccessfullyVectorIndexToShardedSSTable()
+    {
+        // UCS is configured to use 2 static shards
+        createTable(String.format("CREATE TABLE %%s (k int PRIMARY KEY, v vector<float, %d>) WITH compaction = {\n" +
+                                  "    'class': 'UnifiedCompactionStrategy',\n" +
+                                  "    'base_shard_count': '2',\n" +
+                                  "    'min_sstable_size' : '0MiB', \n" +
+                                  "    'sstable_growth' : '1'\n" +
+                                  "}", word2vec.dimension()));
+        createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
+
+
+        int vectorCount = 100;
+        List<float[]> vectors = IntStream.range(0, vectorCount).mapToObj(s -> randomVector()).collect(Collectors.toList());
+
+        int pk = 0;
+        for (float[] vector : vectors)
+            execute("INSERT INTO %s (k, v) VALUES (?," + vectorString(vector) + ")", pk++);
+
+        flush();
+    }
+
     private UntypedResultSet search(String stringValue, float[] queryVector, int limit)
     {
         UntypedResultSet result = execute("SELECT * FROM %s WHERE str_val = '" + stringValue + "' ORDER BY val ann of " + Arrays.toString(queryVector) + " LIMIT " + limit);