Support multiple slices over RocksDBPartitions

Instagram · Mar 21, 2018 · 9c8ada2 · 9c8ada2
1 parent b5a1bb5
commit 9c8ada2
Showing 1 changed file with 45 additions and 2 deletions.
diff --git a/src/java/org/apache/cassandra/rocksdb/RocksDBPartition.java b/src/java/org/apache/cassandra/rocksdb/RocksDBPartition.java
@@ -20,6 +20,7 @@
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 
 import org.apache.commons.lang.NotImplementedException;
@@ -147,8 +148,50 @@ public UnfilteredRowIterator unfilteredIterator()
 
     public UnfilteredRowIterator unfilteredIterator(ColumnFilter columns, Slices slices, boolean reversed)
     {
-        //TODO: support multiple slices
-        return sliceIterator(slices.get(0), columns, (reversed ? PartitionIterOrder.REVERSED : PartitionIterOrder.NORMAL));
+        if(slices.size() == 1) {
+            return sliceIterator(slices.get(0), columns, reversed
+                    ? PartitionIterOrder.REVERSED
+                    : PartitionIterOrder.NORMAL);
+        }
+
+        return new SlicesIterator(slices, columns, reversed);
+    }
+
+    private class SlicesIterator extends AbstractUnfilteredRowIterator {
+        private final Iterator<Slice> slices;
+        private final PartitionIterOrder iterOrder;
+        private final ColumnFilter columnFilter;
+
+        private Iterator<Unfiltered> currentSlice;
+
+        private SlicesIterator(Slices slices, ColumnFilter columnFilter, boolean isReverseOrder) {
+            super(RocksDBPartition.this.metadata, RocksDBPartition.this.partitionKey, DeletionTime.LIVE,
+                    RocksDBPartition.this.metadata.partitionColumns(), null, isReverseOrder, EncodingStats.NO_STATS);
+            this.slices = slices.iterator();
+            this.iterOrder = isReverseOrder ? PartitionIterOrder.REVERSED : PartitionIterOrder.NORMAL;
+            this.columnFilter = columnFilter;
+        }
+
+        protected Unfiltered computeNext()
+        {
+            while (true)
+            {
+                if (currentSlice == null)
+                {
+                    if(slices.hasNext()) {
+                        currentSlice = sliceIterator(slices.next(), columnFilter, iterOrder);
+                    }
+                    else {
+                        return endOfData();
+                    }
+                }
+
+                if (currentSlice.hasNext())
+                    return currentSlice.next();
+
+                currentSlice = null;
+            }
+        }
     }