natas_hw
/
cassandra
mirror of https://github.com/apache/cassandra
0
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

optimizing write path

This commit is contained in:
Pranav Shenoy 2026-07-02 19:44:29 -07:00
parent 3a5152e004
commit bdcd52fd18
10 changed files with 498 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

87
src/java/org/apache/cassandra/index/sai/disk/RowMapping.java
View File

@ -24,11 +24,13 @@ import javax.annotation.concurrent.NotThreadSafe;
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 @@ public class RowMapping
};
}
/**
* 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.
*/

47
src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java
View File

@ -28,7 +28,6 @@ import com.google.common.base.Stopwatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.cassandra.config.CassandraRelevantProperties;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.index.sai.disk.PerColumnIndexWriter;
import org.apache.cassandra.index.sai.disk.RowMapping;
@ -36,14 +35,11 @@ import org.apache.cassandra.index.sai.disk.format.IndexComponent;
import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
import org.apache.cassandra.index.sai.disk.v1.bbtree.NumericIndexWriter;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentTrieBuffer;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentWriter;
import org.apache.cassandra.index.sai.disk.v1.trie.LiteralIndexWriter;
import org.apache.cassandra.index.sai.memory.MemtableIndex;
import org.apache.cassandra.index.sai.memory.MemtableTermsIterator;
import org.apache.cassandra.index.sai.metrics.IndexMetrics;
import org.apache.cassandra.index.sai.postings.PostingList;
import org.apache.cassandra.index.sai.utils.IndexEntry;
import org.apache.cassandra.index.sai.utils.IndexIdentifier;
import org.apache.cassandra.index.sai.utils.IndexTermType;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
@ -59,7 +55,6 @@ public class MemtableIndexWriter implements PerColumnIndexWriter
{
private static final Logger logger = LoggerFactory.getLogger(MemtableIndexWriter.class);
private static final int NO_ROWS = -1;
private static final int MAX_RECURSIVE_TERM_LENGTH = 128;
private final IndexDescriptor indexDescriptor;
private final IndexTermType indexTermType;
@ -186,26 +181,36 @@ public class MemtableIndexWriter implements PerColumnIndexWriter
{
SegmentMetadata.ComponentMetadataMap indexMetas;
long numRows;
// getMin/MaxSSTableRowId() are declared on MemtableTermsIterator, not the TermsIterator interface,
// so this must be typed as the concrete class.
MemtableTermsIterator boundsSource = terms; // V1 path uses the incoming iterator for row-id bounds
if (indexTermType.isLiteral() && literalPrefixEnabled)
{
int skip = CassandraRelevantProperties.SAI_POSTINGS_SKIP.getInt();
SegmentTrieBuffer buffer = new SegmentTrieBuffer(depth -> depth % skip == 0);
// Fast path: TrieMemoryIndex already accumulated prefix postings at intermediate nodes during
// insert (via putSingleton with the prefixAtDepth policy). Resolve PrimaryKey -> sstableRowId for
// both the exact and prefix sections directly, with no SegmentTrieBuffer rebuild.
final Iterator<Pair<ByteComparable, LongArrayList>> v2Iterator = rowMapping.mergeV2(memtable);
while (terms.hasNext())
if (v2Iterator.hasNext())
{
IndexEntry entry = terms.next();
try (PostingList postings = entry.postingList)
try (MemtableTermsIterator v2Terms = new MemtableTermsIterator(memtable.getMinTerm(),
memtable.getMaxTerm(),
v2Iterator,
true))
{
long rowId;
while ((rowId = postings.nextPosting()) != PostingList.END_OF_STREAM)
buffer.add(entry.term, MAX_RECURSIVE_TERM_LENGTH, (int) rowId);
LiteralIndexWriter writer = new LiteralIndexWriter(indexDescriptor, indexIdentifier);
indexMetas = writer.writeCompleteSegment(v2Terms, true);
numRows = writer.getNumberOfRows();
boundsSource = v2Terms;
}
}
LiteralIndexWriter writer = new LiteralIndexWriter(indexDescriptor, indexIdentifier);
indexMetas = writer.writeCompleteSegment(buffer.iterator(), true);
numRows = writer.getNumberOfRows();
else
{
// Every primary key was deleted before flush; emit empty metadata so numRows == 0 is handled below.
indexMetas = new SegmentMetadata.ComponentMetadataMap();
numRows = 0;
}
}
else
{
@ -218,8 +223,6 @@ public class MemtableIndexWriter implements PerColumnIndexWriter
numRows = writer.getNumberOfRows();
}
// If no rows were written we need to delete any created column index components
// so that the index is correctly identified as being empty (only having a completion marker)
if (numRows == 0)
{
indexDescriptor.deleteColumnIndex(indexTermType, indexIdentifier);
@ -229,9 +232,9 @@ public class MemtableIndexWriter implements PerColumnIndexWriter
// During index memtable flush, the data is sorted based on terms.
SegmentMetadata metadata = new SegmentMetadata(0,
numRows,
terms.getMinSSTableRowId(), terms.getMaxSSTableRowId(),
minKey, maxKey,
terms.getMinTerm(), terms.getMaxTerm(),
boundsSource.getMinSSTableRowId(), boundsSource.getMaxSSTableRowId(),
minKey, maxKey,
boundsSource.getMinTerm(), boundsSource.getMaxTerm(),
indexMetas);
try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerIndexOutput(IndexComponent.META, indexIdentifier)))

10
src/java/org/apache/cassandra/index/sai/memory/MemoryIndex.java
View File

@ -20,6 +20,7 @@ package org.apache.cassandra.index.sai.memory;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.Iterator;
import java.util.function.Function;
@ -65,6 +66,15 @@ public abstract class MemoryIndex implements MemtableOrdering
*/
public abstract Iterator<Pair<ByteComparable, PrimaryKeys>> iterator();
/**
* Returns a sectioned iterator for V2 prefix flush.
* Default returns an empty iterator; overridden in {@link TrieMemoryIndex}.
*/
public Iterator<Pair<ByteComparable, SectionedPrimaryKeys>> iteratorSectioned()
{
return Collections.emptyIterator();
}
public abstract SegmentMetadata.ComponentMetadataMap writeDirect(IndexDescriptor indexDescriptor,
IndexIdentifier indexIdentifier,
Function<PrimaryKey, Integer> postingTransformer) throws IOException;

5
src/java/org/apache/cassandra/index/sai/memory/MemtableIndex.java
View File

@ -113,6 +113,11 @@ public class MemtableIndex implements MemtableOrdering
return memoryIndex.iterator();
}
public Iterator<Pair<ByteComparable, SectionedPrimaryKeys>> iteratorSectioned()
{
return memoryIndex.iteratorSectioned();
}
public SegmentMetadata.ComponentMetadataMap writeDirect(IndexDescriptor indexDescriptor,
IndexIdentifier indexIdentifier,
Function<PrimaryKey, Integer> postingTransformer) throws IOException

36
src/java/org/apache/cassandra/index/sai/memory/MemtableTermsIterator.java
View File

@ -35,23 +35,40 @@ import org.apache.cassandra.utils.bytecomparable.ByteComparable;
*/
public class MemtableTermsIterator implements TermsIterator
{
// Number of header longs prepended to each V2 posting list: [exactCount, totalCount].
// Mirrors org.apache.cassandra.index.sai.disk.v1.segment.PackedLongValuesList#FILTER_TYPES
// (which is package-private and not accessible from this package).
private static final int V2_HEADER_SIZE = 2;
private final ByteBuffer minTerm;
private final ByteBuffer maxTerm;
private final Iterator<Pair<ByteComparable, LongArrayList>> iterator;
private final boolean isV2;
private Pair<ByteComparable, LongArrayList> current;
private long maxSSTableRowId = -1;
private long minSSTableRowId = Long.MAX_VALUE;
/** V1 constructor: LongArrayList contains raw sorted row IDs only. */
public MemtableTermsIterator(ByteBuffer minTerm,
ByteBuffer maxTerm,
Iterator<Pair<ByteComparable, LongArrayList>> iterator)
{
this(minTerm, maxTerm, iterator, false);
}
/** V2 constructor: LongArrayList contains [exactCount, totalCount, rowId...]. */
public MemtableTermsIterator(ByteBuffer minTerm,
ByteBuffer maxTerm,
Iterator<Pair<ByteComparable, LongArrayList>> iterator,
boolean isV2)
{
Preconditions.checkArgument(iterator != null);
this.minTerm = minTerm;
this.maxTerm = maxTerm;
this.iterator = iterator;
this.isV2 = isV2;
}
@Override
@ -98,11 +115,20 @@ public class MemtableTermsIterator implements TermsIterator
assert list.size() > 0;
final long minSegmentRowID = list.get(0);
final long maxSegmentRowID = list.get(list.size() - 1);
minSSTableRowId = Math.min(minSSTableRowId, minSegmentRowID);
maxSSTableRowId = Math.max(maxSSTableRowId, maxSegmentRowID);
if (isV2)
{
for (int i = V2_HEADER_SIZE; i < list.size(); i++)
{
long rowId = list.get(i);
minSSTableRowId = Math.min(minSSTableRowId, rowId);
maxSSTableRowId = Math.max(maxSSTableRowId, rowId);
}
}
else
{
minSSTableRowId = Math.min(minSSTableRowId, list.get(0));
maxSSTableRowId = Math.max(maxSSTableRowId, list.get(list.size() - 1));
}
final Iterator<LongCursor> it = list.iterator();

85
src/java/org/apache/cassandra/index/sai/memory/SectionedPrimaryKeys.java Normal file
View File

@ -0,0 +1,85 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.cassandra.index.sai.memory;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.PrimaryKeys;
/**
* Trie node payload for {@link TrieMemoryIndex} when the SAI prefix feature is enabled.
* <p>
* Holds two sets of {@link PrimaryKeys}:
* <ul>
* <li>{@link #exact} — primary keys of rows whose indexed term is exactly the term at this node
* (populated via {@link org.apache.cassandra.db.tries.InMemoryTrie.UpsertTransformer#apply}).</li>
* <li>{@link #prefix} — primary keys of rows in this node's subtree, accumulated at eligible
* intermediate depths via {@link org.apache.cassandra.db.tries.InMemoryTrie.UpsertTransformer#applyIntermediate}.
* This set is only populated when {@code prefixAtDepth.test(depth)} is true, following
* the same eligibility rule as {@link org.apache.cassandra.config.CassandraRelevantProperties#SAI_POSTINGS_SKIP}.
* </li>
* </ul>
*
* At flush time the {@code prefix} set is converted to on-disk combined postings sections.
* During in-memory search only {@code exact} is consulted.
*/
public class SectionedPrimaryKeys
{
private final PrimaryKeys exact = new PrimaryKeys();
private PrimaryKeys prefix; // lazily allocated; null when no intermediate accumulation
/** Add a primary key to the exact section (terminal-node insert). */
public long addExact(PrimaryKey key)
{
return exact.add(key);
}
/** Add a primary key to the prefix section (intermediate-node insert). */
public long addPrefix(PrimaryKey key)
{
if (prefix == null)
prefix = new PrimaryKeys();
return prefix.add(key);
}
/** Primary keys whose indexed term exactly equals this trie node's key. Never null. */
public PrimaryKeys exact()
{
return exact;
}
/**
* Primary keys accumulated from this subtree for the prefix combined section.
* Returns an empty {@link PrimaryKeys} when no intermediate accumulation has occurred.
*/
public PrimaryKeys prefix()
{
return prefix == null ? PrimaryKeys.EMPTY : prefix;
}
/** Returns true if neither exact nor prefix contains any keys. */
public boolean isEmpty()
{
return exact.isEmpty() && (prefix == null || prefix.isEmpty());
}
/** Approximate heap overhead of this object and its contents. */
public long unsharedHeapSize()
{
return exact.unsharedHeapSize() + (prefix == null ? 0 : prefix.unsharedHeapSize());
}
}

106
src/java/org/apache/cassandra/index/sai/memory/TrieMemoryIndex.java
View File

@ -28,16 +28,19 @@ import java.util.SortedSet;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.LongAdder;
import java.util.function.Function;
import java.util.function.IntPredicate;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.cassandra.config.CassandraRelevantProperties;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.memtable.TrieMemtable;
import org.apache.cassandra.db.tries.InMemoryTrie;
import org.apache.cassandra.db.tries.Trie;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
@ -66,8 +69,11 @@ public class TrieMemoryIndex extends MemoryIndex
private static final int MAX_RECURSIVE_KEY_LENGTH = 128;
private static final int MINIMUM_PRIORITY_QUEUE_SIZE = 128;
private final InMemoryTrie<PrimaryKeys> data;
private final PrimaryKeysReducer primaryKeysReducer;
private final InMemoryTrie<SectionedPrimaryKeys> data;
private final SectionedPrimaryKeysReducer sectionedReducer;
/** Non-null when prefix SAI is enabled; null for plain (V1) indexes. */
@Nullable
private final IntPredicate prefixAtDepth;
private ByteBuffer minTerm;
private ByteBuffer maxTerm;
@ -81,7 +87,19 @@ public class TrieMemoryIndex extends MemoryIndex
{
super(index);
this.data = new InMemoryTrie<>(TrieMemtable.BUFFER_TYPE);
this.primaryKeysReducer = new PrimaryKeysReducer();
boolean prefixEnabled = index.termType().isLiteral()
&& index.indexWriterConfig().isLiteralPrefixEnabled();
if (prefixEnabled)
{
int skip = CassandraRelevantProperties.SAI_POSTINGS_SKIP.getInt();
this.prefixAtDepth = depth -> depth % skip == 0;
}
else
{
this.prefixAtDepth = null;
}
this.sectionedReducer = new SectionedPrimaryKeysReducer();
}
/**
@ -100,7 +118,7 @@ public class TrieMemoryIndex extends MemoryIndex
: index.keyFactory().create(key);
final long initialSizeOnHeap = data.sizeOnHeap();
final long initialSizeOffHeap = data.sizeOffHeap();
final long reducerHeapSize = primaryKeysReducer.heapAllocations();
final long reducerHeapSize = sectionedReducer.heapAllocations();
if (index.hasAnalyzer())
{
@ -124,7 +142,7 @@ public class TrieMemoryIndex extends MemoryIndex
}
long onHeap = data.sizeOnHeap();
long offHeap = data.sizeOffHeap();
long heapAllocations = primaryKeysReducer.heapAllocations();
long heapAllocations = sectionedReducer.heapAllocations();
return (onHeap - initialSizeOnHeap) + (offHeap - initialSizeOffHeap) + (heapAllocations - reducerHeapSize);
}
@ -178,7 +196,7 @@ public class TrieMemoryIndex extends MemoryIndex
@Override
public Iterator<Pair<ByteComparable, PrimaryKeys>> iterator()
{
Iterator<Map.Entry<ByteComparable, PrimaryKeys>> iterator = data.entrySet().iterator();
Iterator<Map.Entry<ByteComparable, SectionedPrimaryKeys>> iterator = data.entrySet().iterator();
return new Iterator<>()
{
@Override
@ -190,7 +208,34 @@ public class TrieMemoryIndex extends MemoryIndex
@Override
public Pair<ByteComparable, PrimaryKeys> next()
{
Map.Entry<ByteComparable, PrimaryKeys> entry = iterator.next();
Map.Entry<ByteComparable, SectionedPrimaryKeys> entry = iterator.next();
return Pair.create(entry.getKey(), entry.getValue().exact());
}
};
}
/**
* Returns an iterator over every trie node that has any payload — both exact-only leaf
* nodes and intermediate nodes that accumulated prefix postings. Used exclusively by
* {@link org.apache.cassandra.index.sai.disk.RowMapping#mergeV2} during memtable flush
* when the SAI prefix feature is enabled.
*/
@Override
public Iterator<Pair<ByteComparable, SectionedPrimaryKeys>> iteratorSectioned()
{
Iterator<Map.Entry<ByteComparable, SectionedPrimaryKeys>> iterator = data.entrySet().iterator();
return new Iterator<>()
{
@Override
public boolean hasNext()
{
return iterator.hasNext();
}
@Override
public Pair<ByteComparable, SectionedPrimaryKeys> next()
{
Map.Entry<ByteComparable, SectionedPrimaryKeys> entry = iterator.next();
return Pair.create(entry.getKey(), entry.getValue());
}
};
@ -232,14 +277,8 @@ public class TrieMemoryIndex extends MemoryIndex
try
{
if (term.limit() <= MAX_RECURSIVE_KEY_LENGTH)
{
data.putRecursive(comparableBytes, primaryKey, primaryKeysReducer);
}
else
{
data.apply(Trie.singleton(comparableBytes, primaryKey), primaryKeysReducer);
}
boolean useRecursive = term.limit() <= MAX_RECURSIVE_KEY_LENGTH;
data.putSingleton(comparableBytes, primaryKey, sectionedReducer, useRecursive, prefixAtDepth);
}
catch (InMemoryTrie.SpaceExhaustedException e)
{
@ -265,9 +304,10 @@ public class TrieMemoryIndex extends MemoryIndex
{
ByteComparable comparableMatch = expression.lower() == null ? ByteComparable.EMPTY
: asComparableBytes(expression.lower().value.encoded);
PrimaryKeys primaryKeys = data.get(comparableMatch);
return primaryKeys == null ? KeyRangeIterator.empty()
: new FilteringInMemoryKeyRangeIterator(primaryKeys.keys(), keyRange);
SectionedPrimaryKeys sectioned = data.get(comparableMatch);
PrimaryKeys primaryKeys = sectioned == null ? null : sectioned.exact();
return primaryKeys == null || primaryKeys.isEmpty() ? KeyRangeIterator.empty()
: new FilteringInMemoryKeyRangeIterator(primaryKeys.keys(), keyRange);
}
@Override
@ -357,10 +397,10 @@ public class TrieMemoryIndex extends MemoryIndex
}
Collector cd = new Collector(keyRange, lastPriorityQueueSize.get());
Iterator<PrimaryKeys> values = data.subtrie(lowerBound, lowerInclusive, upperBound, upperInclusive).valueIterator();
Iterator<SectionedPrimaryKeys> values = data.subtrie(lowerBound, lowerInclusive, upperBound, upperInclusive).valueIterator();
while (values.hasNext())
cd.processContent(values.next());
cd.processContent(values.next().exact());
if (cd.mergedKeys.isEmpty())
return KeyRangeIterator.empty();
@ -380,10 +420,10 @@ public class TrieMemoryIndex extends MemoryIndex
ByteComparable upperBound = successor == null ? null : asComparableBytes(successor);
Collector cd = new Collector(keyRange, lastPriorityQueueSize.get());
Iterator<PrimaryKeys> values = data.subtrie(lowerBound, true, upperBound, false).valueIterator();
Iterator<SectionedPrimaryKeys> values = data.subtrie(lowerBound, true, upperBound, false).valueIterator();
while (values.hasNext())
cd.processContent(values.next());
cd.processContent(values.next().exact());
if (cd.mergedKeys.isEmpty())
return KeyRangeIterator.empty();
@ -412,19 +452,31 @@ public class TrieMemoryIndex extends MemoryIndex
return ByteBuffer.wrap(successor);
}
private static class PrimaryKeysReducer implements InMemoryTrie.UpsertTransformer<PrimaryKeys, PrimaryKey>
private static class SectionedPrimaryKeysReducer implements InMemoryTrie.UpsertTransformer<SectionedPrimaryKeys, PrimaryKey>
{
private final LongAdder heapAllocations = new LongAdder();
@Override
public PrimaryKeys apply(PrimaryKeys existing, PrimaryKey neww)
public SectionedPrimaryKeys apply(SectionedPrimaryKeys existing, PrimaryKey key)
{
if (existing == null)
{
existing = new PrimaryKeys();
existing = new SectionedPrimaryKeys();
heapAllocations.add(existing.unsharedHeapSize());
}
heapAllocations.add(existing.add(neww));
heapAllocations.add(existing.addExact(key));
return existing;
}
@Override
public SectionedPrimaryKeys applyIntermediate(SectionedPrimaryKeys existing, PrimaryKey key)
{
if (existing == null)
{
existing = new SectionedPrimaryKeys();
heapAllocations.add(existing.unsharedHeapSize());
}
heapAllocations.add(existing.addPrefix(key));
return existing;
}

10
src/java/org/apache/cassandra/index/sai/utils/PrimaryKeys.java
View File

@ -33,6 +33,16 @@ import org.apache.cassandra.utils.ObjectSizes;
@ThreadSafe
public class PrimaryKeys implements Iterable<PrimaryKey>
{
/** Immutable empty sentinel. */
public static final PrimaryKeys EMPTY = new PrimaryKeys()
{
@Override
public long add(PrimaryKey key)
{
throw new UnsupportedOperationException();
}
};
private static final long EMPTY_SIZE = ObjectSizes.measure(new PrimaryKeys());
// from https://github.com/gaul/java-collection-overhead
private static final long SET_ENTRY_OVERHEAD = 36;

38
test/unit/org/apache/cassandra/index/sai/disk/v1/segment/PrefixSectionOptimizationTest.java
View File

@ -26,6 +26,7 @@ import org.apache.cassandra.index.sai.SAITester;
import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig;
import static org.apache.cassandra.index.sai.SAITester.getRandom;
import static org.junit.Assert.assertEquals;
/**
* Tests for prefix search optimization when query prefix lands at different depths
@ -351,4 +352,41 @@ public class PrefixSectionOptimizationTest extends SAITester
assertRowCount(execute("SELECT * FROM %s WHERE value LIKE ? ALLOW FILTERING", "zzz%"), 80);
}
/**
* Verifies that the memtable prefix flush fast path (which resolves prefix postings directly from
* the {@link org.apache.cassandra.index.sai.memory.TrieMemoryIndex} via
* {@code RowMapping.mergeV2}, eliminating the intermediate {@code SegmentTrieBuffer} rebuild)
* produces query results identical to the pre-flush in-memory path for the same data.
*/
@Test
public void testMemtablePrefixFlushMatchesMemtablePath() throws Throwable
{
createTable("CREATE TABLE %s (pk int PRIMARY KEY, value text)");
createIndex(String.format("CREATE INDEX ON %%s(value) USING 'sai' WITH OPTIONS = {'%s': 'true'}",
IndexWriterConfig.ENABLE_LITERAL_PREFIX_SAI));
// Ten groups "prod00".."prod09", 30 rows each: "prod0X" reaches depth 6 (eligible for a prefix
// section) while each group individually is below the 64-row minimum, exercising both the
// exact and (pruned) prefix sections of the V2 flush path.
for (int i = 0; i < 300; i++)
execute("INSERT INTO %s (pk, value) VALUES (?, ?)", i, String.format("prod%02d_item_%04d", i % 10, i));
// Query while the data still lives only in the memtable (in-memory search path).
long memtableGroupCount = execute("SELECT pk FROM %s WHERE value LIKE ?", "prod05%").size();
long memtableAllCount = execute("SELECT pk FROM %s WHERE value LIKE ?", "prod%").size();
assertEquals(30, memtableGroupCount);
assertEquals(300, memtableAllCount);
// Flush through the optimized prefix path, then re-run the same queries against the SSTable.
flush();
waitForTableIndexesQueryable();
assertRowCount(execute("SELECT pk FROM %s WHERE value LIKE ?", "prod05%"), (int) memtableGroupCount);
assertRowCount(execute("SELECT pk FROM %s WHERE value LIKE ?", "prod%"), (int) memtableAllCount);
// A single group is below the prefix-section minimum, so its section is pruned but exact
// postings must still return every row.
assertRowCount(execute("SELECT pk FROM %s WHERE value = ?", String.format("prod00_item_%04d", 0)), 1);
}
}

128
test/unit/org/apache/cassandra/index/sai/memory/SectionedPrimaryKeysTest.java Normal file
View File

@ -0,0 +1,128 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.cassandra.index.sai.memory;
import org.junit.Before;
import org.junit.Test;
import org.apache.cassandra.dht.Murmur3Partitioner;
import org.apache.cassandra.index.sai.SAITester;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.PrimaryKeys;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;
/**
* Unit tests for {@link SectionedPrimaryKeys}, the {@link TrieMemoryIndex} node payload used when
* the SAI prefix feature is enabled. Verifies that the exact and prefix sections are populated and
* read independently and that the lazily-allocated prefix section reports the shared empty sentinel
* until it is first written.
*/
public class SectionedPrimaryKeysTest
{
private PrimaryKey.Factory keyFactory;
@Before
public void setup()
{
keyFactory = new PrimaryKey.Factory(Murmur3Partitioner.instance, SAITester.EMPTY_COMPARATOR);
}
private PrimaryKey key(long token)
{
return keyFactory.create(new Murmur3Partitioner.LongToken(token));
}
@Test
public void emptyPayloadReportsEmptyAndSharesEmptySentinel()
{
SectionedPrimaryKeys sectioned = new SectionedPrimaryKeys();
assertTrue(sectioned.isEmpty());
assertNotNull(sectioned.exact());
assertTrue(sectioned.exact().isEmpty());
// prefix is lazily allocated: before any addPrefix it must return the shared EMPTY sentinel.
assertSame(PrimaryKeys.EMPTY, sectioned.prefix());
assertTrue(sectioned.prefix().isEmpty());
}
@Test
public void addExactPopulatesOnlyExactSection()
{
SectionedPrimaryKeys sectioned = new SectionedPrimaryKeys();
sectioned.addExact(key(1));
sectioned.addExact(key(2));
assertFalse(sectioned.isEmpty());
assertEquals(2, sectioned.exact().size());
// The prefix section must remain untouched (still the EMPTY sentinel).
assertSame(PrimaryKeys.EMPTY, sectioned.prefix());
assertTrue(sectioned.prefix().isEmpty());
}
@Test
public void addPrefixPopulatesOnlyPrefixSection()
{
SectionedPrimaryKeys sectioned = new SectionedPrimaryKeys();
sectioned.addPrefix(key(1));
sectioned.addPrefix(key(2));
sectioned.addPrefix(key(3));
assertFalse(sectioned.isEmpty());
assertTrue(sectioned.exact().isEmpty());
assertEquals(3, sectioned.prefix().size());
}
@Test
public void exactAndPrefixSectionsAreIndependent()
{
SectionedPrimaryKeys sectioned = new SectionedPrimaryKeys();
sectioned.addExact(key(10));
sectioned.addPrefix(key(20));
sectioned.addPrefix(key(30));
assertEquals(1, sectioned.exact().size());
assertEquals(2, sectioned.prefix().size());
assertFalse(sectioned.isEmpty());
}
@Test
public void unsharedHeapSizeGrowsWhenPrefixSectionIsAllocated()
{
SectionedPrimaryKeys sectioned = new SectionedPrimaryKeys();
// Only the exact set is allocated up front; the prefix set is lazy (null), so adding exact
// keys does not allocate a new set and the unshared overhead stays at the single-set size.
long beforePrefix = sectioned.unsharedHeapSize();
sectioned.addExact(key(1));
assertEquals("adding an exact key must not allocate the lazy prefix set",
beforePrefix, sectioned.unsharedHeapSize());
// The first addPrefix lazily allocates the second PrimaryKeys, increasing the unshared overhead.
sectioned.addPrefix(key(2));
assertTrue("allocating the prefix section should increase the unshared heap size",
sectioned.unsharedHeapSize() > beforePrefix);
}
}