Fix merging iterator issue dain#95

pcmind · Jan 13, 2018 · d42f639 · d42f639
1 parent f2c0083
commit d42f639
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 0 deletions.
diff --git a/leveldb/src/main/java/org/iq80/leveldb/util/MergingIterator.java b/leveldb/src/main/java/org/iq80/leveldb/util/MergingIterator.java
@@ -62,6 +62,7 @@ protected void seekInternal(InternalKey targetKey)
 
     private void resetPriorityQueue(Comparator<InternalKey> comparator)
     {
+        priorityQueue.clear();
         int i = 1;
         for (InternalIterator level : levels) {
             if (level.hasNext()) {

diff --git a/leveldb/src/test/java/org/iq80/leveldb/util/MergingIteratorTest.java b/leveldb/src/test/java/org/iq80/leveldb/util/MergingIteratorTest.java
@@ -0,0 +1,83 @@
+package org.iq80.leveldb.util;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Maps;
+import org.iq80.leveldb.impl.InternalKey;
+import org.iq80.leveldb.impl.InternalKeyComparator;
+import org.iq80.leveldb.impl.ValueType;
+import org.iq80.leveldb.table.BytewiseComparator;
+import org.testng.annotations.Test;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+
+public class MergingIteratorTest
+{
+    @Test
+    public void testMergeSeek() throws Exception {
+        InternalKeyComparator comparator = new InternalKeyComparator(new BytewiseComparator());
+        InternalKey key1 = new InternalKey(Slices.wrappedBuffer(new byte[]{'1'}), 1, ValueType.VALUE);
+        InternalKey key2 = new InternalKey(Slices.wrappedBuffer(new byte[]{'3'}), 2, ValueType.VALUE);
+        InternalKey key3 = new InternalKey(Slices.wrappedBuffer(new byte[]{'2'}), 3, ValueType.VALUE);
+        InternalKey key4 = new InternalKey(Slices.wrappedBuffer(new byte[]{'4'}), 4, ValueType.VALUE);
+        ImmutableList<Map.Entry<InternalKey, Slice>> of1 = ImmutableList
+                .of(
+                        Maps.immutableEntry(key1, Slices.EMPTY_SLICE),
+                        Maps.immutableEntry(key2, Slices.EMPTY_SLICE)
+                );
+        ImmutableList<Map.Entry<InternalKey, Slice>> of2 = ImmutableList
+                .of(
+                        Maps.immutableEntry(key3, Slices.EMPTY_SLICE),
+                        Maps.immutableEntry(key4, Slices.EMPTY_SLICE)
+                );
+        ImmutableList<InternalIterator> of = ImmutableList.of(new InternalKeySliceAbstractSeekingIterator(of1, comparator), new InternalKeySliceAbstractSeekingIterator(of2, comparator));
+        MergingIterator mergingIterator = new MergingIterator(of, comparator);
+        assertEquals(mergingIterator.next().getKey(), key1);
+        assertEquals(mergingIterator.next().getKey(), key3);
+        mergingIterator.seekToFirst();
+        assertEquals(mergingIterator.next().getKey(), key1);
+        assertEquals(mergingIterator.next().getKey(), key3);
+        assertEquals(mergingIterator.next().getKey(), key2);
+        assertEquals(mergingIterator.next().getKey(), key4);
+        assertFalse(mergingIterator.hasNext());
+        mergingIterator.seekInternal(key2);
+        assertEquals(mergingIterator.next().getKey(), key2);
+        assertEquals(mergingIterator.next().getKey(), key4);
+        assertFalse(mergingIterator.hasNext());
+    }
+
+    private static class InternalKeySliceAbstractSeekingIterator extends AbstractSeekingIterator<InternalKey, Slice> implements InternalIterator {
+        int index = 0;
+        private final List<Map.Entry<InternalKey, Slice>> entries;
+        private final InternalKeyComparator comparator;
+
+        public InternalKeySliceAbstractSeekingIterator(List<Map.Entry<InternalKey, Slice>> entries, InternalKeyComparator comparator) {
+            this.entries = entries;
+            this.comparator = comparator;
+        }
+
+        @Override
+        protected void seekToFirstInternal() {
+            index = 0;
+        }
+
+        @Override
+        protected void seekInternal(InternalKey targetKey) {
+            index = entries.size();
+            for (int i = 0; i < entries.size(); i++) {
+                Map.Entry<InternalKey, Slice> entry = entries.get(i);
+                if (comparator.compare(entry.getKey(), targetKey) >= 0) {
+                    index = i;
+                }
+            }
+        }
+
+        @Override
+        protected Map.Entry<InternalKey, Slice> getNextElement() {
+            return index < entries.size() ? entries.get(index++) : null;
+        }
+    }
+}