From 349e6208172ce305a0dc753b3f4ffa9435274abe Mon Sep 17 00:00:00 2001 From: Andy Arthur Date: Thu, 8 Sep 2022 14:17:31 -0700 Subject: [PATCH] removed ItemSearchFn, simplified cursor usage --- go/store/prolly/tree/map.go | 46 ++-------- go/store/prolly/tree/node_cursor.go | 111 +++++++++-------------- go/store/prolly/tree/node_cursor_test.go | 27 ++---- 3 files changed, 56 insertions(+), 128 deletions(-) diff --git a/go/store/prolly/tree/map.go b/go/store/prolly/tree/map.go index ef9a0e85dd..c281ed6a12 100644 --- a/go/store/prolly/tree/map.go +++ b/go/store/prolly/tree/map.go @@ -80,12 +80,12 @@ func DiffKeyRangeOrderedTrees[K, V ~[]byte, O Ordering[K]]( return err } } else { - fromStart, err = NewCursorAtItem(ctx, from.NodeStore, from.Root, Item(start), from.searchNode) + fromStart, err = NewCursorAtKey(ctx, from.NodeStore, from.Root, start, from.Order) if err != nil { return err } - toStart, err = NewCursorAtItem(ctx, to.NodeStore, to.Root, Item(start), to.searchNode) + toStart, err = NewCursorAtKey(ctx, to.NodeStore, to.Root, start, to.Order) if err != nil { return err } @@ -102,12 +102,12 @@ func DiffKeyRangeOrderedTrees[K, V ~[]byte, O Ordering[K]]( return err } } else { - fromStop, err = NewCursorAtItem(ctx, from.NodeStore, from.Root, Item(stop), from.searchNode) + fromStop, err = NewCursorAtKey(ctx, from.NodeStore, from.Root, stop, from.Order) if err != nil { return err } - toStop, err = NewCursorAtItem(ctx, to.NodeStore, to.Root, Item(stop), to.searchNode) + toStop, err = NewCursorAtKey(ctx, to.NodeStore, to.Root, stop, to.Order) if err != nil { return err } @@ -180,7 +180,7 @@ func (t StaticMap[K, V, O]) WalkNodes(ctx context.Context, cb NodeCb) error { } func (t StaticMap[K, V, O]) Get(ctx context.Context, query K, cb KeyValueFn[K, V]) (err error) { - cur, err := NewLeafCursorAtItem(ctx, t.NodeStore, t.Root, Item(query), t.searchNode) + cur, err := NewLeafCursorAtKey(ctx, t.NodeStore, t.Root, query, t.Order) if err != nil { return err } @@ -200,7 +200,7 @@ func (t StaticMap[K, V, O]) Get(ctx context.Context, query K, cb KeyValueFn[K, V } func (t StaticMap[K, V, O]) Has(ctx context.Context, query K) (ok bool, err error) { - cur, err := NewLeafCursorAtItem(ctx, t.NodeStore, t.Root, Item(query), t.searchNode) + cur, err := NewLeafCursorAtKey(ctx, t.NodeStore, t.Root, query, t.Order) if err != nil { return false, err } @@ -389,7 +389,7 @@ func (t StaticMap[K, V, O]) getKeyRangeCursors(ctx context.Context, startInclusi return nil, nil, err } } else { - lo, err = NewCursorAtItem(ctx, t.NodeStore, t.Root, Item(startInclusive), t.searchNode) + lo, err = NewCursorAtKey(ctx, t.NodeStore, t.Root, startInclusive, t.Order) if err != nil { return nil, nil, err } @@ -401,7 +401,7 @@ func (t StaticMap[K, V, O]) getKeyRangeCursors(ctx context.Context, startInclusi return nil, nil, err } } else { - hi, err = NewCursorAtItem(ctx, t.NodeStore, t.Root, Item(stopExclusive), t.searchNode) + hi, err = NewCursorAtKey(ctx, t.NodeStore, t.Root, stopExclusive, t.Order) if err != nil { return nil, nil, err } @@ -410,35 +410,9 @@ func (t StaticMap[K, V, O]) getKeyRangeCursors(ctx context.Context, startInclusi return } -// searchNode returns the smallest index where nd[i] >= query -// Adapted from search.Sort to inline comparison. -func (t StaticMap[K, V, O]) searchNode(query Item, nd Node) int { - n := int(nd.Count()) - // Define f(-1) == false and f(n) == true. - // Invariant: f(i-1) == false, f(j) == true. - i, j := 0, n - for i < j { - h := int(uint(i+j) >> 1) // avoid overflow when computing h - less := t.Order.Compare(K(query), K(nd.GetKey(h))) <= 0 - // i ≤ h < j - if !less { - i = h + 1 // preserves f(i-1) == false - } else { - j = h // preserves f(j) == true - } - } - // i == j, f(i-1) == false, and - // f(j) (= f(i)) == true => answer is i. - return i -} - -func (t StaticMap[K, V, O]) CompareItems(left, right Item) int { - return t.Order.Compare(K(left), K(right)) -} - // getOrdinalForKey returns the smallest ordinal position at which the key >= |query|. func (t StaticMap[K, V, O]) GetOrdinalForKey(ctx context.Context, query K) (uint64, error) { - cur, err := NewCursorAtItem(ctx, t.NodeStore, t.Root, Item(query), t.searchNode) + cur, err := NewCursorAtKey(ctx, t.NodeStore, t.Root, query, t.Order) if err != nil { return 0, err } @@ -446,8 +420,6 @@ func (t StaticMap[K, V, O]) GetOrdinalForKey(ctx context.Context, query K) (uint return GetOrdinalOfCursor(cur) } -var _ ItemSearchFn = StaticMap[Item, Item, Ordering[Item]]{}.searchNode - type OrderedTreeIter[K, V ~[]byte] struct { // current tuple location curr *Cursor diff --git a/go/store/prolly/tree/node_cursor.go b/go/store/prolly/tree/node_cursor.go index b1c46d3617..a139fd308c 100644 --- a/go/store/prolly/tree/node_cursor.go +++ b/go/store/prolly/tree/node_cursor.go @@ -43,31 +43,6 @@ type Ordering[K ~[]byte] interface { Compare(left, right K) int } -// SearchForKey returns a SearchFn for |key|. -func SearchForKey[K ~[]byte, O Ordering[K]](key K, order O) SearchFn { - return func(nd Node) (idx int) { - n := int(nd.Count()) - // Define f(-1) == false and f(n) == true. - // Invariant: f(i-1) == false, f(j) == true. - i, j := 0, n - for i < j { - h := int(uint(i+j) >> 1) // avoid overflow when computing h - less := order.Compare(key, K(nd.GetKey(h))) <= 0 - // i ≤ h < j - if !less { - i = h + 1 // preserves f(i-1) == false - } else { - j = h // preserves f(j) == true - } - } - // i == j, f(i-1) == false, and - // f(j) (= f(i)) == true => answer is i. - return i - } -} - -type ItemSearchFn func(item Item, nd Node) (idx int) - func NewCursorAtStart(ctx context.Context, ns NodeStore, nd Node) (cur *Cursor, err error) { cur = &Cursor{nd: nd, nrw: ns} var leaf bool @@ -240,49 +215,29 @@ func NewCursorFromSearchFn(ctx context.Context, ns NodeStore, nd Node, search Se return } -func NewCursorAtItem(ctx context.Context, ns NodeStore, nd Node, item Item, search ItemSearchFn) (cur *Cursor, err error) { - cur = &Cursor{nd: nd, nrw: ns} +func NewLeafCursorAtKey[K ~[]byte, O Ordering[K]](ctx context.Context, ns NodeStore, nd Node, key K, order O) (Cursor, error) { + cur := Cursor{nd: nd, nrw: ns} + for { + // binary search |cur.nd| for |key| + i, j := 0, cur.nd.Count() + for i < j { + h := int(uint(i+j) >> 1) + cmp := order.Compare(key, K(cur.nd.GetKey(h))) + if cmp > 0 { + i = h + 1 + } else { + j = h + } + } + cur.idx = i - cur.idx = search(item, cur.nd) - var leaf bool - leaf, err = cur.isLeaf() - if err != nil { - return nil, err - } - for !leaf { - - // stay in bounds for internal nodes - cur.keepInBounds() - - nd, err = fetchChild(ctx, ns, cur.CurrentRef()) + leaf, err := cur.isLeaf() if err != nil { return cur, err + } else if leaf { + break // done } - parent := cur - cur = &Cursor{nd: nd, parent: parent, nrw: ns} - - cur.idx = search(item, cur.nd) - leaf, err = cur.isLeaf() - if err != nil { - return nil, err - } - } - - return -} - -func NewLeafCursorAtItem(ctx context.Context, ns NodeStore, nd Node, item Item, search ItemSearchFn) (cur Cursor, err error) { - cur = Cursor{nd: nd, parent: nil, nrw: ns} - - cur.idx = search(item, cur.nd) - var leaf bool - leaf, err = cur.isLeaf() - if err != nil { - return cur, err - } - for !leaf { - // stay in bounds for internal nodes cur.keepInBounds() @@ -291,17 +246,33 @@ func NewLeafCursorAtItem(ctx context.Context, ns NodeStore, nd Node, item Item, if err != nil { return cur, err } - - cur.idx = search(item, cur.nd) - leaf, err = cur.isLeaf() - if err != nil { - return cur, err - } } - return cur, nil } +// SearchForKey returns a SearchFn for |key|. +func SearchForKey[K ~[]byte, O Ordering[K]](key K, order O) SearchFn { + return func(nd Node) (idx int) { + n := int(nd.Count()) + // Define f(-1) == false and f(n) == true. + // Invariant: f(i-1) == false, f(j) == true. + i, j := 0, n + for i < j { + h := int(uint(i+j) >> 1) // avoid overflow when computing h + less := order.Compare(key, K(nd.GetKey(h))) <= 0 + // i ≤ h < j + if !less { + i = h + 1 // preserves f(i-1) == false + } else { + j = h // preserves f(j) == true + } + } + // i == j, f(i-1) == false, and + // f(j) (= f(i)) == true => answer is i. + return i + } +} + type LeafSpan struct { Leaves []Node LocalStart int diff --git a/go/store/prolly/tree/node_cursor_test.go b/go/store/prolly/tree/node_cursor_test.go index 726b888a39..ead4515ee6 100644 --- a/go/store/prolly/tree/node_cursor_test.go +++ b/go/store/prolly/tree/node_cursor_test.go @@ -17,7 +17,6 @@ package tree import ( "context" "fmt" - "sort" "testing" "github.com/stretchr/testify/assert" @@ -79,7 +78,7 @@ func testNewCursorAtItem(t *testing.T, count int) { ctx := context.Background() for i := range items { key, value := items[i][0], items[i][1] - cur, err := NewCursorAtItem(ctx, ns, root, key, searchTestTree) + cur, err := NewCursorAtKey(ctx, ns, root, val.Tuple(key), keyDesc) require.NoError(t, err) assert.Equal(t, key, cur.CurrentKey()) assert.Equal(t, value, cur.CurrentValue()) @@ -89,18 +88,11 @@ func testNewCursorAtItem(t *testing.T, count int) { } func testGetOrdinalOfCursor(t *testing.T, count int) { - tuples, d := AscendingUintTuples(count) - - search := func(item Item, nd Node) (idx int) { - return sort.Search(int(nd.count), func(i int) bool { - l, r := val.Tuple(item), val.Tuple(nd.GetKey(i)) - return d.Compare(l, r) <= 0 - }) - } + tuples, desc := AscendingUintTuples(count) ctx := context.Background() ns := NewTestNodeStore() - serializer := message.NewProllyMapSerializer(d, ns.Pool()) + serializer := message.NewProllyMapSerializer(desc, ns.Pool()) chkr, err := newEmptyChunker(ctx, ns, serializer) require.NoError(t, err) @@ -112,7 +104,7 @@ func testGetOrdinalOfCursor(t *testing.T, count int) { assert.NoError(t, err) for i := 0; i < len(tuples); i++ { - curr, err := NewCursorAtItem(ctx, ns, nd, Item(tuples[i][0]), search) + curr, err := NewCursorAtKey(ctx, ns, nd, tuples[i][0], desc) require.NoError(t, err) ord, err := GetOrdinalOfCursor(curr) @@ -121,11 +113,11 @@ func testGetOrdinalOfCursor(t *testing.T, count int) { assert.Equal(t, uint64(i), ord) } - b := val.NewTupleBuilder(d) + b := val.NewTupleBuilder(desc) b.PutUint32(0, uint32(len(tuples))) aboveItem := b.Build(sharedPool) - curr, err := NewCursorAtItem(ctx, ns, nd, Item(aboveItem), search) + curr, err := NewCursorAtKey(ctx, ns, nd, aboveItem, desc) require.NoError(t, err) ord, err := GetOrdinalOfCursor(curr) @@ -171,13 +163,6 @@ var valDesc = val.NewTupleDescriptor( val.Type{Enc: val.Int64Enc, Nullable: true}, ) -func searchTestTree(item Item, nd Node) int { - return sort.Search(int(nd.count), func(i int) bool { - l, r := val.Tuple(item), val.Tuple(nd.GetKey(i)) - return keyDesc.Compare(l, r) <= 0 - }) -} - func randomTupleItemPairs(count int, ns NodeStore) (items [][2]Item) { tups := RandomTuplePairs(count, keyDesc, valDesc, ns) items = make([][2]Item, count)