diff --git a/src/java/org/apache/cassandra/index/sai/disk/RowMapping.java b/src/java/org/apache/cassandra/index/sai/disk/RowMapping.java index 5bf64360af..6ec1095478 100644 --- a/src/java/org/apache/cassandra/index/sai/disk/RowMapping.java +++ b/src/java/org/apache/cassandra/index/sai/disk/RowMapping.java @@ -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> merge(MemtableIndex index) { return Collections.emptyIterator(); } + @Override + public Iterator> 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. + *

+ * 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: + * 

+     *   [exactCount, totalCount, exact_rowId_0, ..., prefix_rowId_0, ...]
+     *
+ * consumed verbatim by + * {@link org.apache.cassandra.index.sai.disk.v1.trie.LiteralIndexWriter#writeCompleteSegment}. + *

+ * Early prefix pruning: 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. + *

+ * 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> mergeV2(MemtableIndex index) + { + assert complete : "RowMapping is not built."; + + final int minimumLeaves = CassandraRelevantProperties.SAI_MINIMUM_POSTINGS_LEAVES.getInt(); + Iterator> iterator = index.iteratorSectioned(); + return new AbstractGuavaIterator<>() + { + @Override + protected Pair computeNext() + { + while (iterator.hasNext()) + { + Pair 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. */ diff --git a/src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java b/src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java index 2181d79a73..7fd69278b6 100644 --- a/src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java +++ b/src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java @@ -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> 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))) diff --git a/src/java/org/apache/cassandra/index/sai/memory/MemoryIndex.java b/src/java/org/apache/cassandra/index/sai/memory/MemoryIndex.java index 27190be4cf..c710735af0 100644 --- a/src/java/org/apache/cassandra/index/sai/memory/MemoryIndex.java +++ b/src/java/org/apache/cassandra/index/sai/memory/MemoryIndex.java @@ -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> iterator(); + /** + * Returns a sectioned iterator for V2 prefix flush. + * Default returns an empty iterator; overridden in {@link TrieMemoryIndex}. + */ + public Iterator> iteratorSectioned() + { + return Collections.emptyIterator(); + } + public abstract SegmentMetadata.ComponentMetadataMap writeDirect(IndexDescriptor indexDescriptor, IndexIdentifier indexIdentifier, Function postingTransformer) throws IOException; diff --git a/src/java/org/apache/cassandra/index/sai/memory/MemtableIndex.java b/src/java/org/apache/cassandra/index/sai/memory/MemtableIndex.java index ab24fff087..2cad0eeea7 100644 --- a/src/java/org/apache/cassandra/index/sai/memory/MemtableIndex.java +++ b/src/java/org/apache/cassandra/index/sai/memory/MemtableIndex.java @@ -113,6 +113,11 @@ public class MemtableIndex implements MemtableOrdering return memoryIndex.iterator(); } + public Iterator> iteratorSectioned() + { + return memoryIndex.iteratorSectioned(); + } + public SegmentMetadata.ComponentMetadataMap writeDirect(IndexDescriptor indexDescriptor, IndexIdentifier indexIdentifier, Function postingTransformer) throws IOException diff --git a/src/java/org/apache/cassandra/index/sai/memory/MemtableTermsIterator.java b/src/java/org/apache/cassandra/index/sai/memory/MemtableTermsIterator.java index 41e6e798de..01a121136a 100644 --- a/src/java/org/apache/cassandra/index/sai/memory/MemtableTermsIterator.java +++ b/src/java/org/apache/cassandra/index/sai/memory/MemtableTermsIterator.java @@ -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> iterator; + private final boolean isV2; private Pair 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> iterator) + { + this(minTerm, maxTerm, iterator, false); + } + + /** V2 constructor: LongArrayList contains [exactCount, totalCount, rowId...]. */ + public MemtableTermsIterator(ByteBuffer minTerm, + ByteBuffer maxTerm, + Iterator> 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 it = list.iterator(); diff --git a/src/java/org/apache/cassandra/index/sai/memory/SectionedPrimaryKeys.java b/src/java/org/apache/cassandra/index/sai/memory/SectionedPrimaryKeys.java new file mode 100644 index 0000000000..37c49e289b --- /dev/null +++ b/src/java/org/apache/cassandra/index/sai/memory/SectionedPrimaryKeys.java @@ -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. + *

+ * Holds two sets of {@link PrimaryKeys}: + *

+ * + * 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()); + } +} diff --git a/src/java/org/apache/cassandra/index/sai/memory/TrieMemoryIndex.java b/src/java/org/apache/cassandra/index/sai/memory/TrieMemoryIndex.java index 5c53a056d4..d3cd817797 100644 --- a/src/java/org/apache/cassandra/index/sai/memory/TrieMemoryIndex.java +++ b/src/java/org/apache/cassandra/index/sai/memory/TrieMemoryIndex.java @@ -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 data; - private final PrimaryKeysReducer primaryKeysReducer; + private final InMemoryTrie 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> iterator() { - Iterator> iterator = data.entrySet().iterator(); + Iterator> iterator = data.entrySet().iterator(); return new Iterator<>() { @Override @@ -190,7 +208,34 @@ public class TrieMemoryIndex extends MemoryIndex @Override public Pair next() { - Map.Entry entry = iterator.next(); + Map.Entry 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> iteratorSectioned() + { + Iterator> iterator = data.entrySet().iterator(); + return new Iterator<>() + { + @Override + public boolean hasNext() + { + return iterator.hasNext(); + } + + @Override + public Pair next() + { + Map.Entry 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 values = data.subtrie(lowerBound, lowerInclusive, upperBound, upperInclusive).valueIterator(); + Iterator 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 values = data.subtrie(lowerBound, true, upperBound, false).valueIterator(); + Iterator 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 + private static class SectionedPrimaryKeysReducer implements InMemoryTrie.UpsertTransformer { 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; } diff --git a/src/java/org/apache/cassandra/index/sai/utils/PrimaryKeys.java b/src/java/org/apache/cassandra/index/sai/utils/PrimaryKeys.java index 3ba7c08af0..834dd5d9ba 100644 --- a/src/java/org/apache/cassandra/index/sai/utils/PrimaryKeys.java +++ b/src/java/org/apache/cassandra/index/sai/utils/PrimaryKeys.java @@ -33,6 +33,16 @@ import org.apache.cassandra.utils.ObjectSizes; @ThreadSafe public class PrimaryKeys implements Iterable { + /** 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; diff --git a/test/unit/org/apache/cassandra/index/sai/disk/v1/segment/PrefixSectionOptimizationTest.java b/test/unit/org/apache/cassandra/index/sai/disk/v1/segment/PrefixSectionOptimizationTest.java index 8b58dadcbc..3621e4034a 100644 --- a/test/unit/org/apache/cassandra/index/sai/disk/v1/segment/PrefixSectionOptimizationTest.java +++ b/test/unit/org/apache/cassandra/index/sai/disk/v1/segment/PrefixSectionOptimizationTest.java @@ -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); + } } diff --git a/test/unit/org/apache/cassandra/index/sai/memory/SectionedPrimaryKeysTest.java b/test/unit/org/apache/cassandra/index/sai/memory/SectionedPrimaryKeysTest.java new file mode 100644 index 0000000000..d4d7b3eb56 --- /dev/null +++ b/test/unit/org/apache/cassandra/index/sai/memory/SectionedPrimaryKeysTest.java @@ -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); + } +}