CNDB-13462: Add support for updating terms in SAI's TrieMemoryIndex (#1200)

Add support for updating a row's terms in the `TrieMemoryIndex`.

There are two general cases for row to term mappings: 1:one and 1:many.

In the 1 to 1 case, update is trivial. We check before insertion that we
have a new value, then we insert the new value and remove the old value.

In the 1 to many case, the solution is somewhat bespoke, but is well
documnted and encapsulated. The main challenge
comes in the form of when to add/remove terms without increasing the
overall time complexity of the update operation. In this design, a
majority of the complexity is in the PrimaryKeys object. There, we
monitor seen keys and use whether we've seen the key to decide whether
to reset the frequency counter and then subsequently to decide whether
to leave the key or remove it (the term could be in both the old and the
new values).

This whole design hinges on the fact that we synchronize additions to
the trie.

The one drawback in terms of performance is that collections are added
and removed within the sync block. This change is necessary to deal with
encoding nuances unless we want to largely refactor the class.

The motivation for needing updates is to ensure that we have unique
primary keys in the postings list for scalar values and for a given map
key. Supporting updates means that the order of terms within an index is
internally consistent.

Author: michaeljmarshall

src/java/org/apache/cassandra/db/memtable/AbstractAllocatorMemtable.java

@@ -183,12 +183,12 @@ public void addMemoryUsageTo(MemoryUsage stats)
 
     public void markExtraOnHeapUsed(long additionalSpace, OpOrder.Group opGroup)
     {
-        getAllocator().onHeap().allocate(additionalSpace, opGroup);
+        getAllocator().onHeap().adjust(additionalSpace, opGroup);
     }
 
     public void markExtraOffHeapUsed(long additionalSpace, OpOrder.Group opGroup)
     {
-        getAllocator().offHeap().allocate(additionalSpace, opGroup);
+        getAllocator().offHeap().adjust(additionalSpace, opGroup);
     }
 
     @Override

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

@@ -402,22 +402,26 @@ private boolean validateCumulativeAnalyzedTermLimit(DecoratedKey key, AbstractAn
 
     public void update(DecoratedKey key, Row oldRow, Row newRow, Memtable memtable, OpOrder.Group opGroup)
     {
-        if (!isVector())
-        {
-            index(key, newRow, memtable, opGroup);
-            return;
-        }
-
         MemtableIndex target = liveMemtables.get(memtable);
         if (target == null)
             return;
-        // Use 0 for nowInSecs to get the value from the oldRow regardless of its liveness status. To get to this point,
+
+        // Use 0 for nowInSecs to get the value(s) from the oldRow regardless of its liveness status. To get to this point,
         // C* has already determined this is the current represntation of the oldRow in the memtable, and that means
         // we need to add the newValue to the index and remove the oldValue from it, even if it has already expired via
         // TTL.
-        ByteBuffer oldValue = getValueOf(key, oldRow, 0);
-        ByteBuffer newValue = getValueOf(key, newRow, FBUtilities.nowInSeconds());
-        target.update(key, oldRow.clustering(), oldValue, newValue, memtable, opGroup);
+        if (isNonFrozenCollection())
+        {
+            Iterator<ByteBuffer> oldValues = getValuesOf(oldRow, 0);
+            Iterator<ByteBuffer> newValues = getValuesOf(newRow, FBUtilities.nowInSeconds());
+            target.update(key, oldRow.clustering(), oldValues, newValues, memtable, opGroup);
+        }
+        else
+        {
+            ByteBuffer oldValue = getValueOf(key, oldRow, 0);
+            ByteBuffer newValue = getValueOf(key, newRow, FBUtilities.nowInSeconds());
+            target.update(key, oldRow.clustering(), oldValue, newValue, memtable, opGroup);
+        }
     }
 
     public void renewMemtable(Memtable renewed)
@@ -720,7 +724,10 @@ public boolean supports(Operator op)
         if (op.isLike() || op == Operator.LIKE) return false;
         // Analyzed columns store the indexed result, so we are unable to compute raw equality.
         // The only supported operators are ANALYZER_MATCHES and BM25.
-        if (op == Operator.ANALYZER_MATCHES || op == Operator.BM25) return isAnalyzed;
+        if (op == Operator.ANALYZER_MATCHES) return isAnalyzed;
+        // BM25 frequency calculations only work on non-collection columns because it assumes a 1:1 mapping from PrK
+        // to frequency, but collections have mulitple documents.
+        if (op == Operator.BM25) return isAnalyzed && !isCollection();
 
         // If the column is analyzed and the operator is EQ, we need to check if the analyzer supports it.
         if (op == Operator.EQ && isAnalyzed && !analyzerFactory.supportsEquals())

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

@@ -174,6 +174,12 @@ public void update(DecoratedKey key, Clustering clustering, ByteBuffer oldValue,
         }
     }
 
+    @Override
+    public void update(DecoratedKey key, Clustering clustering, Iterator<ByteBuffer> oldValues, Iterator<ByteBuffer> newValues, Memtable memtable, OpOrder.Group opGroup)
+    {
+        throw new UnsupportedOperationException("Vector index does not support multi-value updates");
+    }
+
     private void updateKeyBounds(PrimaryKey primaryKey) {
         if (minimumKey == null)
             minimumKey = primaryKey;

src/java/org/apache/cassandra/index/sai/memory/MemoryIndex.java

@@ -52,6 +52,28 @@ public abstract void add(DecoratedKey key,
                              LongConsumer onHeapAllocationsTracker,
                              LongConsumer offHeapAllocationsTracker);
 
+    /**
+     * Update the index value for the given key and clustering by removing the old value and adding the new value.
+     * This is meant to be used when the indexed column is any type other than a non-frozen collection.
+     */
+    public abstract void update(DecoratedKey key,
+                                Clustering clustering,
+                                ByteBuffer oldValue,
+                                ByteBuffer newValue,
+                                LongConsumer onHeapAllocationsTracker,
+                                LongConsumer offHeapAllocationsTracker);
+
+    /**
+     * Update the index value for the given key and clustering by removing the old values and adding the new values.
+     * This is meant to be used when the indexed column is a non-frozen collection.
+     */
+    public abstract void update(DecoratedKey key,
+                                Clustering clustering,
+                                Iterator<ByteBuffer> oldValues,
+                                Iterator<ByteBuffer> newValues,
+                                LongConsumer onHeapAllocationsTracker,
+                                LongConsumer offHeapAllocationsTracker);
+
     public abstract CloseableIterator<PrimaryKeyWithSortKey> orderBy(Orderer orderer, Expression slice);
 
     public abstract KeyRangeIterator search(Expression expression, AbstractBounds<PartitionPosition> keyRange);

src/java/org/apache/cassandra/index/sai/memory/MemtableIndex.java

@@ -70,10 +70,8 @@ public interface MemtableIndex extends MemtableOrdering
 
     void index(DecoratedKey key, Clustering clustering, ByteBuffer value, Memtable memtable, OpOrder.Group opGroup);
 
-    default void update(DecoratedKey key, Clustering clustering, ByteBuffer oldValue, ByteBuffer newValue, Memtable memtable, OpOrder.Group opGroup)
-    {
-        throw new UnsupportedOperationException();
-    }
+    void update(DecoratedKey key, Clustering clustering, ByteBuffer oldValue, ByteBuffer newValue, Memtable memtable, OpOrder.Group opGroup);
+    void update(DecoratedKey key, Clustering clustering, Iterator<ByteBuffer> oldValues, Iterator<ByteBuffer> newValues, Memtable memtable, OpOrder.Group opGroup);
 
     KeyRangeIterator search(QueryContext queryContext, Expression expression, AbstractBounds<PartitionPosition> keyRange, int limit);