From d9f5b4e6d29f4d1c7c5c468cf74b68cb4ffc3ee0 Mon Sep 17 00:00:00 2001
From: Josh Snyder <joshs@netflix.com>
Date: Sat, 11 Nov 2017 07:11:09 +0000
Subject: [PATCH] Support multiple slices over RocksDBPartitions

---
 .../cassandra/rocksdb/RocksDBPartition.java   | 47 ++++++++++++++++++-
 1 file changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/java/org/apache/cassandra/rocksdb/RocksDBPartition.java b/src/java/org/apache/cassandra/rocksdb/RocksDBPartition.java
index e32cd109d978..96820fd13927 100644
--- a/src/java/org/apache/cassandra/rocksdb/RocksDBPartition.java
+++ b/src/java/org/apache/cassandra/rocksdb/RocksDBPartition.java
@@ -21,6 +21,7 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.List;
 
 import org.apache.commons.lang.NotImplementedException;
@@ -150,8 +151,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;
+            }
+        }
     }