CNDB-13027 Fix V5VectorPostingsWriter to handle omitted ordinals corr…

…ectly
datastax · Feb 19, 2025 · 8947516 · 8947516
1 parent cf3e012
commit 8947516
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 2 deletions.
diff --git a/src/java/org/apache/cassandra/index/sai/disk/v5/V5VectorPostingsWriter.java b/src/java/org/apache/cassandra/index/sai/disk/v5/V5VectorPostingsWriter.java
@@ -281,7 +281,7 @@ public void writeGenericOrdinalToRowIdMapping(SequentialWriter writer,
             int postingListSize;
             if (originalOrdinal == OrdinalMapper.OMITTED)
             {
-                assert remappedPostings.structure == Structure.ONE_TO_MANY;
+                assert remappedPostings.structure == Structure.ZERO_OR_ONE_TO_MANY;
                 postingListSize = 0;
             }
             else
@@ -298,7 +298,7 @@ public void writeGenericOrdinalToRowIdMapping(SequentialWriter writer,
             int originalOrdinal = newToOldMapper.applyAsInt(i);
             if (originalOrdinal == OrdinalMapper.OMITTED)
             {
-                assert remappedPostings.structure == Structure.ONE_TO_MANY;
+                assert remappedPostings.structure == Structure.ZERO_OR_ONE_TO_MANY;
                 writer.writeInt(0);
                 continue;
             }
@@ -321,6 +321,9 @@ public void writeGenericRowIdMapping(SequentialWriter writer,
         var rowIdToOrdinalMap = new Int2IntHashMap(remappedPostings.maxNewOrdinal, 0.65f, OrdinalMapper.OMITTED);
         for (int i = 0; i <= remappedPostings.maxNewOrdinal; i++) {
             int ord = remappedPostings.ordinalMapper.newToOld(i);
+            if (ord == OrdinalMapper.OMITTED)
+                continue;
+
             var rowIds = postingsMap.get(vectorValues.getVector(ord)).getRowIds();
             for (int r = 0; r < rowIds.size(); r++)
             {

diff --git a/test/unit/org/apache/cassandra/index/sai/cql/VectorCompactionTest.java b/test/unit/org/apache/cassandra/index/sai/cql/VectorCompactionTest.java
@@ -110,6 +110,69 @@ public void testOneToManyCompactionTooManyHoles()
         testOneToManyCompactionHolesInternal(MIN_PQ_ROWS, sstables);
     }
 
+    @Test
+    public void testZeroOrOneToManyCompaction()
+    {
+        for (int sstables = 2; sstables <= 3; sstables++)
+        {
+            testZeroOrOneToManyCompactionInternal(10, sstables);
+            testZeroOrOneToManyCompactionInternal(MIN_PQ_ROWS, sstables);
+        }
+    }
+
+    public void testZeroOrOneToManyCompactionInternal(int vectorsPerSstable, int sstables)
+    {
+        createTable();
+        disableCompaction();
+
+        insertZeroOrOneToManyRows(vectorsPerSstable, sstables);
+
+        validateQueries();
+        compact();
+        validateQueries();
+    }
+
+    private void insertZeroOrOneToManyRows(int vectorsPerSstable, int sstables)
+    {
+        var R = getRandom();
+        double duplicateChance = R.nextDouble() * 0.2;
+        int j = 0;
+        boolean nullInserted = false;
+
+        for (int i = 0; i < sstables; i++)
+        {
+            var vectorsInserted = new ArrayList<Vector<Float>>();
+            var duplicateExists = false;
+            while (vectorsInserted.size() < vectorsPerSstable || !duplicateExists)
+            {
+                if (!nullInserted && vectorsInserted.size() == vectorsPerSstable/2)
+                {
+                    // Insert one null vector in the middle
+                    execute("INSERT INTO %s (pk, v) VALUES (?, null)", j++);
+                    nullInserted = true;
+                    continue;
+                }
+
+                Vector<Float> v;
+                if (R.nextDouble() < duplicateChance && !vectorsInserted.isEmpty())
+                {
+                    // insert a duplicate
+                    v = vectorsInserted.get(R.nextIntBetween(0, vectorsInserted.size() - 1));
+                    duplicateExists = true;
+                }
+                else
+                {
+                    // insert a new random vector
+                    v = randomVectorBoxed(2);
+                    vectorsInserted.add(v);
+                }
+                assert v != null;
+                execute("INSERT INTO %s (pk, v) VALUES (?, ?)", j++, v);
+            }
+            flush();
+        }
+    }
+
     public void testOneToManyCompactionHolesInternal(int vectorsPerSstable, int sstables)
     {
         createTable();