Sai prefixv2

src/java/org/apache/cassandra/db/tries/InMemoryTrie.java

@@ -22,6 +22,7 @@
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 import java.util.concurrent.atomic.AtomicReferenceArray;
+import java.util.function.IntPredicate;
 
 import com.google.common.annotations.VisibleForTesting;
 
@@ -809,6 +810,17 @@ public interface UpsertTransformer<T, U>
          * @return The combined value to use. Cannot be null.
          */
         T apply(T existing, U update);
+
+        /**
+         * Called when intermediate-node content is applied during a recursive put driven by an insertion policy
+         * (see {@link #putSingleton(ByteComparable, Object, UpsertTransformer, boolean, IntPredicate)}).
+         * The default implementation delegates to {@link #apply(Object, Object)}; transformers that need to
+         * distinguish intermediate (prefix) applications from terminal (exact) ones should override it.
+         */
+        default T applyIntermediate(T existing, U update)
+        {
+            return apply(existing, update);
+        }
     }
 
     /**
@@ -899,6 +911,69 @@ public <R> void putRecursive(ByteComparable key, R value, final UpsertTransforme
             root = newRoot;
     }
 
+    /**
+     * A version of {@link #putSingleton(ByteComparable, Object, UpsertTransformer, boolean)} which, in addition to
+     * applying the value at the terminal node, also applies it as intermediate content at every depth selected by
+     * {@code accumulateIntermediateAtDepth}. Intermediate applications go through
+     * {@link UpsertTransformer#applyIntermediate}, while the terminal application uses {@link UpsertTransformer#apply}.
+     * <p>
+     * Depth is 1-based on the key bytes (the first byte is depth 1); the empty prefix (depth 0, the root) is never
+     * accumulated. A null policy is equivalent to {@link #putSingleton(ByteComparable, Object, UpsertTransformer, boolean)}.
+     */
+    public <R> void putSingleton(ByteComparable key,
+                                 R value,
+                                 UpsertTransformer<T, ? super R> transformer,
+                                 boolean useRecursive,
+                                 IntPredicate accumulateIntermediateAtDepth) throws SpaceExhaustedException
+    {
+        if (accumulateIntermediateAtDepth == null)
+        {
+            putSingleton(key, value, transformer, useRecursive);
+            return;
+        }
+        putRecursiveWithPolicy(key, value, transformer, accumulateIntermediateAtDepth);
+    }
+
+    @SuppressWarnings("unchecked")
+    public <R> void putRecursiveWithPolicy(ByteComparable key,
+                                           R value,
+                                           UpsertTransformer<T, ? super R> transformer,
+                                           IntPredicate accumulateIntermediateAtDepth) throws SpaceExhaustedException
+    {
+        int newRoot = putRecursiveWithPolicy(root, key.asComparableBytes(BYTE_COMPARABLE_VERSION), value,
+                                             (UpsertTransformer<T, R>) transformer, accumulateIntermediateAtDepth, 0);
+        if (newRoot != root)
+            root = newRoot;
+    }
+
+    private <R> int putRecursiveWithPolicy(int node, ByteSource key, R value, final UpsertTransformer<T, R> transformer,
+                                           IntPredicate accumulateIntermediateAtDepth, int depth) throws SpaceExhaustedException
+    {
+        int transition = key.next();
+        if (transition == ByteSource.END_OF_STREAM)
+            return applyContent(node, value, transformer, false);
+
+        boolean appliedHere = false;
+        if (depth >= 1 && accumulateIntermediateAtDepth.test(depth))
+        {
+            node = applyContent(node, value, transformer, true);
+            appliedHere = true;
+        }
+
+        int child = getChild(node, transition);
+
+        int newChild = putRecursiveWithPolicy(child, key, value, transformer, accumulateIntermediateAtDepth, depth + 1);
+        if (newChild == child && !appliedHere)
+            return node;
+
+        int skippedContent = followContentTransition(node);
+        int attachedChild = !isNull(skippedContent)
+                            ? attachChild(skippedContent, transition, newChild)  // Single path, no copying required
+                            : expandOrCreateChainNode(transition, newChild);
+
+        return preserveContent(node, skippedContent, attachedChild);
+    }
+
     private <R> int putRecursive(int node, ByteSource key, R value, final UpsertTransformer<T, R> transformer) throws SpaceExhaustedException
     {
         int transition = key.next();
@@ -920,25 +995,35 @@ private <R> int putRecursive(int node, ByteSource key, R value, final UpsertTran
     }
 
     private <R> int applyContent(int node, R value, UpsertTransformer<T, R> transformer) throws SpaceExhaustedException
+    {
+        return applyContent(node, value, transformer, false);
+    }
+
+    private <R> int applyContent(int node, R value, UpsertTransformer<T, R> transformer, boolean intermediate) throws SpaceExhaustedException
     {
         if (isNull(node))
-            return ~addContent(transformer.apply(null, value));
+            return ~addContent(combine(transformer, null, value, intermediate));
 
         if (isLeaf(node))
         {
             int contentIndex = ~node;
-            setContent(contentIndex, transformer.apply(getContent(contentIndex), value));
+            setContent(contentIndex, combine(transformer, getContent(contentIndex), value, intermediate));
             return node;
         }
 
         if (offset(node) == PREFIX_OFFSET)
         {
             int contentIndex = getInt(node + PREFIX_CONTENT_OFFSET);
-            setContent(contentIndex, transformer.apply(getContent(contentIndex), value));
+            setContent(contentIndex, combine(transformer, getContent(contentIndex), value, intermediate));
             return node;
         }
         else
-            return createPrefixNode(addContent(transformer.apply(null, value)), node, false);
+            return createPrefixNode(addContent(combine(transformer, null, value, intermediate)), node, false);
+    }
+
+    private <R> T combine(UpsertTransformer<T, R> transformer, T existing, R value, boolean intermediate)
+    {
+        return intermediate ? transformer.applyIntermediate(existing, value) : transformer.apply(existing, value);
     }
 
     /**

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

@@ -157,6 +157,7 @@ public class StorageAttachedIndex implements Index
                                                                      IndexWriterConfig.CONSTRUCTION_BEAM_WIDTH,
                                                                      IndexWriterConfig.SIMILARITY_FUNCTION,
                                                                      IndexWriterConfig.OPTIMIZE_FOR,
+                                                                     IndexWriterConfig.ENABLE_LITERAL_PREFIX_SAI,
                                                                      NonTokenizingOptions.CASE_SENSITIVE,
                                                                      NonTokenizingOptions.NORMALIZE,
                                                                      NonTokenizingOptions.ASCII);
@@ -279,6 +280,17 @@ public static Map<String, String> validateOptions(Map<String, String> options, T
         AbstractAnalyzer.fromOptions(indexTermType, analysisOptions);
         IndexWriterConfig config = IndexWriterConfig.fromOptions(null, indexTermType, options);
 
+        if (options.containsKey(IndexWriterConfig.ENABLE_LITERAL_PREFIX_SAI))
+        {
+            String val = options.get(IndexWriterConfig.ENABLE_LITERAL_PREFIX_SAI);
+            if (!"true".equalsIgnoreCase(val) && !"false".equalsIgnoreCase(val))
+                throw new InvalidRequestException(
+                    IndexWriterConfig.ENABLE_LITERAL_PREFIX_SAI + " must be 'true' or 'false', got: " + val);
+            if (!indexTermType.isLiteral())
+                throw new InvalidRequestException(
+                    IndexWriterConfig.ENABLE_LITERAL_PREFIX_SAI + " is only supported on string/literal columns");
+        }
+
         // If we are indexing map entries we need to validate the subtypes
         if (indexTermType.isComposite())
         {
@@ -450,7 +462,16 @@ public boolean dependsOn(ColumnMetadata column)
     @Override
     public boolean supportsExpression(ColumnMetadata column, Operator operator)
     {
-        return dependsOn(column) && indexTermType.supports(operator);
+        if (!dependsOn(column))
+            return false;
+
+        // LIKE is initially parsed as the generic operator (the specific variant is resolved from the bound value),
+        // so advertise support for both the generic LIKE and LIKE_PREFIX. Only prefix-enabled literal indexes qualify;
+        // other LIKE variants (suffix/contains/matches) are rejected at execution time.
+        if (operator == Operator.LIKE || operator == Operator.LIKE_PREFIX)
+            return indexTermType.isLiteral() && indexWriterConfig.isLiteralPrefixEnabled();
+
+        return indexTermType.supports(operator);
     }
 
     @Override

src/java/org/apache/cassandra/index/sai/disk/RowMapping.java

@@ -24,11 +24,13 @@
 
 import com.carrotsearch.hppc.LongArrayList;
 
+import org.apache.cassandra.config.CassandraRelevantProperties;
 import org.apache.cassandra.db.compaction.OperationType;
 import org.apache.cassandra.db.rows.RangeTombstoneMarker;
 import org.apache.cassandra.db.rows.Row;
 import org.apache.cassandra.db.tries.InMemoryTrie;
 import org.apache.cassandra.index.sai.memory.MemtableIndex;
+import org.apache.cassandra.index.sai.memory.SectionedPrimaryKeys;
 import org.apache.cassandra.index.sai.utils.PrimaryKey;
 import org.apache.cassandra.index.sai.utils.PrimaryKeys;
 import org.apache.cassandra.io.compress.BufferType;
@@ -54,6 +56,9 @@ public class RowMapping
         @Override
         public Iterator<Pair<ByteComparable, LongArrayList>> merge(MemtableIndex index) { return Collections.emptyIterator(); }
 
+        @Override
+        public Iterator<Pair<ByteComparable, LongArrayList>> mergeV2(MemtableIndex index) { return Collections.emptyIterator(); }
+
         @Override
         public void complete() {}
 
@@ -136,6 +141,88 @@ protected Pair<ByteComparable, LongArrayList> computeNext()
         };
     }
 
+    /**
+     * V2 version of {@link #merge} for indexes with the SAI prefix feature enabled.
+     * <p>
+     * Iterates the sectioned payload of a {@link org.apache.cassandra.index.sai.memory.TrieMemoryIndex}
+     * (via {@link MemtableIndex#iteratorSectioned()}) and produces a {@link LongArrayList} per term in
+     * V2 wire format:
+     * <pre>
+     *   [exactCount, totalCount, exact_rowId_0, ..., prefix_rowId_0, ...]
+     * </pre>
+     * consumed verbatim by
+     * {@link org.apache.cassandra.index.sai.disk.v1.trie.LiteralIndexWriter#writeCompleteSegment}.
+     * <p>
+     * <b>Early prefix pruning:</b> if a node's raw prefix primary-key count is below
+     * {@link CassandraRelevantProperties#SAI_MINIMUM_POSTINGS_LEAVES}, the prefix section is
+     * dropped before any {@link #rowMapping} lookups are performed. The in-memory count is always
+     * >= the resolved row-id count (deleted rows are absent from the row mapping), so this is a safe
+     * over-approximation of the check {@link org.apache.cassandra.index.sai.disk.v1.trie.LiteralIndexWriter}
+     * would perform anyway.
+     * <p>
+     * Nodes whose exact and (pruned) prefix sections are both empty after resolution are skipped.
+     *
+     * @param index a {@link MemtableIndex} whose backing memory index was constructed with prefix
+     * accumulation enabled
+     * @return iterator of (term, v2-posting-list) pairs in trie-sorted order
+     */
+    public Iterator<Pair<ByteComparable, LongArrayList>> mergeV2(MemtableIndex index)
+    {
+        assert complete : "RowMapping is not built.";
+
+        final int minimumLeaves = CassandraRelevantProperties.SAI_MINIMUM_POSTINGS_LEAVES.getInt();
+        Iterator<Pair<ByteComparable, SectionedPrimaryKeys>> iterator = index.iteratorSectioned();
+        return new AbstractGuavaIterator<>()
+        {
+            @Override
+            protected Pair<ByteComparable, LongArrayList> computeNext()
+            {
+                while (iterator.hasNext())
+                {
+                    Pair<ByteComparable, SectionedPrimaryKeys> pair = iterator.next();
+                    SectionedPrimaryKeys sectioned = pair.right;
+
+                    LongArrayList exactIds = resolveRowIds(sectioned.exact());
+
+                    LongArrayList prefixIds;
+                    if (sectioned.prefix().size() >= minimumLeaves)
+                        prefixIds = resolveRowIds(sectioned.prefix());
+                    else
+                        prefixIds = new LongArrayList();
+
+                    if (exactIds.isEmpty() && prefixIds.isEmpty())
+                        continue;
+
+                    int exactCount = exactIds.size();
+                    int totalCount = exactCount + prefixIds.size();
+
+                    LongArrayList v2 = new LongArrayList(2 + totalCount);
+                    v2.add(exactCount);
+                    v2.add(totalCount);
+                    for (int i = 0; i < exactIds.size(); i++)
+                        v2.add(exactIds.get(i));
+                    for (int i = 0; i < prefixIds.size(); i++)
+                        v2.add(prefixIds.get(i));
+
+                    return Pair.create(pair.left, v2);
+                }
+                return endOfData();
+            }
+
+            private LongArrayList resolveRowIds(PrimaryKeys keys)
+            {
+                LongArrayList ids = new LongArrayList();
+                for (PrimaryKey pk : keys)
+                {
+                    Long sstableRowId = rowMapping.get(pk);
+                    if (sstableRowId != null)
+                        ids.add((long) sstableRowId);
+                }
+                return ids;
+            }
+        };
+    }
+
     /**
      * Complete building in memory RowMapping, mark it as immutable.
      */

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

@@ -52,13 +52,15 @@ public class IndexWriterConfig
 
     public static final String OPTIMIZE_FOR = "optimize_for";
     private static final OptimizeFor DEFAULT_OPTIMIZE_FOR = OptimizeFor.LATENCY;
+
+    public static final String ENABLE_LITERAL_PREFIX_SAI = "enable_literal_prefix_sai";
     private static final String validOptimizeFor = Arrays.stream(OptimizeFor.values())
                                                          .map(Enum::name)
                                                          .collect(Collectors.joining(", "));
 
     public static final int MAX_TOP_K = SAI_VECTOR_SEARCH_MAX_TOP_K.getInt();
 
-    private static final IndexWriterConfig EMPTY_CONFIG = new IndexWriterConfig(-1, -1, null, null);
+    private static final IndexWriterConfig EMPTY_CONFIG = new IndexWriterConfig(-1, -1, null, null, false);
 
     // The maximum number of outgoing connections a node can have in a graph.
     private final int maximumNodeConnections;
@@ -71,15 +73,19 @@ public class IndexWriterConfig
 
     private final OptimizeFor optimizeFor;
 
+    private final boolean literalPrefixEnabled;
+
     public IndexWriterConfig(int maximumNodeConnections,
                              int constructionBeamWidth,
                              VectorSimilarityFunction similarityFunction,
-                             OptimizeFor optimizerFor)
+                             OptimizeFor optimizerFor,
+                             boolean literalPrefixEnabled)
     {
         this.maximumNodeConnections = maximumNodeConnections;
         this.constructionBeamWidth = constructionBeamWidth;
         this.similarityFunction = similarityFunction;
         this.optimizeFor = optimizerFor;
+        this.literalPrefixEnabled = literalPrefixEnabled;
     }
 
     public int getMaximumNodeConnections()
@@ -102,6 +108,11 @@ public OptimizeFor getOptimizeFor()
         return optimizeFor;
     }
 
+    public boolean isLiteralPrefixEnabled()
+    {
+        return literalPrefixEnabled;
+    }
+
     public static IndexWriterConfig fromOptions(String indexName, IndexTermType indexTermType, Map<String, String> options)
     {
         int maximumNodeConnections = DEFAULT_MAXIMUM_NODE_CONNECTIONS;
@@ -178,7 +189,8 @@ public static IndexWriterConfig fromOptions(String indexName, IndexTermType inde
                 }
             }
         }
-        return new IndexWriterConfig(maximumNodeConnections, queueSize, similarityFunction, optimizeFor);
+        boolean literalPrefixEnabled = "true".equalsIgnoreCase(options.get(ENABLE_LITERAL_PREFIX_SAI));
+        return new IndexWriterConfig(maximumNodeConnections, queueSize, similarityFunction, optimizeFor, literalPrefixEnabled);
     }
 
     public static IndexWriterConfig emptyConfig()