Clean up KeyRangeIterator classes

* replace KeyRangeConcatIterator's PriorityQeueu with List
* remove KeyRangeIterator.current and simplify
* remove injected exception and tests - not relevant to the current implementation
* expand randomized testing
* inline getCurrent() -> peek(); rename getCount to getMaxKeys
* redefine skipTo contract to not return a value (which saves unnecessary work when skipTo is called multiple times in a row)
* calling hasNext in skipTo is a pessimization; if the iterator is in DONE state, then skipTo will see it and avoid further effort; if it is not, then we are computing a next value that we're just going to throw away
* fix SingleNodeQueryFailureTest - tests now multi and single index
* rationalize/standardize the way we release SSTableIndexes in QueryController

patch by Ekaterina Dimitrova; reviewed by Caleb Rackliffe, Ekaterina Dimitrova for CASSANDRA-19428

Co-authored-by: Caleb Rackliffe <calebrackliffe@gmail.com>
Co-authored-by: Jonathan Ellis <jbellis@apache.org>
Co-authored-by: Piotr Kolaczkowski <pkolaczk@datastax.com>
Co-authored-by: Michael Marshall <michael.marshall@datastax.com>

CHANGES.txt

@@ -1,4 +1,5 @@
 5.0-beta2
+ * Clean up KeyRangeIterator classes (CASSANDRA-19428)
  * Warn clients about possible consistency violations for filtering queries against multiple mutable columns (CASSANDRA-19489)
  * Align buffer with commitlog segment size (CASSANDRA-19471)
  * Ensure SAI indexes empty byte buffers for types that allow them as a valid input (CASSANDRA-19461)

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

@@ -41,17 +41,12 @@ public class IndexSearchResultIterator extends KeyRangeIterator
 {
     private static final Logger logger = LoggerFactory.getLogger(IndexSearchResultIterator.class);
 
-    private final QueryContext context;
     private final KeyRangeIterator union;
-    private final Collection<SSTableIndex> referencedIndexes;
 
-    private IndexSearchResultIterator(KeyRangeIterator union, Collection<SSTableIndex> referencedIndexes, QueryContext queryContext)
+    private IndexSearchResultIterator(KeyRangeIterator union, Runnable onClose)
     {
-        super(union.getMinimum(), union.getMaximum(), union.getCount());
+        super(union.getMinimum(), union.getMaximum(), union.getMaxKeys(), onClose);
-
         this.union = union;
-        this.referencedIndexes = referencedIndexes;
-        this.context = queryContext;
     }
 
     /**
@@ -62,7 +57,8 @@ public class IndexSearchResultIterator extends KeyRangeIterator
                                                   Collection<SSTableIndex> sstableIndexes,
                                                   AbstractBounds<PartitionPosition> keyRange,
                                                   QueryContext queryContext,
-                                                  boolean includeMemtables)
+                                                  boolean includeMemtables,
+                                                  Runnable onClose)
     {
         List<KeyRangeIterator> subIterators = new ArrayList<>(sstableIndexes.size() + (includeMemtables ? 1 : 0));
 
@@ -97,24 +93,25 @@ public class IndexSearchResultIterator extends KeyRangeIterator
             }
         }
 
-        KeyRangeIterator union = KeyRangeUnionIterator.build(subIterators);
+        KeyRangeIterator union = KeyRangeUnionIterator.build(subIterators, () -> {});
-        return new IndexSearchResultIterator(union, sstableIndexes, queryContext);
+        return new IndexSearchResultIterator(union, onClose);
     }
 
     public static IndexSearchResultIterator build(List<KeyRangeIterator> sstableIntersections,
                                                   KeyRangeIterator memtableResults,
                                                   Set<SSTableIndex> referencedIndexes,
-                                                  QueryContext queryContext)
+                                                  QueryContext queryContext,
+                                                  Runnable onClose)
     {
         queryContext.sstablesHit += referencedIndexes
                                     .stream()
                                     .map(SSTableIndex::getSSTable).collect(Collectors.toSet()).size();
         queryContext.checkpoint();
-        KeyRangeIterator union = KeyRangeUnionIterator.builder(sstableIntersections.size() + 1)
+        KeyRangeIterator union = KeyRangeUnionIterator.builder(sstableIntersections.size() + 1, () -> {})
                                                       .add(sstableIntersections)
                                                       .add(memtableResults)
                                                       .build();
-        return new IndexSearchResultIterator(union, referencedIndexes, queryContext);
+        return new IndexSearchResultIterator(union, onClose);
     }
 
     protected PrimaryKey computeNext()
@@ -127,22 +124,10 @@ public class IndexSearchResultIterator extends KeyRangeIterator
         union.skipTo(nextKey);
     }
 
+    @Override
     public void close()
     {
+        super.close();
         FileUtils.closeQuietly(union);
-        referencedIndexes.forEach(IndexSearchResultIterator::releaseQuietly);
-        referencedIndexes.clear();
-    }
-
-    private static void releaseQuietly(SSTableIndex index)
-    {
-        try
-        {
-            index.release();
-        }
-        catch (Throwable e)
-        {
-            logger.error(index.getIndexIdentifier().logMessage(String.format("Failed to release index on SSTable %s", index.getSSTable())), e);
-        }
     }
 }

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

@@ -77,7 +77,7 @@ public class PostingListRangeIterator extends KeyRangeIterator
                                     PrimaryKeyMap primaryKeyMap,
                                     IndexSegmentSearcherContext searcherContext)
     {
-        super(searcherContext.minimumKey, searcherContext.maximumKey, searcherContext.count());
+        super(searcherContext.minimumKey, searcherContext.maximumKey, searcherContext.count(), () -> {});
 
         this.indexIdentifier = indexIdentifier;
         this.primaryKeyMap = primaryKeyMap;

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

@@ -18,9 +18,7 @@
 package org.apache.cassandra.index.sai.iterators;
 
 import java.util.ArrayList;
-import java.util.Comparator;
 import java.util.List;
-import java.util.PriorityQueue;
 
 import com.google.common.annotations.VisibleForTesting;
 
@@ -28,90 +26,97 @@ import org.apache.cassandra.index.sai.utils.PrimaryKey;
 import org.apache.cassandra.io.util.FileUtils;
 
 /**
- * {@link KeyRangeConcatIterator} takes a list of sorted range iterator and concatenates them, leaving duplicates in
+ * {@link KeyRangeConcatIterator} takes a list of sorted range iterators and concatenates them, leaving duplicates in
  * place, to produce a new stably sorted iterator. Duplicates are eliminated later in
  * {@link org.apache.cassandra.index.sai.plan.StorageAttachedIndexSearcher}
  * as results from multiple SSTable indexes and their respective segments are consumed.
- *
+ * <p>
  * ex. (1, 2, 3) + (3, 3, 4, 5) -> (1, 2, 3, 3, 3, 4, 5)
  * ex. (1, 2, 2, 3) + (3, 4, 4, 6, 6, 7) -> (1, 2, 2, 3, 3, 4, 4, 6, 6, 7)
- *
- * TODO Investigate removing the use of PriorityQueue from this class <a href="https://issues.apache.org/jira/browse/CASSANDRA-18165">CASSANDRA-18165</a>
  */
 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> ranges;
-    private final List<KeyRangeIterator> toRelease;
 
-    protected KeyRangeConcatIterator(KeyRangeIterator.Builder.Statistics statistics, PriorityQueue<KeyRangeIterator> ranges)
+    private int current;
+
+    protected KeyRangeConcatIterator(KeyRangeIterator.Builder.Statistics statistics, List<KeyRangeIterator> ranges, Runnable onClose)
     {
-        super(statistics);
+        super(statistics, onClose);
 
+        if (ranges.isEmpty())
+            throw new IllegalArgumentException("Cannot concatenate empty list of ranges");
+
+        this.current = 0;
         this.ranges = ranges;
-        this.toRelease = new ArrayList<>(ranges);
     }
 
     @Override
     protected void performSkipTo(PrimaryKey nextKey)
     {
-        while (!ranges.isEmpty())
+        while (current < ranges.size())
         {
-            if (ranges.peek().getCurrent().compareTo(nextKey) >= 0)
+            KeyRangeIterator currentIterator = ranges.get(current);
+
+            if (currentIterator.hasNext() && currentIterator.peek().compareTo(nextKey) >= 0)
                 break;
 
-            KeyRangeIterator head = ranges.poll();
+            if (currentIterator.getMaximum().compareTo(nextKey) >= 0)
-
-            if (head.getMaximum().compareTo(nextKey) >= 0)
             {
-                head.skipTo(nextKey);
+                currentIterator.skipTo(nextKey);
-                ranges.add(head);
                 break;
             }
-        }
-    }
 
-    @Override
+            current++;
-    public void close()
+        }
-    {
-        // due to lazy key fetching, we cannot close iterator immediately
-        FileUtils.closeQuietly(toRelease);
     }
 
     @Override
     protected PrimaryKey computeNext()
     {
-        while (!ranges.isEmpty())
+        while (current < ranges.size())
         {
-            KeyRangeIterator current = ranges.poll();
+            KeyRangeIterator currentIterator = ranges.get(current);
-            if (current.hasNext())
-            {
-                PrimaryKey next = current.next();
-                // hasNext will update RangeIterator's current which is used to sort in PQ
-                if (current.hasNext())
-                    ranges.add(current);
 
-                return next;
+            if (currentIterator.hasNext())
-            }
+                return currentIterator.next();
+
+            current++;
         }
 
         return endOfData();
     }
 
+    @Override
+    public void close()
+    {
+        super.close();
+
+        // due to lazy key fetching, we cannot close iterator immediately
+        FileUtils.closeQuietly(ranges);
+    }
+
     public static Builder builder(int size)
     {
-        return new Builder(size);
+        return builder(size, () -> {});
+    }
+
+    public static Builder builder(int size, Runnable onClose)
+    {
+        return new Builder(size, onClose);
     }
 
     @VisibleForTesting
     public static class Builder extends KeyRangeIterator.Builder
     {
-        private final PriorityQueue<KeyRangeIterator> ranges;
+        // We can use a list because the iterators are already in order
+        private final List<KeyRangeIterator> ranges;
 
-        Builder(int size)
+        Builder(int size, Runnable onClose)
         {
-            super(new ConcatStatistics());
+            super(new ConcatStatistics(), onClose);
-            ranges = new PriorityQueue<>(size, Comparator.comparing(KeyRangeIterator::getCurrent));
+            this.ranges = new ArrayList<>(size);
         }
 
         @Override
@@ -120,7 +125,7 @@ public class KeyRangeConcatIterator extends KeyRangeIterator
             if (range == null)
                 return this;
 
-            if (range.getCount() > 0)
+            if (range.getMaxKeys() > 0)
                 ranges.add(range);
             else
                 FileUtils.closeQuietly(range);
@@ -138,16 +143,26 @@ public class KeyRangeConcatIterator extends KeyRangeIterator
         @Override
         public void cleanup()
         {
+            super.cleanup();
             FileUtils.closeQuietly(ranges);
         }
 
         @Override
         protected KeyRangeIterator buildIterator()
         {
+            if (rangeCount() == 0)
+            {
+                onClose.run();
+                return empty();
+            }
             if (rangeCount() == 1)
-                return ranges.poll();
+            {
+                KeyRangeIterator single = ranges.get(0);
+                single.setOnClose(onClose);
+                return single;
+            }
 
-            return new KeyRangeConcatIterator(statistics, ranges);
+            return new KeyRangeConcatIterator(statistics, ranges, onClose);
         }
     }
 
@@ -157,7 +172,7 @@ public class KeyRangeConcatIterator extends KeyRangeIterator
         public void update(KeyRangeIterator range)
         {
             // range iterators should be sorted, but previous max must not be greater than next min.
-            if (range.getCount() > 0)
+            if (range.getMaxKeys() > 0)
             {
                 if (count == 0)
                 {
@@ -169,7 +184,7 @@ public class KeyRangeConcatIterator extends KeyRangeIterator
                 }
 
                 max = range.getMaximum();
-                count += range.getCount();
+                count += range.getMaxKeys();
             }
         }
     }

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

@@ -38,7 +38,7 @@ import javax.annotation.Nullable;
  * A simple intersection iterator that makes no real attempts at optimising the iteration apart from
  * initially sorting the ranges. This implementation also supports an intersection limit via
  * {@code CassandraRelevantProperties.SAI_INTERSECTION_CLAUSE_LIMIT} which limits the number of ranges that will 
- * be included in the intersection.
+ * be included in the intersection. This currently defaults to 2.
  * <p> 
  * Intersection only works for ranges that are compatible according to {@link PrimaryKey.Kind#isIntersectable(Kind)}.
  */
@@ -52,21 +52,20 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
     }
 
     private final List<KeyRangeIterator> ranges;
+    private PrimaryKey highestKey;
 
-    private KeyRangeIntersectionIterator(Builder.Statistics statistics, List<KeyRangeIterator> ranges)
+    private KeyRangeIntersectionIterator(Builder.Statistics statistics, List<KeyRangeIterator> ranges, Runnable onClose)
     {
-        super(statistics);
+        super(statistics, onClose);
         this.ranges = ranges;
+        this.highestKey = null;
     }
 
     @Override
     protected PrimaryKey computeNext()
     {
-        // Advance one iterator to the next key and remember the key as the highest seen so far.
+        if (highestKey == null)
-        // It can become null when we reach the end of the iterator.
+            highestKey = computeHighestKey();
-        // If there are both static and non-static keys being iterated here, we advance a non-static one,
-        // regardless of the order of ranges in the ranges list.
-        PrimaryKey highestKey = advanceOneRange();
 
         outer:
         // After advancing one iterator, we must try to advance all the other iterators that got behind,
@@ -81,14 +80,17 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
             // Once this inner loop finishes normally, all iterators are guaranteed to be at the same value.
             for (KeyRangeIterator range : ranges)
             {
-                if (range.getCurrent().compareTo(highestKey) < 0)
+                if (!range.hasNext())
+                    return endOfData();
+
+                if (range.peek().compareTo(highestKey) < 0)
                 {
                     // If we advance a STATIC key, then we must advance it to the same partition as the highestKey.
                     // Advancing a STATIC key to a WIDE key directly (without throwing away the clustering) would
                     // go too far, as WIDE keys are stored after STATIC in the posting list.
-                    PrimaryKey nextKey = range.getCurrent().kind() == Kind.STATIC
+                    PrimaryKey nextKey = range.peek().kind() == Kind.STATIC
-                                         ? nextOrNull(range, highestKey.toStatic())
+                                         ? skipAndPeek(range, highestKey.toStatic())
-                                         : nextOrNull(range, highestKey);
+                                         : skipAndPeek(range, highestKey);
 
                     // We use strict comparison here, since it orders WIDE primary keys after STATIC primary keys
                     // in the same partition. When WIDE keys are present, we want to return them rather than STATIC
@@ -111,7 +113,17 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
 
             // If we get here, all iterators have been advanced to the same key. When STATIC and WIDE keys are
             // mixed, this means WIDE keys point to exactly the same row, and STATIC keys the same partition.
-            return highestKey;
+            PrimaryKey result = highestKey;
+
+            // Advance one iterator to the next key and remember the key as the highest seen so far.
+            // It can become null when we reach the end of the iterator.
+            // If there are both static and non-static keys being iterated here, we advance a non-static one,
+            // regardless of the order of ranges in the ranges list.
+            highestKey = advanceOneRange();
+
+            // If we get here, all iterators have been advanced to the same key. When STATIC and WIDE keys are
+            // mixed, this means WIDE keys point to exactly the same row, and STATIC keys the same partition.
+            return result;
         }
 
         return endOfData();
@@ -125,30 +137,53 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
      */
     private @Nullable PrimaryKey advanceOneRange()
     {
-        PrimaryKey highestKey = getCurrent();
-
         for (KeyRangeIterator range : ranges)
-            if (range.getCurrent().kind() != Kind.STATIC)
+            if (range.peek().kind() != Kind.STATIC)
-                return nextOrNull(range, highestKey);
+            {
+                range.next();
+                return range.hasNext() ? range.peek() : null;
+            }
         
         for (KeyRangeIterator range : ranges)
-            if (range.getCurrent().kind() == Kind.STATIC)
+            if (range.peek().kind() == Kind.STATIC)
-                return nextOrNull(range, highestKey);
+            {
+                range.next();
+                return range.hasNext() ? range.peek() : null;
+            }
 
         throw new IllegalStateException("There should be at least one range to advance!");
     }
 
+    private @Nullable PrimaryKey computeHighestKey()
+    {
+        PrimaryKey max = getMinimum();
+        for (KeyRangeIterator range : ranges)
+        {
+            if (!range.hasNext())
+                return null;
+            if (range.peek().compareToStrict(max) > 0)
+                max = range.peek();
+        }
+        return max;
+    }
+
     @Override
     protected void performSkipTo(PrimaryKey nextKey)
     {
+        // Resist the temptation to call range.hasNext before skipTo: this is a pessimisation, hasNext will invoke
+        // computeNext under the hood, which is an expensive operation to produce a value that we plan to throw away.
+        // Instead, it is the responsibility of the child iterators to make skipTo fast when the iterator is exhausted.
         for (KeyRangeIterator range : ranges)
-            if (range.hasNext())
+            range.skipTo(nextKey);
-                range.skipTo(nextKey);
+
+        // Force recomputing the highest key on the next call to computeNext()
+        highestKey = null;
     }
 
     @Override
     public void close()
     {
+        super.close();
         FileUtils.closeQuietly(ranges);
     }
 
@@ -156,21 +191,26 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
      * Fetches the next available item from the iterator, such that the item is not lower than the given key.
      * If no such items are available, returns null.
      */
-    private PrimaryKey nextOrNull(KeyRangeIterator iterator, PrimaryKey minKey)
+    private PrimaryKey skipAndPeek(KeyRangeIterator iterator, PrimaryKey minKey)
     {
         iterator.skipTo(minKey);
-        return iterator.hasNext() ? iterator.next() : null;
+        return iterator.hasNext() ? iterator.peek() : null;
     }
 
-    public static Builder builder(int size)
+    public static Builder builder(int size, int limit)
     {
-        return new Builder(size);
+        return builder(size, limit, () -> {});
+    }
+
+    public static Builder builder(int size, Runnable onClose)
+    {
+        return new Builder(size, onClose);
     }
 
     @VisibleForTesting
-    public static Builder builder(int size, int limit)
+    public static Builder builder(int size, int limit, Runnable onClose)
     {
-        return new Builder(size, limit);
+        return new Builder(size, limit, onClose);
     }
 
     @VisibleForTesting
@@ -187,14 +227,14 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
 
         protected final List<KeyRangeIterator> rangeIterators;
 
-        Builder(int size)
+        Builder(int size, Runnable onClose)
         {
-            this(size, CassandraRelevantProperties.SAI_INTERSECTION_CLAUSE_LIMIT.getInt());
+            this(size, CassandraRelevantProperties.SAI_INTERSECTION_CLAUSE_LIMIT.getInt(), onClose);
         }
 
-        Builder(int size, int limit)
+        Builder(int size, int limit, Runnable onClose)
         {
-            super(new IntersectionStatistics());
+            super(new IntersectionStatistics(), onClose);
             rangeIterators = new ArrayList<>(size);
             this.limit = limit;
         }
@@ -205,7 +245,7 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
             if (range == null)
                 return this;
 
-            if (range.getCount() > 0)
+            if (range.getMaxKeys() > 0)
                 rangeIterators.add(range);
             else
                 FileUtils.closeQuietly(range);
@@ -224,13 +264,14 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
         @Override
         public void cleanup()
         {
+            super.cleanup();
             FileUtils.closeQuietly(rangeIterators);
         }
 
         @Override
         protected KeyRangeIterator buildIterator()
         {
-            rangeIterators.sort(Comparator.comparingLong(KeyRangeIterator::getCount));
+            rangeIterators.sort(Comparator.comparingLong(KeyRangeIterator::getMaxKeys));
             int initialSize = rangeIterators.size();
             // all ranges will be included
             if (limit >= rangeIterators.size() || limit <= 0)
@@ -248,7 +289,7 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
                 Tracing.trace("Selecting {} {} of {} out of {} indexes",
                               rangeIterators.size(),
                               rangeIterators.size() > 1 ? "indexes with cardinalities" : "index with cardinality",
-                              rangeIterators.stream().map(KeyRangeIterator::getCount).map(Object::toString).collect(Collectors.joining(", ")),
+                              rangeIterators.stream().map(KeyRangeIterator::getMaxKeys).map(Object::toString).collect(Collectors.joining(", ")),
                               initialSize);
 
             return buildIterator(selectiveStatistics, rangeIterators);
@@ -266,18 +307,27 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
             if (isDisjoint)
             {
                 FileUtils.closeQuietly(ranges);
+                onClose.run();
                 return KeyRangeIterator.empty();
             }
 
             if (ranges.size() == 1)
-                return ranges.get(0);
+            {
+                KeyRangeIterator single = ranges.get(0);
+                single.setOnClose(onClose);
+                return single;
+            }
 
             // Make sure intersection is supported on the ranges provided:
             PrimaryKey.Kind firstKind = null;
             
             for (KeyRangeIterator range : ranges)
             {
-                PrimaryKey key = range.getCurrent();
+                PrimaryKey key;
+                if(range.hasNext())
+                    key = range.peek();
+                else
+                    key = range.getMaximum();
 
                 if (key != null)
                     if (firstKind == null)
@@ -286,7 +336,7 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
                         throw new IllegalArgumentException("Cannot intersect " + firstKind + " and " + key.kind() + " ranges!");
             }
 
-            return new KeyRangeIntersectionIterator(statistics, ranges);
+            return new KeyRangeIntersectionIterator(statistics, ranges, onClose);
         }
 
         private void updateStatistics(Statistics statistics, KeyRangeIterator range)
@@ -310,11 +360,11 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
             if (empty)
             {
                 empty = false;
-                count = range.getCount();
+                count = range.getMaxKeys();
             }
             else
             {
-                count = Math.min(count, range.getCount());
+                count = Math.min(count, range.getMaxKeys());
             }
         }
     }
@@ -322,7 +372,7 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
     @VisibleForTesting
     protected static boolean isDisjoint(KeyRangeIterator a, KeyRangeIterator b)
     {
-        return isDisjointInternal(a.getCurrent(), a.getMaximum(), b);
+        return isDisjointInternal(a.peek(), a.getMaximum(), b);
     }
 
     /**
@@ -344,6 +394,6 @@ public class KeyRangeIntersectionIterator extends KeyRangeIterator
      */
     private static boolean isDisjointInternal(PrimaryKey min, PrimaryKey max, KeyRangeIterator b)
     {
-        return min == null || max == null || b.getCount() == 0 || min.compareTo(b.getMaximum()) > 0 || b.getCurrent().compareTo(max) > 0;
+        return min == null || max == null || b.getMaxKeys() == 0 || min.compareTo(b.getMaximum()) > 0 || (b.hasNext() && b.peek().compareTo(max) > 0);
     }
 }

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

@@ -26,40 +26,55 @@ import com.google.common.collect.Iterables;
 import org.apache.cassandra.index.sai.utils.PrimaryKey;
 import org.apache.cassandra.utils.AbstractGuavaIterator;
 
+import javax.annotation.concurrent.NotThreadSafe;
+
 /**
  * An abstract implementation of {@link AbstractGuavaIterator} that supports the building and management of
  * concatanation, union and intersection iterators.
- *
+ * <p>
+ * Range iterators contain primary keys, in sorted order, with no duplicates.  They also
+ * know their minimum and maximum keys, and an upper bound on the number of keys they contain.
+ * <p>
  * Only certain methods are designed to be overriden.  The others are marked private or final.
  */
+@NotThreadSafe
 public abstract class KeyRangeIterator extends AbstractGuavaIterator<PrimaryKey> implements Closeable
 {
     private final PrimaryKey min, max;
     private final long count;
-    private PrimaryKey current;
+    private Runnable onClose;
 
-    protected KeyRangeIterator(Builder.Statistics statistics)
+    protected KeyRangeIterator(Builder.Statistics statistics, Runnable onClose)
     {
-        this(statistics.min, statistics.max, statistics.count);
+        this(statistics.min, statistics.max, statistics.count, onClose);
     }
 
-    public KeyRangeIterator(KeyRangeIterator range)
+    public KeyRangeIterator(KeyRangeIterator range, Runnable onClose)
     {
         this(range == null ? null : range.min,
              range == null ? null : range.max,
-             range == null ? -1 : range.count);
+             range == null ? -1 : range.count,
+             onClose);
     }
 
     public KeyRangeIterator(PrimaryKey min, PrimaryKey max, long count)
+    {
+        this(min, max, count, () -> {});
+    }
+
+    public KeyRangeIterator(PrimaryKey min, PrimaryKey max, long count, Runnable onClose)
     {
         boolean isComplete = min != null && max != null && count != 0;
         boolean isEmpty = min == null && max == null && (count == 0 || count == -1);
         Preconditions.checkArgument(isComplete || isEmpty, "Range: [%s,%s], Count: %d", min, max, count);
 
+        if (isEmpty)
+          endOfData();
+
         this.min = min;
-        this.current = min;
         this.max = max;
         this.count = count;
+        this.onClose = onClose;
     }
 
     public final PrimaryKey getMinimum()
@@ -67,75 +82,63 @@ public abstract class KeyRangeIterator extends AbstractGuavaIterator<PrimaryKey>
         return min;
     }
 
-    public final PrimaryKey getCurrent()
-    {
-        return current;
-    }
-
     public final PrimaryKey getMaximum()
     {
         return max;
     }
 
-    public final long getCount()
+    /**
+     * @return an upper bound on the number of keys that can be returned by this iterator.
+     */
+    public final long getMaxKeys()
     {
         return count;
     }
 
     /**
-     * When called, this iterators current position should
+     * When called, this iterator's current position will
      * be skipped forwards until finding either:
-     *   1) an element equal to or bigger than next
+     *   1) an element equal to or bigger than nextKey
      *   2) the end of the iterator
      *
      * @param nextKey value to skip the iterator forward until matching
-     *
-     * @return The key skipped to, which will be the key returned by the
-     * next call to {@link #next()}, i.e., we are "peeking" at the next key as part of the skip.
      */
-    public final PrimaryKey skipTo(PrimaryKey nextKey)
+    public final void skipTo(PrimaryKey nextKey)
     {
-        if (min == null || max == null)
+        if (state == State.DONE)
-            return endOfData();
+            return;
 
-        // In the case of deferred iterators the current value may not accurately
+        if (state == State.READY && next.compareTo(nextKey) >= 0)
-        // reflect the next value, so we need to check that as well
+            return;
-        if (current.compareTo(nextKey) >= 0)
-        {
-            next = next == null ? recomputeNext() : next;
-            if (next == null)
-                return endOfData();
-            else if (next.compareTo(nextKey) >= 0)
-                return next;
-        }
 
         if (max.compareTo(nextKey) < 0)
-            return endOfData();
+        {
+            endOfData();
+            return;
+        }
 
         performSkipTo(nextKey);
-        return recomputeNext();
+        state = State.NOT_READY;
     }
 
     /**
      * Skip to nextKey.
-     *
+     * <p>
      * That is, implementations should set up the iterator state such that
      * calling computeNext() will return nextKey if present,
      * or the first one after it if not present.
      */
     protected abstract void performSkipTo(PrimaryKey nextKey);
 
-    private PrimaryKey recomputeNext()
+    public void setOnClose(Runnable onClose)
     {
-        return tryToComputeNext() ? peek() : endOfData();
+        this.onClose = onClose;
     }
 
     @Override
-    protected final boolean tryToComputeNext()
+    public void close()
     {
-        boolean hasNext = super.tryToComputeNext();
+        onClose.run();
-        current = hasNext ? next : getMaximum();
-        return hasNext;
     }
 
     public static KeyRangeIterator empty()
@@ -146,7 +149,7 @@ public abstract class KeyRangeIterator extends AbstractGuavaIterator<PrimaryKey>
     private static class EmptyRangeIterator extends KeyRangeIterator
     {
         static final KeyRangeIterator instance = new EmptyRangeIterator();
-        EmptyRangeIterator() { super(null, null, 0); }
+        EmptyRangeIterator() { super(null, null, 0, () -> {}); }
         public PrimaryKey computeNext() { return endOfData(); }
         protected void performSkipTo(PrimaryKey nextKey) { }
         public void close() { }
@@ -156,10 +159,12 @@ public abstract class KeyRangeIterator extends AbstractGuavaIterator<PrimaryKey>
     public static abstract class Builder
     {
         protected final Statistics statistics;
+        protected final Runnable onClose;
 
-        public Builder(Statistics statistics)
+        public Builder(Statistics statistics, Runnable onClose)
         {
             this.statistics = statistics;
+            this.onClose = onClose;
         }
 
         public PrimaryKey getMinimum()
@@ -189,16 +194,24 @@ public abstract class KeyRangeIterator extends AbstractGuavaIterator<PrimaryKey>
         public final KeyRangeIterator build()
         {
             if (rangeCount() == 0)
+            {
+                onClose.run();
                 return empty();
+            }
             else
+            {
                 return buildIterator();
+            }
         }
 
         public abstract Builder add(KeyRangeIterator range);
 
         public abstract int rangeCount();
 
-        public abstract void cleanup();
+        public void cleanup()
+        {
+            onClose.run();
+        }
 
         protected abstract KeyRangeIterator buildIterator();
 

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

@@ -42,7 +42,7 @@ public class KeyRangeOrderingIterator extends KeyRangeIterator
 
     public KeyRangeOrderingIterator(KeyRangeIterator input, int chunkSize, Function<List<PrimaryKey>, KeyRangeIterator> nextRangeFunction)
     {
-        super(input);
+        super(input, () -> {});
         this.input = input;
         this.chunkSize = chunkSize;
         this.nextRangeFunction = nextRangeFunction;

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

@@ -33,9 +33,9 @@ public class KeyRangeUnionIterator extends KeyRangeIterator
     private final List<KeyRangeIterator> ranges;
     private final List<KeyRangeIterator> candidates;
 
-    private KeyRangeUnionIterator(Builder.Statistics statistics, List<KeyRangeIterator> ranges)
+    private KeyRangeUnionIterator(Builder.Statistics statistics, List<KeyRangeIterator> ranges, Runnable onClose)
     {
-        super(statistics);
+        super(statistics, onClose);
         this.ranges = ranges;
         this.candidates = new ArrayList<>(ranges.size());
     }
@@ -43,38 +43,42 @@ public class KeyRangeUnionIterator extends KeyRangeIterator
     @Override
     public PrimaryKey computeNext()
     {
+        // the design is to find the next best value from all the ranges,
+        // and then advance all the ranges that have the same value.
         candidates.clear();
         PrimaryKey candidateKey = null;
         for (KeyRangeIterator range : ranges)
         {
-            if (range.hasNext())
+            if (!range.hasNext())
+                continue;
+
+            if (candidateKey == null)
             {
-                if (candidateKey == null)
+                candidateKey = range.peek();
+                candidates.add(range);
+            }
+            else
+            {
+                PrimaryKey peeked = range.peek();
+    
+                int cmp = candidateKey.compareTo(peeked);
+
+                if (cmp == 0)
                 {
-                    candidateKey = range.peek();
+                    // Replace any existing candidate key if this one is STATIC:
+                    if (peeked.kind() == PrimaryKey.Kind.STATIC)
+                        candidateKey = peeked;
+
                     candidates.add(range);
                 }
-                else
+                else if (cmp > 0)
                 {
-                    PrimaryKey peeked = range.peek();
+                    // we found a new best candidate, throw away the old ones
-    
+                    candidates.clear();
-                    int cmp = candidateKey.compareTo(peeked);
+                    candidateKey = peeked;
-
+                    candidates.add(range);
-                    if (cmp == 0)
-                    {
-                        // Replace any existing candidate key if this one is STATIC:  
-                        if (peeked.kind() == PrimaryKey.Kind.STATIC)
-                            candidateKey = peeked;
-
-                        candidates.add(range);
-                    }
-                    else if (cmp > 0)
-                    {
-                        candidates.clear();
-                        candidateKey = peeked;
-                        candidates.add(range);
-                    }
                 }
+                // else, existing candidate is less than the next in this range
             }
         }
         if (candidates.isEmpty())
@@ -106,18 +110,29 @@ public class KeyRangeUnionIterator extends KeyRangeIterator
     @Override
     public void close()
     {
+        super.close();
+
         // Due to lazy key fetching, we cannot close iterator immediately
         FileUtils.closeQuietly(ranges);
     }
 
     public static Builder builder(int size)
     {
-        return new Builder(size);
+        return builder(size, () -> {});
     }
 
+    public static Builder builder(int size, Runnable onClose)
+    {
+        return new Builder(size, onClose);
+    }
+
+    public static KeyRangeIterator build(List<KeyRangeIterator> keys, Runnable onClose)
+    {
+        return new Builder(keys.size(), onClose).add(keys).build();
+    }
     public static KeyRangeIterator build(List<KeyRangeIterator> keys)
     {
-        return new Builder(keys.size()).add(keys).build();
+        return build(keys, () -> {});
     }
 
     @VisibleForTesting
@@ -125,9 +140,9 @@ public class KeyRangeUnionIterator extends KeyRangeIterator
     {
         protected final List<KeyRangeIterator> rangeIterators;
 
-        Builder(int size)
+        Builder(int size, Runnable onClose)
         {
-            super(new UnionStatistics());
+            super(new UnionStatistics(), onClose);
             this.rangeIterators = new ArrayList<>(size);
         }
 
@@ -137,7 +152,7 @@ public class KeyRangeUnionIterator extends KeyRangeIterator
             if (range == null)
                 return this;
 
-            if (range.getCount() > 0)
+            if (range.getMaxKeys() > 0)
             {
                 rangeIterators.add(range);
                 statistics.update(range);
@@ -159,6 +174,7 @@ public class KeyRangeUnionIterator extends KeyRangeIterator
         @Override
         public void cleanup()
         {
+            super.cleanup();
             FileUtils.closeQuietly(rangeIterators);
         }
 
@@ -166,9 +182,13 @@ public class KeyRangeUnionIterator extends KeyRangeIterator
         protected KeyRangeIterator buildIterator()
         {
             if (rangeCount() == 1)
-                return rangeIterators.get(0);
+            {
+                KeyRangeIterator single = rangeIterators.get(0);
+                single.setOnClose(onClose);
+                return single;
+            }
 
-            return new KeyRangeUnionIterator(statistics, rangeIterators);
+            return new KeyRangeUnionIterator(statistics, rangeIterators, onClose);
         }
     }
 
@@ -179,7 +199,7 @@ public class KeyRangeUnionIterator extends KeyRangeIterator
         {
             min = nullSafeMin(min, range.getMinimum());
             max = nullSafeMax(max, range.getMaximum());
-            count += range.getCount();
+            count += range.getMaxKeys();
         }
     }
 }

src/java/org/apache/cassandra/index/sai/memory/InMemoryKeyRangeIterator.java

@@ -37,7 +37,7 @@ public class InMemoryKeyRangeIterator extends KeyRangeIterator
      */
     public InMemoryKeyRangeIterator(SortedSet<PrimaryKey> keys)
     {
-        super(keys.first(), keys.last(), keys.size());
+        super(keys.first(), keys.last(), keys.size(), () -> {});
         this.keys = new PriorityQueue<>(keys);
         this.uniqueKeys = true;
     }
@@ -48,7 +48,7 @@ public class InMemoryKeyRangeIterator extends KeyRangeIterator
      */
     public InMemoryKeyRangeIterator(PrimaryKey min, PrimaryKey max, PriorityQueue<PrimaryKey> keys)
     {
-        super(min, max, keys.size());
+        super(min, max, keys.size(), () -> {});
         this.keys = keys;
         this.uniqueKeys = false;
     }

src/java/org/apache/cassandra/index/sai/plan/QueryController.java

@@ -175,11 +175,11 @@ public class QueryController
      * the {@link SSTableIndex}s that will satisfy the expression.
      * <p>
      * Each (expression, SSTable indexes) pair is then passed to
-     * {@link IndexSearchResultIterator#build(Expression, Collection, AbstractBounds, QueryContext, boolean)}
+     * {@link IndexSearchResultIterator#build(Expression, Collection, AbstractBounds, QueryContext, boolean, Runnable)}
      * to search the in-memory index associated with the expression and the SSTable indexes, the results of
      * which are unioned and returned.
      * <p>
-     * The results from each call to {@link IndexSearchResultIterator#build(Expression, Collection, AbstractBounds, QueryContext, boolean)}
+     * The results from each call to {@link IndexSearchResultIterator#build(Expression, Collection, AbstractBounds, QueryContext, boolean, Runnable)}
      * are added to a {@link KeyRangeIntersectionIterator} and returned if strict filtering is allowed.
      * <p>
      * If strict filtering is not allowed, indexes are split into two groups according to the repaired status of their 
@@ -195,10 +195,11 @@ public class QueryController
         // VSTODO move ANN out of expressions and into its own abstraction? That will help get generic ORDER BY support
         expressions = expressions.stream().filter(e -> e.getIndexOperator() != Expression.IndexOperator.ANN).collect(Collectors.toList());
 
-        KeyRangeIterator.Builder builder = command.rowFilter().isStrict()
-                                           ? KeyRangeIntersectionIterator.builder(expressions.size())
-                                           : KeyRangeUnionIterator.builder(expressions.size());
         QueryViewBuilder.QueryView queryView = new QueryViewBuilder(expressions, mergeRange).build();
+        Runnable onClose = () -> queryView.referencedIndexes.forEach(SSTableIndex::releaseQuietly);
+        KeyRangeIterator.Builder builder = command.rowFilter().isStrict()
+                                           ? KeyRangeIntersectionIterator.builder(expressions.size(), onClose)
+                                           : KeyRangeUnionIterator.builder(expressions.size(), onClose);
 
         try
         {
@@ -210,11 +211,11 @@ public class QueryController
                 // This usually means we are making this local index query in the context of a user query that reads 
                 // from a single replica and thus can safely perform local intersections.
                 for (Pair<Expression, Collection<SSTableIndex>> queryViewPair : queryView.view)
-                    builder.add(IndexSearchResultIterator.build(queryViewPair.left, queryViewPair.right, mergeRange, queryContext, true));
+                    builder.add(IndexSearchResultIterator.build(queryViewPair.left, queryViewPair.right, mergeRange, queryContext, true, () -> {}));
             }
             else
             {
-                KeyRangeIterator.Builder repairedBuilder = KeyRangeIntersectionIterator.builder(expressions.size());
+                KeyRangeIterator.Builder repairedBuilder = KeyRangeIntersectionIterator.builder(expressions.size(), () -> {});
 
                 for (Pair<Expression, Collection<SSTableIndex>> queryViewPair : queryView.view)
                 {
@@ -232,10 +233,10 @@ public class QueryController
 
                     // Always build an iterator for the un-repaired set, given this must include Memtable indexes...  
                     IndexSearchResultIterator unrepairedIterator =
-                            IndexSearchResultIterator.build(queryViewPair.left, unrepaired, mergeRange, queryContext, true);
+                            IndexSearchResultIterator.build(queryViewPair.left, unrepaired, mergeRange, queryContext, true, () -> {});
 
                     // ...but ignore it if our combined results are empty.
-                    if (unrepairedIterator.getCount() > 0)
+                    if (unrepairedIterator.getMaxKeys() > 0)
                     {
                         builder.add(unrepairedIterator);
                         queryContext.hasUnrepairedMatches = true;
@@ -248,7 +249,7 @@ public class QueryController
 
                     // ...then only add an iterator to the repaired intersection if repaired SSTable indexes exist. 
                     if (!repaired.isEmpty())
-                        repairedBuilder.add(IndexSearchResultIterator.build(queryViewPair.left, repaired, mergeRange, queryContext, false));
+                        repairedBuilder.add(IndexSearchResultIterator.build(queryViewPair.left, repaired, mergeRange, queryContext, false, () -> {}));
                 }
 
                 if (repairedBuilder.rangeCount() > 0)
@@ -259,7 +260,6 @@ public class QueryController
         {
             // all sstable indexes in view have been referenced, need to clean up when exception is thrown
             builder.cleanup();
-            queryView.referencedIndexes.forEach(SSTableIndex::releaseQuietly);
             throw t;
         }
         return builder;
@@ -315,6 +315,7 @@ public class QueryController
         KeyRangeIterator memtableResults = index.memtableIndexManager().searchMemtableIndexes(queryContext, planExpression, mergeRange);
 
         QueryViewBuilder.QueryView queryView = new QueryViewBuilder(Collections.singleton(planExpression), mergeRange).build();
+        Runnable onClose = () -> queryView.referencedIndexes.forEach(SSTableIndex::releaseQuietly);
 
         try
         {
@@ -322,12 +323,13 @@ public class QueryController
                                                                    .stream()
                                                                    .map(this::createRowIdIterator)
                                                                    .collect(Collectors.toList());
-            return IndexSearchResultIterator.build(sstableIntersections, memtableResults, queryView.referencedIndexes, queryContext);
+
+            return IndexSearchResultIterator.build(sstableIntersections, memtableResults, queryView.referencedIndexes, queryContext, onClose);
         }
         catch (Throwable t)
         {
             // all sstable indexes in view have been referenced, need to clean up when exception is thrown
-            queryView.referencedIndexes.forEach(SSTableIndex::release);
+            onClose.run();
             throw t;
         }
     }
@@ -353,6 +355,7 @@ public class QueryController
         // search memtable before referencing sstable indexes; otherwise we may miss newly flushed memtable index
         KeyRangeIterator memtableResults = index.memtableIndexManager().limitToTopResults(queryContext, sourceKeys, planExpression);
         QueryViewBuilder.QueryView queryView = new QueryViewBuilder(Collections.singleton(planExpression), mergeRange).build();
+        Runnable onClose = () -> queryView.referencedIndexes.forEach(SSTableIndex::releaseQuietly);
 
         try
         {
@@ -371,12 +374,12 @@ public class QueryController
                                                                    })
                                                                    .collect(Collectors.toList());
 
-            return IndexSearchResultIterator.build(sstableIntersections, memtableResults, queryView.referencedIndexes, queryContext);
+            return IndexSearchResultIterator.build(sstableIntersections, memtableResults, queryView.referencedIndexes, queryContext, onClose);
         }
         catch (Throwable t)
         {
             // all sstable indexes in view have been referenced, need to clean up when exception is thrown
-            queryView.referencedIndexes.forEach(SSTableIndex::release);
+            onClose.run();
             throw t;
         }
     }

src/java/org/apache/cassandra/utils/AbstractGuavaIterator.java

@@ -24,6 +24,8 @@ import java.util.NoSuchElementException;
 
 import com.google.common.collect.PeekingIterator;
 
+import javax.annotation.concurrent.NotThreadSafe;
+
 import static com.google.common.base.Preconditions.checkState;
 
 /**
@@ -31,9 +33,10 @@ import static com.google.common.base.Preconditions.checkState;
  * is that the next variable is now protected so that the KeyRangeIterator.skipTo
  * method can avoid early state changed.
  */
+@NotThreadSafe
 public abstract class AbstractGuavaIterator<T> implements PeekingIterator<T>
 {
-    private State state = State.NOT_READY;
+    protected State state = State.NOT_READY;
 
     /** Constructor for use by subclasses. */
     protected AbstractGuavaIterator() {}

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

@@ -37,13 +37,14 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 public class SingleNodeQueryFailureTest extends SAITester
 {
-    private static final String CREATE_TABLE_TEMPLATE = "CREATE TABLE %s (id text PRIMARY KEY, v1 text) WITH " +
+    private static final String CREATE_TABLE_TEMPLATE = "CREATE TABLE %s (id text PRIMARY KEY, v1 int, v2 text) WITH " +
                                                         "compaction = {'class' : 'SizeTieredCompactionStrategy', 'enabled' : false }";
 
     @Before
     public void setup()
     {
         requireNetwork();
+        setupTableAndIndexes();
     }
 
     @After
@@ -52,47 +53,57 @@ public class SingleNodeQueryFailureTest extends SAITester
         Injections.deleteAll();
     }
 
+    // Single Index Tests
+    @Test
+    public void testFailedRangeIteratorOnSingleIndexQuery() throws Throwable
+    {
+        testFailedQuery("range_iterator_single", PostingListRangeIterator.class, "getNextRowId", true);
+    }
+
+    @Test
+    public void testFailedTermsReaderOnSingleIndexQuery() throws Throwable
+    {
+        testFailedQuery("terms_reader_single", LiteralIndexSegmentTermsReader.TermQuery.class, "lookupPostingsOffset", true);
+    }
+
+    // Multi Index Tests
     @Test
     public void testFailedRangeIteratorOnMultiIndexesQuery() throws Throwable
     {
-        testFailedMultiIndexesQuery("range_iterator", PostingListRangeIterator.class, "getNextRowId");
+        testFailedQuery("range_iterator_multi", PostingListRangeIterator.class, "getNextRowId", false);
     }
 
     @Test
     public void testFailedTermsReaderOnMultiIndexesQuery() throws Throwable
     {
-        testFailedMultiIndexesQuery("terms_reader", LiteralIndexSegmentTermsReader.TermQuery.class, "lookupPostingsOffset");
+        testFailedQuery("terms_reader_multi", LiteralIndexSegmentTermsReader.TermQuery.class, "lookupPostingsOffset", false);
     }
 
-    private void testFailedMultiIndexesQuery(String name, Class<?> targetClass, String targetMethod) throws Throwable
+    private void setupTableAndIndexes()
     {
-        Injection injection = Injections.newCustom(name)
-                                        .add(newInvokePoint().onClass(targetClass).onMethod(targetMethod))
-                                        .add(newActionBuilder().actions().doThrow(RuntimeException.class, quote("Injected failure!")))
-                                        .build();
-
         createTable(CREATE_TABLE_TEMPLATE);
         createIndex(String.format(CREATE_INDEX_TEMPLATE, "v1"));
+        createIndex(String.format(CREATE_INDEX_TEMPLATE, "v2"));
 
-        execute("INSERT INTO %s (id, v1) VALUES ('1', '0')");
+        execute("INSERT INTO %s (id, v1, v2) VALUES ('1', 0, '0')");
         flush();
-        execute("INSERT INTO %s (id, v1) VALUES ('2', '1')");
+        execute("INSERT INTO %s (id, v1, v2) VALUES ('2', 1, '1')");
         flush();
-        execute("INSERT INTO %s (id, v1) VALUES ('3', '2')");
+        execute("INSERT INTO %s (id, v1, v2) VALUES ('3', 2, '2')");
         flush();
+    }
+
+    private void testFailedQuery(String name, Class<?> targetClass, String targetMethod, boolean isSingleIndexTest) throws Throwable
+    {
+        Injection injection = Injections.newCustom(name)
+                .add(newInvokePoint().onClass(targetClass).onMethod(targetMethod))
+                .add(newActionBuilder().actions().doThrow(RuntimeException.class, quote("Injected failure!")))
+                .build();
 
         try
         {
             Injections.inject(injection);
-
+            performTestQueries(isSingleIndexTest);
-            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v1 = '0'"))
-                    .isInstanceOf(ReadFailureException.class);
-
-            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v1 = '1'"))
-                    .isInstanceOf(ReadFailureException.class);
-
-            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v1 = '2'"))
-                    .isInstanceOf(ReadFailureException.class);
         }
         catch (Exception e)
         {
@@ -103,8 +114,49 @@ public class SingleNodeQueryFailureTest extends SAITester
             injection.disable();
         }
 
-        Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v1 = '0'").all().size());
+        verifyResults(isSingleIndexTest);
-        Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v1 = '1'").all().size());
-        Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v1 = '2'").all().size());
     }
-}
+
+    private void performTestQueries(boolean isSingleIndexTest)
+    {
+        if (isSingleIndexTest)
+        {
+            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v2 = '0'"))
+                    .isInstanceOf(ReadFailureException.class);
+
+            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v2 = '1'"))
+                    .isInstanceOf(ReadFailureException.class);
+
+            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v2 = '2'"))
+                    .isInstanceOf(ReadFailureException.class);
+        }
+        else
+        {
+            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v1 < 1 and v2 = '0'"))
+                    .isInstanceOf(ReadFailureException.class);
+
+            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v1 >= 1 and v2 = '1'"))
+                    .isInstanceOf(ReadFailureException.class);
+
+            assertThatThrownBy(() -> executeNet("SELECT id FROM %s WHERE v1 >= 2 and v2 = '2'"))
+                    .isInstanceOf(ReadFailureException.class);
+        }
+    }
+
+    private void verifyResults(boolean isSingleIndexTest)
+    {
+        if (isSingleIndexTest)
+        {
+            Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v2 = '0'").all().size());
+            Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v2 = '1'").all().size());
+            Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v2 = '2'").all().size());
+        }
+        else
+        {
+            Assert.assertEquals(3, executeNet("SELECT id FROM %s WHERE v1 >= 0").all().size());
+            Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v2 = '0'").all().size());
+            Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v2 = '1'").all().size());
+            Assert.assertEquals(1, executeNet("SELECT id FROM %s WHERE v2 = '2'").all().size());
+        }
+    }
+}

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

@@ -152,7 +152,6 @@ public class BalancedTreeIndexSearcherTest extends SAIRandomizedTester
                                                                        .add(Operator.EQ, rawType.decompose(rawValueProducer.apply(EQ_TEST_LOWER_BOUND_INCLUSIVE)))
                                                              , null, mock(QueryContext.class)))
         {
-            assertEquals(results.getMinimum(), results.getCurrent());
             assertTrue(results.hasNext());
 
             assertEquals(0L, results.next().token().getLongValue());
@@ -186,7 +185,6 @@ public class BalancedTreeIndexSearcherTest extends SAIRandomizedTester
                                                                        .add(Operator.LTE, rawType.decompose(rawValueProducer.apply((short)7))),
                                                              null, mock(QueryContext.class)))
         {
-            assertEquals(results.getMinimum(), results.getCurrent());
             assertTrue(results.hasNext());
 
             List<Long> actualTokenList = Lists.newArrayList(Iterators.transform(results, key -> key.token().getLongValue()));

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

@@ -114,7 +114,6 @@ public class InvertedIndexSearcherTest extends SAIRandomizedTester
             {
                 try (KeyRangeIterator results = searcher.search(Expression.create(index).add(Operator.EQ, wrap(termsEnum.get(t).left)), null, context))
                 {
-                    assertEquals(results.getMinimum(), results.getCurrent());
                     assertTrue(results.hasNext());
 
                     for (int p = 0; p < numPostings; ++p)
@@ -129,7 +128,6 @@ public class InvertedIndexSearcherTest extends SAIRandomizedTester
 
                 try (KeyRangeIterator results = searcher.search(Expression.create(index).add(Operator.EQ, wrap(termsEnum.get(t).left)), null, context))
                 {
-                    assertEquals(results.getMinimum(), results.getCurrent());
                     assertTrue(results.hasNext());
 
                     // test skipping to the last block

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

@@ -18,11 +18,14 @@
 package org.apache.cassandra.index.sai.iterators;
 
 import java.util.Arrays;
+import java.util.Set;
 import java.util.function.BiFunction;
+import java.util.stream.Collectors;
+
+import org.junit.Assert;
 
 import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
-
+import org.apache.cassandra.utils.Pair;
-import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 public class AbstractKeyRangeIteratorTester extends SAIRandomizedTester
 {
@@ -36,11 +39,6 @@ public class AbstractKeyRangeIteratorTester extends SAIRandomizedTester
         return Arrays.stream(intArray).mapToLong(i -> i).toArray();
     }
 
-    void assertOnError(KeyRangeIterator range)
-    {
-        assertThatThrownBy(() -> LongIterator.convert(range)).isInstanceOf(RuntimeException.class);
-    }
-
     final KeyRangeIterator buildIntersection(KeyRangeIterator... ranges)
     {
         return KeyRangeIntersectionIterator.builder(16, Integer.MAX_VALUE).add(Arrays.asList(ranges)).build();
@@ -66,62 +64,77 @@ public class AbstractKeyRangeIteratorTester extends SAIRandomizedTester
         return KeyRangeUnionIterator.builder(ranges.length).add(Arrays.asList(ranges)).build();
     }
 
-    final KeyRangeIterator buildUnion(long[]... ranges)
+    static KeyRangeIterator buildConcat(KeyRangeIterator... ranges)
-    {
-        return buildUnion(toRangeIterator(ranges));
-    }
-
-    final KeyRangeIterator buildConcat(KeyRangeIterator... ranges)
     {
         return KeyRangeConcatIterator.builder(ranges.length).add(Arrays.asList(ranges)).build();
     }
-
+    private static KeyRangeIterator[] toRangeIterator(long[]... ranges)
-    final KeyRangeIterator buildConcat(long[]... ranges)
     {
-        return buildConcat(toRangeIterator(ranges));
+        return Arrays.stream(ranges).map(AbstractKeyRangeIteratorTester::build).toArray(KeyRangeIterator[]::new);
     }
 
-    private KeyRangeIterator[] toRangeIterator(long[]... ranges)
+    protected static LongIterator build(long... tokens)
     {
-        return Arrays.stream(ranges).map(this::build).toArray(KeyRangeIterator[]::new);
+        return new LongIterator(tokens);
-    }
-
-    protected LongIterator build(long... tokens)
-    {
-        return build(tokens, false);
-    }
-
-    protected LongIterator build(long[] tokensA, boolean onErrorA)
-    {
-        LongIterator rangeA = new LongIterator(tokensA);
-
-        if (onErrorA)
-            rangeA.throwsException();
-
-        return rangeA;
-    }
-
-    protected KeyRangeIterator buildOnError(BiFunction<KeyRangeIterator, KeyRangeIterator, KeyRangeIterator> builder, long[] tokensA, long[] tokensB)
-    {
-        return build(builder, tokensA, true, tokensB, true);
-    }
-
-    protected KeyRangeIterator buildOnErrorA(BiFunction<KeyRangeIterator, KeyRangeIterator, KeyRangeIterator> builder, long[] tokensA, long[] tokensB)
-    {
-        return build(builder, tokensA, true, tokensB, false);
-    }
-
-    protected KeyRangeIterator buildOnErrorB(BiFunction<KeyRangeIterator, KeyRangeIterator, KeyRangeIterator> builder, long[] tokensA, long[] tokensB)
-    {
-        return build(builder, tokensA, false, tokensB, true);
     }
 
     protected KeyRangeIterator build(BiFunction<KeyRangeIterator, KeyRangeIterator, KeyRangeIterator> builder,
                                      long[] tokensA,
-                                     boolean onErrorA,
+                                     long[] tokensB)
-                                     long[] tokensB,
-                                     boolean onErrorB)
     {
-        return builder.apply(build(tokensA, onErrorA), build(tokensB, onErrorB));
+        return builder.apply(build(tokensA), build(tokensB));
+    }
+
+    static void validateWithSkipping(KeyRangeIterator ri, long[] totalOrdering)
+    {
+        int count = 0;
+        while (ri.hasNext())
+        {
+            // make sure hasNext plays nice with skipTo
+            if (nextBoolean())
+                ri.hasNext();
+
+            // skipping to the same element should also be a no-op
+            if (nextBoolean())
+                ri.skipTo(LongIterator.fromToken(totalOrdering[count]));
+
+            // skip a few elements
+            if (nextDouble() < 0.1)
+            {
+                int n = nextInt(1, 3);
+                if (count + n < totalOrdering.length)
+                {
+                    count += n;
+                    ri.skipTo(LongIterator.fromToken(totalOrdering[count]));
+                }
+            }
+            Assert.assertEquals(totalOrdering[count++], ri.next().token().getLongValue());
+        }
+        Assert.assertEquals(totalOrdering.length, count);
+    }
+
+    static Set<Long> toSet(long[] tokens)
+    {
+        return Arrays.stream(tokens).boxed().collect(Collectors.toSet());
+    }
+
+    /**
+     * @return a random {Concat,Intersection, Union} iterator, and a long[] of the elements in the iterator.
+     *         elements will range from 0..1024.
+     */
+    static Pair<KeyRangeIterator, long[]> createRandomIterator()
+    {
+        var n = between(0, 3);
+        switch (n)
+        {
+            case 0:
+                return KeyRangeConcatIteratorTest.createRandom();
+            case 1:
+                return KeyRangeIntersectionIteratorTest.createRandom(nextInt(1, 16));
+            case 2:
+                return KeyRangeUnionIteratorTest.createRandom(nextInt(1, 16));
+            default:
+                throw new AssertionError();
+        }
     }
 }

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

@@ -17,6 +17,7 @@
  */
 package org.apache.cassandra.index.sai.iterators;
 
+import java.util.ArrayList;
 import java.util.function.Supplier;
 import java.util.stream.IntStream;
 
@@ -24,6 +25,7 @@ import org.junit.Test;
 
 import org.apache.cassandra.dht.Murmur3Partitioner;
 import org.apache.cassandra.index.sai.utils.PrimaryKey;
+import org.apache.cassandra.utils.Pair;
 
 import static org.apache.cassandra.index.sai.iterators.LongIterator.convert;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -109,7 +111,7 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
         assertNotNull(keyIterator);
         assertEquals(1L, keyIterator.getMinimum().token().getLongValue());
         assertEquals(9L, keyIterator.getMaximum().token().getLongValue());
-        assertEquals(9L, keyIterator.getCount());
+        assertEquals(9L, keyIterator.getMaxKeys());
 
         for (long i = 1; i < 10; i++)
         {
@@ -196,7 +198,7 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, keyIterator.getMinimum().token().getLongValue());
         assertEquals(19L, keyIterator.getMaximum().token().getLongValue());
         assertTrue(keyIterator.hasNext());
-        assertEquals(10, keyIterator.getCount());
+        assertEquals(10, keyIterator.getMaxKeys());
     }
 
     @Test
@@ -211,7 +213,7 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, keyIterator.getMinimum().token().getLongValue());
         assertEquals(10L, keyIterator.getMaximum().token().getLongValue());
         assertTrue(keyIterator.hasNext());
-        assertEquals(1, keyIterator.getCount());
+        assertEquals(1, keyIterator.getMaxKeys());
     }
 
     @Test
@@ -226,7 +228,7 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, keyIterator.getMinimum().token().getLongValue());
         assertEquals(19L, keyIterator.getMaximum().token().getLongValue());
         assertTrue(keyIterator.hasNext());
-        assertEquals(10, keyIterator.getCount());
+        assertEquals(10, keyIterator.getMaxKeys());
     }
 
     @Test
@@ -241,7 +243,7 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, keyIterator.getMinimum().token().getLongValue());
         assertEquals(10L, keyIterator.getMaximum().token().getLongValue());
         assertTrue(keyIterator.hasNext());
-        assertEquals(1, keyIterator.getCount());
+        assertEquals(1, keyIterator.getMaxKeys());
     }
 
     @Test
@@ -257,7 +259,7 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, keyIterator.getMinimum().token().getLongValue());
         assertEquals(19L, keyIterator.getMaximum().token().getLongValue());
         assertTrue(keyIterator.hasNext());
-        assertEquals(10, keyIterator.getCount());
+        assertEquals(10, keyIterator.getMaxKeys());
     }
 
     @Test
@@ -273,7 +275,7 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, keyIterator.getMinimum().token().getLongValue());
         assertEquals(19L, keyIterator.getMaximum().token().getLongValue());
         assertTrue(keyIterator.hasNext());
-        assertEquals(10, keyIterator.getCount());
+        assertEquals(10, keyIterator.getMaxKeys());
     }
 
     @Test
@@ -318,37 +320,6 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(convert(1L, 3L, 5L, 7L, 9L), convert(buildIntersection(concatA, concatB)));
     }
 
-    @Test
-    public void testConcatOnError()
-    {
-        assertOnError(buildOnErrorA(this::buildConcat, arr(1L, 2L, 3L), arr(4L, 5L, 6L)));
-        assertOnError(buildOnErrorB(this::buildConcat, arr( 1L, 2L, 3L), arr(4L)));
-    }
-
-    @Test
-    public void testConcatOfUnionsOnError()
-    {
-        KeyRangeIterator unionA = buildUnion(arr(1L, 2L, 3L), arr(4L));
-        KeyRangeIterator unionB = buildOnErrorB(this::buildUnion, arr(6L), arr(8L, 9L));
-        assertOnError(buildConcat(unionA, unionB));
-
-        unionA = buildOnErrorA(this::buildUnion, arr( 1L, 2L, 3L), arr( 4L));
-        unionB = buildUnion(arr( 5L), arr( 5L, 6L));
-        assertOnError(buildConcat(unionA, unionB));
-    }
-
-    @Test
-    public void testConcatOfIntersectionsOnError()
-    {
-        KeyRangeIterator intersectionA = buildOnErrorA(this::buildIntersection, arr(1L, 2L, 3L), arr(2L, 3L, 4L));
-        KeyRangeIterator intersectionB = buildIntersection(arr(6L, 7L, 8L), arr(7L, 8L, 9L));
-        assertOnError(buildConcat(intersectionA, intersectionB));
-
-        intersectionA = buildIntersection(arr( 1L, 2L, 3L), arr( 2L, 3L, 4L));
-        intersectionB = buildOnErrorB(this::buildIntersection, arr( 6L, 7L, 8L, 9L, 10L), arr(  7L, 8L, 9L));
-        assertOnError(buildConcat(intersectionA, intersectionB));
-    }
-
     @Test
     public void testDuplicatedElementsInTheSameIterator()
     {
@@ -424,4 +395,40 @@ public class KeyRangeConcatIteratorTest extends AbstractKeyRangeIteratorTester
                              primaryKeyFactory.create(new Murmur3Partitioner.LongToken(max)),
                              primaryKeyFactory.create(new Murmur3Partitioner.LongToken(min)));
     }
+
+    @Test
+    public void testRandom()
+    {
+        for (int testIteration = 0; testIteration < 16; testIteration++)
+        {
+            var p = createRandom();
+            validateWithSkipping(p.left, p.right);
+        }
+    }
+
+    static Pair<KeyRangeIterator, long[]> createRandom()
+    {
+        var ranges = new ArrayList<KeyRangeIterator>();
+        var current = new ArrayList<Long>();
+        var allValues = new ArrayList<Long>();
+        int maxValue = 1024;
+        for (int i = 0; i < maxValue; i++)
+        {
+            allValues.add((long) i);
+            current.add((long) i);
+            if (nextDouble() < 0.05)
+            {
+                ranges.add(build(current.stream().mapToLong(Long::longValue).toArray()));
+                current.clear();
+            }
+            if (nextDouble() < 0.1)
+                i += nextInt(5);
+        }
+        ranges.add(build(current.stream().mapToLong(Long::longValue).toArray()));
+
+        long[] totalOrdered = allValues.stream().mapToLong(Long::longValue).toArray();
+        KeyRangeIterator it = buildConcat(ranges.toArray(KeyRangeIterator[]::new));
+        assertEquals(totalOrdered.length, it.getMaxKeys());
+        return Pair.create(it, totalOrdered);
+    }
 }

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

@@ -18,15 +18,18 @@
 package org.apache.cassandra.index.sai.iterators;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Set;
+import java.util.stream.IntStream;
 
+import org.junit.Assert;
 import org.junit.Test;
 
 import com.carrotsearch.hppc.LongHashSet;
 import com.carrotsearch.hppc.LongSet;
 import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.utils.Pair;
 
 import static org.apache.cassandra.index.sai.iterators.LongIterator.convert;
 import static org.junit.Assert.assertEquals;
@@ -115,23 +118,23 @@ public class KeyRangeIntersectionIteratorTest extends AbstractKeyRangeIteratorTe
         assertNotNull(range);
 
         // first let's skipTo something before range
-        assertEquals(4L, range.skipTo(LongIterator.fromToken(3L)).token().getLongValue());
+        range.skipTo(LongIterator.fromToken(3L));
-        assertEquals(4L, range.getCurrent().token().getLongValue());
+        Assert.assertEquals(4L, range.peek().token().getLongValue());
 
         // now let's skip right to the send value
-        assertEquals(6L, range.skipTo(LongIterator.fromToken(5L)).token().getLongValue());
+        range.skipTo(LongIterator.fromToken(5L));
-        assertEquals(6L, range.getCurrent().token().getLongValue());
+        Assert.assertEquals(6L, range.peek().token().getLongValue());
 
         // now right to the element
-        assertEquals(7L, range.skipTo(LongIterator.fromToken(7L)).token().getLongValue());
+        range.skipTo(LongIterator.fromToken(7L));
-        assertEquals(7L, range.getCurrent().token().getLongValue());
+        Assert.assertEquals(7L, range.peek().token().getLongValue());
         assertEquals(7L, range.next().token().getLongValue());
 
         assertTrue(range.hasNext());
-        assertEquals(10L, range.getCurrent().token().getLongValue());
+        Assert.assertEquals(10L, range.peek().token().getLongValue());
 
         // now right after the last element
-        assertNull(range.skipTo(LongIterator.fromToken(11L)));
+        range.skipTo(LongIterator.fromToken(11L));
         assertFalse(range.hasNext());
     }
 
@@ -152,7 +155,7 @@ public class KeyRangeIntersectionIteratorTest extends AbstractKeyRangeIteratorTe
         assertNotNull(tokens);
         assertEquals(7L, tokens.getMinimum().token().getLongValue());
         assertEquals(9L, tokens.getMaximum().token().getLongValue());
-        assertEquals(3L, tokens.getCount());
+        assertEquals(3L, tokens.getMaxKeys());
 
         assertEquals(convert(9L), convert(builder.build()));
     }
@@ -198,7 +201,7 @@ public class KeyRangeIntersectionIteratorTest extends AbstractKeyRangeIteratorTe
         FileUtils.closeQuietly(tokens);
 
         KeyRangeIterator emptyTokens = KeyRangeIntersectionIterator.builder(16, Integer.MAX_VALUE).build();
-        assertEquals(0, emptyTokens.getCount());
+        assertEquals(0, emptyTokens.getMaxKeys());
 
         builder = KeyRangeIntersectionIterator.builder(16, Integer.MAX_VALUE);
         assertEquals(0L, builder.add((KeyRangeIterator) null).rangeCount());
@@ -218,7 +221,7 @@ public class KeyRangeIntersectionIteratorTest extends AbstractKeyRangeIteratorTe
         builder.add(single);
         assertEquals(1L, builder.rangeCount());
         range = builder.build();
-        assertEquals(0, range.getCount());
+        assertEquals(0, range.getMaxKeys());
 
         // disjoint case
         builder = KeyRangeIntersectionIterator.builder(16, Integer.MAX_VALUE);
@@ -298,7 +301,7 @@ public class KeyRangeIntersectionIteratorTest extends AbstractKeyRangeIteratorTe
         assertNull(range.getMinimum());
         assertNull(range.getMaximum());
         assertFalse(range.hasNext());
-        assertEquals(0, range.getCount());
+        assertEquals(0, range.getMaxKeys());
     }
 
     @Test
@@ -343,48 +346,38 @@ public class KeyRangeIntersectionIteratorTest extends AbstractKeyRangeIteratorTe
     {
         for (int attempt = 0; attempt < 16; attempt++)
         {
-            final int maxRanges = nextInt(2, 16);
+            var p = createRandom(nextInt(2, 16));
-
+            validateWithSkipping(p.left, p.right);
-            // generate randomize ranges
-            long[][] ranges = new long[maxRanges][];
-            for (int i = 0; i < ranges.length; i++)
-            {
-                int rangeSize = nextInt(16, 512);
-                LongSet range = new LongHashSet(rangeSize);
-
-                for (int j = 0; j < rangeSize; j++)
-                    range.add(nextLong(0, 100));
-
-                ranges[i] = range.toArray();
-                Arrays.sort(ranges[i]);
-            }
-
-            List<Long> expected = new ArrayList<>();
-            // determine unique tokens which intersect every range
-            for (long token : ranges[0])
-            {
-                boolean intersectsAll = true;
-                for (int i = 1; i < ranges.length; i++)
-                {
-                    if (Arrays.binarySearch(ranges[i], token) < 0)
-                    {
-                        intersectsAll = false;
-                        break;
-                    }
-                }
-
-                if (intersectsAll)
-                    expected.add(token);
-            }
-
-            KeyRangeIterator.Builder builder = KeyRangeIntersectionIterator.builder(16, Integer.MAX_VALUE);
-            for (long[] range : ranges)
-                builder.add(new LongIterator(range));
-
-            assertEquals(expected, convert(builder.build()));
         }
     }
 
+    /**
+     * @return a long[][] of random elements, and a long[] of the intersection of those elements
+     */
+    static Pair<KeyRangeIterator, long[]> createRandom(int nRanges)
+    {
+        // generate randomized ranges
+        long[][] ranges = new long[nRanges][];
+        for (int i = 0; i < ranges.length; i++)
+        {
+            int rangeSize = nextInt(16, 512);
+            LongSet range = new LongHashSet(rangeSize);
+
+            for (int j = 0; j < rangeSize; j++)
+                range.add(nextInt(1024));
+
+            ranges[i] = range.toArray();
+            Arrays.sort(ranges[i]);
+        }
+        var builder = KeyRangeIntersectionIterator.builder(16, Integer.MAX_VALUE);
+        for (long[] range : ranges)
+            builder.add(new LongIterator(range));
+
+        Set<Long> expectedSet = toSet(ranges[0]);
+        IntStream.range(1, ranges.length).forEach(i -> expectedSet.retainAll(toSet(ranges[i])));
+        return Pair.create(builder.build(), expectedSet.stream().mapToLong(Long::longValue).sorted().toArray());
+    }
+
     // SAI specific tests
     @Test
     public void testSelectiveIntersection()
@@ -403,34 +396,4 @@ public class KeyRangeIntersectionIteratorTest extends AbstractKeyRangeIteratorTe
 
         assertEquals(convert(2L, 4L, 6L), convert(intersection));
     }
-
-    @Test
-    public void testIntersectionOfUnionsOnError()
-    {
-        // intersection of two unions
-        KeyRangeIterator unionA = buildOnErrorB(this::buildUnion, arr(1L, 2L, 3L), arr(5L, 6L, 7L));
-        KeyRangeIterator unionB = buildUnion(arr(2L, 4L, 6L), arr(5L, 7L, 9L));
-        assertOnError(buildIntersection(unionA, unionB));
-
-        // intersection of union and intersection
-        KeyRangeIterator unionC = buildOnErrorB(this::buildUnion, arr(2L, 4L, 6L), arr(5L, 6L, 9L));
-        KeyRangeIterator intersectionA = buildIntersection(arr(3L, 4L, 6L, 9L), arr(2L, 3L, 6L, 9L));
-        assertOnError(buildIntersection(unionC, intersectionA));
-    }
-
-    @Test
-    public void testIntersectionOfIntersectionsOnError()
-    {
-        KeyRangeIterator intersectionA = buildIntersection(arr(1L, 2L, 3L, 6L), arr(2L, 3L, 6L));
-        KeyRangeIterator intersectionB = buildOnErrorA(this::buildIntersection, arr(2L, 4L, 6L), arr(5L, 6L, 7L, 9L));
-        assertOnError(buildIntersection(intersectionA, intersectionB));
-
-        intersectionA = buildOnErrorB(this::buildIntersection, arr(1L, 2L, 3L, 4L, 5L), arr(2L, 3L, 4L));
-        intersectionB = buildIntersection(arr(1L, 2L, 3L, 4L, 6L), arr(2L, 3L, 4L, 7L, 9L));
-        assertOnError(buildIntersection(intersectionA, intersectionB));
-
-        intersectionA = buildOnError(this::buildIntersection, arr(1L, 2L, 3L, 5L), arr(3L, 4L));
-        intersectionB = buildIntersection(arr(1L, 2L, 3L, 4L, 6L), arr(2L, 3L, 4L, 7L, 9L));
-        assertOnError(buildIntersection(intersectionA, intersectionB));
-    }
 }

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

@@ -19,12 +19,15 @@ package org.apache.cassandra.index.sai.iterators;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import org.junit.Assert;
 import org.junit.Test;
 
 import org.apache.cassandra.io.util.FileUtils;
+import org.apache.cassandra.utils.Pair;
 
 import static org.apache.cassandra.index.sai.iterators.LongIterator.convert;
 import static org.junit.Assert.assertEquals;
@@ -106,50 +109,38 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
     @Test
     public void testRandomSequences()
     {
-        long[][] values = new long[getRandom().nextIntBetween(1, 20)][];
+        for (int testIteration = 0; testIteration < 16; testIteration++)
-        int numTests = getRandom().nextIntBetween(10, 20);
-
-        for (int tests = 0; tests < numTests; tests++)
         {
-            KeyRangeUnionIterator.Builder builder = KeyRangeUnionIterator.builder(16);
+            var p = createRandom(nextInt(1, 20));
-            int totalCount = 0;
+            validateWithSkipping(p.left, p.right);
-
-            for (int i = 0; i < values.length; i++)
-            {
-                long[] part = new long[getRandom().nextIntBetween(1, 500)];
-                for (int j = 0; j < part.length; j++)
-                    part[j] = getRandom().nextLong();
-
-                // all the parts have to be sorted to mimic SSTable
-                Arrays.sort(part);
-
-                values[i] = part;
-                builder.add(new LongIterator(part));
-                totalCount += part.length;
-            }
-
-            long[] totalOrdering = new long[totalCount];
-            int index = 0;
-
-            for (long[] part : values)
-            {
-                for (long value : part)
-                    totalOrdering[index++] = value;
-            }
-
-            Arrays.sort(totalOrdering);
-
-            int count = 0;
-            KeyRangeIterator tokens = builder.build();
-
-            Assert.assertNotNull(tokens);
-            while (tokens.hasNext())
-                assertEquals(totalOrdering[count++], tokens.next().token().getLongValue());
-
-            assertEquals(totalCount, count);
         }
     }
 
+    static Pair<KeyRangeIterator, long[]> createRandom(int nRanges)
+    {
+        long[][] values = new long[nRanges][];
+        KeyRangeUnionIterator.Builder builder = KeyRangeUnionIterator.builder(10);
+        var allValues = new HashSet<Long>();
+        // add a random number of random values
+        for (int i = 0; i < values.length; i++)
+        {
+            int partLength = nextInt(1, 500);
+            var part = new HashSet<Long>(partLength);
+            for (int j = 0; j < partLength; j++)
+            {
+                long m = nextLong(0, 1024);
+                part.add(m);
+                allValues.add(m);
+            }
+
+            // all the parts have to be sorted to mimic SSTable
+            builder.add(new LongIterator(part.stream().mapToLong(Long::longValue).sorted().toArray()));
+        }
+        long[] totalOrdering = allValues.stream().mapToLong(Long::longValue).sorted().toArray();
+        KeyRangeIterator tokens = builder.build();
+        return Pair.create(tokens, totalOrdering);
+    }
+
     @Test
     public void testMinMaxAndCount()
     {
@@ -167,7 +158,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         Assert.assertNotNull(tokens);
         assertEquals(1L, tokens.getMinimum().token().getLongValue());
         assertEquals(9L, tokens.getMaximum().token().getLongValue());
-        assertEquals(9L, tokens.getCount());
+        assertEquals(9L, tokens.getMaxKeys());
 
         for (long i = 1; i < 10; i++)
         {
@@ -202,7 +193,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(4L, builder.rangeIterators.get(1).getMinimum().token().getLongValue());
         assertEquals(7L, builder.rangeIterators.get(2).getMinimum().token().getLongValue());
 
-        KeyRangeIterator tokens = KeyRangeUnionIterator.build(new ArrayList<KeyRangeIterator>()
+        KeyRangeIterator tokens = KeyRangeUnionIterator.build(new ArrayList<>()
         {{
             add(new LongIterator(new long[]{1L, 2L, 4L}));
             add(new LongIterator(new long[]{3L, 5L, 6L}));
@@ -213,7 +204,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         FileUtils.closeQuietly(tokens);
 
         KeyRangeIterator emptyTokens = KeyRangeUnionIterator.builder(16).build();
-        assertEquals(0, emptyTokens.getCount());
+        assertEquals(0, emptyTokens.getMaxKeys());
 
         builder = KeyRangeUnionIterator.builder(16);
         assertEquals(0L, builder.add((KeyRangeIterator) null).rangeCount());
@@ -281,27 +272,27 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
             for (int i = 0; i <= 3; i++)
             {
                 Assert.assertTrue(tokens.hasNext());
-                assertEquals(i, tokens.getCurrent().token().getLongValue());
+                assertEquals(i, tokens.peek().token().getLongValue());
                 assertEquals(i, tokens.next().token().getLongValue());
             }
         }
 
         try (LongIterator tokens = new LongIterator(new long[] { 0L, 1L, 3L, 5L }))
         {
-            assertEquals(3L, tokens.skipTo(LongIterator.fromToken(2L)).token().getLongValue());
+            tokens.skipTo(LongIterator.fromToken(2L));
             Assert.assertTrue(tokens.hasNext());
-            assertEquals(3L, tokens.getCurrent().token().getLongValue());
+            assertEquals(3L, tokens.peek().token().getLongValue());
             assertEquals(3L, tokens.next().token().getLongValue());
 
-            assertEquals(5L, tokens.skipTo(LongIterator.fromToken(5L)).token().getLongValue());
+            tokens.skipTo(LongIterator.fromToken(5L));
             Assert.assertTrue(tokens.hasNext());
-            assertEquals(5L, tokens.getCurrent().token().getLongValue());
+            assertEquals(5L, tokens.peek().token().getLongValue());
             assertEquals(5L, tokens.next().token().getLongValue());
         }
 
         try (LongIterator empty = LongIterator.newEmptyIterator())
         {
-            Assert.assertNull(empty.skipTo(LongIterator.fromToken(3L)));
+            empty.skipTo(LongIterator.fromToken(3L));
             Assert.assertFalse(empty.hasNext());
         }
     }
@@ -320,7 +311,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, range.getMinimum().token().getLongValue());
         assertEquals(19L, range.getMaximum().token().getLongValue());
         Assert.assertTrue(range.hasNext());
-        assertEquals(10, range.getCount());
+        assertEquals(10, range.getMaxKeys());
 
         builder = KeyRangeUnionIterator.builder(16);
         builder.add(LongIterator.newEmptyIterator());
@@ -329,7 +320,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, range.getMinimum().token().getLongValue());
         assertEquals(10L, range.getMaximum().token().getLongValue());
         Assert.assertTrue(range.hasNext());
-        assertEquals(1, range.getCount());
+        assertEquals(1, range.getMaxKeys());
 
         // non-empty, then empty
         builder = KeyRangeUnionIterator.builder(16);
@@ -340,7 +331,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10, range.getMinimum().token().getLongValue());
         assertEquals(19, range.getMaximum().token().getLongValue());
         Assert.assertTrue(range.hasNext());
-        assertEquals(10, range.getCount());
+        assertEquals(10, range.getMaxKeys());
 
         builder = KeyRangeUnionIterator.builder(16);
         builder.add(new LongIterator(new long[] {10}));
@@ -349,7 +340,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, range.getMinimum().token().getLongValue());
         assertEquals(10L, range.getMaximum().token().getLongValue());
         Assert.assertTrue(range.hasNext());
-        assertEquals(1, range.getCount());
+        assertEquals(1, range.getMaxKeys());
 
         // empty, then non-empty then empty again
         builder = KeyRangeUnionIterator.builder(16);
@@ -361,7 +352,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, range.getMinimum().token().getLongValue());
         assertEquals(19L, range.getMaximum().token().getLongValue());
         Assert.assertTrue(range.hasNext());
-        assertEquals(10, range.getCount());
+        assertEquals(10, range.getMaxKeys());
 
         // non-empty, empty, then non-empty again
         builder = KeyRangeUnionIterator.builder(16);
@@ -374,7 +365,7 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         assertEquals(10L, range.getMinimum().token().getLongValue());
         assertEquals(19L, range.getMaximum().token().getLongValue());
         Assert.assertTrue(range.hasNext());
-        assertEquals(10, range.getCount());
+        assertEquals(10, range.getMaxKeys());
     }
 
     // SAI specific tests
@@ -397,54 +388,28 @@ public class KeyRangeUnionIteratorTest extends AbstractKeyRangeIteratorTester
         assertSame(KeyRangeUnionIterator.class, union.getClass());
 
         // union of one intersected intersection and one non-intersected intersection
-        intersectionA = buildIntersection(arr(1L, 2L, 3L), arr(2L, 3L, 4L ));
+        intersectionA = buildIntersection(arr(1L, 2L, 3L), arr(2L, 3L, 4L));
-        intersectionB = buildIntersection(arr(6L, 7L, 8L), arr(10L ));
+        intersectionB = buildIntersection(arr(6L, 7L, 8L), arr(10L));
 
         union = buildUnion(intersectionA, intersectionB);
         assertEquals(convert(2L, 3L), convert(union));
     }
 
     @Test
-    public void testUnionOnError()
+    public void testUnionOfRandom()
     {
-        assertOnError(buildOnError(this::buildUnion, arr(1L, 3L, 4L ), arr(7L, 8L)));
+        for (int testIteration = 0; testIteration < 16; testIteration++)
-        assertOnError(buildOnErrorA(this::buildUnion, arr(1L, 3L, 4L ), arr(4L, 5L)));
+        {
-        assertOnError(buildOnErrorB(this::buildUnion, arr(1L), arr(2)));
+            var allValues = new HashSet<Long>();
-    }
+            var builder = KeyRangeUnionIterator.builder(10);
-
+            for (int i = 0; i < nextInt(2, 3); i++)
-    @Test
+            {
-    public void testUnionOfIntersectionsOnError()
+                var p = createRandomIterator();
-    {
+                builder.add(p.left);
-        KeyRangeIterator intersectionA = buildIntersection(arr(1L, 2L, 3L, 6L), arr(2L, 3L, 6L));
+                allValues.addAll(Arrays.stream(p.right).boxed().collect(Collectors.toList()));
-        KeyRangeIterator intersectionB = buildOnErrorA(this::buildIntersection, arr(2L, 4L, 6L), arr(5L, 6L, 7L, 9L));
+            }
-        assertOnError(buildUnion(intersectionA, intersectionB));
+            long[] totalOrdered = allValues.stream().mapToLong(Long::longValue).sorted().toArray();
-
+            validateWithSkipping(builder.build(), totalOrdered);
-        intersectionA = buildOnErrorB(this::buildIntersection, arr(1L, 2L, 3L, 4L, 5L), arr(2L, 3L, 5L));
+        }
-        intersectionB = buildIntersection(arr(2L, 4L, 5L), arr(5L, 6L, 7L));
-        assertOnError(buildUnion(intersectionA, intersectionB));
-    }
-
-    @Test
-    public void testUnionOfUnionsOnError()
-    {
-        KeyRangeIterator unionA = buildUnion(arr(1L, 2L, 3L, 6L), arr(6L, 7L, 8L));
-        KeyRangeIterator unionB = buildOnErrorA(this::buildUnion, arr(2L, 4L, 6L), arr (6L, 7L, 9L));
-        assertOnError(buildUnion(unionA, unionB));
-
-        unionA = buildOnErrorB(this::buildUnion, arr(1L, 2L, 3L), arr(3L, 7L, 8L));
-        unionB = buildUnion(arr(2L, 4L, 5L), arr (5L, 7L, 9L));
-        assertOnError(buildUnion(unionA, unionB));
-    }
-
-    @Test
-    public void testUnionOfMergingOnError()
-    {
-        KeyRangeIterator mergingA = buildConcat(arr(1L, 2L, 3L, 6L), arr(6L, 7L, 8L));
-        KeyRangeIterator mergingB = buildOnErrorA(this::buildConcat, arr(2L, 4L, 6L), arr (6L, 7L, 9L));
-        assertOnError(buildUnion(mergingA, mergingB));
-
-        mergingA = buildOnErrorB(this::buildConcat, arr(1L, 2L, 3L), arr(3L, 7L, 8L));
-        mergingB = buildConcat(arr(2L, 4L, 5L), arr (5L, 7L, 9L));
-        assertOnError(buildUnion(mergingA, mergingB));
     }
 }

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

@@ -29,11 +29,6 @@ public class LongIterator extends KeyRangeIterator
     private final List<PrimaryKey> keys;
     private int currentIdx = 0;
 
-    /**
-     * whether LongIterator should throw exception during iteration.
-     */
-    private boolean shouldThrow = false;
-
     public static LongIterator newEmptyIterator()
     {
         return new LongIterator();
@@ -54,18 +49,9 @@ public class LongIterator extends KeyRangeIterator
             this.keys.add(fromToken(token));
     }
 
-    public void throwsException()
-    {
-        this.shouldThrow = true;
-    }
-
     @Override
     protected PrimaryKey computeNext()
     {
-        // throws exception if it's last element or chosen 1 out of n
-        if (shouldThrow && (currentIdx >= keys.size() - 1 || SAITester.getRandom().nextInt(keys.size()) == 0))
-            throw new RuntimeException("injected exception");
-
         if (currentIdx >= keys.size())
             return endOfData();
 
@@ -75,14 +61,11 @@ public class LongIterator extends KeyRangeIterator
     @Override
     protected void performSkipTo(PrimaryKey nextKey)
     {
-        for (int i = currentIdx == 0 ? 0 : currentIdx - 1; i < keys.size(); i++)
+        for ( ; currentIdx < keys.size(); currentIdx++)
         {
-            PrimaryKey token = keys.get(i);
+            PrimaryKey token = keys.get(currentIdx);
             if (token.compareTo(nextKey) >= 0)
-            {
-                currentIdx = i;
                 break;
-            }
         }
     }
 

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

@@ -129,6 +129,16 @@ public class SAIRandomizedTester extends SAITester
         return getRandom().nextIntBetween(min, max - 1);
     }
 
+    public static double nextDouble()
+    {
+        return getRandom().nextDouble();
+    }
+
+    public static boolean nextBoolean()
+    {
+        return getRandom().nextBoolean();
+    }
+
     public static long between(long min, long max)
     {
         return randomLongBetween(min, max);