Merge pull request #574 from mastersenthilkumar/master

Fix for KeyIterator doesn't loop through all keys

core/src/main/java/me/prettyprint/cassandra/service/KeyIterator.java

@@ -35,6 +35,7 @@ public class KeyIterator<K> implements Iterable<K> {
   private K lastReadValue = null;
   private K endKey;
   private boolean firstRun = true;
+  private int rowCount = MAX_ROW_COUNT_DEFAULT;
 
   private Iterator<K> keyIterator = new Iterator<K>() {
     @Override
@@ -99,12 +100,17 @@ public KeyIterator(Keyspace keyspace, String columnFamily, Serializer<K> seriali
       .setRowCount(maxRowCount);
 
     endKey = end;
+    if(maxRowCount < Integer.MAX_VALUE) {
+      rowCount = maxRowCount+1; //to compensate the first entry skip (except in first run)
+    }
     runQuery(start, end);
   }
 
   private void runQuery(K start, K end) {
     query.setKeys(start, end);
-
+    if(!firstRun) {
+        query.setRowCount(rowCount);
+    }
     rowsIterator = null;
     QueryResult<OrderedRows<K, String, String>> result = query.execute();
     OrderedRows<K, String, String> rows = (result != null) ? result.get() : null;

core/src/test/java/me/prettyprint/cassandra/service/KeyIteratorTest.java

@@ -54,6 +54,12 @@ public void testIterator() {
     assertStringKeys(5, "k5", null);
     assertStringKeys(9, null, null);
     assertStringKeys(7, null, "k7");
+
+    assertKeys(5, "k5", 1);
+    assertKeys(9, null, 2);
+    assertKeys(9, null, 5);
+    assertKeys(8, "k2", 7);
+    assertKeys(7, "k3", 10);
   }
 
   private void assertKeys(int expected, String start, String end) {
@@ -76,4 +82,14 @@ private void assertStringKeys(int expected, String start, String end) {
     assertEquals(expected, tot);
   }
 
+  private void assertKeys(int expected, String start, int count) {
+    Iterable<String> it = new KeyIterator.Builder<String>(keyspace, CF, se).start(start).maxRowCount(count).build();
+
+    int tot = 0;
+    for (String key : it)
+      tot++;
+
+    assertEquals(expected, tot);
+  }
+
 }