Repository: petermattis/pebble (go)

Issue:
batch: iterator ignores bounds
A `flushableBatchIter` is created from a `flushableBatch` which is in turn created whenever a `Batch` is too large to insert into a `memTable`.

However, a `BatchIter` supports lower and upper bounds but a `flushableBatchIter` does not. This is a bug that needs to be fixed.

Pull request description:
Fix lower and upper bound for flushableBatchIter
Fix `flushableBatchIter` by adding support for lower and upper bounds.

Fixes #148.

Patch:
diff --git a/batch.go b/batch.go
index 0988639..e38eb3a 100644
--- a/batch.go
+++ b/batch.go
@@ -937,6 +937,8 @@ func (b *flushableBatch) newIter(o *IterOptions) internalIterator {
 		offsets: b.offsets,
 		cmp:     b.cmp,
 		index:   -1,
+		lower:   o.GetLowerBound(),
+		upper:   o.GetUpperBound(),
 	}
 }
 
@@ -997,6 +999,8 @@ type flushableBatchIter struct {
 	index   int
 	key     InternalKey
 	err     error
+	lower   []byte
+	upper   []byte
 }
 
 // flushableBatchIter implements the internalIterator interface.
@@ -1011,6 +1015,10 @@ func (i *flushableBatchIter) SeekGE(key []byte) (*InternalKey, []byte) {
 		return nil, nil
 	}
 	i.key = i.getKey(i.index)
+	if i.upper != nil && i.cmp(i.key.UserKey, i.upper) >= 0 {
+		i.index = len(i.offsets)
+		return nil, nil
+	}
 	return &i.key, i.Value()
 }
 
@@ -1028,6 +1036,10 @@ func (i *flushableBatchIter) SeekLT(key []byte) (*InternalKey, []byte) {
 		return nil, nil
 	}
 	i.key = i.getKey(i.index)
+	if i.lower != nil && i.cmp(i.key.UserKey, i.lower) < 0 {
+		i.index = -1
+		return nil, nil
+	}
 	return &i.key, i.Value()
 }
 
@@ -1037,6 +1049,10 @@ func (i *flushableBatchIter) First() (*InternalKey, []byte) {
 	}
 	i.index = 0
 	i.key = i.getKey(i.index)
+	if i.upper != nil && i.cmp(i.key.UserKey, i.upper) >= 0 {
+		i.index = len(i.offsets)
+		return nil, nil
+	}
 	return &i.key, i.Value()
 }
 
@@ -1046,6 +1062,10 @@ func (i *flushableBatchIter) Last() (*InternalKey, []byte) {
 	}
 	i.index = len(i.offsets) - 1
 	i.key = i.getKey(i.index)
+	if i.lower != nil && i.cmp(i.key.UserKey, i.lower) < 0 {
+		i.index = -1
+		return nil, nil
+	}
 	return &i.key, i.Value()
 }
 
@@ -1058,6 +1078,10 @@ func (i *flushableBatchIter) Next() (*InternalKey, []byte) {
 		return nil, nil
 	}
 	i.key = i.getKey(i.index)
+	if i.upper != nil && i.cmp(i.key.UserKey, i.upper) >= 0 {
+		i.index = len(i.offsets)
+		return nil, nil
+	}
 	return &i.key, i.Value()
 }
 
@@ -1070,6 +1094,10 @@ func (i *flushableBatchIter) Prev() (*InternalKey, []byte) {
 		return nil, nil
 	}
 	i.key = i.getKey(i.index)
+	if i.lower != nil && i.cmp(i.key.UserKey, i.lower) < 0 {
+		i.index = -1
+		return nil, nil
+	}
 	return &i.key, i.Value()
 }
 
@@ -1106,6 +1134,6 @@ func (i *flushableBatchIter) Close() error {
 }
 
 func (i *flushableBatchIter) SetBounds(lower, upper []byte) {
-	// This should not be called as bounds are not used for this iterator.
-	panic("TODO(ryan): pebble: SetBounds unimplemented")
+	i.lower = lower
+	i.upper = upper
 }


Test patch:
diff --git a/batch_test.go b/batch_test.go
index cf0f1a4..b7243d0 100644
--- a/batch_test.go
+++ b/batch_test.go
@@ -285,7 +285,7 @@ func TestFlushableBatchIter(t *testing.T) {
 	})
 }
 
-func TestFlushableBatchSeqNum(t *testing.T) {
+func TestFlushableBatch(t *testing.T) {
 	var b *flushableBatch
 	datadriven.RunTest(t, "testdata/flushable_batch", func(d *datadriven.TestData) string {
 		switch d.Cmd {
@@ -307,6 +307,26 @@ func TestFlushableBatchSeqNum(t *testing.T) {
 			b = newFlushableBatch(batch, DefaultComparer)
 			return ""
 
+		case "iter":
+			var opts IterOptions
+			for _, arg := range d.CmdArgs {
+				if len(arg.Vals) != 1 {
+					return fmt.Sprintf("%s: %s=<value>", d.Cmd, arg.Key)
+				}
+				switch arg.Key {
+				case "lower":
+					opts.LowerBound = []byte(arg.Vals[0])
+				case "upper":
+					opts.UpperBound = []byte(arg.Vals[0])
+				default:
+					return fmt.Sprintf("%s: unknown arg: %s", d.Cmd, arg.Key)
+				}
+			}
+
+			iter := b.newIter(&opts)
+			defer iter.Close()
+			return runInternalIterCmd(d, iter)
+
 		case "dump":
 			if len(d.CmdArgs) != 1 || len(d.CmdArgs[0].Vals) != 1 || d.CmdArgs[0].Key != "seq" {
 				return fmt.Sprintf("dump seq=<value>\n")
diff --git a/testdata/flushable_batch b/testdata/flushable_batch
index 315750c..2f84099 100644
--- a/testdata/flushable_batch
+++ b/testdata/flushable_batch
@@ -36,3 +36,107 @@ dump seq=1000
 a#1002,1:1
 b#1001,1:1
 c#1000,1:1
+
+iter
+first
+next
+next
+next
+----
+a:1
+b:1
+c:1
+.
+
+iter
+last
+prev
+prev
+prev
+----
+c:1
+b:1
+a:1
+.
+
+iter lower=b
+seek-ge b
+prev
+----
+b:1
+.
+
+iter lower=c
+seek-lt b
+----
+.
+
+iter lower=e
+last
+----
+.
+
+iter upper=1
+first
+----
+.
+
+iter upper=b
+first
+next
+----
+a:1
+.
+
+iter upper=b
+seek-ge b
+----
+.
+
+iter
+set-bounds lower=b
+seek-ge b
+prev
+----
+b:1
+.
+
+iter
+set-bounds lower=c
+seek-lt b
+----
+.
+
+iter
+set-bounds lower=e
+last
+----
+.
+
+iter
+set-bounds upper=1
+first
+----
+.
+
+iter
+set-bounds upper=b
+first
+next
+----
+a:1
+.
+
+iter
+set-bounds upper=b
+seek-ge b
+----
+.
+
+iter
+first
+set-bounds upper=b
+seek-ge b
+----
+a:1
+.