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

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>
This commit is contained in:
Ekaterina Dimitrova 2024-03-04 13:46:25 -05:00
parent 973aad7b68
commit b0150e86fb
20 changed files with 604 additions and 525 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGES.txt
View File

@ -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)

39
src/java/org/apache/cassandra/index/sai/disk/IndexSearchResultIterator.java
View File

@ -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);
return new IndexSearchResultIterator(union, sstableIndexes, queryContext);
KeyRangeIterator union = KeyRangeUnionIterator.build(subIterators, () -> {});
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);
}
}
}

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

@ -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;

105
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeConcatIterator.java
View File

@ -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> toRelease;
private final List<KeyRangeIterator> ranges;
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 (head.getMaximum().compareTo(nextKey) >= 0)
if (currentIterator.getMaximum().compareTo(nextKey) >= 0)
{
head.skipTo(nextKey);
ranges.add(head);
currentIterator.skipTo(nextKey);
break;
}
}
}
@Override
public void close()
{
// due to lazy key fetching, we cannot close iterator immediately
FileUtils.closeQuietly(toRelease);
current++;
}
}
@Override
protected PrimaryKey computeNext()
{
while (!ranges.isEmpty())
while (current < ranges.size())
{
KeyRangeIterator current = ranges.poll();
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);
KeyRangeIterator currentIterator = ranges.get(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());
ranges = new PriorityQueue<>(size, Comparator.comparing(KeyRangeIterator::getCurrent));
super(new ConcatStatistics(), onClose);
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();
}
}
}

134
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeIntersectionIterator.java
View File

@ -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.
// 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.
PrimaryKey highestKey = advanceOneRange();
if (highestKey == null)
highestKey = computeHighestKey();
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
? nextOrNull(range, highestKey.toStatic())
: nextOrNull(range, highestKey);
PrimaryKey nextKey = range.peek().kind() == Kind.STATIC
? skipAndPeek(range, highestKey.toStatic())
: 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.peek().kind() != Kind.STATIC)
{
range.next();
return range.hasNext() ? range.peek() : null;
}
for (KeyRangeIterator range : ranges)
if (range.getCurrent().kind() != Kind.STATIC)
return nextOrNull(range, highestKey);
for (KeyRangeIterator range : ranges)
if (range.getCurrent().kind() == Kind.STATIC)
return nextOrNull(range, highestKey);
if (range.peek().kind() == Kind.STATIC)
{
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);
}
}

99
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeIterator.java
View File

@ -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)
return endOfData();
if (state == State.DONE)
return;
// In the case of deferred iterators the current value may not accurately
// reflect the next value, so we need to check that as well
if (current.compareTo(nextKey) >= 0)
{
next = next == null ? recomputeNext() : next;
if (next == null)
return endOfData();
else if (next.compareTo(nextKey) >= 0)
return next;
}
if (state == State.READY && next.compareTo(nextKey) >= 0)
return;
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();
current = hasNext ? next : getMaximum();
return hasNext;
onClose.run();
}
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();

2
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeOrderingIterator.java
View File

@ -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;

84
src/java/org/apache/cassandra/index/sai/iterators/KeyRangeUnionIterator.java
View File

@ -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();
int cmp = candidateKey.compareTo(peeked);
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);
}
// we found a new best candidate, throw away the old ones
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();
}
}
}

4
src/java/org/apache/cassandra/index/sai/memory/InMemoryKeyRangeIterator.java
View File

@ -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;
}

33
src/java/org/apache/cassandra/index/sai/plan/QueryController.java
View File

@ -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;
}
}

5
src/java/org/apache/cassandra/utils/AbstractGuavaIterator.java
View File

@ -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() {}

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

@ -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);
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);
performTestQueries(isSingleIndexTest);
}
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());
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());
verifyResults(isSingleIndexTest);
}
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());
}
}
}

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

@ -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()));

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

@ -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

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

@ -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 static org.assertj.core.api.Assertions.assertThatThrownBy;
import org.apache.cassandra.utils.Pair;
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)
{
return buildUnion(toRangeIterator(ranges));
}
final KeyRangeIterator buildConcat(KeyRangeIterator... ranges)
static KeyRangeIterator buildConcat(KeyRangeIterator... ranges)
{
return KeyRangeConcatIterator.builder(ranges.length).add(Arrays.asList(ranges)).build();
}
final KeyRangeIterator buildConcat(long[]... ranges)
private static KeyRangeIterator[] toRangeIterator(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);
}
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);
return new LongIterator(tokens);
}
protected KeyRangeIterator build(BiFunction<KeyRangeIterator, KeyRangeIterator, KeyRangeIterator> builder,
long[] tokensA,
boolean onErrorA,
long[] tokensB,
boolean onErrorB)
long[] tokensB)
{
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();
}
}
}

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

@ -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);
}
}

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

@ -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());
assertEquals(4L, range.getCurrent().token().getLongValue());
range.skipTo(LongIterator.fromToken(3L));
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());
assertEquals(6L, range.getCurrent().token().getLongValue());
range.skipTo(LongIterator.fromToken(5L));
Assert.assertEquals(6L, range.peek().token().getLongValue());
// now right to the element
assertEquals(7L, range.skipTo(LongIterator.fromToken(7L)).token().getLongValue());
assertEquals(7L, range.getCurrent().token().getLongValue());
range.skipTo(LongIterator.fromToken(7L));
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);
// 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()));
var p = createRandom(nextInt(2, 16));
validateWithSkipping(p.left, p.right);
}
}
/**
* @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));
}
}

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

@ -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)][];
int numTests = getRandom().nextIntBetween(10, 20);
for (int tests = 0; tests < numTests; tests++)
for (int testIteration = 0; testIteration < 16; testIteration++)
{
KeyRangeUnionIterator.Builder builder = KeyRangeUnionIterator.builder(16);
int totalCount = 0;
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);
var p = createRandom(nextInt(1, 20));
validateWithSkipping(p.left, p.right);
}
}
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 ));
intersectionB = buildIntersection(arr(6L, 7L, 8L), arr(10L ));
intersectionA = buildIntersection(arr(1L, 2L, 3L), arr(2L, 3L, 4L));
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)));
assertOnError(buildOnErrorA(this::buildUnion, arr(1L, 3L, 4L ), arr(4L, 5L)));
assertOnError(buildOnErrorB(this::buildUnion, arr(1L), arr(2)));
}
@Test
public void testUnionOfIntersectionsOnError()
{
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(buildUnion(intersectionA, intersectionB));
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));
for (int testIteration = 0; testIteration < 16; testIteration++)
{
var allValues = new HashSet<Long>();
var builder = KeyRangeUnionIterator.builder(10);
for (int i = 0; i < nextInt(2, 3); i++)
{
var p = createRandomIterator();
builder.add(p.left);
allValues.addAll(Arrays.stream(p.right).boxed().collect(Collectors.toList()));
}
long[] totalOrdered = allValues.stream().mapToLong(Long::longValue).sorted().toArray();
validateWithSkipping(builder.build(), totalOrdered);
}
}
}

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

@ -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;
}
}
}

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

@ -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);