Merge branch 'cassandra-5.0' into trunk

* cassandra-5.0:
  Reduce size of per-SSTable index components for SAI

.build/parent-pom-template.xml

@@ -660,13 +660,13 @@
       <dependency>
         <groupId>org.openjdk.jmh</groupId>
         <artifactId>jmh-core</artifactId>
-        <version>1.36</version>
+        <version>1.37</version>
         <scope>test</scope>
       </dependency>
       <dependency>
         <groupId>org.openjdk.jmh</groupId>
         <artifactId>jmh-generator-annprocess</artifactId>
-        <version>1.36</version>
+        <version>1.37</version>
         <scope>test</scope>
       </dependency>
       <dependency>

CHANGES.txt

@@ -1,5 +1,6 @@
 5.1
 Merged from 5.0:
+ * Reduce size of per-SSTable index components for SAI (CASSANDRA-18673)
  * Remove unnecessary Netty dependencies after upgrade to version 4.1.96 (CASSANDRA-18729)
  * Prevent InaccessibleObjectException when the Leak Detector is traversing objects (CASSANDRA-18708)
  * Remove legacy command line options from cassandra-stress (CASSANDRA-18529)

src/java/org/apache/cassandra/config/CassandraRelevantProperties.java

@@ -437,6 +437,16 @@ public enum CassandraRelevantProperties
      */
     SAI_POSTINGS_SKIP("cassandra.sai.postings_skip", "3"),
 
+    /**
+     * Used to determine the block size and block mask for the clustering sorted terms.
+     */
+    SAI_SORTED_TERMS_CLUSTERING_BLOCK_SHIFT("cassandra.sai.sorted_terms_clustering_block_shift", "4"),
+
+    /**
+     * Used to determine the block size and block mask for the partition sorted terms.
+     */
+    SAI_SORTED_TERMS_PARTITION_BLOCK_SHIFT("cassandra.sai.sorted_terms_partition_block_shift", "4"),
+
     SAI_TEST_BALANCED_TREE_DEBUG_ENABLED("cassandra.sai.test.balanced_tree_debug_enabled", "false"),
     SAI_TEST_DISABLE_TIMEOUT("cassandra.sai.test.disable.timeout", "false"),
 

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

@@ -131,7 +131,7 @@ public class SSTableContext extends SharedCloseableImpl
      */
     public int openFilesPerSSTable()
     {
-        return indexDescriptor.version.onDiskFormat().openFilesPerSSTableIndex();
+        return indexDescriptor.version.onDiskFormat().openFilesPerSSTableIndex(indexDescriptor.hasClustering());
     }
 
     @Override

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

@@ -219,10 +219,10 @@ public class StorageAttachedIndexGroup implements Index.Group, INotificationCons
         return getComponents(indexes);
     }
 
-    static Set<Component> getComponents(Collection<StorageAttachedIndex> indices)
+    private Set<Component> getComponents(Collection<StorageAttachedIndex> indices)
     {
         Set<Component> components = Version.LATEST.onDiskFormat()
-                                                  .perSSTableIndexComponents()
+                                                  .perSSTableIndexComponents(baseCfs.metadata.get().comparator.size() > 0)
                                                   .stream()
                                                   .map(Version.LATEST::makePerSSTableComponent)
                                                   .collect(Collectors.toSet());

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

@@ -19,6 +19,7 @@ package org.apache.cassandra.index.sai.disk;
 
 import java.io.IOException;
 
+import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.index.sai.utils.PrimaryKey;
 
 /**
@@ -28,6 +29,9 @@ public interface PerSSTableIndexWriter
 {
     PerSSTableIndexWriter NONE = (key) -> {};
 
+    default void startPartition(DecoratedKey decoratedKey) throws IOException
+    {}
+
     void nextRow(PrimaryKey primaryKey) throws IOException;
 
     default void complete() throws IOException

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

@@ -38,7 +38,7 @@ import org.apache.cassandra.utils.bytecomparable.ByteComparable;
 /**
  * In memory representation of {@link PrimaryKey} to row ID mappings which only contains
  * {@link Row} regardless of whether it's live or deleted. ({@link RangeTombstoneMarker} is not included.)
- *
+ * <p>
  * While this inherits the threading behaviour of {@link InMemoryTrie} of single-writer / multiple-reader,
  * since it is only used by {@link StorageAttachedIndexWriter}, which is not threadsafe, we can consider
  * this class not threadsafe as well.
@@ -102,7 +102,7 @@ public class RowMapping
         assert complete : "RowMapping is not built.";
 
         Iterator<Pair<ByteComparable, PrimaryKeys>> iterator = index.iterator();
-        return new AbstractGuavaIterator<Pair<ByteComparable, LongArrayList>>()
+        return new AbstractGuavaIterator<>()
         {
             @Override
             protected Pair<ByteComparable, LongArrayList> computeNext()
@@ -116,9 +116,7 @@ public class RowMapping
 
                     while (primaryKeys.hasNext())
                     {
-                        PrimaryKey primaryKey = primaryKeys.next();
-                        ByteComparable byteComparable = primaryKey::asComparableBytes;
-                        Long sstableRowId = rowMapping.get(byteComparable);
+                        Long sstableRowId = rowMapping.get(primaryKeys.next());
 
                         // The in-memory index does not handle deletions, so it is possible to
                         // have a primary key in the index that doesn't exist in the row mapping
@@ -157,8 +155,7 @@ public class RowMapping
     {
         assert !complete : "Cannot modify built RowMapping.";
 
-        ByteComparable byteComparable = key::asComparableBytes;
-        rowMapping.putSingleton(byteComparable, sstableRowId, OVERWRITE_TRANSFORMER);
+        rowMapping.putSingleton(key, sstableRowId, OVERWRITE_TRANSFORMER);
 
         maxSSTableRowId = Math.max(maxSSTableRowId, sstableRowId);
 

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

@@ -105,6 +105,16 @@ public class StorageAttachedIndexWriter implements SSTableFlushObserver
         if (aborted) return;
         
         currentKey = key;
+
+        try
+        {
+            perSSTableWriter.startPartition(key);
+        }
+        catch (Throwable t)
+        {
+            logger.error(indexDescriptor.logMessage("Failed to record a partition during an index build"), t);
+            abort(t, true);
+        }
     }
 
     @Override

src/java/org/apache/cassandra/index/sai/disk/format/IndexComponent.java

@@ -68,19 +68,29 @@ public enum IndexComponent
     TOKEN_VALUES("TokenValues"),
 
     /**
-     * An on-disk trie containing the primary keys used for looking up the rowId from a partition key
+     * An on-disk block packed index containing the starting and ending rowIds for each partition.
      */
-    PRIMARY_KEY_TRIE("PrimaryKeyTrie"),
+    PARTITION_SIZES("PartitionSizes"),
 
     /**
-     * Prefix-compressed blocks of primary keys used for rowId to partition key lookups
+     * Prefix-compressed blocks of partition keys used for rowId to partition key lookups
      */
-    PRIMARY_KEY_BLOCKS("PrimaryKeyBlocks"),
+    PARTITION_KEY_BLOCKS("PartitionKeyBlocks"),
 
     /**
-     * Encoded sequence of offsets to primary key blocks
+     * Encoded sequence of offsets to partition key blocks
      */
-    PRIMARY_KEY_BLOCK_OFFSETS("PrimaryKeyBlockOffsets"),
+    PARTITION_KEY_BLOCK_OFFSETS("PartitionKeyBlockOffsets"),
+
+    /**
+     * Prefix-compressed blocks of clustering keys used for rowId to clustering key lookups
+     */
+    CLUSTERING_KEY_BLOCKS("ClusteringKeyBlocks"),
+
+    /**
+     * Encoded sequence of offsets to clustering key blocks
+     */
+    CLUSTERING_KEY_BLOCK_OFFSETS("ClusteringKeyBlockOffsets"),
 
     /**
      * Metadata for per-SSTable on-disk components.

src/java/org/apache/cassandra/index/sai/disk/format/IndexDescriptor.java

@@ -49,6 +49,7 @@ import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.io.util.File;
 import org.apache.cassandra.io.util.FileHandle;
 import org.apache.cassandra.utils.FBUtilities;
+import org.apache.cassandra.utils.Throwables;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.IOUtils;
 
@@ -104,6 +105,11 @@ public class IndexDescriptor
                                    sstable.metadata().comparator);
     }
 
+    public boolean hasClustering()
+    {
+        return clusteringComparator.size() > 0;
+    }
+
     public String componentName(IndexComponent indexComponent)
     {
         return version.fileNameFormatter().format(indexComponent, null);
@@ -251,36 +257,66 @@ public class IndexDescriptor
         return writer;
     }
 
-    public FileHandle createPerSSTableFileHandle(IndexComponent indexComponent)
+    public FileHandle createPerSSTableFileHandle(IndexComponent indexComponent, Throwables.DiscreteAction<?> cleanup)
     {
-        final File file = fileFor(indexComponent);
-
-        if (logger.isTraceEnabled())
+        try
         {
-            logger.trace(logMessage("Opening {} file handle for {} ({})"),
-                         file, FBUtilities.prettyPrintMemory(file.length()));
-        }
+            final File file = fileFor(indexComponent);
 
-        return new FileHandle.Builder(file).mmapped(true).complete();
+            if (logger.isTraceEnabled())
+            {
+                logger.trace(logMessage("Opening {} file handle for {} ({})"),
+                             file, FBUtilities.prettyPrintMemory(file.length()));
+            }
+
+            return new FileHandle.Builder(file).mmapped(true).complete();
+        }
+        catch (Throwable t)
+        {
+            throw handleFileHandleCleanup(t, cleanup);
+        }
     }
 
-    public FileHandle createPerIndexFileHandle(IndexComponent indexComponent, IndexContext indexContext)
+    public FileHandle createPerIndexFileHandle(IndexComponent indexComponent, IndexContext indexContext, Throwables.DiscreteAction<?> cleanup)
     {
-        final File file = fileFor(indexComponent, indexContext);
-
-        if (logger.isTraceEnabled())
+        try
         {
-            logger.trace(indexContext.logMessage("Opening file handle for {} ({})"),
-                         file, FBUtilities.prettyPrintMemory(file.length()));
-        }
+            final File file = fileFor(indexComponent, indexContext);
 
-        return new FileHandle.Builder(file).mmapped(true).complete();
+            if (logger.isTraceEnabled())
+            {
+                logger.trace(indexContext.logMessage("Opening file handle for {} ({})"),
+                             file, FBUtilities.prettyPrintMemory(file.length()));
+            }
+
+            return new FileHandle.Builder(file).mmapped(true).complete();
+        }
+        catch (Throwable t)
+        {
+            throw handleFileHandleCleanup(t, cleanup);
+        }
+    }
+
+    private RuntimeException handleFileHandleCleanup(Throwable t, Throwables.DiscreteAction<?> cleanup)
+    {
+        if (cleanup != null)
+        {
+            try
+            {
+                cleanup.perform();
+            }
+            catch (Exception e)
+            {
+                return Throwables.unchecked(Throwables.merge(t, e));
+            }
+        }
+        return Throwables.unchecked(t);
     }
 
     public Set<Component> getLivePerSSTableComponents()
     {
         return version.onDiskFormat()
-                      .perSSTableIndexComponents()
+                      .perSSTableIndexComponents(hasClustering())
                       .stream()
                       .filter(c -> fileFor(c).exists())
                       .map(version::makePerSSTableComponent)
@@ -300,7 +336,7 @@ public class IndexDescriptor
     public long sizeOnDiskOfPerSSTableComponents()
     {
         return version.onDiskFormat()
-                      .perSSTableIndexComponents()
+                      .perSSTableIndexComponents(hasClustering())
                       .stream()
                       .map(this::fileFor)
                       .filter(File::exists)
@@ -380,7 +416,7 @@ public class IndexDescriptor
     public void deletePerSSTableIndexComponents()
     {
         version.onDiskFormat()
-               .perSSTableIndexComponents()
+               .perSSTableIndexComponents(hasClustering())
                .stream()
                .map(this::fileFor)
                .filter(File::exists)

src/java/org/apache/cassandra/index/sai/disk/format/OnDiskFormat.java

@@ -136,8 +136,10 @@ public interface OnDiskFormat
      * Returns the set of {@link IndexComponent} for the per-SSTable part of an index.
      * This is a complete set of components that could exist on-disk. It does not imply that the
      * components currently exist on-disk.
+
+     * @param hasClustering true if the SSTable forms part of a table using clustering columns
      */
-    Set<IndexComponent> perSSTableIndexComponents();
+    Set<IndexComponent> perSSTableIndexComponents(boolean hasClustering);
 
     /**
      * Returns the set of {@link IndexComponent} for the per-column part of an index.
@@ -152,8 +154,10 @@ public interface OnDiskFormat
      * Return the number of open per-SSTable files that can be open during a query.
      * This is a static indication of the files that can be held open by an index
      * for queries. It is not a dynamic calculation.
+     *
+     * @param hasClustering true if the SSTable forms part of a table using clustering columns
      */
-    int openFilesPerSSTableIndex();
+    int openFilesPerSSTableIndex(boolean hasClustering);
 
     /**
      * Return the number of open per-column index files that can be open during a query.

src/java/org/apache/cassandra/index/sai/disk/io/IndexFileUtils.java

@@ -36,12 +36,12 @@ import org.apache.lucene.store.IndexInput;
 public class IndexFileUtils
 {
     @VisibleForTesting
-    protected static final SequentialWriterOption DEFAULT_WRITER_OPTION = SequentialWriterOption.newBuilder()
-                                                                                                .trickleFsync(DatabaseDescriptor.getTrickleFsync())
-                                                                                                .trickleFsyncByteInterval(DatabaseDescriptor.getTrickleFsyncIntervalInKiB() * 1024)
-                                                                                                .bufferType(BufferType.OFF_HEAP)
-                                                                                                .finishOnClose(true)
-                                                                                                .build();
+    public static final SequentialWriterOption DEFAULT_WRITER_OPTION = SequentialWriterOption.newBuilder()
+                                                                                             .trickleFsync(DatabaseDescriptor.getTrickleFsync())
+                                                                                             .trickleFsyncByteInterval(DatabaseDescriptor.getTrickleFsyncIntervalInKiB() * 1024)
+                                                                                             .bufferType(BufferType.OFF_HEAP)
+                                                                                             .finishOnClose(true)
+                                                                                             .build();
 
     public static final IndexFileUtils instance = new IndexFileUtils(DEFAULT_WRITER_OPTION);
 

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

@@ -37,6 +37,14 @@ public interface LongArray extends Closeable
      */
     long length();
 
+    /**
+     * Using the given value returns the first index corresponding to the value.
+     *
+     * @param value Value to lookup, and it must not be smaller than previous value
+     * @return The index of the given value or negative value if target value is greater than all values
+     */
+    long indexOf(long value);
+
     @Override
     default void close() throws IOException { }
 
@@ -66,6 +74,13 @@ public interface LongArray extends Closeable
             return longArray.length();
         }
 
+        @Override
+        public long indexOf(long value)
+        {
+            open();
+            return longArray.indexOf(value);
+        }
+
         @Override
         public void close() throws IOException
         {

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

@@ -46,13 +46,13 @@ public class PerColumnIndexFiles implements Closeable
         this.indexContext = indexContext;
         if (indexContext.isLiteral())
         {
-            files.put(IndexComponent.POSTING_LISTS, indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext));
-            files.put(IndexComponent.TERMS_DATA, indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext));
+            files.put(IndexComponent.POSTING_LISTS, indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, this::close));
+            files.put(IndexComponent.TERMS_DATA, indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext, this::close));
         }
         else
         {
-            files.put(IndexComponent.BALANCED_TREE, indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext));
-            files.put(IndexComponent.POSTING_LISTS, indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext));
+            files.put(IndexComponent.BALANCED_TREE, indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext, this::close));
+            files.put(IndexComponent.POSTING_LISTS, indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, this::close));
         }
     }
 

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

@@ -1,201 +0,0 @@
-/*
- * 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.disk.v1;
-
-import java.io.IOException;
-import java.nio.ByteBuffer;
-import java.util.Arrays;
-import javax.annotation.concurrent.NotThreadSafe;
-import javax.annotation.concurrent.ThreadSafe;
-
-import org.apache.cassandra.db.BufferDecoratedKey;
-import org.apache.cassandra.db.Clustering;
-import org.apache.cassandra.db.ClusteringComparator;
-import org.apache.cassandra.db.DecoratedKey;
-import org.apache.cassandra.db.marshal.ByteBufferAccessor;
-import org.apache.cassandra.dht.IPartitioner;
-import org.apache.cassandra.dht.Token;
-import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
-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.bitpack.BlockPackedReader;
-import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesMeta;
-import org.apache.cassandra.index.sai.disk.v1.sortedterms.SortedTermsMeta;
-import org.apache.cassandra.index.sai.disk.v1.sortedterms.SortedTermsReader;
-import org.apache.cassandra.index.sai.disk.v1.trie.TriePrefixSearcher;
-import org.apache.cassandra.index.sai.utils.PrimaryKey;
-import org.apache.cassandra.io.sstable.format.SSTableReader;
-import org.apache.cassandra.io.util.FileHandle;
-import org.apache.cassandra.io.util.FileUtils;
-import org.apache.cassandra.utils.Throwables;
-import org.apache.cassandra.utils.bytecomparable.ByteComparable;
-import org.apache.cassandra.utils.bytecomparable.ByteSource;
-import org.apache.cassandra.utils.bytecomparable.ByteSourceInverse;
-
-/**
- * A row-aware {@link PrimaryKeyMap}
- *
- * This uses the following on-disk structures:
- * <ul>
- *     <li>Block-packed structure for rowId to token lookups using {@link BlockPackedReader}.
- *     Uses component {@link IndexComponent#TOKEN_VALUES} </li>
- *     <li>A sorted-terms structure for rowId to {@link PrimaryKey} and {@link PrimaryKey} to rowId lookups using
- *     {@link SortedTermsReader} and {@link TriePrefixSearcher}. Uses components {@link IndexComponent#PRIMARY_KEY_TRIE},
- *     {@link IndexComponent#PRIMARY_KEY_BLOCKS} and {@link IndexComponent#PRIMARY_KEY_BLOCK_OFFSETS}</li>
- * </ul>
- *
- * While the {@link RowAwarePrimaryKeyMapFactory} is threadsafe, individual instances of the {@link RowAwarePrimaryKeyMap}
- * are not.
- */
-@NotThreadSafe
-public class RowAwarePrimaryKeyMap implements PrimaryKeyMap
-{
-    @ThreadSafe
-    public static class RowAwarePrimaryKeyMapFactory implements Factory
-    {
-        private final LongArray.Factory tokenReaderFactory;
-        private final SortedTermsReader sortedTermsReader;
-        private final SortedTermsMeta sortedTermsMeta;
-        private FileHandle tokensFile = null;
-        private FileHandle primaryKeyBlockOffsetsFile = null;
-        private FileHandle primaryKeyBlocksFile = null;
-        private FileHandle primaryKeyTrieFile = null;
-        private final IPartitioner partitioner;
-        private final ClusteringComparator clusteringComparator;
-        private final PrimaryKey.Factory primaryKeyFactory;
-
-        public RowAwarePrimaryKeyMapFactory(IndexDescriptor indexDescriptor, SSTableReader sstable)
-        {
-            try
-            {
-                MetadataSource metadataSource = MetadataSource.loadGroupMetadata(indexDescriptor);
-                NumericValuesMeta tokensMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.TOKEN_VALUES)));
-                NumericValuesMeta blockOffsetsMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS)));
-                this.sortedTermsMeta = new SortedTermsMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCKS)));
-                this.tokensFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES);
-                this.tokenReaderFactory = new BlockPackedReader(tokensFile, tokensMeta);
-                this.primaryKeyBlockOffsetsFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS);
-                this.primaryKeyBlocksFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_BLOCKS);
-                this.primaryKeyTrieFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_TRIE);
-                this.sortedTermsReader = new SortedTermsReader(primaryKeyBlocksFile, primaryKeyBlockOffsetsFile, sortedTermsMeta, blockOffsetsMeta);
-                this.partitioner = sstable.metadata().partitioner;
-                this.primaryKeyFactory = indexDescriptor.primaryKeyFactory;
-                this.clusteringComparator = indexDescriptor.clusteringComparator;
-            }
-            catch (Throwable t)
-            {
-                throw Throwables.unchecked(Throwables.close(t, Arrays.asList(tokensFile, primaryKeyBlocksFile, primaryKeyBlockOffsetsFile, primaryKeyTrieFile)));
-            }
-        }
-
-        @Override
-        @SuppressWarnings({"resource", "RedundantSuppression"})
-        public PrimaryKeyMap newPerSSTablePrimaryKeyMap() throws IOException
-        {
-            LongArray rowIdToToken = new LongArray.DeferredLongArray(tokenReaderFactory::open);
-            return new RowAwarePrimaryKeyMap(rowIdToToken,
-                                             new TriePrefixSearcher(primaryKeyTrieFile.instantiateRebufferer(null), sortedTermsMeta.trieFilePointer),
-                                             sortedTermsReader.openCursor(),
-                                             partitioner,
-                                             primaryKeyFactory,
-                                             clusteringComparator);
-        }
-
-        @Override
-        public void close()
-        {
-            FileUtils.closeQuietly(Arrays.asList(tokensFile, primaryKeyBlocksFile, primaryKeyBlockOffsetsFile, primaryKeyTrieFile));
-        }
-    }
-
-    private final LongArray rowIdToToken;
-    private final TriePrefixSearcher triePrefixSearcher;
-    private final SortedTermsReader.Cursor cursor;
-    private final IPartitioner partitioner;
-    private final PrimaryKey.Factory primaryKeyFactory;
-    private final ClusteringComparator clusteringComparator;
-    private final ByteBuffer tokenBuffer = ByteBuffer.allocate(Long.BYTES);
-
-    private RowAwarePrimaryKeyMap(LongArray rowIdToToken,
-                                  TriePrefixSearcher triePrefixSearcher,
-                                  SortedTermsReader.Cursor cursor,
-                                  IPartitioner partitioner,
-                                  PrimaryKey.Factory primaryKeyFactory,
-                                  ClusteringComparator clusteringComparator)
-    {
-        this.rowIdToToken = rowIdToToken;
-        this.triePrefixSearcher = triePrefixSearcher;
-        this.cursor = cursor;
-        this.partitioner = partitioner;
-        this.primaryKeyFactory = primaryKeyFactory;
-        this.clusteringComparator = clusteringComparator;
-    }
-
-    @Override
-    public PrimaryKey primaryKeyFromRowId(long sstableRowId)
-    {
-        tokenBuffer.putLong(rowIdToToken.get(sstableRowId));
-        tokenBuffer.rewind();
-        return primaryKeyFactory.createDeferred(partitioner.getTokenFactory().fromByteArray(tokenBuffer), () -> supplier(sstableRowId));
-    }
-
-    @Override
-    public long rowIdFromPrimaryKey(PrimaryKey key)
-    {
-        return triePrefixSearcher.prefixSearch(key.asComparableBytes(ByteComparable.Version.OSS50));
-    }
-
-    @Override
-    public void close()
-    {
-        FileUtils.closeQuietly(Arrays.asList(triePrefixSearcher, cursor, rowIdToToken));
-    }
-
-    private PrimaryKey supplier(long sstableRowId)
-    {
-        try
-        {
-            cursor.seekToPointId(sstableRowId);
-            ByteSource.Peekable peekable = ByteSource.peekable(cursor.term().asComparableBytes(ByteComparable.Version.OSS50));
-
-            Token token = partitioner.getTokenFactory().fromComparableBytes(ByteSourceInverse.nextComponentSource(peekable),
-                                                                            ByteComparable.Version.OSS50);
-            byte[] keyBytes = ByteSourceInverse.getUnescapedBytes(ByteSourceInverse.nextComponentSource(peekable));
-
-            if (keyBytes == null)
-                return primaryKeyFactory.createTokenOnly(token);
-
-            DecoratedKey partitionKey = new BufferDecoratedKey(token, ByteBuffer.wrap(keyBytes));
-
-            Clustering<?> clustering = clusteringComparator.size() == 0
-                                       ? Clustering.EMPTY
-                                       : clusteringComparator.clusteringFromByteComparable(ByteBufferAccessor.instance,
-                                                                                           v -> ByteSourceInverse.nextComponentSource(peekable));
-
-            if (clustering == null)
-                clustering = Clustering.EMPTY;
-
-            return primaryKeyFactory.create(partitionKey, clustering);
-        }
-        catch (IOException e)
-        {
-            throw Throwables.cleaned(e);
-        }
-    }
-}

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

@@ -23,14 +23,17 @@ import java.io.IOException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.cassandra.config.CassandraRelevantProperties;
+import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.index.sai.disk.PerSSTableIndexWriter;
 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.io.IndexOutputWriter;
 import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesWriter;
-import org.apache.cassandra.index.sai.disk.v1.sortedterms.SortedTermsWriter;
+import org.apache.cassandra.index.sai.disk.v1.keystore.KeyStoreWriter;
 import org.apache.cassandra.index.sai.utils.PrimaryKey;
-import org.apache.lucene.util.IOUtils;
+import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.utils.bytecomparable.ByteSource;
 
 public class SSTableComponentsWriter implements PerSSTableIndexWriter
 {
@@ -38,34 +41,61 @@ public class SSTableComponentsWriter implements PerSSTableIndexWriter
 
     private final IndexDescriptor indexDescriptor;
     private final MetadataWriter metadataWriter;
+    private final NumericValuesWriter partitionSizeWriter;
     private final NumericValuesWriter tokenWriter;
-    private final SortedTermsWriter sortedTermsWriter;
+    private final KeyStoreWriter partitionKeysWriter;
+    private final KeyStoreWriter clusteringKeysWriter;
+
+    private long partitionId = -1;
 
     @SuppressWarnings({"resource", "RedundantSuppression"})
     public SSTableComponentsWriter(IndexDescriptor indexDescriptor) throws IOException
     {
         this.indexDescriptor = indexDescriptor;
         this.metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META));
-        this.tokenWriter = new NumericValuesWriter(indexDescriptor.componentName(IndexComponent.TOKEN_VALUES),
-                                                   indexDescriptor.openPerSSTableOutput(IndexComponent.TOKEN_VALUES),
-                                                   metadataWriter, false);
-        IndexOutputWriter primaryKeyTrieWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_TRIE);
-        IndexOutputWriter primaryKeyBlocksWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCKS);
-        NumericValuesWriter primaryKeyBlockOffsetWriter = new NumericValuesWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                     indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                     metadataWriter, true);
-        this.sortedTermsWriter = new SortedTermsWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCKS),
-                                                       metadataWriter,
-                                                       primaryKeyBlocksWriter,
-                                                       primaryKeyBlockOffsetWriter,
-                                                       primaryKeyTrieWriter);
+        this.tokenWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.TOKEN_VALUES, metadataWriter, false);
+        this.partitionSizeWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.PARTITION_SIZES, metadataWriter, true);
+        IndexOutputWriter partitionKeyBlocksWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PARTITION_KEY_BLOCKS);
+        NumericValuesWriter partitionKeyBlockOffsetWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.PARTITION_KEY_BLOCK_OFFSETS, metadataWriter, true);
+        this.partitionKeysWriter = new KeyStoreWriter(indexDescriptor.componentName(IndexComponent.PARTITION_KEY_BLOCKS),
+                                                      metadataWriter,
+                                                      partitionKeyBlocksWriter,
+                                                      partitionKeyBlockOffsetWriter,
+                                                      CassandraRelevantProperties.SAI_SORTED_TERMS_PARTITION_BLOCK_SHIFT.getInt(),
+                                                      false);
+        if (indexDescriptor.hasClustering())
+        {
+            IndexOutputWriter clusteringKeyBlocksWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.CLUSTERING_KEY_BLOCKS);
+            NumericValuesWriter clusteringKeyBlockOffsetWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.CLUSTERING_KEY_BLOCK_OFFSETS, metadataWriter, true);
+            this.clusteringKeysWriter = new KeyStoreWriter(indexDescriptor.componentName(IndexComponent.CLUSTERING_KEY_BLOCKS),
+                                                           metadataWriter,
+                                                           clusteringKeyBlocksWriter,
+                                                           clusteringKeyBlockOffsetWriter,
+                                                           CassandraRelevantProperties.SAI_SORTED_TERMS_CLUSTERING_BLOCK_SHIFT.getInt(),
+                                                           true);
+        }
+        else
+        {
+            this.clusteringKeysWriter = null;
+        }
+    }
+
+    @Override
+    public void startPartition(DecoratedKey partitionKey) throws IOException
+    {
+        partitionId++;
+        partitionKeysWriter.add(v -> ByteSource.of(partitionKey.getKey(), v));
+        if (indexDescriptor.hasClustering())
+            clusteringKeysWriter.startPartition();
     }
 
     @Override
     public void nextRow(PrimaryKey primaryKey) throws IOException
     {
         tokenWriter.add(primaryKey.token().getLongValue());
-        sortedTermsWriter.add(primaryKey::asComparableBytes);
+        partitionSizeWriter.add(partitionId);
+        if (indexDescriptor.hasClustering())
+            clusteringKeysWriter.add(indexDescriptor.clusteringComparator.asByteComparable(primaryKey.clustering()));
     }
 
     @Override
@@ -77,7 +107,7 @@ public class SSTableComponentsWriter implements PerSSTableIndexWriter
         }
         finally
         {
-            IOUtils.close(tokenWriter, sortedTermsWriter, metadataWriter);
+            FileUtils.close(tokenWriter, partitionSizeWriter, partitionKeysWriter, clusteringKeysWriter, metadataWriter);
         }
     }
 

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

@@ -0,0 +1,208 @@
+/*
+ * 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.disk.v1;
+
+import org.apache.cassandra.db.BufferDecoratedKey;
+import org.apache.cassandra.db.Clustering;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+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.bitpack.BlockPackedReader;
+import org.apache.cassandra.index.sai.disk.v1.bitpack.MonotonicBlockPackedReader;
+import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesMeta;
+import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookupMeta;
+import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookup;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.io.util.FileHandle;
+import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.utils.Throwables;
+import org.apache.cassandra.utils.bytecomparable.ByteComparable;
+import org.apache.cassandra.utils.bytecomparable.ByteSource;
+import org.apache.cassandra.utils.bytecomparable.ByteSourceInverse;
+
+import javax.annotation.concurrent.NotThreadSafe;
+import javax.annotation.concurrent.ThreadSafe;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+
+/**
+ * A {@link PrimaryKeyMap} for skinny tables (those with no clustering columns).
+ * <p>
+ * This uses the following on-disk structures:
+ * <ul>
+ *     <li>A block-packed structure for rowId to token lookups using {@link BlockPackedReader}.
+ *     Uses the {@link IndexComponent#TOKEN_VALUES} component</li>
+ *     <li>A monotonic block packed structure for rowId to partitionId lookups using {@link MonotonicBlockPackedReader}.
+ *     Uses the {@link IndexComponent#PARTITION_SIZES} component</li>
+ *     <li>A key store for rowId to {@link PrimaryKey} and {@link PrimaryKey} to rowId lookups using
+ *     {@link KeyLookup}. Uses the {@link IndexComponent#PARTITION_KEY_BLOCKS} and
+ *     {@link IndexComponent#PARTITION_KEY_BLOCK_OFFSETS} components</li>
+ * </ul>
+ *
+ * While the {@link Factory} is threadsafe, individual instances of the {@link SkinnyPrimaryKeyMap}
+ * are not.
+ */
+@NotThreadSafe
+public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
+{
+    @ThreadSafe
+    public static class Factory implements PrimaryKeyMap.Factory
+    {
+        protected final MetadataSource metadataSource;
+        protected final LongArray.Factory tokenReaderFactory;
+        protected final LongArray.Factory partitionReaderFactory;
+        protected final KeyLookup partitionKeyReader;
+        protected final IPartitioner partitioner;
+        protected final PrimaryKey.Factory primaryKeyFactory;
+
+        private final FileHandle tokensFile;
+        private final FileHandle partitionsFile;
+        private final FileHandle partitionKeyBlockOffsetsFile;
+        private final FileHandle partitionKeyBlocksFile;
+
+        public Factory(IndexDescriptor indexDescriptor, SSTableReader sstable)
+        {
+            this.tokensFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES, this::close);
+            this.partitionsFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PARTITION_SIZES, this::close);
+            this.partitionKeyBlockOffsetsFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PARTITION_KEY_BLOCK_OFFSETS, this::close);
+            this.partitionKeyBlocksFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PARTITION_KEY_BLOCKS, this::close);
+            try
+            {
+                this.metadataSource = MetadataSource.loadGroupMetadata(indexDescriptor);
+                NumericValuesMeta tokensMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.TOKEN_VALUES)));
+                this.tokenReaderFactory = new BlockPackedReader(tokensFile, tokensMeta);
+                NumericValuesMeta partitionsMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PARTITION_SIZES)));
+                this.partitionReaderFactory = new MonotonicBlockPackedReader(partitionsFile, partitionsMeta);
+                NumericValuesMeta partitionKeyBlockOffsetsMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PARTITION_KEY_BLOCK_OFFSETS)));
+                KeyLookupMeta partitionKeysMeta = new KeyLookupMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PARTITION_KEY_BLOCKS)));
+                this.partitionKeyReader = new KeyLookup(partitionKeyBlocksFile, partitionKeyBlockOffsetsFile, partitionKeysMeta, partitionKeyBlockOffsetsMeta);
+                this.partitioner = sstable.metadata().partitioner;
+                this.primaryKeyFactory = indexDescriptor.primaryKeyFactory;
+            }
+            catch (Throwable t)
+            {
+                throw Throwables.unchecked(t);
+            }
+        }
+
+        @Override
+        @SuppressWarnings({"resource", "RedundantSuppression"})
+        public PrimaryKeyMap newPerSSTablePrimaryKeyMap() throws IOException
+        {
+            LongArray rowIdToToken = new LongArray.DeferredLongArray(tokenReaderFactory::open);
+            LongArray rowIdToPartitionId = new LongArray.DeferredLongArray(partitionReaderFactory::open);
+            return new SkinnyPrimaryKeyMap(rowIdToToken,
+                                           rowIdToPartitionId,
+                                           partitionKeyReader.openCursor(),
+                                           partitioner,
+                                           primaryKeyFactory);
+        }
+
+        @Override
+        public void close()
+        {
+            FileUtils.closeQuietly(Arrays.asList(tokensFile, partitionsFile, partitionKeyBlocksFile, partitionKeyBlockOffsetsFile));
+        }
+    }
+
+    protected final LongArray tokenArray;
+    protected final LongArray partitionArray;
+    protected final KeyLookup.Cursor partitionKeyCursor;
+    protected final IPartitioner partitioner;
+    protected final PrimaryKey.Factory primaryKeyFactory;
+    protected final ByteBuffer tokenBuffer = ByteBuffer.allocate(Long.BYTES);
+
+    protected SkinnyPrimaryKeyMap(LongArray tokenArray,
+                                  LongArray partitionArray,
+                                  KeyLookup.Cursor partitionKeyCursor,
+                                  IPartitioner partitioner,
+                                  PrimaryKey.Factory primaryKeyFactory)
+    {
+        this.tokenArray = tokenArray;
+        this.partitionArray = partitionArray;
+        this.partitionKeyCursor = partitionKeyCursor;
+        this.partitioner = partitioner;
+        this.primaryKeyFactory = primaryKeyFactory;
+    }
+
+    @Override
+    public PrimaryKey primaryKeyFromRowId(long sstableRowId)
+    {
+        tokenBuffer.putLong(tokenArray.get(sstableRowId));
+        tokenBuffer.rewind();
+        return primaryKeyFactory.createDeferred(partitioner.getTokenFactory().fromByteArray(tokenBuffer), () -> supplier(sstableRowId));
+    }
+
+    @Override
+    public long rowIdFromPrimaryKey(PrimaryKey primaryKey)
+    {
+        long rowId = tokenArray.indexOf(primaryKey.token().getLongValue());
+        // If the key is token only, the token is out of range, we are at the end of our keys, or we have skipped a token
+        // we can return straight away.
+        if (primaryKey.isTokenOnly() || rowId < 0 || rowId + 1 == tokenArray.length() || tokenArray.get(rowId) != primaryKey.token().getLongValue())
+            return rowId;
+        // Otherwise we need to check for token collision.
+        return tokenCollisionDetection(primaryKey, rowId);
+    }
+
+    @Override
+    public void close()
+    {
+        FileUtils.closeQuietly(Arrays.asList(partitionKeyCursor, tokenArray, partitionArray));
+    }
+
+    // Look for token collision by if the ajacent token in the token array matches the
+    // current token. If we find a collision we need to compare the partition key instead.
+    protected long tokenCollisionDetection(PrimaryKey primaryKey, long rowId)
+    {
+        // Look for collisions while we haven't reached the end of the tokens and the tokens don't collide
+        while (rowId + 1 < tokenArray.length() && primaryKey.token().getLongValue() == tokenArray.get(rowId + 1))
+        {
+            // If we had a collision then see if the partition key for this row is >= to the lookup partition key
+            if (readPartitionKey(rowId).compareTo(primaryKey.partitionKey()) >= 0)
+                return rowId;
+
+            rowId++;
+        }
+        // Note: We would normally expect to get here without going into the while loop
+        return rowId;
+    }
+
+    protected PrimaryKey supplier(long sstableRowId)
+    {
+        return primaryKeyFactory.create(readPartitionKey(sstableRowId), Clustering.EMPTY);
+    }
+
+    protected DecoratedKey readPartitionKey(long sstableRowId)
+    {
+        long partitionId = partitionArray.get(sstableRowId);
+        ByteSource.Peekable peekable = ByteSource.peekable(partitionKeyCursor.seekToPointId(partitionId).asComparableBytes(ByteComparable.Version.OSS50));
+
+        byte[] keyBytes = ByteSourceInverse.getUnescapedBytes(peekable);
+
+        assert keyBytes != null : "Primary key from map did not contain a partition key";
+
+        ByteBuffer decoratedKey = ByteBuffer.wrap(keyBytes);
+        return new BufferDecoratedKey(partitioner.getToken(decoratedKey), decoratedKey);
+    }
+}

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

@@ -58,12 +58,23 @@ public class V1OnDiskFormat implements OnDiskFormat
     private static final Logger logger = LoggerFactory.getLogger(V1OnDiskFormat.class);
 
     @VisibleForTesting
-    public static final Set<IndexComponent> PER_SSTABLE_COMPONENTS = EnumSet.of(IndexComponent.GROUP_COMPLETION_MARKER,
-                                                                                IndexComponent.GROUP_META,
-                                                                                IndexComponent.TOKEN_VALUES,
-                                                                                IndexComponent.PRIMARY_KEY_TRIE,
-                                                                                IndexComponent.PRIMARY_KEY_BLOCKS,
-                                                                                IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS);
+    public static final Set<IndexComponent> SKINNY_PER_SSTABLE_COMPONENTS = EnumSet.of(IndexComponent.GROUP_COMPLETION_MARKER,
+                                                                                       IndexComponent.GROUP_META,
+                                                                                       IndexComponent.TOKEN_VALUES,
+                                                                                       IndexComponent.PARTITION_SIZES,
+                                                                                       IndexComponent.PARTITION_KEY_BLOCKS,
+                                                                                       IndexComponent.PARTITION_KEY_BLOCK_OFFSETS);
+
+    @VisibleForTesting
+    public static final Set<IndexComponent> WIDE_PER_SSTABLE_COMPONENTS = EnumSet.of(IndexComponent.GROUP_COMPLETION_MARKER,
+                                                                                     IndexComponent.GROUP_META,
+                                                                                     IndexComponent.TOKEN_VALUES,
+                                                                                     IndexComponent.PARTITION_SIZES,
+                                                                                     IndexComponent.PARTITION_KEY_BLOCKS,
+                                                                                     IndexComponent.PARTITION_KEY_BLOCK_OFFSETS,
+                                                                                     IndexComponent.CLUSTERING_KEY_BLOCKS,
+                                                                                     IndexComponent.CLUSTERING_KEY_BLOCK_OFFSETS);
+
     @VisibleForTesting
     public static final Set<IndexComponent> LITERAL_COMPONENTS = EnumSet.of(IndexComponent.COLUMN_COMPLETION_MARKER,
                                                                             IndexComponent.META,
@@ -112,7 +123,8 @@ public class V1OnDiskFormat implements OnDiskFormat
     @Override
     public PrimaryKeyMap.Factory newPrimaryKeyMapFactory(IndexDescriptor indexDescriptor, SSTableReader sstable)
     {
-        return new RowAwarePrimaryKeyMap.RowAwarePrimaryKeyMapFactory(indexDescriptor, sstable);
+        return indexDescriptor.hasClustering() ? new WidePrimaryKeyMap.Factory(indexDescriptor, sstable)
+                                               : new SkinnyPrimaryKeyMap.Factory(indexDescriptor, sstable);
     }
 
     @Override
@@ -165,7 +177,7 @@ public class V1OnDiskFormat implements OnDiskFormat
     @Override
     public void validatePerSSTableIndexComponents(IndexDescriptor indexDescriptor, boolean checksum)
     {
-        for (IndexComponent indexComponent : perSSTableIndexComponents())
+        for (IndexComponent indexComponent : perSSTableIndexComponents(indexDescriptor.hasClustering()))
         {
             if (isNotBuildCompletionMarker(indexComponent))
             {
@@ -224,9 +236,9 @@ public class V1OnDiskFormat implements OnDiskFormat
     }
 
     @Override
-    public Set<IndexComponent> perSSTableIndexComponents()
+    public Set<IndexComponent> perSSTableIndexComponents(boolean hasClustering)
     {
-        return PER_SSTABLE_COMPONENTS;
+        return hasClustering ? WIDE_PER_SSTABLE_COMPONENTS : SKINNY_PER_SSTABLE_COMPONENTS;
     }
 
     @Override
@@ -236,11 +248,14 @@ public class V1OnDiskFormat implements OnDiskFormat
     }
 
     @Override
-    public int openFilesPerSSTableIndex()
+    public int openFilesPerSSTableIndex(boolean hasClustering)
     {
-        // For the V1 format there are always 4 open files per SSTable - token values, primary key trie,
-        // primary key blocks, primary key block offsets
-        return 4;
+        // For the V1 format the number of open files depends on whether the table has clustering. For wide tables
+        // the number of open files will be 6 per SSTable - token values, partition sizes index, partition key blocks,
+        // partition key block offsets, clustering key blocks & clustering key block offsets and for skinny tables
+        // the number of files will be 4 per SSTable - token values, partition key sizes, partition key blocks &
+        // partition key block offsets.
+        return hasClustering ? 6 : 4;
     }
 
     @Override

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

@@ -0,0 +1,179 @@
+/*
+ * 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.disk.v1;
+
+import org.apache.cassandra.db.Clustering;
+import org.apache.cassandra.db.ClusteringComparator;
+import org.apache.cassandra.db.marshal.ByteBufferAccessor;
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+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.bitpack.NumericValuesMeta;
+import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookupMeta;
+import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookup;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.io.util.FileHandle;
+import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.utils.Throwables;
+import org.apache.cassandra.utils.bytecomparable.ByteComparable;
+import org.apache.cassandra.utils.bytecomparable.ByteSource;
+
+import javax.annotation.concurrent.NotThreadSafe;
+import javax.annotation.concurrent.ThreadSafe;
+import java.io.IOException;
+import java.util.Arrays;
+
+/**
+ * An extension of the {@link SkinnyPrimaryKeyMap} for wide tables (those with clustering columns).
+ * <p>
+ * This used the following additional on-disk structures to the {@link SkinnyPrimaryKeyMap}
+ * <ul>
+ *     <li>A key store for rowId to {@link Clustering} and {@link Clustering} to rowId lookups using
+ *     {@link KeyLookup}. Uses the {@link IndexComponent#CLUSTERING_KEY_BLOCKS} and
+ *     {@link IndexComponent#CLUSTERING_KEY_BLOCK_OFFSETS} components</li>
+ * </ul>
+ * While the {@link Factory} is threadsafe, individual instances of the {@link WidePrimaryKeyMap}
+ * are not.
+ */
+@NotThreadSafe
+public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
+{
+    @ThreadSafe
+    public static class Factory extends SkinnyPrimaryKeyMap.Factory
+    {
+        private final ClusteringComparator clusteringComparator;
+        private final KeyLookup clusteringKeyReader;
+        private final FileHandle clusteringKeyBlockOffsetsFile;
+        private final FileHandle clustingingKeyBlocksFile;
+
+        public Factory(IndexDescriptor indexDescriptor, SSTableReader sstable)
+        {
+            super(indexDescriptor, sstable);
+
+            this.clusteringKeyBlockOffsetsFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.CLUSTERING_KEY_BLOCK_OFFSETS, this::close);
+            this.clustingingKeyBlocksFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.CLUSTERING_KEY_BLOCKS, this::close);
+
+            try
+            {
+                this.clusteringComparator = indexDescriptor.clusteringComparator;
+                NumericValuesMeta clusteringKeyBlockOffsetsMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.CLUSTERING_KEY_BLOCK_OFFSETS)));
+                KeyLookupMeta clusteringKeyMeta = new KeyLookupMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.CLUSTERING_KEY_BLOCKS)));
+                this.clusteringKeyReader = new KeyLookup(clustingingKeyBlocksFile, clusteringKeyBlockOffsetsFile, clusteringKeyMeta, clusteringKeyBlockOffsetsMeta);
+            }
+            catch (Throwable t)
+            {
+                throw Throwables.unchecked(t);
+            }
+        }
+
+        @Override
+        @SuppressWarnings({ "resource", "RedundantSuppression" })
+        public PrimaryKeyMap newPerSSTablePrimaryKeyMap() throws IOException
+        {
+            LongArray rowIdToToken = new LongArray.DeferredLongArray(tokenReaderFactory::open);
+            LongArray partitionIdToToken = new LongArray.DeferredLongArray(partitionReaderFactory::open);
+
+            return new WidePrimaryKeyMap(rowIdToToken,
+                                         partitionIdToToken,
+                                         partitionKeyReader.openCursor(),
+                                         clusteringKeyReader.openCursor(),
+                                         partitioner,
+                                         primaryKeyFactory,
+                                         clusteringComparator);
+        }
+
+        @Override
+        public void close()
+        {
+            super.close();
+            FileUtils.closeQuietly(Arrays.asList(clustingingKeyBlocksFile, clusteringKeyBlockOffsetsFile));
+        }
+    }
+
+    private final ClusteringComparator clusteringComparator;
+    private final KeyLookup.Cursor clusteringKeyCursor;
+
+    private WidePrimaryKeyMap(LongArray tokenArray,
+                              LongArray partitionArray,
+                              KeyLookup.Cursor partitionKeyCursor,
+                              KeyLookup.Cursor clusteringKeyCursor,
+                              IPartitioner partitioner,
+                              PrimaryKey.Factory primaryKeyFactory,
+                              ClusteringComparator clusteringComparator)
+    {
+        super(tokenArray, partitionArray, partitionKeyCursor, partitioner, primaryKeyFactory);
+
+        this.clusteringComparator = clusteringComparator;
+        this.clusteringKeyCursor = clusteringKeyCursor;
+    }
+
+    @Override
+    public long rowIdFromPrimaryKey(PrimaryKey primaryKey)
+    {
+        long rowId = tokenArray.indexOf(primaryKey.token().getLongValue());
+
+        // If the key only has a token (initial range skip in the query), the token is out of range,
+        // or we have skipped a token, return the rowId from the token array.
+        if (primaryKey.isTokenOnly() || rowId < 0 || tokenArray.get(rowId) != primaryKey.token().getLongValue())
+            return rowId;
+
+        rowId = tokenCollisionDetection(primaryKey, rowId);
+
+        // Search the key store for the key in the same partition
+        return clusteringKeyCursor.clusteredSeekToKey(clusteringComparator.asByteComparable(primaryKey.clustering()), rowId, startOfNextPartition(rowId));
+    }
+
+    @Override
+    public void close()
+    {
+        super.close();
+        FileUtils.closeQuietly(clusteringKeyCursor);
+    }
+
+    @Override
+    protected PrimaryKey supplier(long sstableRowId)
+    {
+        return primaryKeyFactory.create(readPartitionKey(sstableRowId), readClusteringKey(sstableRowId));
+    }
+
+    private Clustering<?> readClusteringKey(long sstableRowId)
+    {
+        ByteSource.Peekable peekable = ByteSource.peekable(clusteringKeyCursor.seekToPointId(sstableRowId)
+                                                                              .asComparableBytes(ByteComparable.Version.OSS50));
+
+        Clustering<?> clustering = clusteringComparator.clusteringFromByteComparable(ByteBufferAccessor.instance, v -> peekable);
+
+        if (clustering == null)
+            clustering = Clustering.EMPTY;
+
+        return clustering;
+    }
+
+    // Returns the rowId of the next partition or the number of rows if supplied rowId is in the last partition
+    private long startOfNextPartition(long rowId)
+    {
+        long partitionId = partitionArray.get(rowId);
+        long nextPartitionRowId = partitionArray.indexOf(++partitionId);
+        if (nextPartitionRowId == -1)
+            nextPartitionRowId = partitionArray.length();
+        return nextPartitionRowId;
+    }
+}

src/java/org/apache/cassandra/index/sai/disk/v1/bbtree/NumericIndexWriter.java

@@ -148,7 +148,7 @@ public class NumericIndexWriter
             components.put(IndexComponent.BALANCED_TREE, treePosition, treeOffset, treeLength, attributes);
         }
 
-        try (BlockBalancedTreeWalker reader = new BlockBalancedTreeWalker(indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext), treePosition);
+        try (BlockBalancedTreeWalker reader = new BlockBalancedTreeWalker(indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext, null), treePosition);
              IndexOutputWriter postingsOutput = indexDescriptor.openPerIndexOutput(IndexComponent.POSTING_LISTS, indexContext, true))
         {
             long postingsOffset = postingsOutput.getFilePointer();

src/java/org/apache/cassandra/index/sai/disk/v1/bitpack/AbstractBlockPackedReader.java

@@ -34,6 +34,9 @@ public abstract class AbstractBlockPackedReader implements LongArray
     private final byte[] blockBitsPerValue;
     private final SeekingRandomAccessInput input;
 
+    private long previousValue = Long.MIN_VALUE;
+    private long lastIndex; // the last index visited by token -> row ID searches
+
     AbstractBlockPackedReader(IndexInput indexInput, byte[] blockBitsPerValue, int blockShift, int blockMask, long valueCount)
     {
         this.blockShift = blockShift;
@@ -67,5 +70,168 @@ public abstract class AbstractBlockPackedReader implements LongArray
         return valueCount;
     }
 
+    @Override
+    public long indexOf(long value)
+    {
+        // already out of range
+        if (lastIndex >= valueCount)
+            return -1;
+
+        // We keep track previous returned value in lastIndex, so searching backward will not return correct result.
+        // Also it's logically wrong to search backward during token iteration in PostingListRangeIterator.
+        if (value < previousValue)
+            throw new IllegalArgumentException(String.format("%d is smaller than prev token value %d", value, previousValue));
+        previousValue = value;
+
+        int blockIndex = binarySearchBlockMinValues(value);
+
+        // We need to check next block's min value on an exact match.
+        boolean exactMatch = blockIndex >= 0;
+
+        if (blockIndex < 0)
+        {
+            // A non-exact match, which is the negative index of the first value greater than the target.
+            // For example, searching for 4 against min values [3,3,5,7] produces -2, which we convert to 2.
+            blockIndex = -blockIndex;
+        }
+
+        if (blockIndex > 0)
+        {
+            // Start at the previous block, because there could be duplicate values in the previous block.
+            // For example, with block 1: [1,2,3,3] & block 2: [3,3,5,7], binary search for 3 would find
+            // block 2, but we need to start from block 1 and search both.
+            // In case non-exact match, we need to pivot left as target is less than next block's min.
+            blockIndex--;
+        }
+
+        // Find the global (not block-specific) index of the target token, which is equivalent to its row ID:
+        lastIndex = findBlockRowID(value, blockIndex, exactMatch);
+        return lastIndex >= valueCount ? -1 : lastIndex;
+    }
+
+    /**
+     *
+     * @return a positive block index for an exact match, or a negative one for a non-exact match
+     */
+    private int binarySearchBlockMinValues(long targetValue)
+    {
+        int high = Math.toIntExact(blockBitsPerValue.length) - 1;
+
+        // Assume here that we'll never move backward through the blocks:
+        int low = Math.toIntExact(lastIndex >> blockShift);
+
+        // Short-circuit the search if the target is in current block:
+        if (low + 1 <= high)
+        {
+            long cmp = Long.compare(targetValue, delta(low + 1, 0));
+
+            if (cmp == 0)
+            {
+                // We have an exact match, so return the index of the next block, which means we'll start
+                // searching from the current one and also inspect the first value of the next block.
+                return low + 1;
+            }
+            else if (cmp < 0)
+            {
+                // We're in the same block. Indicate a non-exact match, and this value will be both
+                // negated and then decremented to wind up at the current value of "low" here.
+                return -low - 1;
+            }
+
+            // The target is greater than the next block's min value, so advance to that
+            // block before starting the usual search...
+            low++;
+        }
+
+        while (low <= high)
+        {
+            int mid = low + ((high - low) >> 1);
+
+            long midVal = delta(mid, 0);
+
+            if (midVal < targetValue)
+            {
+                low = mid + 1;
+            }
+            else if (midVal > targetValue)
+            {
+                high = mid - 1;
+            }
+            else
+            {
+                // target found, but we need to check for duplicates
+                if (mid > 0 && delta(mid - 1, 0) == targetValue)
+                {
+                    // there are duplicates, pivot left
+                    high = mid - 1;
+                }
+                else
+                {
+                    // no duplicates
+                    return mid;
+                }
+            }
+        }
+
+        return -low; // no exact match found
+    }
+
+    private long findBlockRowID(long targetValue, long blockIdx, boolean exactMatch)
+    {
+        // Calculate the global offset for the selected block:
+        long offset = blockIdx << blockShift;
+
+        // Resume from previous index if it's larger than offset
+        long low = Math.max(lastIndex, offset);
+
+        // The high is either the last local index in the block, or something smaller if the block isn't full:
+        long high = Math.min(offset + blockMask + (exactMatch ? 1 : 0), valueCount - 1);
+
+        return binarySearchBlock(targetValue, low, high);
+    }
+
+    /**
+     * binary search target value between low and high.
+     *
+     * @return index if exact match is found, or *positive* insertion point if no exact match is found.
+     */
+    private long binarySearchBlock(long target, long low, long high)
+    {
+        while (low <= high)
+        {
+            long mid = low + ((high - low) >> 1);
+
+            long midVal = get(mid);
+
+            if (midVal < target)
+            {
+                low = mid + 1;
+                // future rowId cannot be smaller than mid as long as next token not smaller than current token.
+                lastIndex = mid;
+            }
+            else if (midVal > target)
+            {
+                high = mid - 1;
+            }
+            else
+            {
+                // target found, but we need to check for duplicates
+                if (mid > 0 && get(mid - 1) == target)
+                {
+                    // there are duplicates, pivot left
+                    high = mid - 1;
+                }
+                else
+                {
+                    // exact match and no duplicates
+                    return mid;
+                }
+            }
+        }
+
+        // target not found
+        return low;
+    }
+
     abstract long delta(int block, int idx);
 }

src/java/org/apache/cassandra/index/sai/disk/v1/bitpack/NumericValuesWriter.java

@@ -34,44 +34,32 @@ public class NumericValuesWriter implements Closeable
     public static final int MONOTONIC_BLOCK_SIZE = 16384;
     public static final int BLOCK_SIZE = 128;
 
-    private final IndexOutput output;
+    private final IndexOutput indexOutput;
     private final AbstractBlockPackedWriter writer;
     private final MetadataWriter metadataWriter;
     private final String componentName;
     private final int blockSize;
     private long count = 0;
 
-    public NumericValuesWriter(String componentName,
-                               IndexOutput indexOutput,
+    public NumericValuesWriter(IndexDescriptor indexDescriptor,
+                               IndexComponent indexComponent,
                                MetadataWriter metadataWriter,
                                boolean monotonic) throws IOException
     {
-        this(componentName, indexOutput, metadataWriter, monotonic, monotonic ? MONOTONIC_BLOCK_SIZE : BLOCK_SIZE);
+        this(indexDescriptor, indexComponent, metadataWriter, monotonic, monotonic ? MONOTONIC_BLOCK_SIZE : BLOCK_SIZE);
     }
 
     public NumericValuesWriter(IndexDescriptor indexDescriptor,
-                               IndexComponent component,
+                               IndexComponent indexComponent,
                                MetadataWriter metadataWriter,
                                boolean monotonic,
                                int blockSize) throws IOException
     {
-        this(indexDescriptor.componentName(component),
-             indexDescriptor.openPerSSTableOutput(component),
-             metadataWriter,
-             monotonic,
-             blockSize);
-    }
-
-    private NumericValuesWriter(String componentName,
-                                IndexOutput indexOutput,
-                                MetadataWriter metadataWriter,
-                                boolean monotonic, int blockSize) throws IOException
-    {
+        this.componentName = indexDescriptor.componentName(indexComponent);
+        this.indexOutput = indexDescriptor.openPerSSTableOutput(indexComponent);
         SAICodecUtils.writeHeader(indexOutput);
         this.writer = monotonic ? new MonotonicBlockPackedWriter(indexOutput, blockSize)
                                 : new BlockPackedWriter(indexOutput, blockSize);
-        this.output = indexOutput;
-        this.componentName = componentName;
         this.metadataWriter = metadataWriter;
         this.blockSize = blockSize;
     }
@@ -82,13 +70,13 @@ public class NumericValuesWriter implements Closeable
         try (IndexOutput o = metadataWriter.builder(componentName))
         {
             long fp = writer.finish();
-            SAICodecUtils.writeFooter(output);
+            SAICodecUtils.writeFooter(indexOutput);
 
             NumericValuesMeta.write(o, count, blockSize, fp);
         }
         finally
         {
-            output.close();
+            indexOutput.close();
         }
     }
 

src/java/org/apache/cassandra/index/sai/disk/v1/keystore/KeyLookup.java

@@ -0,0 +1,377 @@
+/*
+ * 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.disk.v1.keystore;
+
+import java.io.IOException;
+import javax.annotation.Nonnull;
+import javax.annotation.concurrent.NotThreadSafe;
+
+import com.google.common.annotations.VisibleForTesting;
+
+import org.apache.cassandra.index.sai.disk.io.IndexInputReader;
+import org.apache.cassandra.index.sai.disk.v1.LongArray;
+import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils;
+import org.apache.cassandra.index.sai.disk.v1.bitpack.MonotonicBlockPackedReader;
+import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesMeta;
+import org.apache.cassandra.io.util.FileHandle;
+import org.apache.cassandra.utils.FastByteOperations;
+import org.apache.cassandra.utils.Throwables;
+import org.apache.cassandra.utils.bytecomparable.ByteComparable;
+import org.apache.cassandra.utils.bytecomparable.ByteSource;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.BytesRefBuilder;
+
+/**
+ * Provides read access to an on-disk sequence of partition or clustering keys written by {@link KeyStoreWriter}.
+ * <p>
+ * Care has been taken to make this structure as efficient as possible.
+ * Reading keys does not require allocating data heap buffers per each read operation.
+ * Only one key at a time is loaded to memory.
+ * Low complexity algorithms are used – a lookup of the key by point id is constant time,
+ * and a lookup of the point id by the key is logarithmic.
+ * <p>
+ * Because the blocks are prefix compressed, random access applies only to the locating the whole block.
+ * In order to jump to a concrete key inside the block, the block keys are iterated from the block beginning.
+ *
+ * @see KeyStoreWriter
+ */
+@NotThreadSafe
+public class KeyLookup
+{
+    public static final String INDEX_OUT_OF_BOUNDS = "The target point id [%d] cannot be less than 0 or greater than or equal to the key count [%d]";
+
+    private final FileHandle keysFileHandle;
+    private final KeyLookupMeta keyLookupMeta;
+    private final LongArray.Factory keyBlockOffsetsFactory;
+
+    /**
+     * Creates a new reader based on its data components.
+     * <p>
+     * It does not own the components, so you must close them separately after you're done with the reader.
+     * @param keysFileHandle handle to the file with a sequence of prefix-compressed blocks
+     *                       each storing a fixed number of keys
+     * @param keysBlockOffsets handle to the file containing an encoded sequence of the file offsets pointing to the blocks
+     * @param keyLookupMeta metadata object created earlier by the writer
+     * @param keyBlockOffsetsMeta metadata object for the block offsets
+     */
+    public KeyLookup(@Nonnull FileHandle keysFileHandle,
+                     @Nonnull FileHandle keysBlockOffsets,
+                     @Nonnull KeyLookupMeta keyLookupMeta,
+                     @Nonnull NumericValuesMeta keyBlockOffsetsMeta) throws IOException
+    {
+        this.keysFileHandle = keysFileHandle;
+        this.keyLookupMeta = keyLookupMeta;
+        this.keyBlockOffsetsFactory = new MonotonicBlockPackedReader(keysBlockOffsets, keyBlockOffsetsMeta);
+    }
+
+    /**
+     * Opens a cursor over the keys stored in the keys file.
+     * <p>
+     * This will read the first key into the key buffer and point to the first point in the keys file.
+     * <p>
+     * The cursor is to be used in a single thread.
+     * The cursor is valid as long this object hasn't been closed.
+     * You must close the cursor when you no longer need it.
+     */
+    public @Nonnull Cursor openCursor() throws IOException
+    {
+        return new Cursor(keysFileHandle, keyBlockOffsetsFactory);
+    }
+
+    /**
+     * Allows reading the keys from the keys file.
+     * Can quickly seek to a random key by point id.
+     * <p>
+     * This object is stateful and not thread safe.
+     * It maintains a position to the current key as well as a buffer that can hold one key.
+     */
+    @NotThreadSafe
+    public class Cursor implements AutoCloseable
+    {
+        private final IndexInputReader keysInput;
+        private final int blockShift;
+        private final int blockMask;
+        private final boolean clustering;
+        private final long keysFilePointer;
+        private final LongArray blockOffsets;
+
+        // The key the cursor currently points to. Initially empty.
+        private final BytesRef currentKey;
+
+        // A temporary buffer used to hold the key at the start of the next block.
+        private final BytesRef nextBlockKey;
+
+        // The point id the cursor currently points to.
+        private long currentPointId;
+        private long currentBlockIndex;
+
+        Cursor(FileHandle keysFileHandle, LongArray.Factory blockOffsetsFactory) throws IOException
+        {
+            this.keysInput = IndexInputReader.create(keysFileHandle);
+            SAICodecUtils.validate(this.keysInput);
+            this.blockShift = this.keysInput.readVInt();
+            this.blockMask = (1 << this.blockShift) - 1;
+            this.clustering = this.keysInput.readByte() == 1;
+            this.keysFilePointer = this.keysInput.getFilePointer();
+            this.blockOffsets = new LongArray.DeferredLongArray(blockOffsetsFactory::open);
+            this.currentKey = new BytesRef(keyLookupMeta.maxKeyLength);
+            this.nextBlockKey = new BytesRef(keyLookupMeta.maxKeyLength);
+            keysInput.seek(keysFilePointer);
+            readKey(currentPointId, currentKey);
+        }
+
+        /**
+         * Positions the cursor on the target point id and reads the key at the target to the current key buffer.
+         * <p>
+         * It is allowed to position the cursor before the first item or after the last item;
+         * in these cases the internal buffer is cleared.
+         *
+         * @param pointId point id to lookup
+         * @return The {@link ByteComparable} containing the key
+         * @throws IndexOutOfBoundsException if the target point id is less than -1 or greater than the number of keys
+         */
+        public @Nonnull ByteComparable seekToPointId(long pointId)
+        {
+            if (pointId < 0 || pointId >= keyLookupMeta.keyCount)
+                throw new IndexOutOfBoundsException(String.format(INDEX_OUT_OF_BOUNDS, pointId, keyLookupMeta.keyCount));
+
+            if (pointId != currentPointId)
+            {
+                long blockIndex = pointId >>> blockShift;
+                // We need to reset the block if the block index has changed or the pointId < currentPointId.
+                // We can read forward in the same block without a reset, but we can't read backwards, and token
+                // collision can result in us moving backwards.
+                if (blockIndex != currentBlockIndex || pointId < currentPointId)
+                {
+                    currentBlockIndex = blockIndex;
+                    resetToCurrentBlock();
+                }
+            }
+            while (currentPointId < pointId)
+            {
+                currentPointId++;
+                readCurrentKey();
+                updateCurrentBlockIndex(currentPointId);
+            }
+
+            return ByteComparable.fixedLength(currentKey.bytes, currentKey.offset, currentKey.length);
+        }
+
+        /**
+         * Finds the pointId for a clustering key within a range of pointIds. The start and end of the range must not
+         * exceed the number of keys available. The keys within the range are expected to be in lexographical order.
+         * <p>
+         * If the key is not in the block containing the start of the range a binary search is done to find
+         * the block containing the search key. That block is then searched to return the pointId that corresponds
+         * to the key that is either equal to or next highest to the search key.
+         *
+         * @param key The key to seek for with the partition
+         * @param startingPointId the inclusive starting point for the partition
+         * @param endingPointId the exclusive ending point for the partition.
+         *                      Note: this can be equal to the number of keys if this is the last partition
+         * @return a {@code long} representing the pointId of the key that is >= to the key passed to the method, or
+         *         -1 if the key passed is > all the keys.
+         */
+        public long clusteredSeekToKey(ByteComparable key, long startingPointId, long endingPointId)
+        {
+            assert clustering : "Cannot do a clustered seek to a key on non-clustered keys";
+
+            BytesRef skipKey = asBytesRef(key);
+
+            updateCurrentBlockIndex(startingPointId);
+            resetToCurrentBlock();
+
+            if (compareKeys(currentKey, skipKey) == 0)
+                return startingPointId;
+
+            if (notInCurrentBlock(startingPointId, skipKey))
+            {
+                long split = (endingPointId - startingPointId) >>> blockShift;
+                long splitPointId = startingPointId;
+                while (split > 0)
+                {
+                    updateCurrentBlockIndex(Math.min((splitPointId >>> blockShift) + split, blockOffsets.length() - 1));
+                    resetToCurrentBlock();
+
+                    if (currentPointId >= endingPointId)
+                    {
+                        updateCurrentBlockIndex((endingPointId - 1));
+                        resetToCurrentBlock();
+                    }
+
+                    int cmp = compareKeys(currentKey, skipKey);
+
+                    if (cmp == 0)
+                        return currentPointId;
+
+                    if (cmp < 0)
+                        splitPointId = currentPointId;
+
+                    split /= 2;
+                }
+                // After we finish the binary search we need to move the block back till we hit a block that has
+                // a starting key that is less than or equal to the skip key
+                while (currentBlockIndex > 0 && compareKeys(currentKey, skipKey) > 0)
+                {
+                    currentBlockIndex--;
+                    resetToCurrentBlock();
+                }
+            }
+
+            // Depending on where we are in the block we may need to move forwards to the starting point ID
+            while (currentPointId < startingPointId)
+            {
+                currentPointId++;
+                readCurrentKey();
+                updateCurrentBlockIndex(currentPointId);
+            }
+
+            // Move forward to the ending point ID, returning the point ID if we find our key
+            while (currentPointId < endingPointId)
+            {
+                if (compareKeys(currentKey, skipKey) >= 0)
+                    return currentPointId;
+                currentPointId++;
+                if (currentPointId == keyLookupMeta.keyCount)
+                    return -1;
+                readCurrentKey();
+                updateCurrentBlockIndex(currentPointId);
+            }
+            return endingPointId < keyLookupMeta.keyCount ? endingPointId : -1;
+        }
+
+        @VisibleForTesting
+        public void reset() throws IOException
+        {
+            currentPointId = 0;
+            currentBlockIndex = 0;
+            keysInput.seek(keysFilePointer);
+            readCurrentKey();
+        }
+
+
+        @Override
+        public void close()
+        {
+            keysInput.close();
+        }
+
+        private void updateCurrentBlockIndex(long pointId)
+        {
+            currentBlockIndex = pointId >>> blockShift;
+        }
+
+        private boolean notInCurrentBlock(long pointId, BytesRef key)
+        {
+            if (inLastBlock(pointId))
+                return false;
+
+            // Load the starting value of the next block into nextBlockKey.
+            long blockIndex = (pointId >>> blockShift) + 1;
+            long currentFp = keysInput.getFilePointer();
+            keysInput.seek(blockOffsets.get(blockIndex) + keysFilePointer);
+            readKey(blockIndex << blockShift, nextBlockKey);
+            keysInput.seek(currentFp);
+
+            return compareKeys(key, nextBlockKey) >= 0;
+        }
+
+        private boolean inLastBlock(long pointId)
+        {
+            return pointId >>> blockShift == blockOffsets.length() - 1;
+        }
+
+        // Reset currentPointId and currentKey to be at the start of the block
+        // pointed to by currentBlockIndex.
+        private void resetToCurrentBlock()
+        {
+            keysInput.seek(blockOffsets.get(currentBlockIndex) + keysFilePointer);
+            currentPointId = currentBlockIndex << blockShift;
+            readCurrentKey();
+        }
+
+        private void readCurrentKey()
+        {
+            readKey(currentPointId, currentKey);
+        }
+
+        // Read the next key indicated by pointId.
+        //
+        // Note: pointId is only used to determine whether we are at the start of a block. It is
+        // important that resetPosition is called prior to multiple calls to readKey. It is
+        // easy to get out of position.
+        private void readKey(long pointId, BytesRef key)
+        {
+            try
+            {
+                int prefixLength;
+                int suffixLength;
+                if ((pointId & blockMask) == 0L)
+                {
+                    prefixLength = 0;
+                    suffixLength = keysInput.readVInt();
+                }
+                else
+                {
+                    // Read the prefix and suffix lengths following the compression mechanism described
+                    // in the KeyStoreWriterWriter. If the lengths contained in the starting byte are less
+                    // than the 4 bit maximum then nothing further is read. Otherwise, the lengths in the
+                    // following vints are added.
+                    int compressedLengths = Byte.toUnsignedInt(keysInput.readByte());
+                    prefixLength = compressedLengths & 0x0F;
+                    suffixLength = compressedLengths >>> 4;
+                    if (prefixLength == 15)
+                        prefixLength += keysInput.readVInt();
+                    if (suffixLength == 15)
+                        suffixLength += keysInput.readVInt();
+                }
+
+                assert prefixLength + suffixLength <= keyLookupMeta.maxKeyLength;
+                if (prefixLength + suffixLength > 0)
+                {
+                    key.length = prefixLength + suffixLength;
+                    // The currentKey is appended to as the suffix for the current key is
+                    // added to the existing prefix.
+                    keysInput.readBytes(key.bytes, prefixLength, suffixLength);
+                }
+            }
+            catch (IOException e)
+            {
+                throw Throwables.cleaned(e);
+            }
+        }
+
+        private int compareKeys(BytesRef left, BytesRef right)
+        {
+            return FastByteOperations.compareUnsigned(left.bytes, left.offset, left.offset + left.length,
+                                                      right.bytes, right.offset, right.offset + right.length);
+        }
+
+        private BytesRef asBytesRef(ByteComparable source)
+        {
+            BytesRefBuilder builder = new BytesRefBuilder();
+
+            ByteSource byteSource = source.asComparableBytes(ByteComparable.Version.OSS50);
+            int val;
+            while ((val = byteSource.next()) != ByteSource.END_OF_STREAM)
+                builder.append((byte) val);
+            return builder.get();
+        }
+    }
+}

src/java/org/apache/cassandra/index/sai/disk/v1/keystore/KeyLookupMeta.java

@@ -16,7 +16,7 @@
  * limitations under the License.
  */
 
-package org.apache.cassandra.index.sai.disk.v1.sortedterms;
+package org.apache.cassandra.index.sai.disk.v1.keystore;
 
 import java.io.IOException;
 
@@ -24,25 +24,22 @@ import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.IndexOutput;
 
 /**
- * Metadata produced by {@link SortedTermsWriter}, needed by {@link SortedTermsReader}.
+ * Metadata produced by {@link KeyStoreWriter}, needed by {@link KeyLookup}.
  */
-public class SortedTermsMeta
+public class KeyLookupMeta
 {
-    public final long trieFilePointer;
-    public final long termCount;
-    public final int maxTermLength;
+    public final long keyCount;
+    public final int maxKeyLength;
 
-    public SortedTermsMeta(DataInput input) throws IOException
+    public KeyLookupMeta(DataInput input) throws IOException
     {
-        this.trieFilePointer = input.readLong();
-        this.termCount = input.readLong();
-        this.maxTermLength = input.readInt();
+        this.keyCount = input.readLong();
+        this.maxKeyLength = input.readInt();
     }
 
-    public static void write(IndexOutput output, long trieFilePointer, long termCount, int maxTermLength) throws IOException
+    public static void write(IndexOutput output, long keyCount, int maxKeyLength) throws IOException
     {
-        output.writeLong(trieFilePointer);
-        output.writeLong(termCount);
-        output.writeInt(maxTermLength);
+        output.writeLong(keyCount);
+        output.writeInt(maxKeyLength);
     }
 }

src/java/org/apache/cassandra/index/sai/disk/v1/keystore/KeyStoreWriter.java

@@ -0,0 +1,218 @@
+/*
+ * 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.disk.v1.keystore;
+
+import java.io.Closeable;
+import java.io.IOException;
+import javax.annotation.Nonnull;
+import javax.annotation.concurrent.NotThreadSafe;
+
+import org.apache.cassandra.index.sai.disk.v1.MetadataWriter;
+import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils;
+import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesWriter;
+import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.utils.FastByteOperations;
+import org.apache.cassandra.utils.bytecomparable.ByteComparable;
+import org.apache.cassandra.utils.bytecomparable.ByteSource;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.BytesRefBuilder;
+import org.apache.lucene.util.StringHelper;
+
+/**
+ * Writes a sequence of partition keys or clustering keys for use with {@link KeyLookup}.
+ * <p>
+ * Partition keys are written unordered and clustering keys are written in ordered partitions determined by calls to
+ * {@link #startPartition()}. In either case keys can be of varying lengths.
+ * <p>
+ * The {@link #blockShift} field is used to quickly determine the id of the current block
+ * based on a point id or to check if we are exactly at the beginning of the block.
+ * <p>
+ * Keys are organized in blocks of (2 ^ {@link #blockShift}) keys.
+ * <p>
+ * The blocks should not be too small because they allow prefix compression of the keys except the first key in a block.
+ * <p>
+ * The blocks should not be too large because we can't just randomly jump to the key inside the block, but we have to
+ * iterate through all the keys from the start of the block.
+ *
+ * @see KeyLookup
+ */
+@NotThreadSafe
+public class KeyStoreWriter implements Closeable
+{
+    private final int blockShift;
+    private final int blockMask;
+    private final boolean clustering;
+    private final IndexOutput keysOutput;
+    private final NumericValuesWriter offsetsWriter;
+    private final String componentName;
+    private final MetadataWriter metadataWriter;
+
+    private BytesRefBuilder prevKey = new BytesRefBuilder();
+    private BytesRefBuilder tempKey = new BytesRefBuilder();
+
+    private final long bytesStartFP;
+
+    private boolean inPartition = false;
+    private int maxKeyLength = -1;
+    private long pointId = 0;
+
+    /**
+     * Creates a new writer.
+     * <p>
+     * It does not own the components, so you must close the components by yourself
+     * after you're done with the writer.
+     *
+     * @param componentName the component name for the {@link KeyLookupMeta}
+     * @param metadataWriter the {@link MetadataWriter} for storing the {@link KeyLookupMeta}
+     * @param keysOutput where to write the prefix-compressed keys
+     * @param keysBlockOffsets  where to write the offsets of each block of keys
+     * @param blockShift the block shift that is used to determine the block size
+     * @param clustering determines whether the keys will be written as ordered partitions
+     */
+    public KeyStoreWriter(String componentName,
+                          MetadataWriter metadataWriter,
+                          IndexOutput keysOutput,
+                          NumericValuesWriter keysBlockOffsets,
+                          int blockShift,
+                          boolean clustering) throws IOException
+    {
+        this.componentName = componentName;
+        this.metadataWriter = metadataWriter;
+        SAICodecUtils.writeHeader(keysOutput);
+        this.blockShift = blockShift;
+        this.blockMask = (1 << this.blockShift) - 1;
+        this.clustering = clustering;
+        this.keysOutput = keysOutput;
+        this.keysOutput.writeVInt(blockShift);
+        this.keysOutput.writeByte((byte ) (clustering ? 1 : 0));
+        this.bytesStartFP = keysOutput.getFilePointer();
+        this.offsetsWriter = keysBlockOffsets;
+    }
+
+    public void startPartition()
+    {
+        assert clustering : "Cannot start a partition on a non-clustering key store";
+
+        inPartition = false;
+    }
+
+    /**
+     * Appends a key at the end of the sequence.
+     *
+     * @throws IOException if write to disk fails
+     * @throws IllegalArgumentException if the key is not greater than the previous added key
+     */
+    public void add(final @Nonnull ByteComparable key) throws IOException
+    {
+        tempKey.clear();
+        copyBytes(key, tempKey);
+
+        BytesRef keyRef = tempKey.get();
+
+        if (clustering && inPartition)
+        {
+            if (compareKeys(keyRef, prevKey.get()) <= 0)
+                throw new IllegalArgumentException("Clustering keys must be in ascending lexographical order");
+        }
+
+        inPartition = true;
+
+        writeKey(keyRef);
+
+        maxKeyLength = Math.max(maxKeyLength, keyRef.length);
+
+        BytesRefBuilder temp = this.tempKey;
+        this.tempKey = this.prevKey;
+        this.prevKey = temp;
+
+        pointId++;
+    }
+
+    private void writeKey(BytesRef key) throws IOException
+    {
+        if ((pointId & blockMask) == 0)
+        {
+            offsetsWriter.add(keysOutput.getFilePointer() - bytesStartFP);
+
+            keysOutput.writeVInt(key.length);
+            keysOutput.writeBytes(key.bytes, key.offset, key.length);
+        }
+        else
+        {
+            int prefixLength = 0;
+            int suffixLength = 0;
+
+            // If the key is the same as the previous key then we use prefix and suffix lengths of 0.
+            // This means that we store a byte of 0 and don't write any data for the key.
+            if (compareKeys(prevKey.get(), key) != 0)
+            {
+                prefixLength = StringHelper.bytesDifference(prevKey.get(), key);
+                suffixLength = key.length - prefixLength;
+            }
+            // The prefix and suffix lengths are written as a byte followed by up to 2 vints. An attempt is
+            // made to compress the lengths into the byte (if prefix length < 15 and/or suffix length < 15).
+            // If either length exceeds the compressed byte maximum, it is written as a vint following the byte.
+            keysOutput.writeByte((byte) (Math.min(prefixLength, 15) | (Math.min(15, suffixLength) << 4)));
+
+            if (prefixLength + suffixLength > 0)
+            {
+                if (prefixLength >= 15)
+                    keysOutput.writeVInt(prefixLength - 15);
+                if (suffixLength >= 15)
+                    keysOutput.writeVInt(suffixLength - 15);
+
+                keysOutput.writeBytes(key.bytes, key.offset + prefixLength, key.length - prefixLength);
+            }
+        }
+    }
+
+    /**
+     * Flushes any in-memory buffers to the output streams.
+     * Does not close the output streams.
+     * No more writes are allowed.
+     */
+    @Override
+    public void close() throws IOException
+    {
+        try (IndexOutput output = metadataWriter.builder(componentName))
+        {
+            SAICodecUtils.writeFooter(keysOutput);
+            KeyLookupMeta.write(output, pointId, maxKeyLength);
+        }
+        finally
+        {
+            FileUtils.close(offsetsWriter, keysOutput);
+        }
+    }
+
+    private int compareKeys(BytesRef left, BytesRef right)
+    {
+        return FastByteOperations.compareUnsigned(left.bytes, left.offset, left.offset + left.length,
+                                                  right.bytes, right.offset, right.offset + right.length);
+    }
+
+    private void copyBytes(ByteComparable source, BytesRefBuilder dest)
+    {
+        ByteSource byteSource = source.asComparableBytes(ByteComparable.Version.OSS50);
+        int val;
+        while ((val = byteSource.next()) != ByteSource.END_OF_STREAM)
+            dest.append((byte) val);
+    }
+}

src/java/org/apache/cassandra/index/sai/disk/v1/postings/PostingsReader.java

@@ -148,6 +148,12 @@ public class PostingsReader implements OrdinalPostingList
             {
                 return length;
             }
+
+            @Override
+            public long indexOf(long value)
+            {
+                throw new UnsupportedOperationException();
+            }
         }
     }
 

src/java/org/apache/cassandra/index/sai/disk/v1/sortedterms/SortedTermsReader.java

@@ -1,234 +0,0 @@
-/*
- * 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.disk.v1.sortedterms;
-
-import java.io.IOException;
-import javax.annotation.Nonnull;
-import javax.annotation.concurrent.NotThreadSafe;
-
-import com.google.common.annotations.VisibleForTesting;
-
-import org.apache.cassandra.index.sai.disk.io.IndexInputReader;
-import org.apache.cassandra.index.sai.disk.v1.LongArray;
-import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils;
-import org.apache.cassandra.index.sai.disk.v1.bitpack.MonotonicBlockPackedReader;
-import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesMeta;
-import org.apache.cassandra.io.util.FileHandle;
-import org.apache.cassandra.utils.bytecomparable.ByteComparable;
-import org.apache.lucene.util.BytesRef;
-
-import static org.apache.cassandra.index.sai.disk.v1.sortedterms.SortedTermsWriter.TERMS_DICT_BLOCK_MASK;
-import static org.apache.cassandra.index.sai.disk.v1.sortedterms.SortedTermsWriter.TERMS_DICT_BLOCK_SHIFT;
-
-/**
- * Provides read access to a sorted on-disk sequence of terms written by {@link SortedTermsWriter}.
- * <p>
- * Allows constant-time look up of the term at a given point id.
- * <p>
- * Care has been taken to make this structure as efficient as possible.
- * Reading terms does not require allocating data heap buffers per each read operation.
- * Only one term at a time is loaded to memory.
- * Low complexity algorithms are used – a lookup of the term by point id is constant time,
- * and a lookup of the point id by the term is logarithmic.
- * <p>
- * Because the blocks are prefix compressed, random access applies only to the locating the whole block.
- * In order to jump to a concrete term inside the block, the block terms are iterated from the block beginning.
- * Expect random access by {@link Cursor#seekToPointId(long)} to be slower
- * than just moving to the next term with {@link Cursor#advance()}.
- * <p>
- * For documentation of the underlying on-disk data structures, see the package documentation.
- *
- * @see SortedTermsWriter
- * @see org.apache.cassandra.index.sai.disk.v1.sortedterms
- */
-@NotThreadSafe
-public class SortedTermsReader
-{
-    private static final long BEFORE_START = -1;
-
-    private final FileHandle termsData;
-    private final SortedTermsMeta meta;
-    private final LongArray.Factory blockOffsetsFactory;
-
-    /**
-     * Creates a new reader based on its data components.
-     * <p>
-     * It does not own the components, so you must close them separately after you're done with the reader.
-     * @param termsDataFileHandle handle to the file with a sequence of prefix-compressed blocks
-     *                  each storing a fixed number of terms
-     * @param termsDataBlockOffsets handle to the file containing an encoded sequence of the file offsets pointing to the blocks
-     * @param meta metadata object created earlier by the writer
-     * @param blockOffsetsMeta metadata object for the block offsets
-     */
-    public SortedTermsReader(@Nonnull FileHandle termsDataFileHandle,
-                             @Nonnull FileHandle termsDataBlockOffsets,
-                             @Nonnull SortedTermsMeta meta,
-                             @Nonnull NumericValuesMeta blockOffsetsMeta) throws IOException
-    {
-        this.termsData = termsDataFileHandle;
-        this.meta = meta;
-        this.blockOffsetsFactory = new MonotonicBlockPackedReader(termsDataBlockOffsets, blockOffsetsMeta);
-    }
-
-    /**
-     * Opens a cursor over the terms stored in the terms file.
-     * <p>
-     * This does not read any data yet.
-     * The cursor is initially positioned before the first item.
-     * <p>
-     * The cursor is to be used in a single thread.
-     * The cursor is valid as long this object hasn't been closed.
-     * You must close the cursor when you no longer need it.
-     */
-    public @Nonnull Cursor openCursor() throws IOException
-    {
-        return new Cursor(termsData, blockOffsetsFactory);
-    }
-
-    /**
-     * Allows reading the terms from the terms file.
-     * Can quickly seek to a random term by point id.
-     * <p>
-     * This object is stateful and not thread safe.
-     * It maintains a position to the current term as well as a buffer that can hold one term.
-     */
-    @NotThreadSafe
-    public class Cursor implements AutoCloseable
-    {
-        private final IndexInputReader termsInput;
-        private final long termsDataFp;
-        private final LongArray blockOffsets;
-
-        // The term the cursor currently points to. Initially empty.
-        private final BytesRef currentTerm;
-
-        // The point id the cursor currently points to. BEFORE_START means before the first item.
-        private long pointId = BEFORE_START;
-
-        Cursor(FileHandle termsFile, LongArray.Factory blockOffsetsFactory) throws IOException
-        {
-            this.termsInput = IndexInputReader.create(termsFile);
-            SAICodecUtils.validate(this.termsInput);
-            this.termsDataFp = this.termsInput.getFilePointer();
-            this.blockOffsets = new LongArray.DeferredLongArray(blockOffsetsFactory::open);
-            this.currentTerm = new BytesRef(meta.maxTermLength);
-        }
-
-        /**
-         * Returns the current position of the cursor.
-         * Initially, before the first call to {@link #advance}, the cursor is positioned at -1.
-         * After reading all the items, the cursor is positioned at index one
-         * greater than the position of the last item.
-         */
-        public long pointId()
-        {
-            return pointId;
-        }
-
-        /**
-         * Returns the current term data as {@link ByteComparable} referencing the internal term buffer.
-         * The term data stored behind that reference is valid only until the next call to
-         * {@link #advance} or {@link #seekToPointId(long)}.
-         */
-        public @Nonnull ByteComparable term()
-        {
-            return ByteComparable.fixedLength(currentTerm.bytes, currentTerm.offset, currentTerm.length);
-        }
-
-        /**
-         * Positions the cursor on the target point id and reads the term at target to the current term buffer.
-         * <p>
-         * It is allowed to position the cursor before the first item or after the last item;
-         * in these cases the internal buffer is cleared.
-         * <p>
-         * This method has constant complexity.
-         *
-         * @param pointId point id to lookup
-         * @throws IOException if a seek and read from the terms file fails
-         * @throws IndexOutOfBoundsException if the target point id is less than -1 or greater than the number of terms
-         */
-        public void seekToPointId(long pointId) throws IOException
-        {
-            if (pointId < 0 || pointId > meta.termCount)
-                throw new IndexOutOfBoundsException(String.format("The target point id [%s] cannot be less than 0 or " +
-                                                                  "greater than the term count [%s]", pointId, meta.termCount));
-            long blockIndex = pointId >>> TERMS_DICT_BLOCK_SHIFT;
-            long blockAddress = blockOffsets.get(blockIndex);
-            termsInput.seek(blockAddress + termsDataFp);
-            this.pointId = (blockIndex << TERMS_DICT_BLOCK_SHIFT) - 1;
-            while (this.pointId < pointId && advance());
-        }
-
-        @Override
-        public void close()
-        {
-            termsInput.close();
-        }
-
-        /**
-         * Advances the cursor to the next term and reads it into the current term buffer.
-         * <p>
-         * If there are no more available terms, clears the term buffer and the cursor's position will point to the
-         * one behind the last item.
-         * <p>
-         * This method has constant time complexity.
-         *
-         * @return true if the cursor was advanced successfully, false if the end of file was reached
-         * @throws IOException if a read from the terms file fails
-         */
-        @VisibleForTesting
-        protected boolean advance() throws IOException
-        {
-            if (pointId >= meta.termCount || ++pointId >= meta.termCount)
-            {
-                currentTerm.length = 0;
-                return false;
-            }
-
-            int prefixLength;
-            int suffixLength;
-            if ((pointId & TERMS_DICT_BLOCK_MASK) == 0L)
-            {
-                prefixLength = 0;
-                suffixLength = termsInput.readVInt();
-            }
-            else
-            {
-                // Read the prefix and suffix lengths following the compression mechanism described
-                // in the SortedTermsWriter. If the lengths contained in the starting byte are less
-                // than the 4 bit maximum then nothing further is read. Otherwise, the lengths in the
-                // following vints are added.
-                int compressedLengths = Byte.toUnsignedInt(termsInput.readByte());
-                prefixLength = compressedLengths & 0x0F;
-                suffixLength = 1 + (compressedLengths >>> 4);
-                if (prefixLength == 15)
-                    prefixLength += termsInput.readVInt();
-                if (suffixLength == 16)
-                    suffixLength += termsInput.readVInt();
-            }
-
-            assert prefixLength + suffixLength <= meta.maxTermLength;
-            currentTerm.length = prefixLength + suffixLength;
-            // The currentTerm is appended to as the suffix for the current term is
-            // added to the existing prefix.
-            termsInput.readBytes(currentTerm.bytes, prefixLength, suffixLength);
-            return true;
-        }
-    }
-}

src/java/org/apache/cassandra/index/sai/disk/v1/sortedterms/SortedTermsWriter.java

@@ -1,214 +0,0 @@
-/*
- * 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.disk.v1.sortedterms;
-
-import java.io.Closeable;
-import java.io.IOException;
-import java.util.Arrays;
-import javax.annotation.Nonnull;
-import javax.annotation.concurrent.NotThreadSafe;
-
-import com.google.common.base.Preconditions;
-
-import org.apache.cassandra.index.sai.disk.io.IndexOutputWriter;
-import org.apache.cassandra.index.sai.disk.v1.MetadataWriter;
-import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesWriter;
-import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils;
-import org.apache.cassandra.io.tries.IncrementalDeepTrieWriterPageAware;
-import org.apache.cassandra.io.util.FileUtils;
-import org.apache.cassandra.utils.bytecomparable.ByteComparable;
-import org.apache.cassandra.utils.bytecomparable.ByteSource;
-import org.apache.lucene.store.IndexOutput;
-import org.apache.lucene.util.BytesRef;
-import org.apache.lucene.util.BytesRefBuilder;
-import org.apache.lucene.util.StringHelper;
-
-import static org.apache.cassandra.index.sai.disk.v1.trie.TrieTermsDictionaryReader.trieSerializer;
-
-/**
- * Writes an ordered sequence of terms for use with {@link SortedTermsReader}.
- * <p>
- * Terms must be added in lexicographical ascending order.
- * Terms can be of varying lengths.
- * <p>
- * For documentation of the underlying on-disk data structures, see the package documentation.
- * <p>
- * The TERMS_DICT_ constants allow for quickly determining the id of the current block based on a point id
- * or to check if we are exactly at the beginning of the block.
- * Terms data are organized in blocks of (2 ^ {@link #TERMS_DICT_BLOCK_SHIFT}) terms.
- * The blocks should not be too small because they allow prefix compression of
- * the terms except the first term in a block.
- * The blocks should not be too large because we can't just randomly jump to the term inside the block,
- * but we have to iterate through all the terms from the start of the block.
- *
- * @see SortedTermsReader
- * @see org.apache.cassandra.index.sai.disk.v1.sortedterms
- */
-@NotThreadSafe
-public class SortedTermsWriter implements Closeable
-{
-    static final int TERMS_DICT_BLOCK_SHIFT = 4;
-    static final int TERMS_DICT_BLOCK_SIZE = 1 << TERMS_DICT_BLOCK_SHIFT;
-    static final int TERMS_DICT_BLOCK_MASK = TERMS_DICT_BLOCK_SIZE - 1;
-
-    private final IncrementalDeepTrieWriterPageAware<Long> trieWriter;
-    private final IndexOutput trieOutput;
-    private final IndexOutput termsOutput;
-    private final NumericValuesWriter offsetsWriter;
-    private final String componentName;
-    private final MetadataWriter metadataWriter;
-
-    private BytesRefBuilder prevTerm = new BytesRefBuilder();
-    private BytesRefBuilder tempTerm = new BytesRefBuilder();
-
-    private final long bytesStartFP;
-
-    private int maxLength = -1;
-    private long pointId = 0;
-
-    /**
-     * Creates a new writer.
-     * <p>
-     * It does not own the components, so you must close the components by yourself
-     * after you're done with the writer.
-     *
-     * @param componentName the component name for the {@link SortedTermsMeta}
-     * @param metadataWriter the {@link MetadataWriter} for storing the {@link SortedTermsMeta}
-     * @param termsData where to write the prefix-compressed terms data
-     * @param termsDataBlockOffsets  where to write the offsets of each block of terms data
-     * @param trieOutput where to write the trie that maps the terms to point ids
-     */
-    public SortedTermsWriter(String componentName,
-                             MetadataWriter metadataWriter,
-                             IndexOutput termsData,
-                             NumericValuesWriter termsDataBlockOffsets,
-                             IndexOutputWriter trieOutput) throws IOException
-    {
-        this.componentName = componentName;
-        this.metadataWriter = metadataWriter;
-        this.trieOutput = trieOutput;
-        SAICodecUtils.writeHeader(this.trieOutput);
-        this.trieWriter = new IncrementalDeepTrieWriterPageAware<>(trieSerializer, trieOutput.asSequentialWriter());
-        SAICodecUtils.writeHeader(termsData);
-        this.termsOutput = termsData;
-        this.bytesStartFP = termsData.getFilePointer();
-        this.offsetsWriter = termsDataBlockOffsets;
-    }
-
-    /**
-     * Appends a term at the end of the sequence.
-     * Terms must be added in lexicographic order.
-     *
-     * @throws IOException if write to disk fails
-     * @throws IllegalArgumentException if the term is not greater than the previous added term
-     */
-    public void add(final @Nonnull ByteComparable term) throws IOException
-    {
-        tempTerm.clear();
-        copyBytes(term, tempTerm);
-
-        BytesRef termRef = tempTerm.get();
-        BytesRef prevTermRef = this.prevTerm.get();
-
-        Preconditions.checkArgument(prevTermRef.length == 0 || prevTermRef.compareTo(termRef) < 0,
-                                    "Terms must be added in lexicographic ascending order.");
-        writeTermData(termRef);
-        writeTermToTrie(term);
-
-        maxLength = Math.max(maxLength, termRef.length);
-        swapTempWithPrevious();
-        pointId++;
-    }
-
-    private void writeTermToTrie(ByteComparable term) throws IOException
-    {
-        trieWriter.add(term, pointId);
-    }
-
-    private void writeTermData(BytesRef term) throws IOException
-    {
-        if ((pointId & TERMS_DICT_BLOCK_MASK) == 0)
-        {
-            offsetsWriter.add(termsOutput.getFilePointer() - bytesStartFP);
-
-            termsOutput.writeVInt(term.length);
-            termsOutput.writeBytes(term.bytes, term.offset, term.length);
-        }
-        else
-        {
-            int prefixLength = StringHelper.bytesDifference(prevTerm.get(), term);
-            int suffixLength = term.length - prefixLength;
-            assert suffixLength > 0: "terms must be unique";
-
-            // The prefix and suffix lengths are written as a byte followed by up to 2 vints. An attempt is
-            // made to compress the lengths into the byte (if prefix length < 15 and/or suffix length < 16).
-            // If either length exceeds the compressed byte maximum, it is written as a vint following the byte.
-            termsOutput.writeByte((byte) (Math.min(prefixLength, 15) | (Math.min(15, suffixLength - 1) << 4)));
-            if (prefixLength >= 15)
-                termsOutput.writeVInt(prefixLength - 15);
-            if (suffixLength >= 16)
-                termsOutput.writeVInt(suffixLength - 16);
-
-            termsOutput.writeBytes(term.bytes, term.offset + prefixLength, term.length - prefixLength);
-        }
-    }
-
-    /**
-     * Flushes any in-memory buffers to the output streams.
-     * Does not close the output streams.
-     * No more writes are allowed.
-     */
-    @Override
-    public void close() throws IOException
-    {
-        try (IndexOutput output = metadataWriter.builder(componentName))
-        {
-            long trieFilePointer = this.trieWriter.complete();
-            SAICodecUtils.writeFooter(trieOutput);
-            SAICodecUtils.writeFooter(termsOutput);
-            SortedTermsMeta.write(output, trieFilePointer, pointId, maxLength);
-            // Don't close the offsets writer quietly because of the work it does
-            // during its close. We need to propagate the error.
-            offsetsWriter.close();
-        }
-        finally
-        {
-            FileUtils.closeQuietly(Arrays.asList(trieWriter, trieOutput, termsOutput));
-        }
-    }
-
-    private void copyBytes(ByteComparable source, BytesRefBuilder dest)
-    {
-        ByteSource byteSource = source.asComparableBytes(ByteComparable.Version.OSS50);
-        int val;
-        while ((val = byteSource.next()) != ByteSource.END_OF_STREAM)
-            dest.append((byte) val);
-    }
-
-    /**
-     * Swaps {@link #tempTerm} with {@link #prevTerm}.
-     * It is faster to swap the pointers instead of copying the data.
-     */
-    private void swapTempWithPrevious()
-    {
-        BytesRefBuilder temp = this.tempTerm;
-        this.tempTerm = this.prevTerm;
-        this.prevTerm = temp;
-    }
-}

src/java/org/apache/cassandra/index/sai/disk/v1/sortedterms/package-info.java

@@ -1,58 +0,0 @@
-/*
- * 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.
- */
-
-/**
- * Space-efficient on-disk data structure for storing a sorted sequence of terms.
- * Provides efficient lookup of terms by their point id, as well as locating them by contents.
- * <p>
- * All the code in the package uses the following teminology:
- * <ul>
- *     <li>Term: arbitrary data provided by the user as a bunch of bytes. Terms can be of variable length.</li>
- *     <li>Point id: the ordinal position of a term in the sequence, 0-based.</li>
- * </ul>
- *
- * Terms are stored in <code>ByteComparable</code> strictly ascending order.
- * Duplicates are not allowed.
- *
- * <p>
- * The structure is immutable, i.e. cannot be modified nor appended after writing to disk is completed.
- * You build it by adding terms in the ascending order using
- * {@link org.apache.cassandra.index.sai.disk.v1.sortedterms.SortedTermsWriter}.
- * Once saved to disk, you can open it for lookups with
- * {@link org.apache.cassandra.index.sai.disk.v1.sortedterms.SortedTermsReader}.
- *
- * <p>
- * The data structure comprises the following components, each stored in a separate file:
- * <ul>
- *     <li>terms data, organized as a sequence of prefix-compressed blocks each storing
- *     {@link org.apache.cassandra.index.sai.disk.v1.sortedterms.SortedTermsWriter#TERMS_DICT_BLOCK_SIZE} terms</li>
- *     <li>a monotonic list of file offsets of the blocks; this component allows to quickly locate the block
- *     that contains the term with a given point id</li>
- *     <li>a trie indexed by terms, with a long payload for the point id,
- *     to quickly locate the point id of a term by the term contents
- *     </li>
- * </ul>
- * </p>
- *
- * The implementation has been based on code from Lucene version 7.5 {@link org.apache.lucene.index.SortedDocValues}.
- * Prefix compression and bitpacking are used extensively to save space.
- * <p>
- * The sorted terms data structure is used for the storage and reading of {@link org.apache.cassandra.index.sai.utils.PrimaryKey}s
- * by the {@link org.apache.cassandra.index.sai.disk.v1.RowAwarePrimaryKeyMap}.
- */
-package org.apache.cassandra.index.sai.disk.v1.sortedterms;

src/java/org/apache/cassandra/index/sai/disk/v1/trie/TriePrefixSearcher.java

@@ -1,98 +0,0 @@
-/*
- * 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.disk.v1.trie;
-
-import javax.annotation.concurrent.NotThreadSafe;
-
-import org.apache.cassandra.io.tries.ValueIterator;
-import org.apache.cassandra.io.tries.Walker;
-import org.apache.cassandra.io.util.Rebufferer;
-import org.apache.cassandra.io.util.SizedInts;
-import org.apache.cassandra.utils.bytecomparable.ByteSource;
-
-/**
- * Thread-unsafe specialized value -> long trie prefix searcher.
- * <p>
- * This is a specialization of the {@link Walker} that will look for a prefix or
- * a complete term and return the first payload associated with the term.
- * <p>
- * TODO Make generic to handle any payload type.
- * TODO Extend search to return first payload or all sub-payloads as iterator (LIKE support?)
- */
-@NotThreadSafe
-public class TriePrefixSearcher extends ValueIterator<TriePrefixSearcher>
-{
-    public static final long NOT_FOUND = -1L;
-
-    public TriePrefixSearcher(Rebufferer source, long root)
-    {
-        super(source, root, false);
-    }
-
-    public long prefixSearch(ByteSource startStream)
-    {
-        IterationPosition prev = null;
-        int childIndex;
-        long payloadedNode = -1;
-
-        try
-        {
-            // Follow start position while we still have a prefix, stacking path and saving prefixes.
-            go(root);
-            while (true)
-            {
-                int s = startStream.next();
-                childIndex = search(s);
-
-                // For a separator trie the latest payload met along the prefix is a potential match for start
-                if (childIndex == 0 || childIndex == -1)
-                {
-                    if (hasPayload())
-                        payloadedNode = position;
-                }
-                else
-                    payloadedNode = -1;
-
-                if (childIndex < 0)
-                    break;
-
-                prev = new IterationPosition(position, childIndex, 256, prev);
-                go(transition(childIndex));
-            }
-
-            childIndex = -1 - childIndex - 1;
-
-            stack = new IterationPosition(position, childIndex, 256, prev);
-
-            // Advancing now gives us first match if we didn't find one already.
-            if (payloadedNode == -1)
-                advanceNode();
-
-            if (stack == null || !hasPayload())
-                return NOT_FOUND;
-
-            return SizedInts.read(buf, payloadPosition(), payloadFlags());
-        }
-        catch (Throwable t)
-        {
-            super.close();
-            throw t;
-        }
-    }
-}

src/java/org/apache/cassandra/index/sai/iterators/KeyRangeConcatIterator.java

@@ -40,6 +40,7 @@ import org.apache.cassandra.io.util.FileUtils;
  */
 public class KeyRangeConcatIterator extends KeyRangeIterator
 {
+    public static final String MUST_BE_SORTED_ERROR = "RangeIterator must be sorted, previous max: %s, next min: %s";
     private final PriorityQueue<KeyRangeIterator> ranges;
     private final List<KeyRangeIterator> toRelease;
 
@@ -166,7 +167,7 @@ public class KeyRangeConcatIterator extends KeyRangeIterator
                 }
                 else if (count > 0 && max.compareTo(range.getMinimum()) > 0)
                 {
-                    throw new IllegalArgumentException("RangeIterator must be sorted, previous max: " + max + ", next min: " + range.getMinimum());
+                    throw new IllegalArgumentException(String.format(MUST_BE_SORTED_ERROR, max, range.getMinimum()));
                 }
 
                 max = range.getMaximum();

src/java/org/apache/cassandra/index/sai/utils/PrimaryKey.java

@@ -37,7 +37,7 @@ import org.apache.cassandra.utils.bytecomparable.ByteSource;
  * Representation of the primary key for a row consisting of the {@link DecoratedKey} and
  * {@link Clustering} associated with a {@link org.apache.cassandra.db.rows.Row}.
  */
-public interface PrimaryKey extends Comparable<PrimaryKey>
+public interface PrimaryKey extends Comparable<PrimaryKey>, ByteComparable
 {
     class Factory
     {
@@ -50,7 +50,7 @@ public interface PrimaryKey extends Comparable<PrimaryKey>
 
         /**
          * Creates a {@link PrimaryKey} that is represented by a {@link Token}.
-         *
+         * <p>
          * {@link Token} only primary keys are used for defining the partition range
          * of a query.
          */
@@ -58,14 +58,14 @@ public interface PrimaryKey extends Comparable<PrimaryKey>
         {
             assert token != null : "Cannot create a primary key with a null token";
 
-            return new ImmutablePrimaryKey(token, null, null);
+            return new TokenOnlyPrimaryKey(token);
         }
 
         public PrimaryKey createPartitionKeyOnly(DecoratedKey partitionKey)
         {
             assert partitionKey != null : "Cannot create a primary key with a null partition key";
 
-            return new ImmutablePrimaryKey(partitionKey.getToken(), partitionKey, null);
+            return new ImmutablePrimaryKey(partitionKey, null);
         }
 
         /**
@@ -77,7 +77,7 @@ public interface PrimaryKey extends Comparable<PrimaryKey>
             assert partitionKey != null : "Cannot create a primary key with a null partition key";
             assert clustering != null : "Cannot create a primary key with a null clustering";
 
-            return new ImmutablePrimaryKey(partitionKey.getToken(), partitionKey, clustering);
+            return new ImmutablePrimaryKey(partitionKey, clustering);
         }
 
         public PrimaryKey createDeferred(Token token, Supplier<PrimaryKey> primaryKeySupplier)
@@ -90,40 +90,34 @@ public interface PrimaryKey extends Comparable<PrimaryKey>
 
         abstract class AbstractPrimaryKey implements PrimaryKey
         {
+            @Override
+            @SuppressWarnings("ConstantConditions")
             public ByteSource asComparableBytes(ByteComparable.Version version)
             {
-                ByteSource tokenComparable = token().asComparableBytes(version);
-                if (partitionKey() == null)
-                    return ByteSource.withTerminator(version == ByteComparable.Version.LEGACY ? ByteSource.END_OF_STREAM
-                                                                                              : ByteSource.TERMINATOR,
-                                                     tokenComparable,
-                                                     null,
-                                                     null);
                 ByteSource keyComparable = ByteSource.of(partitionKey().getKey(), version);
+                if (clusteringComparator.size() == 0)
+                    return keyComparable;
                 // It is important that the ClusteringComparator.asBytesComparable method is used
                 // to maintain the correct clustering sort order
-                ByteSource clusteringComparable = clusteringComparator == null ||
-                                                  clusteringComparator.size() == 0 ||
-                                                  clustering() == null ||
+                ByteSource clusteringComparable = clustering() == null ||
                                                   clustering().isEmpty() ? null
                                                                          : clusteringComparator.asByteComparable(clustering())
                                                                                                .asComparableBytes(version);
                 return ByteSource.withTerminator(version == ByteComparable.Version.LEGACY ? ByteSource.END_OF_STREAM
                                                                                           : ByteSource.TERMINATOR,
-                                                 tokenComparable,
                                                  keyComparable,
                                                  clusteringComparable);
             }
 
             @Override
+            @SuppressWarnings("ConstantConditions")
             public int compareTo(PrimaryKey o)
             {
                 int cmp = token().compareTo(o.token());
 
                 // If the tokens don't match then we don't need to compare any more of the key.
-                // Otherwise, if it's partition key is null or the other partition key is null
-                // then one or both of the keys are token only so we can only compare tokens
-                if ((cmp != 0) || (partitionKey() == null) || o.partitionKey() == null)
+                // Otherwise, if either of the keys are token only we can only compare tokens
+                if ((cmp != 0) || isTokenOnly() || o.isTokenOnly())
                     return cmp;
 
                 // Next compare the partition keys. If they are not equal or
@@ -151,15 +145,57 @@ public interface PrimaryKey extends Comparable<PrimaryKey>
             }
 
             @Override
+            @SuppressWarnings("ConstantConditions")
             public String toString()
             {
-                return String.format("PrimaryKey: { token: %s, partition: %s, clustering: %s:%s} ",
-                                     token(),
-                                     partitionKey(),
-                                     clustering() == null ? null : clustering().kind(),
-                                     clustering() == null ? null : Arrays.stream(clustering().getBufferArray())
-                                                                         .map(ByteBufferUtil::bytesToHex)
-                                                                         .collect(Collectors.joining(", ")));
+                return isTokenOnly() ? String.format("PrimaryKey: { token: %s }", token())
+                                     : String.format("PrimaryKey: { token: %s, partition: %s, clustering: %s:%s } ",
+                                                     token(),
+                                                     partitionKey(),
+                                                     clustering() == null ? null : clustering().kind(),
+                                                     clustering() == null ? null : Arrays.stream(clustering().getBufferArray())
+                                                                                         .map(ByteBufferUtil::bytesToHex)
+                                                                                         .collect(Collectors.joining(", ")));
+            }
+        }
+
+        class TokenOnlyPrimaryKey extends AbstractPrimaryKey
+        {
+            private final Token token;
+
+            TokenOnlyPrimaryKey(Token token)
+            {
+                this.token = token;
+            }
+
+            @Override
+            public boolean isTokenOnly()
+            {
+                return true;
+            }
+
+            @Override
+            public Token token()
+            {
+                return token;
+            }
+
+            @Override
+            public DecoratedKey partitionKey()
+            {
+                throw new UnsupportedOperationException();
+            }
+
+            @Override
+            public Clustering<?> clustering()
+            {
+                throw new UnsupportedOperationException();
+            }
+
+            @Override
+            public ByteSource asComparableBytes(Version version)
+            {
+                throw new UnsupportedOperationException();
             }
         }
 
@@ -169,9 +205,9 @@ public interface PrimaryKey extends Comparable<PrimaryKey>
             private final DecoratedKey partitionKey;
             private final Clustering<?> clustering;
 
-            ImmutablePrimaryKey(Token token, DecoratedKey partitionKey, Clustering<?> clustering)
+            ImmutablePrimaryKey(DecoratedKey partitionKey, Clustering<?> clustering)
             {
-                this.token = token;
+                this.token = partitionKey.getToken();
                 this.partitionKey = partitionKey;
                 this.clustering = clustering;
             }
@@ -243,6 +279,11 @@ public interface PrimaryKey extends Comparable<PrimaryKey>
         }
     }
 
+    default boolean isTokenOnly()
+    {
+        return false;
+    }
+
     Token token();
 
     @Nullable
@@ -258,6 +299,7 @@ public interface PrimaryKey extends Comparable<PrimaryKey>
      *
      * @return {@code true} if the clustering is empty, otherwise {@code false}
      */
+    @SuppressWarnings("ConstantConditions")
     default boolean hasEmptyClustering()
     {
         return clustering() == null || clustering().isEmpty();

test/distributed/org/apache/cassandra/distributed/test/sai/IndexStreamingTest.java

@@ -66,15 +66,17 @@ public class IndexStreamingTest extends TestBaseImpl
     public boolean isLiteral;
     @Parameterized.Parameter(1)
     public boolean isZeroCopyStreaming;
+    @Parameterized.Parameter(2)
+    public boolean isWide;
 
     @Parameterized.Parameters(name = "isLiteral={0}, isZeroCopyStreaming={1}")
     public static List<Object[]> data()
     {
         List<Object[]> result = new ArrayList<>();
-        result.add(new Object[]{ true, true });
-        result.add(new Object[]{ true, false });
-        result.add(new Object[]{ false, true });
-        result.add(new Object[]{ false, false });
+        for (boolean isLiteral : BOOLEANS)
+            for (boolean isZeroCopyStreaming : BOOLEANS)
+                for (boolean isWide : BOOLEANS)
+                    result.add(new Object[]{ isLiteral, isZeroCopyStreaming, isWide });
         return result;
     }
 
@@ -88,15 +90,14 @@ public class IndexStreamingTest extends TestBaseImpl
                                            .start()))
         {
             cluster.schemaChange(withKeyspace(
-                "CREATE TABLE %s.test (pk int PRIMARY KEY, literal text, numeric int, b blob) WITH compression = { 'enabled' : false };"
+                isWide
+                ? "CREATE TABLE %s.test (pk int, ck int , literal text, numeric int, b blob, PRIMARY KEY(pk, ck)) WITH compression = { 'enabled' : false };"
+                : "CREATE TABLE %s.test (pk int PRIMARY KEY , literal text, numeric int, b blob) WITH compression = { 'enabled' : false };"
             ));
 
-            int num_components = isLiteral ? sstableStreamingComponentsCount() +
-                                             V1OnDiskFormat.PER_SSTABLE_COMPONENTS.size() +
-                                             V1OnDiskFormat.LITERAL_COMPONENTS.size()
-                                           : sstableStreamingComponentsCount() +
-                                             V1OnDiskFormat.PER_SSTABLE_COMPONENTS.size() +
-                                             V1OnDiskFormat.NUMERIC_COMPONENTS.size();
+            int numSSTableComponents = isWide ? V1OnDiskFormat.WIDE_PER_SSTABLE_COMPONENTS.size() : V1OnDiskFormat.SKINNY_PER_SSTABLE_COMPONENTS.size();
+            int numIndexComponents = isLiteral ? V1OnDiskFormat.LITERAL_COMPONENTS.size() : V1OnDiskFormat.NUMERIC_COMPONENTS.size();
+            int numComponents = sstableStreamingComponentsCount() + numSSTableComponents + numIndexComponents;
 
             if (isLiteral)
                 cluster.schemaChange(withKeyspace(
@@ -113,10 +114,13 @@ public class IndexStreamingTest extends TestBaseImpl
             IInvokableInstance first = cluster.get(1);
             IInvokableInstance second = cluster.get(2);
             long sstableCount = 10;
-            long expectedFiles = isZeroCopyStreaming ? sstableCount * num_components : sstableCount;
+            long expectedFiles = isZeroCopyStreaming ? sstableCount * numComponents : sstableCount;
             for (int i = 0; i < sstableCount; i++)
             {
-                first.executeInternal(withKeyspace("insert into %s.test(pk, literal, numeric, b) values (?, ?, ?, ?)"), i, "v" + i, i, BLOB);
+                if (isWide)
+                    first.executeInternal(withKeyspace("insert into %s.test(pk, ck, literal, numeric, b) values (?, ?, ?, ?, ?)"), i, i, "v" + i, i, BLOB);
+                else
+                    first.executeInternal(withKeyspace("insert into %s.test(pk, literal, numeric, b) values (?, ?, ?, ?)"), i, "v" + i, i, BLOB);
                 first.flush(KEYSPACE);
             }
 

test/microbench/org/apache/cassandra/test/microbench/sai/KeyLookupBench.java

@@ -0,0 +1,205 @@
+/*
+ * 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.test.microbench.sai;
+
+import java.nio.ByteBuffer;
+import java.nio.file.Files;
+import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.cassandra.Util;
+import org.apache.cassandra.config.CassandraRelevantProperties;
+import org.apache.cassandra.config.DatabaseDescriptor;
+import org.apache.cassandra.db.Clustering;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.marshal.CompositeType;
+import org.apache.cassandra.db.marshal.LongType;
+import org.apache.cassandra.db.marshal.UTF8Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
+import org.apache.cassandra.index.sai.disk.v1.SSTableComponentsWriter;
+import org.apache.cassandra.index.sai.disk.v1.WidePrimaryKeyMap;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.index.sai.utils.TypeUtil;
+import org.apache.cassandra.io.sstable.Descriptor;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.io.util.File;
+import org.apache.cassandra.schema.TableMetadata;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+import static org.apache.cassandra.index.sai.SAITester.getRandom;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+@BenchmarkMode({Mode.Throughput})
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 3, time = 5)
+@Fork(value = 1, jvmArgsAppend = "-Xmx512M")
+@Threads(1)
+@State(Scope.Benchmark)
+public class KeyLookupBench
+{
+    private static final int rows = 1_000_000;
+
+    static
+    {
+        DatabaseDescriptor.toolInitialization();
+        // Partitioner is not set in client mode.
+        if (DatabaseDescriptor.getPartitioner() == null)
+            DatabaseDescriptor.setPartitionerUnsafe(Murmur3Partitioner.instance);
+    }
+
+    protected TableMetadata metadata;
+    protected IndexDescriptor indexDescriptor;
+
+    private PrimaryKeyMap primaryKeyMap;
+
+    private PrimaryKey primaryKey;
+
+    @Param({"3", "4", "5"})
+    public int partitionBlockShift;
+
+    @Param({"3", "4", "5"})
+    public int clusteringBlockShift;
+
+    @Param({"10", "100", "1000", "10000"})
+    public int partitionSize;
+
+    @Param({"true", "false"})
+    public boolean randomClustering;
+
+    @Setup(Level.Trial)
+    public void trialSetup() throws Exception
+    {
+        String keyspaceName = "ks";
+        String tableName = this.getClass().getSimpleName();
+        metadata = TableMetadata
+                   .builder(keyspaceName, tableName)
+                   .partitioner(Murmur3Partitioner.instance)
+                   .addPartitionKeyColumn("pk1", LongType.instance)
+                   .addPartitionKeyColumn("pk2", LongType.instance)
+                   .addClusteringColumn("ck1", UTF8Type.instance)
+                   .addClusteringColumn("ck2", UTF8Type.instance)
+                   .build();
+
+        Descriptor descriptor = new Descriptor(new File(Files.createTempDirectory("jmh").toFile()),
+                                               metadata.keyspace,
+                                               metadata.name,
+                                               Util.newUUIDGen().get());
+
+        indexDescriptor = IndexDescriptor.create(descriptor, metadata.comparator);
+
+        CassandraRelevantProperties.SAI_SORTED_TERMS_PARTITION_BLOCK_SHIFT.setInt(partitionBlockShift);
+        CassandraRelevantProperties.SAI_SORTED_TERMS_CLUSTERING_BLOCK_SHIFT.setInt(clusteringBlockShift);
+        SSTableComponentsWriter writer = new SSTableComponentsWriter(indexDescriptor);
+
+        PrimaryKey.Factory factory = new PrimaryKey.Factory(metadata.comparator);
+
+        PrimaryKey[] primaryKeys = new PrimaryKey[rows];
+        int partition = 0;
+        int partitionRowCounter = 0;
+        for (int index = 0; index < rows; index++)
+        {
+            primaryKeys[index] = factory.create(makeKey(metadata, (long) partition, (long) partition), makeClustering(metadata));
+            partitionRowCounter++;
+            if (partitionRowCounter == partitionSize)
+            {
+                partition++;
+                partitionRowCounter = 0;
+            }
+        }
+
+        Arrays.sort(primaryKeys);
+
+        DecoratedKey lastKey = null;
+        for (PrimaryKey primaryKey : primaryKeys)
+        {
+            if (lastKey == null || lastKey.compareTo(primaryKey.partitionKey()) < 0)
+            {
+                lastKey = primaryKey.partitionKey();
+                writer.startPartition(lastKey);
+            }
+            writer.nextRow(primaryKey);
+        }
+
+        writer.complete();
+
+        SSTableReader sstableReader = mock(SSTableReader.class);
+        when(sstableReader.metadata()).thenReturn(metadata);
+
+        PrimaryKeyMap.Factory mapFactory = new WidePrimaryKeyMap.Factory(indexDescriptor, sstableReader);
+
+        primaryKeyMap = mapFactory.newPerSSTablePrimaryKeyMap();
+
+        primaryKey = primaryKeys[500000];
+    }
+
+    @Benchmark
+    public long advanceToKey()
+    {
+        return primaryKeyMap.rowIdFromPrimaryKey(primaryKey);
+    }
+
+    private static DecoratedKey makeKey(TableMetadata table, Object...partitionKeys)
+    {
+        ByteBuffer key;
+        if (TypeUtil.isComposite(table.partitionKeyType))
+            key = ((CompositeType)table.partitionKeyType).decompose(partitionKeys);
+        else
+            key = table.partitionKeyType.fromString((String)partitionKeys[0]);
+        return table.partitioner.decorateKey(key);
+    }
+
+    private Clustering<?> makeClustering(TableMetadata table)
+    {
+        Clustering<?> clustering;
+        if (table.comparator.size() == 0)
+            clustering = Clustering.EMPTY;
+        else
+        {
+            ByteBuffer[] values = new ByteBuffer[table.comparator.size()];
+            for (int index = 0; index < table.comparator.size(); index++)
+                values[index] = table.comparator.subtype(index).fromString(makeClusteringString());
+            clustering = Clustering.make(values);
+        }
+        return clustering;
+    }
+
+    private String makeClusteringString()
+    {
+        if (randomClustering)
+            return getRandom().nextTextString(10, 100);
+        else
+            return String.format("%08d", getRandom().nextIntBetween(0, partitionSize));
+    }
+}

test/unit/org/apache/cassandra/index/sai/SAITester.java

@@ -486,7 +486,7 @@ public abstract class SAITester extends CQLTester
     {
         Set<File> indexFiles = indexFiles();
 
-        for (IndexComponent indexComponent : Version.LATEST.onDiskFormat().perSSTableIndexComponents())
+        for (IndexComponent indexComponent : Version.LATEST.onDiskFormat().perSSTableIndexComponents(false))
         {
             Set<File> tableFiles = componentFiles(indexFiles, SSTableFormat.Components.Types.CUSTOM.createComponent(Version.LATEST.fileNameFormatter().format(indexComponent, null)));
             assertEquals(tableFiles.toString(), perSSTableFiles, tableFiles.size());

test/unit/org/apache/cassandra/index/sai/cql/StorageAttachedIndexDDLTest.java

@@ -823,7 +823,7 @@ public class StorageAttachedIndexDDLTest extends SAITester
         IndexContext numericIndexContext = createIndexContext(numericIndexName, Int32Type.instance);
         IndexContext stringIndexContext = createIndexContext(stringIndexName, UTF8Type.instance);
 
-        for (IndexComponent component : Version.LATEST.onDiskFormat().perSSTableIndexComponents())
+        for (IndexComponent component : Version.LATEST.onDiskFormat().perSSTableIndexComponents(false))
             verifyRebuildIndexComponent(numericIndexContext, stringIndexContext, component, null, corruptionType, true, true, rebuild);
 
         for (IndexComponent component : Version.LATEST.onDiskFormat().perColumnIndexComponents(numericIndexContext))
@@ -842,12 +842,14 @@ public class StorageAttachedIndexDDLTest extends SAITester
                                              boolean failedNumericIndex,
                                              boolean rebuild) throws Throwable
     {
-        // The completion markers are valid if they exist on the file system so we only need to test
+        // The completion markers are valid if they exist on the file system, so we only need to test
         // their removal. If we are testing with encryption then we don't want to test any components
         // that are encryptable unless they have been removed because encrypted components aren't
         // checksum validated.
 
-        if (component == IndexComponent.PRIMARY_KEY_TRIE || component == IndexComponent.PRIMARY_KEY_BLOCKS || component == IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS)
+        if (component == IndexComponent.PARTITION_SIZES || component == IndexComponent.PARTITION_KEY_BLOCKS ||
+            component == IndexComponent.PARTITION_KEY_BLOCK_OFFSETS || component == IndexComponent.CLUSTERING_KEY_BLOCKS ||
+            component == IndexComponent.CLUSTERING_KEY_BLOCK_OFFSETS)
             return;
 
         if (((component == IndexComponent.GROUP_COMPLETION_MARKER) ||
@@ -894,8 +896,8 @@ public class StorageAttachedIndexDDLTest extends SAITester
             reloadSSTableIndex();
 
             // Verify the index cannot be read:
-            verifySSTableIndexes(numericIndexContext.getIndexName(), Version.LATEST.onDiskFormat().perSSTableIndexComponents().contains(component) ? 0 : 1, failedNumericIndex ? 0 : 1);
-            verifySSTableIndexes(stringIndexContext.getIndexName(), Version.LATEST.onDiskFormat().perSSTableIndexComponents().contains(component) ? 0 : 1, failedStringIndex ? 0 : 1);
+            verifySSTableIndexes(numericIndexContext.getIndexName(), Version.LATEST.onDiskFormat().perSSTableIndexComponents(false).contains(component) ? 0 : 1, failedNumericIndex ? 0 : 1);
+            verifySSTableIndexes(stringIndexContext.getIndexName(), Version.LATEST.onDiskFormat().perSSTableIndexComponents(false).contains(component) ? 0 : 1, failedStringIndex ? 0 : 1);
 
             try
             {

test/unit/org/apache/cassandra/index/sai/cql/TokenCollisionTest.java

@@ -0,0 +1,120 @@
+/*
+ * 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.cql;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.junit.Test;
+
+import org.apache.cassandra.Util;
+import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.utils.ByteBufferUtil;
+
+import static org.junit.Assert.assertEquals;
+
+public class TokenCollisionTest extends SAITester
+{
+    @Test
+    public void skinnyPartitionTest()
+    {
+        doSkinnyPartitionTest(10, 0);
+    }
+
+    @Test
+    public void skinnyPartitionLastRowTest()
+    {
+        doSkinnyPartitionTest(49, 9);
+    }
+
+    private void doSkinnyPartitionTest(int v1Match, int v2Match)
+    {
+        createTable("CREATE TABLE %s (pk blob, v1 int, v2 int, PRIMARY KEY (pk))");
+        createIndex("CREATE INDEX ON %s(v1) USING 'sai'");
+        createIndex("CREATE INDEX ON %s(v2) USING 'sai'");
+
+        ByteBuffer prefix = ByteBufferUtil.bytes("key");
+        int numRows = 100;
+        int v1Count = 0;
+        int v2Count = 0;
+        List<Object[]> matchingPks = new ArrayList<>();
+        for (int pkCount = 0; pkCount < numRows; pkCount++)
+        {
+            ByteBuffer pk = Util.generateMurmurCollision(prefix, (byte) (pkCount / 64), (byte) (pkCount % 64));
+            if (v1Count == v1Match && v2Count == v2Match)
+                matchingPks.add(row(pk, v1Count, v2Count));
+            execute("INSERT INTO %s (pk, v1, v2) VALUES (?, ?, ?)", pk, v1Count++, v2Count++);
+            if (v1Count == 50)
+                v1Count = 0;
+            if (v2Count == 10)
+                v2Count = 0;
+        }
+        assertEquals(2, matchingPks.size());
+        flush();
+
+        assertRowsIgnoringOrder(execute("SELECT * FROM %s WHERE v1=" + v1Match + " AND v2=" + v2Match), matchingPks.get(0), matchingPks.get(1));
+    }
+
+    @Test
+    public void widePartitionTest()
+    {
+        doWidePartitionTest(100, 10, 0);
+    }
+
+    @Test
+    public void widePartitionLastRowTest()
+    {
+        // Reduce the number of rows so the last row occurs at the first clustering value
+        doWidePartitionTest(97, 46, 6);
+    }
+
+    private void doWidePartitionTest(int numRows, int v1Match, int v2Match)
+    {
+        createTable("CREATE TABLE %s (pk blob, ck int, v1 int, v2 int, PRIMARY KEY (pk, ck))");
+        createIndex("CREATE INDEX ON %s(v1) USING 'sai'");
+        createIndex("CREATE INDEX ON %s(v2) USING 'sai'");
+
+        ByteBuffer prefix = ByteBufferUtil.bytes("key");
+        int pkCount = 0;
+        int ckCount = 0;
+        int v1Count = 0;
+        int v2Count = 0;
+        List<Object[]> matchingRows = new ArrayList<>();
+        for (int i = 0; i < numRows; i++)
+        {
+            ByteBuffer pk = Util.generateMurmurCollision(prefix, (byte) (pkCount / 64), (byte) (pkCount % 64));
+            if (v1Count == v1Match && v2Count == v2Match)
+                matchingRows.add(row(pk, ckCount, v1Count, v2Count));
+            execute("INSERT INTO %s (pk, ck, v1, v2) VALUES (?, ?, ?, ?)", pk, ckCount++, v1Count++, v2Count++);
+            if (ckCount == 8)
+            {
+                ckCount = 0;
+                pkCount++;
+            }
+            if (v1Count == 50)
+                v1Count = 0;
+            if (v2Count == 10)
+                v2Count = 0;
+        }
+        assertEquals(2, matchingRows.size());
+        flush();
+
+        assertRowsIgnoringOrder(execute("SELECT * FROM %s WHERE v1=" + v1Match + " AND v2=" + v2Match), matchingRows.get(0), matchingRows.get(1));
+    }
+}

test/unit/org/apache/cassandra/index/sai/disk/NodeStartupTest.java

@@ -336,7 +336,7 @@ public class NodeStartupTest extends SAITester
             case VALID:
                 break;
             case ALL_EMPTY:
-                Version.LATEST.onDiskFormat().perSSTableIndexComponents().forEach(this::remove);
+                Version.LATEST.onDiskFormat().perSSTableIndexComponents(false).forEach(this::remove);
                 Version.LATEST.onDiskFormat().perColumnIndexComponents(indexContext).forEach(c -> remove(c, indexContext));
                 break;
             case PER_SSTABLE_INCOMPLETE:

test/unit/org/apache/cassandra/index/sai/disk/v1/SegmentFlushTest.java

@@ -184,8 +184,8 @@ public class SegmentFlushTest
 
     private void verifyStringIndex(IndexDescriptor indexDescriptor, IndexContext indexContext, SegmentMetadata segmentMetadata) throws IOException
     {
-        FileHandle termsData = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext);
-        FileHandle postingLists = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext);
+        FileHandle termsData = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext, null);
+        FileHandle postingLists = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null);
 
         try (TermsIterator iterator = new TermsScanner(termsData, postingLists, segmentMetadata.componentMetadatas.get(IndexComponent.TERMS_DATA).root))
         {
@@ -203,7 +203,7 @@ public class SegmentFlushTest
         }
     }
 
-    private void verifyTermPostings(TermsIterator iterator, ByteBuffer expectedTerm, int minSegmentRowId, int maxSegmentRowId) throws IOException
+    private void verifyTermPostings(TermsIterator iterator, ByteBuffer expectedTerm, int minSegmentRowId, int maxSegmentRowId)
     {
         ByteComparable term = iterator.next();
         PostingList postings = iterator.postings();

test/unit/org/apache/cassandra/index/sai/disk/v1/TermsReaderTest.java

@@ -82,8 +82,8 @@ public class TermsReaderTest extends SAIRandomizedTester
             indexMetas = writer.writeCompleteSegment(new MemtableTermsIterator(null, null, termsEnum.iterator()));
         }
 
-        FileHandle termsData = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext);
-        FileHandle postingLists = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext);
+        FileHandle termsData = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext, null);
+        FileHandle postingLists = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null);
 
         try (TermsIterator iterator = new TermsScanner(termsData, postingLists, indexMetas.get(IndexComponent.TERMS_DATA).root))
         {
@@ -109,8 +109,8 @@ public class TermsReaderTest extends SAIRandomizedTester
             indexMetas = writer.writeCompleteSegment(new MemtableTermsIterator(null, null, termsEnum.iterator()));
         }
 
-        FileHandle termsData = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext);
-        FileHandle postingLists = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext);
+        FileHandle termsData = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext, null);
+        FileHandle postingLists = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null);
 
         long termsFooterPointer = Long.parseLong(indexMetas.get(IndexComponent.TERMS_DATA).attributes.get(SAICodecUtils.FOOTER_POINTER));
 

test/unit/org/apache/cassandra/index/sai/disk/v1/WideRowPrimaryKeyTest.java

@@ -0,0 +1,115 @@
+/*
+ * 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.disk.v1;
+
+import java.util.Arrays;
+
+import org.junit.Test;
+
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
+import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
+import org.apache.cassandra.index.sai.utils.AbstractPrimaryKeyTester;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
+
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class WideRowPrimaryKeyTest extends AbstractPrimaryKeyTester
+{
+    @Test
+    public void randomTest() throws Throwable
+    {
+        IndexDescriptor indexDescriptor = newClusteringIndexDescriptor(compositePartitionMultipleClusteringAsc);
+
+        SSTableComponentsWriter writer = new SSTableComponentsWriter(indexDescriptor);
+
+        PrimaryKey.Factory factory = new PrimaryKey.Factory(compositePartitionMultipleClusteringAsc.comparator);
+
+        int rows = nextInt(1000, 10000);
+        PrimaryKey[] keys = new PrimaryKey[rows];
+        int partition = 0;
+        int partitionSize = nextInt(5, 500);
+        int partitionCounter = 0;
+        for (int index = 0; index < rows; index++)
+        {
+            keys[index] = factory.create(makeKey(compositePartitionMultipleClusteringAsc, partition, partition),
+                                                 makeClustering(compositePartitionMultipleClusteringAsc,
+                                                                getRandom().nextTextString(10, 100),
+                                                                getRandom().nextTextString(10, 100)));
+            partitionCounter++;
+            if (partitionCounter == partitionSize)
+            {
+                partition++;
+                partitionCounter = 0;
+                partitionSize = nextInt(5, 500);
+            }
+        }
+
+        Arrays.sort(keys);
+
+        DecoratedKey lastKey = null;
+        for (PrimaryKey primaryKey : keys)
+        {
+            if (lastKey == null || lastKey.compareTo(primaryKey.partitionKey()) < 0)
+            {
+                lastKey = primaryKey.partitionKey();
+                writer.startPartition(lastKey);
+            }
+            writer.nextRow(primaryKey);
+        }
+
+        writer.complete();
+
+        SSTableReader sstableReader = mock(SSTableReader.class);
+        when(sstableReader.metadata()).thenReturn(compositePartitionMultipleClusteringAsc);
+
+        try (PrimaryKeyMap.Factory mapFactory = new WidePrimaryKeyMap.Factory(indexDescriptor, sstableReader);
+             PrimaryKeyMap primaryKeyMap = mapFactory.newPerSSTablePrimaryKeyMap())
+        {
+            for (int key = 0; key < rows; key++)
+            {
+                PrimaryKey test = keys[key];
+
+                test = factory.create(test.partitionKey(),
+                                      makeClustering(compositePartitionMultipleClusteringAsc,
+                                                     getRandom().nextTextString(10, 100),
+                                                     getRandom().nextTextString(10, 100)));
+
+                long rowId = primaryKeyMap.rowIdFromPrimaryKey(test);
+
+                if (rowId >= 0)
+                {
+                    PrimaryKey found = keys[(int) rowId];
+
+                    assertTrue(found.compareTo(test) >= 0);
+
+                    if (rowId > 0)
+                        assertTrue(keys[(int) rowId - 1].compareTo(test) < 0);
+                }
+                else
+                {
+                    assertTrue(test.compareTo(keys[keys.length - 1]) > 0);
+                }
+            }
+        }
+    }
+}

test/unit/org/apache/cassandra/index/sai/disk/v1/bbtree/BlockBalancedTreePostingsWriterTest.java

@@ -78,7 +78,7 @@ public class BlockBalancedTreePostingsWriterTest extends SAIRandomizedTester
             fp = writer.finish(output, leaves, indexContext);
         }
 
-        BlockBalancedTreePostingsIndex postingsIndex = new BlockBalancedTreePostingsIndex(indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext), fp);
+        BlockBalancedTreePostingsIndex postingsIndex = new BlockBalancedTreePostingsIndex(indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null), fp);
         assertEquals(10, postingsIndex.size());
 
         // Internal postings...
@@ -126,7 +126,7 @@ public class BlockBalancedTreePostingsWriterTest extends SAIRandomizedTester
         }
 
         // There is only a single posting list...the leaf posting list.
-        BlockBalancedTreePostingsIndex postingsIndex = new BlockBalancedTreePostingsIndex(indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext), fp);
+        BlockBalancedTreePostingsIndex postingsIndex = new BlockBalancedTreePostingsIndex(indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null), fp);
         assertEquals(1, postingsIndex.size());
     }
 
@@ -148,7 +148,7 @@ public class BlockBalancedTreePostingsWriterTest extends SAIRandomizedTester
         }
 
         // There is only a single posting list...the leaf posting list.
-        BlockBalancedTreePostingsIndex postingsIndex = new BlockBalancedTreePostingsIndex(indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext), fp);
+        BlockBalancedTreePostingsIndex postingsIndex = new BlockBalancedTreePostingsIndex(indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null), fp);
         assertEquals(1, postingsIndex.size());
     }
 

test/unit/org/apache/cassandra/index/sai/disk/v1/bbtree/BlockBalancedTreeReaderTest.java

@@ -334,8 +334,8 @@ public class BlockBalancedTreeReaderTest extends SAIRandomizedTester
         final long postingsPosition = metadata.get(IndexComponent.POSTING_LISTS).root;
         assertThat(postingsPosition, is(greaterThan(0L)));
 
-        FileHandle treeHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext);
-        FileHandle treePostingsHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext);
+        FileHandle treeHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext, null);
+        FileHandle treePostingsHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null);
         return new BlockBalancedTreeReader(indexContext,
                                            treeHandle,
                                            treePosition,

test/unit/org/apache/cassandra/index/sai/disk/v1/bbtree/NumericIndexWriterTest.java

@@ -84,8 +84,8 @@ public class NumericIndexWriterTest extends SAIRandomizedTester
                                                            rowCount);
         indexMetas = writer.writeCompleteSegment(ramBuffer.iterator());
 
-        final FileHandle treeHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext);
-        final FileHandle treePostingsHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext);
+        final FileHandle treeHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext, null);
+        final FileHandle treePostingsHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null);
 
         try (BlockBalancedTreeReader reader = new BlockBalancedTreeReader(indexContext,
                                                                           treeHandle,
@@ -130,8 +130,8 @@ public class NumericIndexWriterTest extends SAIRandomizedTester
                                                            maxSegmentRowId);
         indexMetas = writer.writeCompleteSegment(BlockBalancedTreeIterator.fromTermsIterator(termEnum, Int32Type.instance));
 
-        final FileHandle treeHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext);
-        final FileHandle treePostingsHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext);
+        final FileHandle treeHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.BALANCED_TREE, indexContext, null);
+        final FileHandle treePostingsHandle = indexDescriptor.createPerIndexFileHandle(IndexComponent.POSTING_LISTS, indexContext, null);
 
         try (BlockBalancedTreeReader reader = new BlockBalancedTreeReader(indexContext,
                                                                           treeHandle,

test/unit/org/apache/cassandra/index/sai/disk/v1/bitpack/NumericValuesTest.java

@@ -67,7 +67,7 @@ public class NumericValuesTest extends SAIRandomizedTester
 
         NumericValuesMeta tokensMeta = new NumericValuesMeta(source.get(indexDescriptor.componentName(IndexComponent.TOKEN_VALUES)));
 
-        try (FileHandle fileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES);
+        try (FileHandle fileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES, null);
              LongArray reader = monotonic ? new MonotonicBlockPackedReader(fileHandle, tokensMeta).open()
                                           : new BlockPackedReader(fileHandle, tokensMeta).open())
         {
@@ -82,12 +82,12 @@ public class NumericValuesTest extends SAIRandomizedTester
     {
         final long[] array = new long[64_000];
         final IndexDescriptor indexDescriptor = newIndexDescriptor();
-        writeTokens(monotonic, indexDescriptor, array, prev -> monotonic ? prev + nextInt(100) : nextInt(100));
+        writeTokens(monotonic, indexDescriptor, array, prev -> monotonic ? prev + nextInt(100) : nextLong(0, Long.MAX_VALUE));
 
         final MetadataSource source = MetadataSource.loadGroupMetadata(indexDescriptor);
         NumericValuesMeta tokensMeta = new NumericValuesMeta(source.get(indexDescriptor.componentName(IndexComponent.TOKEN_VALUES)));
 
-        try (FileHandle fileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES);
+        try (FileHandle fileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES, null);
              LongArray reader = (monotonic ? new MonotonicBlockPackedReader(fileHandle, tokensMeta)
                                            : new BlockPackedReader(fileHandle, tokensMeta)).open())
         {

test/unit/org/apache/cassandra/index/sai/disk/v1/keystore/KeyLookupTest.java

@@ -0,0 +1,413 @@
+/*
+ * 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.disk.v1.keystore;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import org.apache.cassandra.db.Clustering;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.db.marshal.UTF8Type;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.index.sai.SAITester;
+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.io.IndexOutputWriter;
+import org.apache.cassandra.index.sai.disk.v1.MetadataSource;
+import org.apache.cassandra.index.sai.disk.v1.MetadataWriter;
+import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils;
+import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesMeta;
+import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesWriter;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
+import org.apache.cassandra.io.util.FileHandle;
+import org.apache.cassandra.utils.bytecomparable.ByteComparable;
+import org.apache.cassandra.utils.bytecomparable.ByteSource;
+import org.apache.cassandra.utils.bytecomparable.ByteSourceInverse;
+import org.apache.lucene.store.IndexInput;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class KeyLookupTest extends SAIRandomizedTester
+{
+    protected IndexDescriptor indexDescriptor;
+
+    @Before
+    public void setup() throws Exception
+    {
+        indexDescriptor = newIndexDescriptor();
+    }
+
+    @Test
+    public void testFileValidation() throws Exception
+    {
+        List<PrimaryKey> primaryKeys = new ArrayList<>();
+
+        for (int x = 0; x < 11; x++)
+        {
+            ByteBuffer buffer = UTF8Type.instance.decompose(Integer.toString(x));
+            DecoratedKey partitionKey = Murmur3Partitioner.instance.decorateKey(buffer);
+            PrimaryKey primaryKey = SAITester.TEST_FACTORY.create(partitionKey, Clustering.EMPTY);
+            primaryKeys.add(primaryKey);
+        }
+
+        primaryKeys.sort(PrimaryKey::compareTo);
+
+        try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META)))
+        {
+            IndexOutputWriter bytesWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PARTITION_KEY_BLOCKS);
+            NumericValuesWriter blockFPWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.PARTITION_KEY_BLOCK_OFFSETS, metadataWriter, true);
+            try (KeyStoreWriter writer = new KeyStoreWriter(indexDescriptor.componentName(IndexComponent.PARTITION_KEY_BLOCKS),
+                                                            metadataWriter,
+                                                            bytesWriter,
+                                                            blockFPWriter,
+                                                            4,
+                                                            false))
+            {
+                primaryKeys.forEach(primaryKey -> {
+                    try
+                    {
+                        writer.add(primaryKey);
+                    }
+                    catch (IOException e)
+                    {
+                        e.printStackTrace();
+                    }
+                });
+            }
+        }
+        assertTrue(validateComponent(IndexComponent.PARTITION_KEY_BLOCKS, true));
+        assertTrue(validateComponent(IndexComponent.PARTITION_KEY_BLOCKS, false));
+        assertTrue(validateComponent(IndexComponent.PARTITION_KEY_BLOCK_OFFSETS, true));
+        assertTrue(validateComponent(IndexComponent.PARTITION_KEY_BLOCK_OFFSETS, false));
+    }
+
+    @Test
+    public void testLongPrefixesAndSuffixes() throws Exception
+    {
+        List<byte[]> keys = new ArrayList<>();
+        writeKeys(writer ->  {
+            // The following writes a set of keys that cover the following conditions:
+
+            // Start value 0
+            byte[] bytes = new byte[20];
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+            // prefix > 15
+            bytes = new byte[20];
+            Arrays.fill(bytes, 16, 20, (byte)1);
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+            // prefix == 15
+            bytes = new byte[20];
+            Arrays.fill(bytes, 15, 20, (byte)1);
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+            // prefix < 15
+            bytes = new byte[20];
+            Arrays.fill(bytes, 14, 20, (byte)1);
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+            // suffix > 16
+            bytes = new byte[20];
+            Arrays.fill(bytes, 0, 4, (byte)1);
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+            // suffix == 16
+            bytes = new byte[20];
+            Arrays.fill(bytes, 0, 5, (byte)1);
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+            // suffix < 16
+            bytes = new byte[20];
+            Arrays.fill(bytes, 0, 6, (byte)1);
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+
+            bytes = new byte[32];
+            Arrays.fill(bytes, 0, 16, (byte)1);
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+            // prefix >= 15 && suffix >= 16
+            bytes = new byte[32];
+            Arrays.fill(bytes, 0, 32, (byte)1);
+            keys.add(bytes);
+            writer.add(ByteComparable.fixedLength(bytes));
+        }, false);
+
+        doTestKeyLookup(keys);
+    }
+
+    @Test
+    public void testNonUniqueKeys() throws Exception
+    {
+        List<byte[]> keys = new ArrayList<>();
+
+        writeKeys(writer ->  {
+            for (int x = 0; x < 4000; x++)
+            {
+                ByteBuffer buffer = Int32Type.instance.decompose(5000);
+                ByteSource byteSource = Int32Type.instance.asComparableBytes(buffer, ByteComparable.Version.OSS50);
+                byte[] bytes = ByteSourceInverse.readBytes(byteSource);
+                keys.add(bytes);
+
+                writer.add(ByteComparable.fixedLength(bytes));
+            }
+        }, false);
+
+        doTestKeyLookup(keys);
+    }
+
+    @Test
+    public void testSeekToPointId() throws Exception
+    {
+        List<byte[]> keys = new ArrayList<>();
+
+        writeKeys(writer -> {
+            for (int x = 0; x < 4000; x++)
+            {
+                ByteBuffer buffer = Int32Type.instance.decompose(x);
+                ByteSource byteSource = Int32Type.instance.asComparableBytes(buffer, ByteComparable.Version.OSS50);
+                byte[] bytes = ByteSourceInverse.readBytes(byteSource);
+                keys.add(bytes);
+
+                writer.add(ByteComparable.fixedLength(bytes));
+            }
+        }, false);
+
+        doTestKeyLookup(keys);
+    }
+
+    @Test
+    public void testSeekToPointIdOutOfRange() throws Exception
+    {
+        writeKeys(writer -> {
+            for (int x = 0; x < 4000; x++)
+            {
+                ByteBuffer buffer = Int32Type.instance.decompose(x);
+                ByteSource byteSource = Int32Type.instance.asComparableBytes(buffer, ByteComparable.Version.OSS50);
+                byte[] bytes = ByteSourceInverse.readBytes(byteSource);
+
+                writer.add(ByteComparable.fixedLength(bytes));
+            }
+        }, false);
+
+        withKeyLookupCursor(cursor -> {
+            assertThatThrownBy(() -> cursor.seekToPointId(-2)).isInstanceOf(IndexOutOfBoundsException.class)
+                                                              .hasMessage(String.format(KeyLookup.INDEX_OUT_OF_BOUNDS, -2, 4000));
+            assertThatThrownBy(() -> cursor.seekToPointId(Long.MAX_VALUE)).isInstanceOf(IndexOutOfBoundsException.class)
+                                                                          .hasMessage(String.format(KeyLookup.INDEX_OUT_OF_BOUNDS, Long.MAX_VALUE, 4000));
+            assertThatThrownBy(() -> cursor.seekToPointId(4000)).isInstanceOf(IndexOutOfBoundsException.class)
+                                                                .hasMessage(String.format(KeyLookup.INDEX_OUT_OF_BOUNDS, 4000, 4000));
+        });
+    }
+
+    @Test
+    public void testSeekToKey() throws Exception
+    {
+        Map<Long, byte[]> keys = new HashMap<>();
+
+        writeKeys(writer -> {
+            long pointId = 0;
+            for (int x = 0; x < 4000; x += 4)
+            {
+                byte[] key = makeKey(x);
+                keys.put(pointId++, key);
+
+                writer.add(ByteComparable.fixedLength(key));
+            }
+        }, true);
+
+        withKeyLookupCursor(cursor -> {
+            assertEquals(0L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(0L)), 0L, 10L));
+            cursor.reset();
+            assertEquals(160L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(160L)), 160L, 170L));
+            cursor.reset();
+            assertEquals(165L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(165L)), 160L, 170L));
+            cursor.reset();
+            assertEquals(175L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(175L)), 160L, 176L));
+            cursor.reset();
+            assertEquals(176L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(176L)), 160L, 177L));
+            cursor.reset();
+            assertEquals(176L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(176L)), 175L, 177L));
+            cursor.reset();
+            assertEquals(176L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(makeKey(701)), 160L, 177L));
+            cursor.reset();
+            assertEquals(504L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(504L)), 200L, 600L));
+            cursor.reset();
+            assertEquals(-1L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(makeKey(4000)), 0L, 1000L));
+            cursor.reset();
+            assertEquals(-1L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(makeKey(4000)), 999L, 1000L));
+            cursor.reset();
+            assertEquals(999L, cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(999L)), 0L, 1000L));
+        });
+    }
+
+    @Test
+    public void seekToKeyOnNonPartitionedTest() throws Throwable
+    {
+        Map<Long, byte[]> keys = new HashMap<>();
+
+        writeKeys(writer -> {
+            long pointId = 0;
+            for (int x = 0; x < 16; x += 4)
+            {
+                byte[] key = makeKey(x);
+                keys.put(pointId++, key);
+
+                writer.add(ByteComparable.fixedLength(key));
+            }
+        }, false);
+
+        withKeyLookupCursor(cursor -> assertThatThrownBy(() -> cursor.clusteredSeekToKey(ByteComparable.fixedLength(keys.get(0L)), 0L, 10L))
+                                        .isInstanceOf(AssertionError.class));
+    }
+
+    @Test
+    public void partitionedKeysMustBeInOrderInPartitions() throws Throwable
+    {
+        writeKeys(writer -> {
+            writer.startPartition();
+            writer.add(ByteComparable.fixedLength(makeKey(0)));
+            writer.add(ByteComparable.fixedLength(makeKey(10)));
+            assertThatThrownBy(() -> writer.add(ByteComparable.fixedLength(makeKey(9)))).isInstanceOf(IllegalArgumentException.class);
+            writer.startPartition();
+            writer.add(ByteComparable.fixedLength(makeKey(9)));
+        }, true);
+    }
+
+    private byte[] makeKey(int value)
+    {
+        ByteBuffer buffer = Int32Type.instance.decompose(value);
+        ByteSource byteSource = Int32Type.instance.asComparableBytes(buffer, ByteComparable.Version.OSS50);
+        return ByteSourceInverse.readBytes(byteSource);
+    }
+
+    private void doTestKeyLookup(List<byte[]> keys) throws Exception
+    {
+        // iterate ascending
+        withKeyLookupCursor(cursor -> {
+            for (int x = 0; x < keys.size(); x++)
+            {
+                ByteComparable key = cursor.seekToPointId(x);
+
+                byte[] bytes = ByteSourceInverse.readBytes(key.asComparableBytes(ByteComparable.Version.OSS50));
+
+                assertArrayEquals(keys.get(x), bytes);
+            }
+        });
+
+        // iterate ascending skipping blocks
+        withKeyLookupCursor(cursor -> {
+            for (int x = 0; x < keys.size(); x += 17)
+            {
+                ByteComparable key = cursor.seekToPointId(x);
+
+                byte[] bytes = ByteSourceInverse.readBytes(key.asComparableBytes(ByteComparable.Version.OSS50));
+
+                assertArrayEquals(keys.get(x), bytes);
+            }
+        });
+
+        withKeyLookupCursor(cursor -> {
+            ByteComparable key = cursor.seekToPointId(7);
+            byte[] bytes = ByteSourceInverse.readBytes(key.asComparableBytes(ByteComparable.Version.OSS50));
+            assertArrayEquals(keys.get(7), bytes);
+
+            key = cursor.seekToPointId(7);
+            bytes = ByteSourceInverse.readBytes(key.asComparableBytes(ByteComparable.Version.OSS50));
+            assertArrayEquals(keys.get(7), bytes);
+        });
+    }
+
+    protected void writeKeys(ThrowingConsumer<KeyStoreWriter> testCode, boolean clustering) throws IOException
+    {
+        try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META)))
+        {
+            IndexOutputWriter bytesWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PARTITION_KEY_BLOCKS);
+            NumericValuesWriter blockFPWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.PARTITION_KEY_BLOCK_OFFSETS, metadataWriter, true);
+            try (KeyStoreWriter writer = new KeyStoreWriter(indexDescriptor.componentName(IndexComponent.PARTITION_KEY_BLOCKS),
+                                                            metadataWriter,
+                                                            bytesWriter,
+                                                            blockFPWriter,
+                                                            4,
+                                                            clustering))
+            {
+                testCode.accept(writer);
+            }
+        }
+    }
+
+    @FunctionalInterface
+    public interface ThrowingConsumer<T>
+    {
+        void accept(T t) throws IOException;
+    }
+
+    private void withKeyLookup(ThrowingConsumer<KeyLookup> testCode) throws IOException
+    {
+        MetadataSource metadataSource = MetadataSource.loadGroupMetadata(indexDescriptor);
+        NumericValuesMeta blockPointersMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PARTITION_KEY_BLOCK_OFFSETS)));
+        KeyLookupMeta keyLookupMeta = new KeyLookupMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PARTITION_KEY_BLOCKS)));
+        try (FileHandle keysData = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PARTITION_KEY_BLOCKS, null);
+             FileHandle blockOffsets = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PARTITION_KEY_BLOCK_OFFSETS, null))
+        {
+            KeyLookup reader = new KeyLookup(keysData, blockOffsets, keyLookupMeta, blockPointersMeta);
+            testCode.accept(reader);
+        }
+    }
+
+    private void withKeyLookupCursor(ThrowingConsumer<KeyLookup.Cursor> testCode) throws IOException
+    {
+        withKeyLookup(reader -> {
+            try (KeyLookup.Cursor cursor = reader.openCursor())
+            {
+                testCode.accept(cursor);
+            }
+        });
+    }
+
+    private boolean validateComponent(IndexComponent indexComponent, boolean checksum)
+    {
+        try (IndexInput input = indexDescriptor.openPerSSTableInput(indexComponent))
+        {
+            if (checksum)
+                SAICodecUtils.validateChecksum(input);
+            else
+                SAICodecUtils.validate(input);
+            return true;
+        }
+        catch (Throwable ignore)
+        {
+            return false;
+        }
+    }
+}

test/unit/org/apache/cassandra/index/sai/disk/v1/sortedterms/SortedTermsTest.java

@@ -1,415 +0,0 @@
-/*
- * 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.disk.v1.sortedterms;
-
-import java.io.IOException;
-import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-
-import org.junit.Test;
-
-import org.apache.cassandra.cql3.CQLTester;
-import org.apache.cassandra.db.Clustering;
-import org.apache.cassandra.db.DecoratedKey;
-import org.apache.cassandra.db.marshal.Int32Type;
-import org.apache.cassandra.db.marshal.UTF8Type;
-import org.apache.cassandra.dht.Murmur3Partitioner;
-import org.apache.cassandra.index.sai.SAITester;
-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.io.IndexOutputWriter;
-import org.apache.cassandra.index.sai.disk.v1.MetadataSource;
-import org.apache.cassandra.index.sai.disk.v1.MetadataWriter;
-import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils;
-import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesMeta;
-import org.apache.cassandra.index.sai.disk.v1.bitpack.NumericValuesWriter;
-import org.apache.cassandra.index.sai.utils.PrimaryKey;
-import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
-import org.apache.cassandra.io.util.FileHandle;
-import org.apache.cassandra.utils.bytecomparable.ByteComparable;
-import org.apache.cassandra.utils.bytecomparable.ByteSource;
-import org.apache.cassandra.utils.bytecomparable.ByteSourceInverse;
-import org.apache.lucene.store.IndexInput;
-
-import static org.junit.Assert.assertArrayEquals;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
-
-public class SortedTermsTest extends SAIRandomizedTester
-{
-    @Test
-    public void testLexicographicException() throws Exception
-    {
-        IndexDescriptor indexDescriptor = newIndexDescriptor();
-        try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META)))
-        {
-            NumericValuesWriter blockFPWriter = new NumericValuesWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                                        indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                                        metadataWriter, true);
-            IndexOutputWriter trieWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_TRIE);
-            IndexOutputWriter bytesWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCKS);
-            try (SortedTermsWriter writer = new SortedTermsWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCKS),
-                                                                  metadataWriter,
-                                                                  bytesWriter,
-                                                                  blockFPWriter,
-                                                                  trieWriter))
-            {
-                ByteBuffer buffer = Int32Type.instance.decompose(99999);
-                ByteSource byteSource = Int32Type.instance.asComparableBytes(buffer, ByteComparable.Version.OSS50);
-                byte[] bytes1 = ByteSourceInverse.readBytes(byteSource);
-
-                writer.add(ByteComparable.fixedLength(bytes1));
-
-                buffer = Int32Type.instance.decompose(444);
-                byteSource = Int32Type.instance.asComparableBytes(buffer, ByteComparable.Version.OSS50);
-                byte[] bytes2 = ByteSourceInverse.readBytes(byteSource);
-
-                assertThrows(IllegalArgumentException.class, () -> writer.add(ByteComparable.fixedLength(bytes2)));
-            }
-        }
-    }
-
-    @Test
-    public void testFileValidation() throws Exception
-    {
-        IndexDescriptor indexDescriptor = newIndexDescriptor();
-
-        List<PrimaryKey> primaryKeys = new ArrayList<>();
-
-        for (int x = 0; x < 11; x++)
-        {
-            ByteBuffer buffer = UTF8Type.instance.decompose(Integer.toString(x));
-            DecoratedKey partitionKey = Murmur3Partitioner.instance.decorateKey(buffer);
-            PrimaryKey primaryKey = SAITester.TEST_FACTORY.create(partitionKey, Clustering.EMPTY);
-            primaryKeys.add(primaryKey);
-        }
-
-        primaryKeys.sort(PrimaryKey::compareTo);
-
-        try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META)))
-        {
-            IndexOutputWriter trieWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_TRIE);
-            IndexOutputWriter bytesWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCKS);
-            NumericValuesWriter blockFPWriter = new NumericValuesWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                                        indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                                        metadataWriter, true);
-            try (SortedTermsWriter writer = new SortedTermsWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCKS),
-                                                                  metadataWriter,
-                                                                  bytesWriter,
-                                                                  blockFPWriter,
-                                                                  trieWriter))
-            {
-                primaryKeys.forEach(primaryKey -> {
-                    try
-                    {
-                        writer.add(primaryKey::asComparableBytes);
-                    }
-                    catch (IOException e)
-                    {
-                        e.printStackTrace();
-                    }
-                });
-            }
-        }
-        assertTrue(validateComponent(indexDescriptor, IndexComponent.PRIMARY_KEY_TRIE, true));
-        assertTrue(validateComponent(indexDescriptor, IndexComponent.PRIMARY_KEY_TRIE, false));
-        assertTrue(validateComponent(indexDescriptor, IndexComponent.PRIMARY_KEY_BLOCKS, true));
-        assertTrue(validateComponent(indexDescriptor, IndexComponent.PRIMARY_KEY_BLOCKS, false));
-        assertTrue(validateComponent(indexDescriptor, IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS, true));
-        assertTrue(validateComponent(indexDescriptor, IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS, false));
-    }
-
-    @Test
-    public void testLongPrefixesAndSuffixes() throws Exception
-    {
-        IndexDescriptor descriptor = newIndexDescriptor();
-
-        List<byte[]> terms = new ArrayList<>();
-        writeLongPrefixAndSuffixTerms(descriptor, terms);
-
-        withSortedTermsCursor(descriptor, cursor ->
-        {
-            int x = 0;
-            while (cursor.advance())
-            {
-                ByteComparable term = cursor.term();
-
-                byte[] bytes = ByteSourceInverse.readBytes(term.asComparableBytes(ByteComparable.Version.OSS50));
-                assertArrayEquals(terms.get(x), bytes);
-                x++;
-            }
-
-            // assert we don't increase the point id beyond one point after the last item
-            assertEquals(cursor.pointId(), terms.size());
-            assertFalse(cursor.advance());
-            assertEquals(cursor.pointId(), terms.size());
-        });
-    }
-
-    @Test
-    public void testAdvance() throws IOException
-    {
-        IndexDescriptor descriptor = newIndexDescriptor();
-
-        List<byte[]> terms = new ArrayList<>();
-        writeTerms(descriptor, terms);
-
-        withSortedTermsCursor(descriptor, cursor ->
-        {
-            int x = 0;
-            while (cursor.advance())
-            {
-                ByteComparable term = cursor.term();
-
-                byte[] bytes = ByteSourceInverse.readBytes(term.asComparableBytes(ByteComparable.Version.OSS50));
-                assertArrayEquals(terms.get(x), bytes);
-
-                x++;
-            }
-
-            // assert we don't increase the point id beyond one point after the last item
-            assertEquals(cursor.pointId(), terms.size());
-            assertFalse(cursor.advance());
-            assertEquals(cursor.pointId(), terms.size());
-        });
-    }
-
-    @Test
-    public void testSeekToPointId() throws Exception
-    {
-        IndexDescriptor descriptor = newIndexDescriptor();
-
-        List<byte[]> terms = new ArrayList<>();
-        writeTerms(descriptor, terms);
-
-        // iterate ascending
-        withSortedTermsCursor(descriptor, cursor ->
-        {
-            for (int x = 0; x < terms.size(); x++)
-            {
-                cursor.seekToPointId(x);
-                ByteComparable term = cursor.term();
-
-                byte[] bytes = ByteSourceInverse.readBytes(term.asComparableBytes(ByteComparable.Version.OSS50));
-                assertArrayEquals(terms.get(x), bytes);
-            }
-        });
-
-        // iterate descending
-        withSortedTermsCursor(descriptor, cursor ->
-        {
-            for (int x = terms.size() - 1; x >= 0; x--)
-            {
-                cursor.seekToPointId(x);
-                ByteComparable term = cursor.term();
-
-                byte[] bytes = ByteSourceInverse.readBytes(term.asComparableBytes(ByteComparable.Version.OSS50));
-                assertArrayEquals(terms.get(x), bytes);
-            }
-        });
-
-        // iterate randomly
-        withSortedTermsCursor(descriptor, cursor ->
-        {
-            for (int x = 0; x < terms.size(); x++)
-            {
-                int target = nextInt(0, terms.size());
-                cursor.seekToPointId(target);
-                ByteComparable term = cursor.term();
-
-                byte[] bytes = ByteSourceInverse.readBytes(term.asComparableBytes(ByteComparable.Version.OSS50));
-                assertArrayEquals(terms.get(target), bytes);
-            }
-        });
-    }
-
-    @Test
-    public void testSeekToPointIdOutOfRange() throws Exception
-    {
-        IndexDescriptor descriptor = newIndexDescriptor();
-
-        List<byte[]> terms = new ArrayList<>();
-        writeTerms(descriptor, terms);
-
-        withSortedTermsCursor(descriptor, cursor ->
-        {
-            assertThrows(IndexOutOfBoundsException.class, () -> cursor.seekToPointId(-2));
-            assertThrows(IndexOutOfBoundsException.class, () -> cursor.seekToPointId(Long.MAX_VALUE));
-        });
-    }
-
-    private void writeTerms(IndexDescriptor indexDescriptor, List<byte[]> terms) throws IOException
-    {
-        try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META)))
-        {
-            IndexOutputWriter trieWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_TRIE);
-            IndexOutputWriter bytesWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCKS);
-            NumericValuesWriter blockFPWriter = new NumericValuesWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                                        indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                                        metadataWriter, true);
-            try (SortedTermsWriter writer = new SortedTermsWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCKS),
-                                                                  metadataWriter,
-                                                                  bytesWriter,
-                                                                  blockFPWriter,
-                                                                  trieWriter))
-            {
-                for (int x = 0; x < 1000 * 4; x++)
-                {
-                    ByteBuffer buffer = Int32Type.instance.decompose(x);
-                    ByteSource byteSource = Int32Type.instance.asComparableBytes(buffer, ByteComparable.Version.OSS50);
-                    byte[] bytes = ByteSourceInverse.readBytes(byteSource);
-                    terms.add(bytes);
-
-                    writer.add(ByteComparable.fixedLength(bytes));
-                }
-            }
-        }
-    }
-
-    private void writeLongPrefixAndSuffixTerms(IndexDescriptor indexDescriptor, List<byte[]> terms) throws IOException
-    {
-        try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META)))
-        {
-            IndexOutputWriter trieWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_TRIE);
-            IndexOutputWriter bytesWriter = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCKS);
-            NumericValuesWriter blockFPWriter = new NumericValuesWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                                        indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS),
-                                                                        metadataWriter, true);
-            try (SortedTermsWriter writer = new SortedTermsWriter(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCKS),
-                                                                  metadataWriter,
-                                                                  bytesWriter,
-                                                                  blockFPWriter,
-                                                                  trieWriter))
-            {
-                // The following writes a lexographically ordered set of terms that cover the following
-                // conditions:
-
-                // Start value 0
-                byte[] bytes = new byte[20];
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-                // prefix > 15
-                bytes = new byte[20];
-                Arrays.fill(bytes, 16, 20, (byte)1);
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-                // prefix == 15
-                bytes = new byte[20];
-                Arrays.fill(bytes, 15, 20, (byte)1);
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-                // prefix < 15
-                bytes = new byte[20];
-                Arrays.fill(bytes, 14, 20, (byte)1);
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-                // suffix > 16
-                bytes = new byte[20];
-                Arrays.fill(bytes, 0, 4, (byte)1);
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-                // suffix == 16
-                bytes = new byte[20];
-                Arrays.fill(bytes, 0, 5, (byte)1);
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-                // suffix < 16
-                bytes = new byte[20];
-                Arrays.fill(bytes, 0, 6, (byte)1);
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-
-                bytes = new byte[32];
-                Arrays.fill(bytes, 0, 16, (byte)1);
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-                // prefix >= 15 && suffix >= 16
-                bytes = new byte[32];
-                Arrays.fill(bytes, 0, 32, (byte)1);
-                terms.add(bytes);
-                writer.add(ByteComparable.fixedLength(bytes));
-            }
-        }
-    }
-
-    @FunctionalInterface
-    public interface ThrowingConsumer<T> {
-        void accept(T t) throws IOException;
-    }
-
-    private void withSortedTermsReader(IndexDescriptor indexDescriptor,
-                                       ThrowingConsumer<SortedTermsReader> testCode) throws IOException
-    {
-        MetadataSource metadataSource = MetadataSource.loadGroupMetadata(indexDescriptor);
-        NumericValuesMeta blockPointersMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS)));
-        SortedTermsMeta sortedTermsMeta = new SortedTermsMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PRIMARY_KEY_BLOCKS)));
-        try (FileHandle termsData = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_BLOCKS);
-             FileHandle blockOffsets = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_BLOCK_OFFSETS))
-        {
-            SortedTermsReader reader = new SortedTermsReader(termsData, blockOffsets, sortedTermsMeta, blockPointersMeta);
-            testCode.accept(reader);
-        }
-    }
-
-    private void withSortedTermsCursor(IndexDescriptor descriptor,
-                                       ThrowingConsumer<SortedTermsReader.Cursor> testCode) throws IOException
-    {
-        withSortedTermsReader(descriptor, reader ->
-        {
-            try (SortedTermsReader.Cursor cursor = reader.openCursor())
-            {
-                testCode.accept(cursor);
-            }
-        });
-    }
-
-    private boolean validateComponent(IndexDescriptor indexDescriptor, IndexComponent indexComponent, boolean checksum)
-    {
-        try (IndexInput input = indexDescriptor.openPerSSTableInput(indexComponent))
-        {
-            if (checksum)
-                SAICodecUtils.validateChecksum(input);
-            else
-                SAICodecUtils.validate(input);
-            return true;
-        }
-        catch (Throwable ignore)
-        {
-            return false;
-        }
-    }
-
-    private static void assertThrows(Class<? extends Throwable> clazz, CQLTester.CheckedFunction function)
-    {
-        try
-        {
-            function.apply();
-            fail("An error should havee been thrown but was not.");
-        }
-        catch (Throwable e)
-        {
-            assertTrue("Unexpected exception type (expected: " + clazz + ", value: " + e.getClass() + ')',
-                       clazz.isAssignableFrom(e.getClass()));
-        }
-    }
-}

test/unit/org/apache/cassandra/index/sai/disk/v1/trie/TriePrefixSearcherTest.java

@@ -1,160 +0,0 @@
-/*
- * 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.disk.v1.trie;
-
-import org.junit.Before;
-import org.junit.Test;
-
-import org.apache.cassandra.db.marshal.UTF8Type;
-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.io.IndexOutputWriter;
-import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
-import org.apache.cassandra.io.tries.IncrementalDeepTrieWriterPageAware;
-import org.apache.cassandra.io.util.FileHandle;
-import org.apache.cassandra.utils.bytecomparable.ByteComparable;
-import org.apache.cassandra.utils.bytecomparable.ByteSource;
-
-import static org.apache.cassandra.index.sai.disk.v1.trie.TrieTermsDictionaryReader.trieSerializer;
-import static org.junit.Assert.assertEquals;
-
-public class TriePrefixSearcherTest extends SAIRandomizedTester
-{
-    private IndexDescriptor indexDescriptor;
-
-    @Before
-    public void createIndexDescriptor() throws Throwable
-    {
-        indexDescriptor =  newIndexDescriptor();
-    }
-
-    @Test
-    public void completeValueTest() throws Throwable
-    {
-        long trieFilePointer = createSimpleTrie(indexDescriptor);
-
-        try (FileHandle trieFileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_TRIE);
-             TriePrefixSearcher searcher = new TriePrefixSearcher(trieFileHandle.instantiateRebufferer(null), trieFilePointer))
-        {
-            assertEquals(1L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abcdef"), ByteComparable.Version.OSS50)));
-            assertEquals(2L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abdefg"), ByteComparable.Version.OSS50)));
-            assertEquals(3L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abdfgh"), ByteComparable.Version.OSS50)));
-        }
-    }
-
-    @Test
-    public void prefixValueTest() throws Throwable
-    {
-        long trieFilePointer = createSimpleTrie(indexDescriptor);
-
-        try (FileHandle trieFileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_TRIE);
-             TriePrefixSearcher searcher = new TriePrefixSearcher(trieFileHandle.instantiateRebufferer(null), trieFilePointer))
-        {
-            assertEquals(1L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString(""), ByteComparable.Version.OSS50)));
-            assertEquals(1L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("a"), ByteComparable.Version.OSS50)));
-            assertEquals(1L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("ab"), ByteComparable.Version.OSS50)));
-            assertEquals(1L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abc"), ByteComparable.Version.OSS50)));
-            assertEquals(2L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abd"), ByteComparable.Version.OSS50)));
-            assertEquals(2L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abde"), ByteComparable.Version.OSS50)));
-            assertEquals(3L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abdf"), ByteComparable.Version.OSS50)));
-        }
-    }
-
-    @Test
-    public void unknownValueTest() throws Throwable
-    {
-        long trieFilePointer = createSimpleTrie(indexDescriptor);
-
-        try (FileHandle trieFileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_TRIE);
-             TriePrefixSearcher searcher = new TriePrefixSearcher(trieFileHandle.instantiateRebufferer(null), trieFilePointer))
-        {
-            assertEquals(-1L, searcher.prefixSearch(UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("b"), ByteComparable.Version.OSS50)));
-        }
-    }
-
-    @Test
-    public void completeMultiPartTest() throws Throwable
-    {
-        long trieFilePointer = createMultiPartTrie(indexDescriptor);
-
-        try (FileHandle trieFileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_TRIE);
-             TriePrefixSearcher searcher = new TriePrefixSearcher(trieFileHandle.instantiateRebufferer(null), trieFilePointer))
-        {
-            assertEquals(1L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc", "def", "ghi")));
-            assertEquals(2L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc", "def", "jkl")));
-            assertEquals(3L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc", "ghi", "jkl")));
-            assertEquals(4L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "def", "ghi", "jkl")));
-        }
-    }
-
-    @Test
-    public void prefixMultiPartTest() throws Throwable
-    {
-        long trieFilePointer = createMultiPartTrie(indexDescriptor);
-
-        try (FileHandle trieFileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PRIMARY_KEY_TRIE);
-             TriePrefixSearcher searcher = new TriePrefixSearcher(trieFileHandle.instantiateRebufferer(null), trieFilePointer))
-        {
-            assertEquals(1L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "ab")));
-            assertEquals(1L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc")));
-            assertEquals(1L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc", "de")));
-            assertEquals(1L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc", "def")));
-            assertEquals(1L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc", "def", "gh")));
-            assertEquals(2L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc", "def", "jk")));
-            assertEquals(3L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "abc", "ghi")));
-            assertEquals(4L, searcher.prefixSearch(createMultiPart(ByteComparable.Version.OSS50, "d")));
-        }
-    }
-
-    private long createSimpleTrie(IndexDescriptor indexDescriptor) throws Throwable
-    {
-        try (IndexOutputWriter trieOutput = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_TRIE);
-        IncrementalDeepTrieWriterPageAware<Long> trieWriter = new IncrementalDeepTrieWriterPageAware<>(trieSerializer, trieOutput.asSequentialWriter()))
-        {
-            trieWriter.add(v -> UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abcdef"), v), 1L);
-            trieWriter.add(v -> UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abdefg"), v), 2L);
-            trieWriter.add(v -> UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abdfgh"), v), 3L);
-            return trieWriter.complete();
-        }
-    }
-
-    private long createMultiPartTrie(IndexDescriptor indexDescriptor) throws Throwable
-    {
-        try (IndexOutputWriter trieOutput = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_TRIE);
-             IncrementalDeepTrieWriterPageAware<Long> trieWriter = new IncrementalDeepTrieWriterPageAware<>(trieSerializer, trieOutput.asSequentialWriter()))
-        {
-            trieWriter.add(v -> createMultiPart(v, "abc", "def", "ghi"), 1L);
-            trieWriter.add(v -> createMultiPart(v, "abc", "def", "jkl"), 2L);
-            trieWriter.add(v -> createMultiPart(v, "abc", "ghi", "jkl"), 3L);
-            trieWriter.add(v -> createMultiPart(v, "def", "ghi", "jkl"), 4L);
-            return trieWriter.complete();
-        }
-    }
-
-    private ByteSource createMultiPart(ByteComparable.Version version, String... parts)
-    {
-        ByteSource [] byteSources = new ByteSource[parts.length];
-        for (int index = 0; index < parts.length; index++)
-            byteSources[index] = UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString(parts[index]), version);
-        return ByteSource.withTerminator(ByteSource.TERMINATOR, byteSources);
-    }
-
-
-
-}

test/unit/org/apache/cassandra/index/sai/disk/v1/trie/TrieTermsDictionaryTest.java

@@ -77,7 +77,7 @@ public class TrieTermsDictionaryTest extends SAIRandomizedTester
             fp = writer.complete(new MutableLong());
         }
 
-        try (FileHandle input = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext);
+        try (FileHandle input = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext, null);
              TrieTermsDictionaryReader reader = new TrieTermsDictionaryReader(input.instantiateRebufferer(null), fp))
         {
             assertEquals(TrieTermsDictionaryReader.NOT_FOUND, reader.exactMatch(asByteComparable("a")));
@@ -104,7 +104,7 @@ public class TrieTermsDictionaryTest extends SAIRandomizedTester
             fp = writer.complete(new MutableLong());
         }
 
-        try (FileHandle input = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext);
+        try (FileHandle input = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext, null);
              TrieTermsIterator iterator = new TrieTermsIterator(input.instantiateRebufferer(null), fp))
         {
             final Iterator<ByteComparable> expected = byteComparables.iterator();
@@ -136,7 +136,7 @@ public class TrieTermsDictionaryTest extends SAIRandomizedTester
             fp = writer.complete(new MutableLong());
         }
 
-        try (FileHandle input = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext);
+        try (FileHandle input = indexDescriptor.createPerIndexFileHandle(IndexComponent.TERMS_DATA, indexContext, null);
              TrieTermsDictionaryReader reader = new TrieTermsDictionaryReader(input.instantiateRebufferer(null), fp))
         {
             final ByteComparable expectedMaxTerm = byteComparables.get(byteComparables.size() - 1);

test/unit/org/apache/cassandra/index/sai/disk/v1/trie/TrieValidationTest.java

@@ -1,93 +0,0 @@
-/*
- * 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.disk.v1.trie;
-
-import org.junit.Before;
-import org.junit.Test;
-
-import org.apache.cassandra.db.marshal.UTF8Type;
-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.io.IndexOutputWriter;
-import org.apache.cassandra.index.sai.disk.v1.SAICodecUtils;
-import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
-import org.apache.cassandra.io.tries.IncrementalDeepTrieWriterPageAware;
-import org.apache.cassandra.utils.bytecomparable.ByteComparable;
-import org.apache.cassandra.utils.bytecomparable.ByteSource;
-import org.apache.lucene.store.IndexInput;
-
-import static org.apache.cassandra.index.sai.disk.v1.trie.TrieTermsDictionaryReader.trieSerializer;
-
-public class TrieValidationTest extends SAIRandomizedTester
-{
-    private IndexDescriptor indexDescriptor;
-
-    @Before
-    public void createIndexDescriptor() throws Throwable
-    {
-        indexDescriptor = newIndexDescriptor();
-    }
-
-    @Test
-    public void testHeaderValidation() throws Throwable
-    {
-        createSimpleTrie(indexDescriptor);
-        try (IndexInput input = indexDescriptor.openPerSSTableInput(IndexComponent.PRIMARY_KEY_TRIE))
-        {
-            SAICodecUtils.validate(input);
-        }
-    }
-
-    @Test
-    public void testChecksumValidation() throws Throwable
-    {
-        createSimpleTrie(indexDescriptor);
-        try (IndexInput input = indexDescriptor.openPerSSTableInput(IndexComponent.PRIMARY_KEY_TRIE))
-        {
-            SAICodecUtils.validateChecksum(input);
-        }
-    }
-
-    private static void createSimpleTrie(IndexDescriptor indexDescriptor) throws Throwable
-    {
-        try (IndexOutputWriter trieOutput = indexDescriptor.openPerSSTableOutput(IndexComponent.PRIMARY_KEY_TRIE);
-             IncrementalDeepTrieWriterPageAware<Long> trieWriter = new IncrementalDeepTrieWriterPageAware<>(trieSerializer, trieOutput.asSequentialWriter()))
-        {
-            SAICodecUtils.writeHeader(trieOutput);
-            trieWriter.add(v -> createMultiPart(v, "abc", "def", "ghi"), 1L);
-            trieWriter.add(v -> createMultiPart(v, "abc", "def", "jkl"), 2L);
-            trieWriter.add(v -> createMultiPart(v, "abc", "ghi", "jkl"), 3L);
-            trieWriter.add(v -> createMultiPart(v, "def", "ghi", "jkl"), 4L);
-            trieWriter.add(v -> UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abcdef"), v), 5L);
-            trieWriter.add(v -> UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abdefg"), v), 6L);
-            trieWriter.add(v -> UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString("abdfgh"), v), 7L);
-            trieWriter.complete();
-            SAICodecUtils.writeFooter(trieOutput);
-        }
-    }
-
-    private static ByteSource createMultiPart(ByteComparable.Version version, String... parts)
-    {
-        ByteSource [] byteSources = new ByteSource[parts.length];
-        for (int index = 0; index < parts.length; index++)
-            byteSources[index] = UTF8Type.instance.asComparableBytes(UTF8Type.instance.fromString(parts[index]), version);
-        return ByteSource.withTerminator(ByteSource.TERMINATOR, byteSources);
-    }
-
-}

test/unit/org/apache/cassandra/index/sai/functional/GroupComponentsTest.java

@@ -57,7 +57,7 @@ public class GroupComponentsTest extends SAITester
         SSTableReader sstable = Iterables.getOnlyElement(cfs.getLiveSSTables());
 
         Set<Component> components = StorageAttachedIndexGroup.getLiveComponents(sstable, getIndexesFromGroup(group));
-        assertEquals(Version.LATEST.onDiskFormat().perSSTableIndexComponents().size() + 1, components.size());
+        assertEquals(Version.LATEST.onDiskFormat().perSSTableIndexComponents(false).size() + 1, components.size());
 
         // index files are released but not removed
         cfs.invalidate(true, false);
@@ -84,7 +84,7 @@ public class GroupComponentsTest extends SAITester
 
         Set<Component> components = StorageAttachedIndexGroup.getLiveComponents(sstables.iterator().next(), getIndexesFromGroup(group));
 
-        assertEquals(Version.LATEST.onDiskFormat().perSSTableIndexComponents().size() + 1, components.size());
+        assertEquals(Version.LATEST.onDiskFormat().perSSTableIndexComponents(false).size() + 1, components.size());
     }
 
     @Test
@@ -105,7 +105,7 @@ public class GroupComponentsTest extends SAITester
 
         Set<Component> components = StorageAttachedIndexGroup.getLiveComponents(sstables.iterator().next(), getIndexesFromGroup(group));
 
-        assertEquals(Version.LATEST.onDiskFormat().perSSTableIndexComponents().size() +
+        assertEquals(Version.LATEST.onDiskFormat().perSSTableIndexComponents(false).size() +
                      Version.LATEST.onDiskFormat().perColumnIndexComponents(indexContext).size(),
                      components.size());
     }

test/unit/org/apache/cassandra/index/sai/iterators/KeyRangeConcatIteratorTest.java

@@ -22,6 +22,9 @@ import java.util.stream.IntStream;
 
 import org.junit.Test;
 
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.index.sai.utils.PrimaryKey;
+
 import static org.apache.cassandra.index.sai.iterators.LongIterator.convert;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.Assert.assertEquals;
@@ -32,26 +35,24 @@ import static org.junit.Assert.assertTrue;
 
 public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
 {
-    private static final String MUST_BE_SORTED_ERROR = "RangeIterator must be sorted, previous max: PrimaryKey: " +
-                                                       "{ token: %d, partition: null, clustering: null:null} , " +
-                                                       "next min: PrimaryKey: { token: %d, partition: null, clustering: null:null} ";
+    PrimaryKey.Factory primaryKeyFactory = new PrimaryKey.Factory(null);
     @Test
     public void testValidation()
     {
         // Iterators being merged via concatanation must not include each other
         assertThatThrownBy(() -> buildConcat(build(1L, 4L), build(2L, 3L))).isInstanceOf(IllegalArgumentException.class)
-                                                                           .hasMessage(MUST_BE_SORTED_ERROR, 4, 2);
+                                                                           .hasMessage(createErrorMessage(4, 2));
 
         // Iterators being merged via concatanation must not overlap
         assertThatThrownBy(() -> buildConcat(build(1L, 4L), build(2L, 5L))).isInstanceOf(IllegalArgumentException.class)
-                                                                           .hasMessage(MUST_BE_SORTED_ERROR, 4, 2);
+                                                                           .hasMessage(createErrorMessage(4, 2));
 
         assertThatThrownBy(() -> buildConcat(build(1L, 4L), build(0L, 3L))).isInstanceOf(IllegalArgumentException.class)
-                                                                           .hasMessage(MUST_BE_SORTED_ERROR, 4, 0);
+                                                                           .hasMessage(createErrorMessage(4, 0));
 
         // Iterators being merged via concatanation must be sorted
         assertThatThrownBy(() -> buildConcat(build(2L, 4L), build(0L, 1L))).isInstanceOf(IllegalArgumentException.class)
-                                                                           .hasMessage(MUST_BE_SORTED_ERROR, 4, 0);
+                                                                           .hasMessage(createErrorMessage(4, 0));
 
         // allow min boundary included
         KeyRangeIterator concat = buildConcat(build(1L, 4L), build(4L, 5L));
@@ -422,4 +423,11 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
     {
         return KeyRangeConcatIterator.builder(16);
     }
+
+    private String createErrorMessage(int max, int min)
+    {
+        return String.format(KeyRangeConcatIterator.MUST_BE_SORTED_ERROR,
+                             primaryKeyFactory.createTokenOnly(new Murmur3Partitioner.LongToken(max)),
+                             primaryKeyFactory.createTokenOnly(new Murmur3Partitioner.LongToken(min)));
+    }
 }

test/unit/org/apache/cassandra/index/sai/metrics/IndexGroupMetricsTest.java

@@ -60,7 +60,7 @@ public class IndexGroupMetricsTest extends AbstractMetricsTest
 
         // with 10 sstable
         int indexopenFileCountWithOnlyNumeric = getOpenIndexFiles();
-        assertEquals(sstables * (Version.LATEST.onDiskFormat().openFilesPerSSTableIndex() +
+        assertEquals(sstables * (Version.LATEST.onDiskFormat().openFilesPerSSTableIndex(false) +
                                  Version.LATEST.onDiskFormat().openFilesPerColumnIndex(v1IndexContext)),
                      indexopenFileCountWithOnlyNumeric);
 
@@ -70,7 +70,7 @@ public class IndexGroupMetricsTest extends AbstractMetricsTest
         // compaction should reduce open files
         compact();
 
-        assertEquals(Version.LATEST.onDiskFormat().openFilesPerSSTableIndex() +
+        assertEquals(Version.LATEST.onDiskFormat().openFilesPerSSTableIndex(false) +
                      Version.LATEST.onDiskFormat().openFilesPerColumnIndex(v1IndexContext),
                      getOpenIndexFiles());
 

test/unit/org/apache/cassandra/index/sai/utils/AbstractPrimaryKeyTester.java

@@ -36,51 +36,51 @@ import static org.junit.Assert.assertTrue;
 
 public class AbstractPrimaryKeyTester extends SAIRandomizedTester
 {
-    static final TableMetadata simplePartition = TableMetadata.builder("test", "test")
+    protected static final TableMetadata simplePartition = TableMetadata.builder("test", "test")
                                                               .partitioner(Murmur3Partitioner.instance)
                                                               .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                               .build();
 
-    static final TableMetadata compositePartition = TableMetadata.builder("test", "test")
+    protected static final TableMetadata compositePartition = TableMetadata.builder("test", "test")
                                                                  .partitioner(Murmur3Partitioner.instance)
                                                                  .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                  .addPartitionKeyColumn("pk2", Int32Type.instance)
                                                                  .build();
 
-    static final TableMetadata simplePartitionSingleClusteringAsc = TableMetadata.builder("test", "test")
+    protected static final TableMetadata simplePartitionSingleClusteringAsc = TableMetadata.builder("test", "test")
                                                                                  .partitioner(Murmur3Partitioner.instance)
                                                                                  .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                  .addClusteringColumn("ck1", UTF8Type.instance)
                                                                                  .build();
 
-    static final TableMetadata simplePartitionMultipleClusteringAsc = TableMetadata.builder("test", "test")
+    protected static final TableMetadata simplePartitionMultipleClusteringAsc = TableMetadata.builder("test", "test")
                                                                                    .partitioner(Murmur3Partitioner.instance)
                                                                                    .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                    .addClusteringColumn("ck1", UTF8Type.instance)
                                                                                    .addClusteringColumn("ck2", UTF8Type.instance)
                                                                                    .build();
 
-    static final TableMetadata simplePartitionSingleClusteringDesc = TableMetadata.builder("test", "test")
+    protected static final TableMetadata simplePartitionSingleClusteringDesc = TableMetadata.builder("test", "test")
                                                                                   .partitioner(Murmur3Partitioner.instance)
                                                                                   .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                   .addClusteringColumn("ck1", ReversedType.getInstance(UTF8Type.instance))
                                                                                   .build();
 
-    static final TableMetadata simplePartitionMultipleClusteringDesc = TableMetadata.builder("test", "test")
+    protected static final TableMetadata simplePartitionMultipleClusteringDesc = TableMetadata.builder("test", "test")
                                                                                     .partitioner(Murmur3Partitioner.instance)
                                                                                     .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                     .addClusteringColumn("ck1", ReversedType.getInstance(UTF8Type.instance))
                                                                                     .addClusteringColumn("ck2", ReversedType.getInstance(UTF8Type.instance))
                                                                                     .build();
 
-    static final TableMetadata compositePartitionSingleClusteringAsc = TableMetadata.builder("test", "test")
+    protected static final TableMetadata compositePartitionSingleClusteringAsc = TableMetadata.builder("test", "test")
                                                                                     .partitioner(Murmur3Partitioner.instance)
                                                                                     .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                     .addPartitionKeyColumn("pk2", Int32Type.instance)
                                                                                     .addClusteringColumn("ck1", UTF8Type.instance)
                                                                                     .build();
 
-    static final TableMetadata compositePartitionMultipleClusteringAsc = TableMetadata.builder("test", "test")
+    protected static final TableMetadata compositePartitionMultipleClusteringAsc = TableMetadata.builder("test", "test")
                                                                                       .partitioner(Murmur3Partitioner.instance)
                                                                                       .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                       .addPartitionKeyColumn("pk2", Int32Type.instance)
@@ -88,14 +88,14 @@ public class AbstractPrimaryKeyTester extends SAIRandomizedTester
                                                                                       .addClusteringColumn("ck2", UTF8Type.instance)
                                                                                       .build();
 
-    static final TableMetadata compositePartitionSingleClusteringDesc = TableMetadata.builder("test", "test")
+    protected static final TableMetadata compositePartitionSingleClusteringDesc = TableMetadata.builder("test", "test")
                                                                                      .partitioner(Murmur3Partitioner.instance)
                                                                                      .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                      .addPartitionKeyColumn("pk2", Int32Type.instance)
                                                                                      .addClusteringColumn("ck1", ReversedType.getInstance(UTF8Type.instance))
                                                                                      .build();
 
-    static final TableMetadata compositePartitionMultipleClusteringDesc = TableMetadata.builder("test", "test")
+    protected static final TableMetadata compositePartitionMultipleClusteringDesc = TableMetadata.builder("test", "test")
                                                                                        .partitioner(Murmur3Partitioner.instance)
                                                                                        .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                        .addPartitionKeyColumn("pk2", Int32Type.instance)
@@ -103,14 +103,14 @@ public class AbstractPrimaryKeyTester extends SAIRandomizedTester
                                                                                        .addClusteringColumn("ck2", ReversedType.getInstance(UTF8Type.instance))
                                                                                        .build();
 
-    static final TableMetadata simplePartitionMultipleClusteringMixed = TableMetadata.builder("test", "test")
+    protected static final TableMetadata simplePartitionMultipleClusteringMixed = TableMetadata.builder("test", "test")
                                                                                      .partitioner(Murmur3Partitioner.instance)
                                                                                      .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                      .addClusteringColumn("ck1", UTF8Type.instance)
                                                                                      .addClusteringColumn("ck2", ReversedType.getInstance(UTF8Type.instance))
                                                                                      .build();
 
-    static final TableMetadata compositePartitionMultipleClusteringMixed = TableMetadata.builder("test", "test")
+    protected static final TableMetadata compositePartitionMultipleClusteringMixed = TableMetadata.builder("test", "test")
                                                                                         .partitioner(Murmur3Partitioner.instance)
                                                                                         .addPartitionKeyColumn("pk1", Int32Type.instance)
                                                                                         .addPartitionKeyColumn("pk2", Int32Type.instance)
@@ -118,14 +118,14 @@ public class AbstractPrimaryKeyTester extends SAIRandomizedTester
                                                                                         .addClusteringColumn("ck2", ReversedType.getInstance(UTF8Type.instance))
                                                                                         .build();
 
-    void assertByteComparison(PrimaryKey a, PrimaryKey b, int expected)
+    protected void assertByteComparison(PrimaryKey a, PrimaryKey b, int expected)
     {
         assertEquals(expected, ByteComparable.compare(a::asComparableBytes,
                                                       b::asComparableBytes,
                                                       ByteComparable.Version.OSS50));
     }
 
-    void assertCompareToAndEquals(PrimaryKey a, PrimaryKey b, int expected)
+    protected void assertCompareToAndEquals(PrimaryKey a, PrimaryKey b, int expected)
     {
         if (expected > 0)
         {
@@ -144,17 +144,17 @@ public class AbstractPrimaryKeyTester extends SAIRandomizedTester
         }
     }
 
-    DecoratedKey makeKey(TableMetadata table, Object...partitionKeys)
+    protected DecoratedKey makeKey(TableMetadata table, Object...partitionKeys)
     {
         ByteBuffer key;
         if (TypeUtil.isComposite(table.partitionKeyType))
             key = ((CompositeType)table.partitionKeyType).decompose(partitionKeys);
         else
-            key = table.partitionKeyType.fromString((String)partitionKeys[0]);
+            key = table.partitionKeyType.decomposeUntyped(partitionKeys[0]);
         return table.partitioner.decorateKey(key);
     }
 
-    Clustering<?> makeClustering(TableMetadata table, String...clusteringKeys)
+    protected Clustering<?> makeClustering(TableMetadata table, Object...clusteringKeys)
     {
         Clustering<?> clustering;
         if (table.comparator.size() == 0)
@@ -163,7 +163,7 @@ public class AbstractPrimaryKeyTester extends SAIRandomizedTester
         {
             ByteBuffer[] values = new ByteBuffer[clusteringKeys.length];
             for (int index = 0; index < table.comparator.size(); index++)
-                values[index] = table.comparator.subtype(index).fromString(clusteringKeys[index]);
+                values[index] = table.comparator.subtype(index).decomposeUntyped(clusteringKeys[index]);
             clustering = Clustering.make(values);
         }
         return clustering;

test/unit/org/apache/cassandra/index/sai/utils/PrimaryKeyTest.java

@@ -33,11 +33,10 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
         int rows = nextInt(10, 100);
         PrimaryKey[] keys = new PrimaryKey[rows];
         for (int index = 0; index < rows; index++)
-            keys[index] = factory.create(makeKey(simplePartition, Integer.toString(index)), Clustering.EMPTY);
+            keys[index] = factory.create(makeKey(simplePartition, index), Clustering.EMPTY);
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -52,7 +51,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -66,7 +64,7 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
         int clustering = 0;
         for (int index = 0; index < rows; index++)
         {
-            keys[index] = factory.create(makeKey(simplePartitionSingleClusteringAsc, Integer.toString(partition)),
+            keys[index] = factory.create(makeKey(simplePartitionSingleClusteringAsc, partition),
                                          makeClustering(simplePartitionSingleClusteringAsc, Integer.toString(clustering++)));
             if (clustering == 5)
             {
@@ -77,7 +75,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -92,7 +89,7 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
         int clustering2 = 0;
         for (int index = 0; index < rows; index++)
         {
-            keys[index] = factory.create(makeKey(simplePartitionMultipleClusteringAsc, Integer.toString(partition)),
+            keys[index] = factory.create(makeKey(simplePartitionMultipleClusteringAsc, partition),
                                          makeClustering(simplePartitionMultipleClusteringAsc, Integer.toString(clustering1), Integer.toString(clustering2++)));
             if (clustering2 == 5)
             {
@@ -108,7 +105,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -122,7 +118,7 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
         int clustering = 0;
         for (int index = 0; index < rows; index++)
         {
-            keys[index] = factory.create(makeKey(simplePartitionSingleClusteringDesc, Integer.toString(partition)),
+            keys[index] = factory.create(makeKey(simplePartitionSingleClusteringDesc, partition),
                                          makeClustering(simplePartitionSingleClusteringDesc, Integer.toString(clustering++)));
             if (clustering == 5)
             {
@@ -133,7 +129,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -148,7 +143,7 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
         int clustering2 = 0;
         for (int index = 0; index < rows; index++)
         {
-            keys[index] = factory.create(makeKey(simplePartitionMultipleClusteringDesc, Integer.toString(partition)),
+            keys[index] = factory.create(makeKey(simplePartitionMultipleClusteringDesc, partition),
                                          makeClustering(simplePartitionMultipleClusteringDesc, Integer.toString(clustering1), Integer.toString(clustering2++)));
             if (clustering2 == 5)
             {
@@ -164,7 +159,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -189,7 +183,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -220,7 +213,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -245,7 +237,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -276,7 +267,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -291,7 +281,7 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
         int clustering2 = 0;
         for (int index = 0; index < rows; index++)
         {
-            keys[index] = factory.create(makeKey(simplePartitionMultipleClusteringMixed, Integer.toString(partition)),
+            keys[index] = factory.create(makeKey(simplePartitionMultipleClusteringMixed, partition),
                                          makeClustering(simplePartitionMultipleClusteringMixed, Integer.toString(clustering1), Integer.toString(clustering2++)));
             if (clustering2 == 5)
             {
@@ -307,7 +297,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -338,7 +327,6 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
 
         Arrays.sort(keys);
 
-        byteComparisonTests(factory, keys);
         compareToAndEqualsTests(factory, keys);
     }
 
@@ -360,22 +348,4 @@ public class PrimaryKeyTest extends AbstractPrimaryKeyTester
             }
         }
     }
-
-    private void byteComparisonTests(PrimaryKey.Factory factory, PrimaryKey... keys)
-    {
-        for (int index = 0; index < keys.length - 1; index++)
-        {
-            PrimaryKey key = keys[index];
-            PrimaryKey tokenOnlyKey = factory.createTokenOnly(key.token());
-            assertByteComparison(tokenOnlyKey, key, -1);
-            assertByteComparison(key, key, 0);
-            assertByteComparison(tokenOnlyKey, tokenOnlyKey, 0);
-
-            for (int comparisonIndex = index + 1; comparisonIndex < keys.length; comparisonIndex++)
-            {
-                assertByteComparison(key, keys[comparisonIndex], -1);
-                assertByteComparison(tokenOnlyKey, keys[comparisonIndex], -1);
-            }
-        }
-    }
 }

test/unit/org/apache/cassandra/index/sai/utils/SAIRandomizedTester.java

@@ -30,13 +30,12 @@ import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.db.marshal.Int32Type;
 import org.apache.cassandra.dht.Murmur3Partitioner;
 import org.apache.cassandra.index.sai.SAITester;
+import org.apache.cassandra.index.sai.disk.io.IndexFileUtils;
 import org.apache.cassandra.index.sai.postings.PostingList;
 import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
-import org.apache.cassandra.io.compress.BufferType;
 import org.apache.cassandra.io.sstable.Descriptor;
 import org.apache.cassandra.io.sstable.SequenceBasedSSTableId;
 import org.apache.cassandra.io.util.File;
-import org.apache.cassandra.io.util.SequentialWriterOption;
 import org.apache.cassandra.schema.TableMetadata;
 
 import static org.junit.Assert.assertEquals;
@@ -71,13 +70,17 @@ public class SAIRandomizedTester extends SAITester
                                                                         randomSimpleString(3, 17),
                                                                         new SequenceBasedSSTableId(getRandom().nextIntBetween(0, 128))),
                                                          metadata,
-                                                         SequentialWriterOption.newBuilder()
-                                                                               .bufferSize(getRandom().nextIntBetween(17, 1 << 13))
-                                                                               .bufferType(getRandom().nextBoolean() ? BufferType.ON_HEAP : BufferType.OFF_HEAP)
-                                                                               .trickleFsync(getRandom().nextBoolean())
-                                                                               .trickleFsyncByteInterval(nextInt(1 << 10, 1 << 16))
-                                                                               .finishOnClose(true)
-                                                                               .build());
+                                                         IndexFileUtils.DEFAULT_WRITER_OPTION);
+    }
+
+    public static IndexDescriptor newClusteringIndexDescriptor(TableMetadata metadata) throws IOException
+    {
+        return indexInputLeakDetector.newIndexDescriptor(new Descriptor(new File(temporaryFolder.newFolder()),
+                                                                        randomSimpleString(5, 13),
+                                                                        randomSimpleString(3, 17),
+                                                                        new SequenceBasedSSTableId(getRandom().nextIntBetween(0, 128))),
+                                                         metadata,
+                                                         IndexFileUtils.DEFAULT_WRITER_OPTION);
     }
 
     public String newIndex()