Merge pull request #78 from Workiva/palm_query

Added query functionality and some benchmarks.

Author: dustinhiatt-wf

btree/palm/action.go

@@ -122,6 +122,41 @@ func newRemoveAction(keys common.Comparators) *removeAction {
 	}
 }
 
+type applyAction struct {
+	fn          func(common.Comparator) bool
+	start, stop common.Comparator
+	completer   *sync.WaitGroup
+}
+
+func (aa *applyAction) operation() operation {
+	return apply
+}
+
+func (aa *applyAction) nodes() []*node {
+	return nil
+}
+
+func (aa *applyAction) addNode(i int64, n *node) {}
+
+func (aa *applyAction) keys() common.Comparators {
+	return nil
+}
+
+func (aa *applyAction) complete() {
+	aa.completer.Done()
+}
+
+func newApplyAction(fn func(common.Comparator) bool, start, stop common.Comparator) *applyAction {
+	aa := &applyAction{
+		fn:        fn,
+		start:     start,
+		stop:      stop,
+		completer: new(sync.WaitGroup),
+	}
+	aa.completer.Add(1)
+	return aa
+}
+
 func minUint64(choices ...uint64) uint64 {
 	min := choices[0]
 	for i := 1; i < len(choices); i++ {

btree/palm/interface.go

@@ -39,7 +39,8 @@ BenchmarkGet-8	 						 2000000	       629 ns/op
 BenchmarkBulkGet-8	    					5000	    223249 ns/op
 BenchmarkDelete-8	  					  500000	      2421 ns/op
 BenchmarkBulkDelete-8	     				 500	   2790461 ns/op
-
+BenchmarkFindQuery-8	 				 1000000	      1166 ns/op
+BenchmarkExecuteQuery-8	   				   10000	   1290732 ns/op
 
 */
 package palm
@@ -58,6 +59,10 @@ type BTree interface {
 	Get(...common.Comparator) common.Comparators
 	// Len returns the number of items in the tree.
 	Len() uint64
+	// Query will return a list of Comparators that fall within the
+	// provided start and stop Comparators.  Start is inclusive while
+	// stop is exclusive, ie [start, stop).
+	Query(start, stop common.Comparator) common.Comparators
 	// Dispose will clean up any resources used by this tree.  This
 	// must be called to prevent a memory leak.
 	Dispose()

btree/palm/node.go

@@ -195,6 +195,7 @@ func (n *node) splitLeaf(i, capacity uint64) (common.Comparator, *node, *node) {
 		keys:   rightKeys,
 		nodes:  newNodes(uint64(cap(n.nodes.list))),
 		isLeaf: true,
+		right:  n.right,
 	}
 	n.right = nn
 	return key, n, nn

btree/palm/tree.go

@@ -32,6 +32,7 @@ const (
 	get operation = iota
 	add
 	remove
+	apply
 )
 
 const multiThreadAt = 1000 // number of keys before we multithread lookups
@@ -89,6 +90,12 @@ func (ptree *ptree) checkAndRun(action action) {
 				} else {
 					ptree.operationRunner(interfaces{action}, false)
 				}
+			case apply:
+				q := action.(*applyAction)
+				n := getParent(ptree.root, q.start)
+				ptree.apply(n, q)
+				q.complete()
+				ptree.reset()
 			}
 		} else {
 			ptree.actions.Put(action)
@@ -154,14 +161,33 @@ func (ptree *ptree) fetchKeys(xns interfaces, inParallel bool) (map[*node][]*key
 				deleteOperations[n] = append(deleteOperations[n], &keyBundle{key: action.keys()[i]})
 			}
 			toComplete = append(toComplete, action)
-		case get:
+		case get, apply:
 			action.complete()
 		}
 	}
 
 	return writeOperations, deleteOperations, toComplete
 }
 
+func (ptree *ptree) apply(n *node, aa *applyAction) {
+	i := n.search(aa.start)
+	if i == n.keys.len() { // nothing to apply against
+		return
+	}
+
+	var k common.Comparator
+	for n != nil {
+		for j := i; j < n.keys.len(); j++ {
+			k = n.keys.byPosition(j)
+			if aa.stop.Compare(k) < 1 || !aa.fn(k) {
+				return
+			}
+		}
+		n = n.right
+		i = 0
+	}
+}
+
 func (ptree *ptree) fetchKeysInSerial(xns interfaces) {
 	for _, ifc := range xns {
 		action := ifc.(action)
@@ -181,6 +207,9 @@ func (ptree *ptree) fetchKeysInSerial(xns interfaces) {
 						action.keys()[i] = k
 					}
 				}
+			case apply:
+				q := action.(*applyAction)
+				ptree.apply(n, q)
 			}
 		}
 	}
@@ -245,6 +274,9 @@ func (ptree *ptree) fetchKeysInParallel(xns []interface{}) {
 							action.keys()[j] = k
 						}
 					}
+				case apply:
+					q := action.(*applyAction)
+					ptree.apply(n, q)
 				}
 			}
 			wg.Done()
@@ -418,6 +450,20 @@ func (ptree *ptree) Len() uint64 {
 	return atomic.LoadUint64(&ptree.number)
 }
 
+// Query will return a list of Comparators that fall within the
+// provided start and stop Comparators.  Start is inclusive while
+// stop is exclusive, ie [start, stop).
+func (ptree *ptree) Query(start, stop common.Comparator) common.Comparators {
+	cmps := make(common.Comparators, 0, 32)
+	aa := newApplyAction(func(cmp common.Comparator) bool {
+		cmps = append(cmps, cmp)
+		return true
+	}, start, stop)
+	ptree.checkAndRun(aa)
+	aa.completer.Wait()
+	return cmps
+}
+
 // Dispose will clean up any resources used by this tree.  This
 // must be called to prevent a memory leak.
 func (ptree *ptree) Dispose() {

btree/palm/tree_test.go

@@ -379,6 +379,57 @@ func TestInsertAndDeletesWithSplits(t *testing.T) {
 	assert.Equal(t, keys2, tree.Get(keys2...))
 }
 
+func TestSimpleQuery(t *testing.T) {
+	tree := newTree(3, 3)
+	defer tree.Dispose()
+	m1 := mockKey(1)
+	tree.Insert(m1)
+
+	result := tree.Query(mockKey(0), mockKey(5))
+	assert.Equal(t, common.Comparators{m1}, result)
+
+	result = tree.Query(mockKey(0), mockKey(1))
+	assert.Len(t, result, 0)
+
+	result = tree.Query(mockKey(2), mockKey(10))
+	assert.Len(t, result, 0)
+
+	result = tree.Query(mockKey(1), mockKey(10))
+	assert.Equal(t, common.Comparators{m1}, result)
+}
+
+func TestMultipleQuery(t *testing.T) {
+	tree := newTree(3, 3)
+	defer tree.Dispose()
+	m1 := mockKey(1)
+	m2 := mockKey(5)
+	tree.Insert(m1, m2)
+
+	result := tree.Query(mockKey(0), mockKey(10))
+	assert.Equal(t, common.Comparators{m1, m2}, result)
+
+	result = tree.Query(mockKey(1), mockKey(5))
+	assert.Equal(t, common.Comparators{m1}, result)
+
+	result = tree.Query(mockKey(6), mockKey(10))
+	assert.Len(t, result, 0)
+
+	result = tree.Query(mockKey(5), mockKey(10))
+	assert.Equal(t, common.Comparators{m2}, result)
+}
+
+func TestCrossNodeQuery(t *testing.T) {
+	tree := newTree(3, 3)
+	defer tree.Dispose()
+	keys := generateKeys(100)
+	tree.Insert(keys...)
+
+	result := tree.Query(mockKey(0), mockKey(len(keys)))
+	if !assert.Equal(t, keys, result) {
+		tree.print(getConsoleLogger())
+	}
+}
+
 func BenchmarkReadAndWrites(b *testing.B) {
 	numItems := 1000
 	keys := make([]common.Comparators, 0, b.N)
@@ -475,9 +526,9 @@ func BenchmarkGet(b *testing.B) {
 }
 
 func BenchmarkBulkGet(b *testing.B) {
-	numItems := 1000
+	numItems := b.N
 	keys := generateRandomKeys(numItems)
-	tree := newTree(16, 8)
+	tree := newTree(8, 8)
 	tree.Insert(keys...)
 
 	b.ResetTimer()
@@ -516,3 +567,29 @@ func BenchmarkBulkDelete(b *testing.B) {
 		trees[i].Delete(keys...)
 	}
 }
+
+func BenchmarkFindQuery(b *testing.B) {
+	numItems := b.N
+	keys := generateKeys(numItems)
+	tree := newTree(8, 8)
+	tree.Insert(keys...)
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		tree.Query(mockKey(numItems/2), mockKey(numItems/2+1))
+	}
+}
+
+func BenchmarkExecuteQuery(b *testing.B) {
+	numItems := b.N
+	keys := generateKeys(numItems)
+	tree := newTree(8, 8)
+	tree.Insert(keys...)
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		tree.Query(mockKey(0), mockKey(numItems))
+	}
+}