Fix SAI intersection queries

- Fix comparison in PostingListRangeIterator for updating skip token
 - Fix binary search in KeyLookup.clusteredSeekToKey
 - Added new on-disk component for storing partition sizes by partition ID

 patch by Mike Adamson; reviewed by Caleb Rackliffe, Mick Semb Wever for CASSANDRA-19011

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

@@ -68,14 +68,19 @@ public enum IndexComponent
 
     // per-sstable components
     /**
-     * Partition key token value for rows including row tombstone and static row. (access key is rowId)
+     * An on-disk block packed index mapping rowIds to token values.
      */
-    TOKEN_VALUES("TokenValues"),
+    ROW_TO_TOKEN("RowToToken"),
 
     /**
-     * An on-disk block packed index containing the starting and ending rowIds for each partition.
+     * An on-disk block packed index mapping rowIds to partitionIds.
      */
-    PARTITION_SIZES("PartitionSizes"),
+    ROW_TO_PARTITION("RowToPartition"),
+
+    /**
+     * An on-disk block packed index mapping partitionIds to the number of rows for the partition.
+     */
+    PARTITION_TO_SIZE("PartitionToSize"),
 
     /**
      * Prefix-compressed blocks of partition keys used for rowId to partition key lookups

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

@@ -42,18 +42,22 @@ public class SSTableComponentsWriter implements PerSSTableIndexWriter
     private final IndexDescriptor indexDescriptor;
     private final MetadataWriter metadataWriter;
     private final NumericValuesWriter partitionSizeWriter;
+    private final NumericValuesWriter partitionRowsWriter;
     private final NumericValuesWriter tokenWriter;
     private final KeyStoreWriter partitionKeysWriter;
     private final KeyStoreWriter clusteringKeysWriter;
 
     private long partitionId = -1;
+    // This is used to record the number of rows in each partition
+    private long partitionRowCount = 0;
 
     public SSTableComponentsWriter(IndexDescriptor indexDescriptor) throws IOException
     {
         this.indexDescriptor = indexDescriptor;
         this.metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META));
-        this.tokenWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.TOKEN_VALUES, metadataWriter, false);
+        this.tokenWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.ROW_TO_TOKEN, metadataWriter, false);
-        this.partitionSizeWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.PARTITION_SIZES, metadataWriter, true);
+        this.partitionRowsWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.ROW_TO_PARTITION, metadataWriter, true);
+        this.partitionSizeWriter = new NumericValuesWriter(indexDescriptor, IndexComponent.PARTITION_TO_SIZE, metadataWriter, false);
         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),
@@ -82,7 +86,11 @@ public class SSTableComponentsWriter implements PerSSTableIndexWriter
     @Override
     public void startPartition(DecoratedKey partitionKey) throws IOException
     {
+        if (partitionId >= 0)
+            partitionSizeWriter.add(partitionRowCount);
+
         partitionId++;
+        partitionRowCount = 0;
         partitionKeysWriter.add(v -> ByteSource.of(partitionKey.getKey(), v));
         if (indexDescriptor.hasClustering())
             clusteringKeysWriter.startPartition();
@@ -92,7 +100,8 @@ public class SSTableComponentsWriter implements PerSSTableIndexWriter
     public void nextRow(PrimaryKey primaryKey) throws IOException
     {
         tokenWriter.add(primaryKey.token().getLongValue());
-        partitionSizeWriter.add(partitionId);
+        partitionRowsWriter.add(partitionId);
+        partitionRowCount++;
         if (indexDescriptor.hasClustering())
             clusteringKeysWriter.add(indexDescriptor.clusteringComparator.asByteComparable(primaryKey.clustering()));
     }
@@ -102,11 +111,12 @@ public class SSTableComponentsWriter implements PerSSTableIndexWriter
     {
         try
         {
+            partitionSizeWriter.add(partitionRowCount);
             indexDescriptor.createComponentOnDisk(IndexComponent.GROUP_COMPLETION_MARKER);
         }
         finally
         {
-            FileUtils.close(tokenWriter, partitionSizeWriter, partitionKeysWriter, clusteringKeysWriter, metadataWriter);
+            FileUtils.close(tokenWriter, partitionSizeWriter, partitionRowsWriter, partitionKeysWriter, clusteringKeysWriter, metadataWriter);
         }
     }
 

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

@@ -34,7 +34,6 @@ 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;
@@ -44,10 +43,10 @@ import org.apache.cassandra.utils.Throwables;
  * <p>
  * This uses the following on-disk structures:
  * <ul>
- *     <li>A block-packed structure for rowId to token lookups using {@link BlockPackedReader}.
+ *     <li>A block-packed structure for rowId to token value lookups using {@link BlockPackedReader}.
- *     Uses the {@link IndexComponent#TOKEN_VALUES} component</li>
+ *     Uses the {@link IndexComponent#ROW_TO_TOKEN} component</li>
  *     <li>A monotonic block packed structure for rowId to partitionId lookups using {@link MonotonicBlockPackedReader}.
- *     Uses the {@link IndexComponent#PARTITION_SIZES} component</li>
+ *     Uses the {@link IndexComponent#ROW_TO_PARTITION} 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>
@@ -63,29 +62,29 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
     public static class Factory implements PrimaryKeyMap.Factory
     {
         protected final MetadataSource metadataSource;
-        protected final LongArray.Factory tokenReaderFactory;
+        protected final LongArray.Factory rowToTokenReaderFactory;
-        protected final LongArray.Factory partitionReaderFactory;
+        protected final LongArray.Factory rowToPartitionReaderFactory;
         protected final KeyLookup partitionKeyReader;
         protected final PrimaryKey.Factory primaryKeyFactory;
 
-        private final FileHandle tokensFile;
+        private final FileHandle rowToTokenFile;
-        private final FileHandle partitionsFile;
+        private final FileHandle rowToPartitionFile;
         private final FileHandle partitionKeyBlockOffsetsFile;
         private final FileHandle partitionKeyBlocksFile;
 
-        public Factory(IndexDescriptor indexDescriptor, SSTableReader sstable)
+        public Factory(IndexDescriptor indexDescriptor)
         {
-            this.tokensFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES, this::close);
+            this.rowToTokenFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.ROW_TO_TOKEN, this::close);
-            this.partitionsFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PARTITION_SIZES, this::close);
+            this.rowToPartitionFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.ROW_TO_PARTITION, 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)));
+                NumericValuesMeta tokensMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.ROW_TO_TOKEN)));
-                this.tokenReaderFactory = new BlockPackedReader(tokensFile, tokensMeta);
+                this.rowToTokenReaderFactory = new BlockPackedReader(rowToTokenFile, tokensMeta);
-                NumericValuesMeta partitionsMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PARTITION_SIZES)));
+                NumericValuesMeta partitionsMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.ROW_TO_PARTITION)));
-                this.partitionReaderFactory = new MonotonicBlockPackedReader(partitionsFile, partitionsMeta);
+                this.rowToPartitionReaderFactory = new MonotonicBlockPackedReader(rowToPartitionFile, 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);
@@ -101,8 +100,8 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
         @SuppressWarnings({"resource", "RedundantSuppression"}) // rowIdToToken, rowIdToPartitionId and cursor are closed by the SkinnyPrimaryKeyMap#close method
         public PrimaryKeyMap newPerSSTablePrimaryKeyMap() throws IOException
         {
-            LongArray rowIdToToken = new LongArray.DeferredLongArray(tokenReaderFactory::open);
+            LongArray rowIdToToken = new LongArray.DeferredLongArray(rowToTokenReaderFactory::open);
-            LongArray rowIdToPartitionId = new LongArray.DeferredLongArray(partitionReaderFactory::open);
+            LongArray rowIdToPartitionId = new LongArray.DeferredLongArray(rowToPartitionReaderFactory::open);
             return new SkinnyPrimaryKeyMap(rowIdToToken,
                                            rowIdToPartitionId,
                                            partitionKeyReader.openCursor(),
@@ -112,22 +111,22 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
         @Override
         public void close()
         {
-            FileUtils.closeQuietly(Arrays.asList(tokensFile, partitionsFile, partitionKeyBlocksFile, partitionKeyBlockOffsetsFile));
+            FileUtils.closeQuietly(Arrays.asList(rowToTokenFile, rowToPartitionFile, partitionKeyBlocksFile, partitionKeyBlockOffsetsFile));
         }
     }
 
-    protected final LongArray tokenArray;
+    protected final LongArray rowIdToTokenArray;
-    protected final LongArray partitionArray;
+    protected final LongArray rowIdToPartitionIdArray;
     protected final KeyLookup.Cursor partitionKeyCursor;
     protected final PrimaryKey.Factory primaryKeyFactory;
 
-    protected SkinnyPrimaryKeyMap(LongArray tokenArray,
+    protected SkinnyPrimaryKeyMap(LongArray rowIdToTokenArray,
-                                  LongArray partitionArray,
+                                  LongArray rowIdToPartitionIdArray,
                                   KeyLookup.Cursor partitionKeyCursor,
                                   PrimaryKey.Factory primaryKeyFactory)
     {
-        this.tokenArray = tokenArray;
+        this.rowIdToTokenArray = rowIdToTokenArray;
-        this.partitionArray = partitionArray;
+        this.rowIdToPartitionIdArray = rowIdToPartitionIdArray;
         this.partitionKeyCursor = partitionKeyCursor;
         this.primaryKeyFactory = primaryKeyFactory;
     }
@@ -141,12 +140,12 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
     @Override
     public long rowIdFromPrimaryKey(PrimaryKey primaryKey)
     {
-        long rowId = tokenArray.indexOf(primaryKey.token().getLongValue());
+        long rowId = rowIdToTokenArray.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.kind() == PrimaryKey.Kind.TOKEN ||
             rowId < 0 ||
-            rowId + 1 == tokenArray.length() || tokenArray.get(rowId) != primaryKey.token().getLongValue())
+            rowId + 1 == rowIdToTokenArray.length() || rowIdToTokenArray.get(rowId) != primaryKey.token().getLongValue())
             return rowId;
         // Otherwise we need to check for token collision.
         return tokenCollisionDetection(primaryKey, rowId);
@@ -155,7 +154,7 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
     @Override
     public long ceiling(Token token)
     {
-        return tokenArray.indexOf(token.getLongValue());
+        return rowIdToTokenArray.indexOf(token.getLongValue());
     }
 
     @Override
@@ -164,14 +163,13 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
         if (token.isMinimum())
             return Long.MIN_VALUE;
 
-        long rowId = tokenArray.indexOf(token.getLongValue());
+        return rowIdToTokenArray.indexOf(token.getLongValue());
-        return rowId < 0 ? rowId : rowId;
     }
 
     @Override
     public void close()
     {
-        FileUtils.closeQuietly(Arrays.asList(partitionKeyCursor, tokenArray, partitionArray));
+        FileUtils.closeQuietly(Arrays.asList(partitionKeyCursor, rowIdToTokenArray, rowIdToPartitionIdArray));
     }
 
     // Look for token collision by if the ajacent token in the token array matches the
@@ -179,7 +177,7 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
     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))
+        while (rowId + 1 < rowIdToTokenArray.length() && primaryKey.token().getLongValue() == rowIdToTokenArray.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)
@@ -193,6 +191,6 @@ public class SkinnyPrimaryKeyMap implements PrimaryKeyMap
 
     protected DecoratedKey readPartitionKey(long sstableRowId)
     {
-        return primaryKeyFactory.partitionKeyFromComparableBytes(partitionKeyCursor.seekToPointId(partitionArray.get(sstableRowId)));
+        return primaryKeyFactory.partitionKeyFromComparableBytes(partitionKeyCursor.seekToPointId(rowIdToPartitionIdArray.get(sstableRowId)));
     }
 }

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

@@ -61,16 +61,17 @@ public class V1OnDiskFormat implements OnDiskFormat
     @VisibleForTesting
     public static final Set<IndexComponent> SKINNY_PER_SSTABLE_COMPONENTS = EnumSet.of(IndexComponent.GROUP_COMPLETION_MARKER,
                                                                                        IndexComponent.GROUP_META,
-                                                                                       IndexComponent.TOKEN_VALUES,
+                                                                                       IndexComponent.ROW_TO_TOKEN,
-                                                                                       IndexComponent.PARTITION_SIZES,
+                                                                                       IndexComponent.ROW_TO_PARTITION,
                                                                                        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.ROW_TO_TOKEN,
-                                                                                     IndexComponent.PARTITION_SIZES,
+                                                                                     IndexComponent.ROW_TO_PARTITION,
+                                                                                     IndexComponent.PARTITION_TO_SIZE,
                                                                                      IndexComponent.PARTITION_KEY_BLOCKS,
                                                                                      IndexComponent.PARTITION_KEY_BLOCK_OFFSETS,
                                                                                      IndexComponent.CLUSTERING_KEY_BLOCKS,
@@ -132,7 +133,7 @@ public class V1OnDiskFormat implements OnDiskFormat
     public PrimaryKeyMap.Factory newPrimaryKeyMapFactory(IndexDescriptor indexDescriptor, SSTableReader sstable)
     {
         return indexDescriptor.hasClustering() ? new WidePrimaryKeyMap.Factory(indexDescriptor, sstable)
-                                               : new SkinnyPrimaryKeyMap.Factory(indexDescriptor, sstable);
+                                               : new SkinnyPrimaryKeyMap.Factory(indexDescriptor);
     }
 
     @Override

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

@@ -29,6 +29,7 @@ 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.keystore.KeyLookupMeta;
 import org.apache.cassandra.index.sai.disk.v1.keystore.KeyLookup;
@@ -43,6 +44,8 @@ import org.apache.cassandra.utils.Throwables;
  * <p>
  * This used the following additional on-disk structures to the {@link SkinnyPrimaryKeyMap}
  * <ul>
+ *     <li>A block-packed structure for partitionId to partition size (number of rows in the partition) lookups using
+ *     {@link BlockPackedReader}. Uses the {@link IndexComponent#PARTITION_TO_SIZE} component</li>
  *     <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>
@@ -58,19 +61,24 @@ public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
     {
         private final ClusteringComparator clusteringComparator;
         private final KeyLookup clusteringKeyReader;
+        private final LongArray.Factory partitionToSizeReaderFactory;
         private final FileHandle clusteringKeyBlockOffsetsFile;
         private final FileHandle clustingingKeyBlocksFile;
+        private final FileHandle partitionToSizeFile;
 
         public Factory(IndexDescriptor indexDescriptor, SSTableReader sstable)
         {
-            super(indexDescriptor, sstable);
+            super(indexDescriptor);
 
             this.clusteringKeyBlockOffsetsFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.CLUSTERING_KEY_BLOCK_OFFSETS, this::close);
             this.clustingingKeyBlocksFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.CLUSTERING_KEY_BLOCKS, this::close);
+            this.partitionToSizeFile = indexDescriptor.createPerSSTableFileHandle(IndexComponent.PARTITION_TO_SIZE, this::close);
 
             try
             {
                 this.clusteringComparator = indexDescriptor.clusteringComparator;
+                NumericValuesMeta partitionSizeMeta = new NumericValuesMeta(metadataSource.get(indexDescriptor.componentName(IndexComponent.PARTITION_TO_SIZE)));
+                this.partitionToSizeReaderFactory = new BlockPackedReader(partitionToSizeFile, partitionSizeMeta);
                 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);
@@ -85,11 +93,13 @@ public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
         @SuppressWarnings({ "resource", "RedundantSuppression" }) // deferred long arrays and cursors are closed in the WidePrimaryKeyMap#close method
         public PrimaryKeyMap newPerSSTablePrimaryKeyMap() throws IOException
         {
-            LongArray rowIdToToken = new LongArray.DeferredLongArray(tokenReaderFactory::open);
+            LongArray rowIdToToken = new LongArray.DeferredLongArray(rowToTokenReaderFactory::open);
-            LongArray partitionIdToToken = new LongArray.DeferredLongArray(partitionReaderFactory::open);
+            LongArray partitionIdToToken = new LongArray.DeferredLongArray(rowToPartitionReaderFactory::open);
+            LongArray partitionIdToSize = new LongArray.DeferredLongArray(partitionToSizeReaderFactory::open);
 
             return new WidePrimaryKeyMap(rowIdToToken,
                                          partitionIdToToken,
+                                         partitionIdToSize,
                                          partitionKeyReader.openCursor(),
                                          clusteringKeyReader.openCursor(),
                                          primaryKeyFactory,
@@ -100,22 +110,25 @@ public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
         public void close()
         {
             super.close();
-            FileUtils.closeQuietly(Arrays.asList(clustingingKeyBlocksFile, clusteringKeyBlockOffsetsFile));
+            FileUtils.closeQuietly(Arrays.asList(clustingingKeyBlocksFile, clusteringKeyBlockOffsetsFile, partitionToSizeFile));
         }
     }
 
+    private final LongArray partitionIdToSizeArray;
     private final ClusteringComparator clusteringComparator;
     private final KeyLookup.Cursor clusteringKeyCursor;
 
-    private WidePrimaryKeyMap(LongArray tokenArray,
+    private WidePrimaryKeyMap(LongArray rowIdToTokenArray,
-                              LongArray partitionArray,
+                              LongArray rowIdToPartitionIdArray,
+                              LongArray partitionIdToSizeArray,
                               KeyLookup.Cursor partitionKeyCursor,
                               KeyLookup.Cursor clusteringKeyCursor,
                               PrimaryKey.Factory primaryKeyFactory,
                               ClusteringComparator clusteringComparator)
     {
-        super(tokenArray, partitionArray, partitionKeyCursor, primaryKeyFactory);
+        super(rowIdToTokenArray, rowIdToPartitionIdArray, partitionKeyCursor, primaryKeyFactory);
 
+        this.partitionIdToSizeArray = partitionIdToSizeArray;
         this.clusteringComparator = clusteringComparator;
         this.clusteringKeyCursor = clusteringKeyCursor;
     }
@@ -129,11 +142,11 @@ public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
     @Override
     public long rowIdFromPrimaryKey(PrimaryKey primaryKey)
     {
-        long rowId = tokenArray.indexOf(primaryKey.token().getLongValue());
+        long rowId = rowIdToTokenArray.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.kind() == PrimaryKey.Kind.TOKEN || rowId < 0 || tokenArray.get(rowId) != primaryKey.token().getLongValue())
+        if (primaryKey.kind() == PrimaryKey.Kind.TOKEN || rowId < 0 || rowIdToTokenArray.get(rowId) != primaryKey.token().getLongValue())
             return rowId;
 
         rowId = tokenCollisionDetection(primaryKey, rowId);
@@ -147,7 +160,7 @@ public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
     {
         if (token.isMinimum())
             return Long.MIN_VALUE;
-        long rowId = tokenArray.indexOf(token.getLongValue());
+        long rowId = rowIdToTokenArray.indexOf(token.getLongValue());
         return rowId < 0 ? rowId : startOfNextPartition(rowId) - 1;
     }
 
@@ -166,10 +179,7 @@ public class WidePrimaryKeyMap extends SkinnyPrimaryKeyMap
     // 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 partitionSize = partitionIdToSizeArray.get(rowIdToPartitionIdArray.get(rowId));
-        long nextPartitionRowId = partitionArray.indexOf(++partitionId);
+        return partitionSize == -1 ? rowIdToPartitionIdArray.length() : rowId + partitionSize;
-        if (nextPartitionRowId == -1)
-            nextPartitionRowId = partitionArray.length();
-        return nextPartitionRowId;
     }
 }

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

@@ -104,6 +104,12 @@ public class MonotonicBlockPackedReader implements LongArray.Factory
             {
                 return blockOffsets.get(block);
             }
+
+            @Override
+            public long indexOf(long value)
+            {
+                throw new UnsupportedOperationException();
+            }
         };
     }
 

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

@@ -192,48 +192,42 @@ public class KeyLookup
         {
             assert clustering : "Cannot do a clustered seek to a key on non-clustered keys";
 
-            BytesRef skipKey = asBytesRef(key);
+            BytesRef searchKey = asBytesRef(key);
 
             updateCurrentBlockIndex(startingPointId);
             resetToCurrentBlock();
 
-            if (compareKeys(currentKey, skipKey) == 0)
+            // We can return immediately if the currentPointId is within the requested partition range and the keys match
-                return startingPointId;
+            if (currentPointId >= startingPointId && currentPointId < endingPointId && compareKeys(currentKey, searchKey) == 0)
+                return currentPointId;
 
-            if (notInCurrentBlock(startingPointId, skipKey))
+            // Now do a binary search over the range if points between [lowSearchId, highSearchId)
+            long lowSearchId = startingPointId;
+            long highSearchId = endingPointId;
+
+            // We will keep going with the binary shift while the search consists of at least one block
+            while ((highSearchId - lowSearchId) >>> blockShift > 0)
             {
-                long split = (endingPointId - startingPointId) >>> blockShift;
+                long midSearchId = lowSearchId + (highSearchId - lowSearchId) / 2;
-                long splitPointId = startingPointId;
+
-                while (split > 0)
+                // See if the searchkey exists in the block containing the midSearchId or is above or below it
+                int position = moveToBlockAndCompareTo(midSearchId, searchKey);
+
+                if (position == 0)
                 {
-                    updateCurrentBlockIndex(Math.min((splitPointId >>> blockShift) + split, blockOffsets.length() - 1));
+                    lowSearchId = currentPointId;
-                    resetToCurrentBlock();
+                    break;
-
-                    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();
                 }
+
+                if (position < 0)
+                    highSearchId = midSearchId;
+                else
+                    lowSearchId = midSearchId;
             }
 
+            updateCurrentBlockIndex(lowSearchId);
+            resetToCurrentBlock();
+
             // Depending on where we are in the block we may need to move forwards to the starting point ID
             while (currentPointId < startingPointId)
             {
@@ -245,11 +239,13 @@ public class KeyLookup
             // Move forward to the ending point ID, returning the point ID if we find our key
             while (currentPointId < endingPointId)
             {
-                if (compareKeys(currentKey, skipKey) >= 0)
+                if (compareKeys(currentKey, searchKey) >= 0)
                     return currentPointId;
+
                 currentPointId++;
                 if (currentPointId == keyLookupMeta.keyCount)
                     return -1;
+
                 readCurrentKey();
                 updateCurrentBlockIndex(currentPointId);
             }
@@ -265,42 +261,43 @@ public class KeyLookup
             readCurrentKey();
         }
 
-
         @Override
         public void close()
         {
             keysInput.close();
         }
 
+        // Move to a block and see if the key is in the block using compareTo logic to indicate the keys position
+        // relative to the block.
+        // Note: It is down to the caller to position the block after a call to this method.
+        private int moveToBlockAndCompareTo(long pointId, BytesRef key)
+        {
+            updateCurrentBlockIndex(pointId);
+            resetToCurrentBlock();
+
+            if (compareKeys(key, currentKey) < 0)
+                return -1;
+
+            // If we are in the last block we will assume for now that the key is in the last block and defer
+            // the final decision to later (if we can't find it).
+            if (currentBlockIndex == blockOffsets.length() -1)
+                return 0;
+
+            // Finish by getting the starting key of the next block and comparing that with the key.
+            keysInput.seek(blockOffsets.get(currentBlockIndex + 1) + keysFilePointer);
+            readKey((currentBlockIndex + 1) << blockShift, nextBlockKey);
+            return compareKeys(key, nextBlockKey) < 0 ? 0 : 1;
+        }
+
         private void updateCurrentBlockIndex(long pointId)
         {
             currentBlockIndex = pointId >>> blockShift;
         }
 
-        private boolean notInCurrentBlock(long pointId, BytesRef key)
+        // Reset currentPointId and currentKey to be at the start of the block pointed to by currentBlockIndex.
-        {
-            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();

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

@@ -67,7 +67,7 @@ public class PostingListRangeIterator extends KeyRangeIterator
     private final long rowIdOffset;
 
     private boolean needsSkipping = false;
-    private PrimaryKey skipToToken = null;
+    private PrimaryKey skipToKey = null;
 
     /**
      * Create a direct PostingListRangeIterator where the underlying PostingList is materialised
@@ -86,28 +86,13 @@ public class PostingListRangeIterator extends KeyRangeIterator
         this.queryContext = searcherContext.context;
     }
 
-    public PostingListRangeIterator(IndexIdentifier indexIdentifier,
-                                    PrimaryKeyMap primaryKeyMap,
-                                    PostingList postingList,
-                                    QueryContext queryContext)
-    {
-        super(primaryKeyMap.primaryKeyFromRowId(postingList.minimum()),
-              primaryKeyMap.primaryKeyFromRowId(postingList.maximum()),
-              postingList.size());
-        this.indexIdentifier = indexIdentifier;
-        this.primaryKeyMap = primaryKeyMap;
-        this.postingList = postingList;
-        this.rowIdOffset = 0;
-        this.queryContext = queryContext;
-    }
-
     @Override
     protected void performSkipTo(PrimaryKey nextKey)
     {
-        if (skipToToken != null && skipToToken.compareTo(nextKey) >= 0)
+        if (skipToKey != null && skipToKey.compareTo(nextKey) > 0)
             return;
 
-        skipToToken = nextKey;
+        skipToKey = nextKey;
         needsSkipping = true;
     }
 
@@ -151,7 +136,7 @@ public class PostingListRangeIterator extends KeyRangeIterator
 
     private boolean exhausted()
     {
-        return needsSkipping && skipToToken.compareTo(getMaximum()) > 0;
+        return needsSkipping && skipToKey.compareTo(getMaximum()) > 0;
     }
 
     /**
@@ -162,7 +147,7 @@ public class PostingListRangeIterator extends KeyRangeIterator
         long segmentRowId;
         if (needsSkipping)
         {
-            long targetRowID = primaryKeyMap.rowIdFromPrimaryKey(skipToToken);
+            long targetRowID = primaryKeyMap.rowIdFromPrimaryKey(skipToKey);
             // skipToToken is larger than max token in token file
             if (targetRowID < 0)
             {

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

@@ -0,0 +1,227 @@
+/*
+ * 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.util.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
+
+@RunWith(Parameterized.class)
+public class RandomIntersectionTest extends SAIRandomizedTester
+{
+    private static final Object[][] EMPTY_ROWS = new Object[][]{};
+
+    @Parameterized.Parameter
+    public String testName;
+
+    @Parameterized.Parameter(1)
+    public boolean partitionRestricted;
+
+    @Parameterized.Parameter(2)
+    public boolean largePartition;
+
+    @Parameterized.Parameter(3)
+    public boolean v1Cardinality;
+
+    @Parameterized.Parameter(4)
+    public boolean v2Cardinality;
+
+    @Parameterized.Parameters(name = "{0}")
+    public static List<Object[]> parameters()
+    {
+        List<Object[]> parameters = new LinkedList<>();
+
+        parameters.add(new Object[]{ "Large partition restricted high high", true, true, true, true });
+        parameters.add(new Object[]{ "Large partition restricted low low", true, true, false, false });
+        parameters.add(new Object[]{ "Large partition restricted high low", true, true, true, false });
+        parameters.add(new Object[]{ "Large partition unrestricted high high", false, true, true, true });
+        parameters.add(new Object[]{ "Large partition unrestricted low low", false, true, false, false });
+        parameters.add(new Object[]{ "Large partition unrestricted high low", false, true, true, false });
+        parameters.add(new Object[]{ "Small partition restricted high high", true, false, true, true });
+        parameters.add(new Object[]{ "Small partition restricted low low", true, false, false, false });
+        parameters.add(new Object[]{ "Small partition restricted high low", true, false, true, false });
+        parameters.add(new Object[]{ "Small partition unrestricted high high", false, false, true, true });
+        parameters.add(new Object[]{ "Small partition unrestricted low low", false, false, false, false });
+        parameters.add(new Object[]{ "Small partition unrestricted high low", false, false, true, false });
+
+        return parameters;
+    }
+
+    private int numRows;
+
+    @Before
+    public void createTableAndIndexes()
+    {
+        createTable("CREATE TABLE %s (pk int, 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'");
+
+        numRows = nextInt(50000, 200000);
+    }
+
+    @Test
+    public void randomIntersectionTest() throws Throwable
+    {
+        if (partitionRestricted)
+            runRestrictedQueries();
+        else
+            runUnrestrictedQueries();
+    }
+
+    private void runRestrictedQueries() throws Throwable
+    {
+        Map<Integer, List<TestRow>> testRowMap = buildAndLoadTestRows();
+
+        beforeAndAfterFlush(() -> {
+            for (int queryCount = 0; queryCount < nextInt(10, 100); queryCount++)
+            {
+                int pk = testRowMap.keySet().stream().skip(nextInt(0, testRowMap.size())).findFirst().orElseThrow();
+                int v1 = nextV1();
+                int v2 = nextV2();
+
+                List<Object[]> expected = testRowMap.get(pk)
+                                                    .stream()
+                                                    .sorted(Comparator.comparingInt(o -> o.ck))
+                                                    .filter(row -> row.v1 > v1 && row.v2 > v2)
+                                                    .map(row -> row(row.ck))
+                                                    .collect(Collectors.toList());
+
+                assertRows(execute("SELECT ck FROM %s WHERE pk = ? AND v1 > ? AND v2 > ?", pk, v1, v2), expected.toArray(EMPTY_ROWS));
+            }
+        });
+    }
+
+    private void runUnrestrictedQueries() throws Throwable
+    {
+        Map<Integer, List<TestRow>> testRowMap = buildAndLoadTestRows();
+
+        beforeAndAfterFlush(() -> {
+            for (int queryCount = 0; queryCount < nextInt(10, 100); queryCount++)
+            {
+                int v1 = nextV1();
+                int v2 = nextV2();
+
+                List<Object[]> expected = testRowMap.values()
+                                                    .stream()
+                                                    .flatMap(Collection::stream)
+                                                    .filter(row -> row.v1 == v1 && row.v2 == v2)
+                                                    .map(row -> row(row.ck))
+                                                    .collect(Collectors.toList());
+
+                assertRowsIgnoringOrder(execute("SELECT ck FROM %s WHERE v1 = ? AND v2 = ?", v1, v2), expected.toArray(EMPTY_ROWS));
+            }
+        });
+    }
+
+    private Map<Integer, List<TestRow>> buildAndLoadTestRows()
+    {
+        Map<Integer, List<TestRow>> testRowMap = new HashMap<>();
+
+        int clusterSize = largePartition ? nextInt(500, 5000) : nextInt(10, 100);
+        int partition = nextInt(0, numRows);
+        List<TestRow> rowList = new ArrayList<>(clusterSize);
+        testRowMap.put(partition, rowList);
+        int clusterCount = 0;
+        for (int index = 0; index < numRows; index++)
+        {
+            TestRow row = new TestRow(partition, nextInt(10, numRows), nextV1(), nextV2());
+            while (rowList.contains(row))
+                row = new TestRow(partition, nextInt(10, numRows), nextV1(), nextV2());
+
+            rowList.add(row);
+            clusterCount++;
+            if (clusterCount == clusterSize)
+            {
+                clusterCount = 0;
+                clusterSize = largePartition ? nextInt(500, 5000) : nextInt(10, 100);
+                partition = nextInt(0, numRows);
+                while (testRowMap.containsKey(partition))
+                    partition = nextInt(0, numRows);
+                rowList = new ArrayList<>(clusterSize);
+                testRowMap.put(partition, rowList);
+            }
+        }
+        testRowMap.values().stream().flatMap(Collection::stream).forEach(row -> execute("INSERT INTO %s (pk, ck, v1, v2) VALUES (?, ?, ?, ?)",
+                                                                                        row.pk, row.ck, row.v1, row.v2));
+        return testRowMap;
+    }
+
+    private int nextV1()
+    {
+        return v1Cardinality ? nextInt(10, numRows/10) : nextInt(10, numRows/1000);
+    }
+
+    private int nextV2()
+    {
+        return v2Cardinality ? nextInt(10, numRows/10) : nextInt(10, numRows/1000);
+    }
+
+    private static class TestRow implements Comparable<TestRow>
+    {
+        final int pk;
+        final int ck;
+        final int v1;
+        final int v2;
+
+        TestRow(int pk, int ck, int v1, int v2)
+        {
+            this.pk = pk;
+            this.ck = ck;
+            this.v1 = v1;
+            this.v2 = v2;
+        }
+
+        @Override
+        public int compareTo(TestRow other)
+        {
+            int cmp = Integer.compare(pk, other.pk);
+            if (cmp != 0)
+                return cmp;
+            return Integer.compare(ck, other.ck);
+        }
+
+        @Override
+        public boolean equals(Object obj)
+        {
+            if (obj instanceof TestRow)
+                return compareTo((TestRow) obj) == 0;
+
+            return false;
+        }
+
+        @Override
+        public int hashCode()
+        {
+            return Objects.hash(pk, ck);
+        }
+    }
+}

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

@@ -23,6 +23,7 @@ package org.apache.cassandra.index.sai.cql;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
+import java.util.EnumSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -122,6 +123,12 @@ public class StorageAttachedIndexDDLTest extends SAITester
                                                                          .add(ActionBuilder.newActionBuilder().actions().doThrow(RuntimeException.class, Expression.quote("Injected failure!")))
                                                                          .build();
 
+    private static final EnumSet<IndexComponent> PER_SSTABLE_COMPONENTS = EnumSet.of(IndexComponent.ROW_TO_PARTITION,
+                                                                                     IndexComponent.PARTITION_TO_SIZE,
+                                                                                     IndexComponent.PARTITION_KEY_BLOCKS,
+                                                                                     IndexComponent.PARTITION_KEY_BLOCK_OFFSETS,
+                                                                                     IndexComponent.CLUSTERING_KEY_BLOCKS,
+                                                                                     IndexComponent.CLUSTERING_KEY_BLOCK_OFFSETS);
     @Before
     public void setup() throws Throwable
     {
@@ -952,9 +959,7 @@ public class StorageAttachedIndexDDLTest extends SAITester
         // that are encryptable unless they have been removed because encrypted components aren't
         // checksum validated.
 
-        if (component == IndexComponent.PARTITION_SIZES || component == IndexComponent.PARTITION_KEY_BLOCKS ||
+        if (PER_SSTABLE_COMPONENTS.contains(component))
-            component == IndexComponent.PARTITION_KEY_BLOCK_OFFSETS || component == IndexComponent.CLUSTERING_KEY_BLOCKS ||
-            component == IndexComponent.CLUSTERING_KEY_BLOCK_OFFSETS)
             return;
 
         if (((component == IndexComponent.GROUP_COMPLETION_MARKER) ||

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

@@ -65,9 +65,9 @@ public class NumericValuesTest extends SAIRandomizedTester
 
         final MetadataSource source = MetadataSource.loadGroupMetadata(indexDescriptor);
 
-        NumericValuesMeta tokensMeta = new NumericValuesMeta(source.get(indexDescriptor.componentName(IndexComponent.TOKEN_VALUES)));
+        NumericValuesMeta tokensMeta = new NumericValuesMeta(source.get(indexDescriptor.componentName(IndexComponent.ROW_TO_TOKEN)));
 
-        try (FileHandle fileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES, null);
+        try (FileHandle fileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.ROW_TO_TOKEN, null);
              LongArray reader = monotonic ? new MonotonicBlockPackedReader(fileHandle, tokensMeta).open()
                                           : new BlockPackedReader(fileHandle, tokensMeta).open())
         {
@@ -85,9 +85,9 @@ public class NumericValuesTest extends SAIRandomizedTester
         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)));
+        NumericValuesMeta tokensMeta = new NumericValuesMeta(source.get(indexDescriptor.componentName(IndexComponent.ROW_TO_TOKEN)));
 
-        try (FileHandle fileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.TOKEN_VALUES, null);
+        try (FileHandle fileHandle = indexDescriptor.createPerSSTableFileHandle(IndexComponent.ROW_TO_TOKEN, null);
              LongArray reader = (monotonic ? new MonotonicBlockPackedReader(fileHandle, tokensMeta)
                                            : new BlockPackedReader(fileHandle, tokensMeta)).open())
         {
@@ -107,7 +107,7 @@ public class NumericValuesTest extends SAIRandomizedTester
         long current = 0;
         try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerSSTableOutput(IndexComponent.GROUP_META));
              final NumericValuesWriter numericWriter = new NumericValuesWriter(indexDescriptor,
-                                                                               IndexComponent.TOKEN_VALUES,
+                                                                               IndexComponent.ROW_TO_TOKEN,
                                                                                metadataWriter,
                                                                                monotonic,
                                                                                blockSize))