Skip to content

Commit

Permalink
SHS-NG M1: Add "max" and "last" to kvstore iterators.
Browse files Browse the repository at this point in the history 
This makes it easier for callers to control the end of iteration,
making it easier to write Scala code that automatically closes
underlying iterator resources. Before, code had to use Scala's
"takeWhile", convert the result to a list, and manually close the
iterators; with these two parameters, that can be avoided in a
bunch of cases, with iterators auto-closing when the last element
is reached.
  • Loading branch information
Marcelo Vanzin committed May 8, 2017
1 parent 7b87021 commit 5197c21
Show file tree
Hide file tree
Showing 3 changed files with 167 additions and 40 deletions.
22 changes: 21 additions & 1 deletion common/kvstore/src/main/java/org/apache/spark/kvstore/KVStoreView.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,9 @@ public abstract class KVStoreView<T> implements Iterable<T> {
boolean ascending = true;
String index = KVIndex.NATURAL_INDEX_NAME;
Object first = null;
Object last = null;
long skip = 0L;
long max = Long.MAX_VALUE;

public KVStoreView(Class<T> type) {
this.type = type;
Expand Down Expand Up @@ -74,7 +76,25 @@ public KVStoreView<T> first(Object value) {
}

/**
* Skips a number of elements in the resulting iterator.
* Stops iteration at the given value of the chosen index.
*/
public KVStoreView<T> last(Object value) {
this.last = value;
return this;
}

/**
* Stops iteration after a number of elements has been retrieved.
*/
public KVStoreView<T> max(long max) {
Preconditions.checkArgument(max > 0L, "max must be positive.");
this.max = max;
return this;
}

/**
* Skips a number of elements at the start of iteration. Skipped elements are not accounted
* when using {@link #max(long)}.
*/
public KVStoreView<T> skip(long n) {
this.skip = n;
Expand Down
50 changes: 44 additions & 6 deletions common/kvstore/src/main/java/org/apache/spark/kvstore/LevelDBIterator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,10 +38,12 @@ class LevelDBIterator<T> implements KVStoreIterator<T> {
private final LevelDBTypeInfo.Index index;
private final byte[] indexKeyPrefix;
private final byte[] end;
private final long max;

private boolean checkedNext;
private T next;
private boolean closed;
private long count;

LevelDBIterator(LevelDB db, KVStoreView<T> params) throws Exception {
this.db = db;
Expand All @@ -51,6 +53,7 @@ class LevelDBIterator<T> implements KVStoreIterator<T> {
this.ti = db.getTypeInfo(type);
this.index = ti.index(params.index);
this.indexKeyPrefix = index.keyPrefix();
this.max = params.max;

byte[] firstKey;
if (params.first != null) {
Expand All @@ -66,14 +69,27 @@ class LevelDBIterator<T> implements KVStoreIterator<T> {
}
it.seek(firstKey);

byte[] end = null;
if (ascending) {
this.end = index.end();
end = params.last != null ? index.end(params.last) : index.end();
} else {
this.end = null;
if (params.last != null) {
end = index.start(params.last);
}
if (it.hasNext()) {
it.next();
// When descending, the caller may have set up the start of iteration at a non-existant
// entry that is guaranteed to be after the desired entry. For example, if you have a
// compound key (a, b) where b is a, integer, you may seek to the end of the elements that
// have the same "a" value by specifying Integer.MAX_VALUE for "b", and that value may not
// exist in the database. So need to check here whether the next value actually belongs to
// the set being returned by the iterator before advancing.
byte[] nextKey = it.peekNext().getKey();
if (compare(nextKey, indexKeyPrefix) <= 0) {
it.next();
}
}
}
this.end = end;

if (params.skip > 0) {
skip(params.skip);
Expand Down Expand Up @@ -147,6 +163,10 @@ public synchronized void close() throws IOException {
}

private T loadNext() {
if (count >= max) {
return null;
}

try {
while (true) {
boolean hasNext = ascending ? it.hasNext() : it.hasPrev();
Expand All @@ -173,11 +193,16 @@ private T loadNext() {
return null;
}

// If there's a known end key and it's found, stop.
if (end != null && Arrays.equals(nextKey, end)) {
return null;
// If there's a known end key and iteration has gone past it, stop.
if (end != null) {
int comp = compare(nextKey, end) * (ascending ? 1 : -1);
if (comp > 0) {
return null;
}
}

count++;

// Next element is part of the iteration, return it.
if (index == null || index.isCopy()) {
return db.serializer.deserialize(nextEntry.getValue(), type);
Expand Down Expand Up @@ -228,4 +253,17 @@ private byte[] stitch(byte[]... comps) {
return dest;
}

private int compare(byte[] a, byte[] b) {
int diff = 0;
int minLen = Math.min(a.length, b.length);
for (int i = 0; i < minLen; i++) {
diff += (a[i] - b[i]);
if (diff != 0) {
return diff;
}
}

return a.length - b.length;
}

}
135 changes: 102 additions & 33 deletions common/kvstore/src/test/java/org/apache/spark/kvstore/DBIteratorSuite.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -152,111 +152,170 @@ public static void cleanup() throws Exception {

@Test
public void naturalIndex() throws Exception {
testIteration(NATURAL_ORDER, view(), null);
testIteration(NATURAL_ORDER, view(), null, null);
}

@Test
public void refIndex() throws Exception {
testIteration(REF_INDEX_ORDER, view().index("id"), null);
testIteration(REF_INDEX_ORDER, view().index("id"), null, null);
}

@Test
public void copyIndex() throws Exception {
testIteration(COPY_INDEX_ORDER, view().index("name"), null);
testIteration(COPY_INDEX_ORDER, view().index("name"), null, null);
}

@Test
public void numericIndex() throws Exception {
testIteration(NUMERIC_INDEX_ORDER, view().index("int"), null);
testIteration(NUMERIC_INDEX_ORDER, view().index("int"), null, null);
}

@Test
public void naturalIndexDescending() throws Exception {
testIteration(NATURAL_ORDER, view().reverse(), null);
testIteration(NATURAL_ORDER, view().reverse(), null, null);
}

@Test
public void refIndexDescending() throws Exception {
testIteration(REF_INDEX_ORDER, view().index("id").reverse(), null);
testIteration(REF_INDEX_ORDER, view().index("id").reverse(), null, null);
}

@Test
public void copyIndexDescending() throws Exception {
testIteration(COPY_INDEX_ORDER, view().index("name").reverse(), null);
testIteration(COPY_INDEX_ORDER, view().index("name").reverse(), null, null);
}

@Test
public void numericIndexDescending() throws Exception {
testIteration(NUMERIC_INDEX_ORDER, view().index("int").reverse(), null);
testIteration(NUMERIC_INDEX_ORDER, view().index("int").reverse(), null, null);
}

@Test
public void naturalIndexWithStart() throws Exception {
CustomType1 first = pickFirst();
testIteration(NATURAL_ORDER, view().first(first.key), first);
CustomType1 first = pickLimit();
testIteration(NATURAL_ORDER, view().first(first.key), first, null);
}

@Test
public void refIndexWithStart() throws Exception {
CustomType1 first = pickFirst();
testIteration(REF_INDEX_ORDER, view().index("id").first(first.id), first);
CustomType1 first = pickLimit();
testIteration(REF_INDEX_ORDER, view().index("id").first(first.id), first, null);
}

@Test
public void copyIndexWithStart() throws Exception {
CustomType1 first = pickFirst();
testIteration(COPY_INDEX_ORDER, view().index("name").first(first.name), first);
CustomType1 first = pickLimit();
testIteration(COPY_INDEX_ORDER, view().index("name").first(first.name), first, null);
}

@Test
public void numericIndexWithStart() throws Exception {
CustomType1 first = pickFirst();
testIteration(NUMERIC_INDEX_ORDER, view().index("int").first(first.num), first);
CustomType1 first = pickLimit();
testIteration(NUMERIC_INDEX_ORDER, view().index("int").first(first.num), first, null);
}

@Test
public void naturalIndexDescendingWithStart() throws Exception {
CustomType1 first = pickFirst();
testIteration(NATURAL_ORDER, view().reverse().first(first.key), first);
CustomType1 first = pickLimit();
testIteration(NATURAL_ORDER, view().reverse().first(first.key), first, null);
}

@Test
public void refIndexDescendingWithStart() throws Exception {
CustomType1 first = pickFirst();
testIteration(REF_INDEX_ORDER, view().reverse().index("id").first(first.id), first);
CustomType1 first = pickLimit();
testIteration(REF_INDEX_ORDER, view().reverse().index("id").first(first.id), first, null);
}

@Test
public void copyIndexDescendingWithStart() throws Exception {
CustomType1 first = pickFirst();
testIteration(COPY_INDEX_ORDER, view().reverse().index("name").first(first.name),
first);
CustomType1 first = pickLimit();
testIteration(COPY_INDEX_ORDER, view().reverse().index("name").first(first.name), first, null);
}

@Test
public void numericIndexDescendingWithStart() throws Exception {
CustomType1 first = pickFirst();
testIteration(NUMERIC_INDEX_ORDER, view().reverse().index("int").first(first.num),
first);
CustomType1 first = pickLimit();
testIteration(NUMERIC_INDEX_ORDER, view().reverse().index("int").first(first.num), first, null);
}

@Test
public void naturalIndexWithSkip() throws Exception {
testIteration(NATURAL_ORDER, view().skip(RND.nextInt(allEntries.size() / 2)), null);
testIteration(NATURAL_ORDER, view().skip(RND.nextInt(allEntries.size() / 2)), null, null);
}

@Test
public void refIndexWithSkip() throws Exception {
testIteration(REF_INDEX_ORDER, view().index("id").skip(RND.nextInt(allEntries.size() / 2)),
null);
null, null);
}

@Test
public void copyIndexWithSkip() throws Exception {
testIteration(COPY_INDEX_ORDER, view().index("name").skip(RND.nextInt(allEntries.size() / 2)),
null);
null, null);
}

@Test
public void naturalIndexWithMax() throws Exception {
testIteration(NATURAL_ORDER, view().max(RND.nextInt(allEntries.size() / 2)), null, null);
}

@Test
public void copyIndexWithMax() throws Exception {
testIteration(COPY_INDEX_ORDER, view().index("name").max(RND.nextInt(allEntries.size() / 2)),
null, null);
}

@Test
public void naturalIndexWithLast() throws Exception {
CustomType1 last = pickLimit();
testIteration(NATURAL_ORDER, view().last(last.key), null, last);
}

@Test
public void refIndexWithLast() throws Exception {
CustomType1 last = pickLimit();
testIteration(REF_INDEX_ORDER, view().index("id").last(last.id), null, last);
}

@Test
public void copyIndexWithLast() throws Exception {
CustomType1 last = pickLimit();
testIteration(COPY_INDEX_ORDER, view().index("name").last(last.name), null, last);
}

@Test
public void numericIndexWithLast() throws Exception {
CustomType1 last = pickLimit();
testIteration(NUMERIC_INDEX_ORDER, view().index("int").last(last.num), null, last);
}

@Test
public void naturalIndexDescendingWithLast() throws Exception {
CustomType1 last = pickLimit();
testIteration(NATURAL_ORDER, view().reverse().last(last.key), null, last);
}

@Test
public void refIndexDescendingWithLast() throws Exception {
CustomType1 last = pickLimit();
testIteration(REF_INDEX_ORDER, view().reverse().index("id").last(last.id), null, last);
}

@Test
public void copyIndexDescendingWithLast() throws Exception {
CustomType1 last = pickLimit();
testIteration(COPY_INDEX_ORDER, view().reverse().index("name").last(last.name),
null, last);
}

@Test
public void numericIndexDescendingWithLast() throws Exception {
CustomType1 last = pickLimit();
testIteration(NUMERIC_INDEX_ORDER, view().reverse().index("int").last(last.num),
null, last);
}

@Test
public void testRefWithIntNaturalKey() throws Exception {
LevelDBSuite.IntKeyType i = new LevelDBSuite.IntKeyType();
Expand All @@ -272,8 +331,8 @@ public void testRefWithIntNaturalKey() throws Exception {
}
}

private CustomType1 pickFirst() {
// Picks a first element that has clashes with other elements in the given index.
private CustomType1 pickLimit() {
// Picks an element that has clashes with other elements in the given index.
return clashingEntries.get(RND.nextInt(clashingEntries.size()));
}

Expand All @@ -297,22 +356,32 @@ private <T extends Comparable<T>> int compareWithFallback(
private void testIteration(
final BaseComparator order,
final KVStoreView<CustomType1> params,
final CustomType1 first) throws Exception {
final CustomType1 first,
final CustomType1 last) throws Exception {
List<CustomType1> indexOrder = sortBy(order.fallback());
if (!params.ascending) {
indexOrder = Lists.reverse(indexOrder);
}

Iterable<CustomType1> expected = indexOrder;
BaseComparator expectedOrder = params.ascending ? order : order.reverse();

if (first != null) {
final BaseComparator expectedOrder = params.ascending ? order : order.reverse();
expected = Iterables.filter(expected, v -> expectedOrder.compare(first, v) <= 0);
}

if (last != null) {
expected = Iterables.filter(expected, v -> expectedOrder.compare(v, last) <= 0);
}

if (params.skip > 0) {
expected = Iterables.skip(expected, (int) params.skip);
}

if (params.max != Long.MAX_VALUE) {
expected = Iterables.limit(expected, (int) params.max);
}

List<CustomType1> actual = collect(params);
compareLists(expected, actual);
}
Expand Down

0 comments on commit 5197c21

Please sign in to comment.