-
Notifications
You must be signed in to change notification settings - Fork 679
/
disk_tree.go
197 lines (168 loc) · 4.83 KB
/
disk_tree.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
// _ _
// __ _____ __ ___ ___ __ _| |_ ___
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
// \ V V / __/ (_| |\ V /| | (_| | || __/
// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
//
// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
//
// CONTACT: hello@weaviate.io
//
package segmentindex
import (
"bytes"
"errors"
"fmt"
"io"
"github.com/weaviate/weaviate/entities/lsmkv"
"github.com/weaviate/weaviate/usecases/byteops"
)
// DiskTree is a read-only wrapper around a marshalled index search tree, which
// can be used for reading, but cannot change the underlying structure. It is
// thus perfectly suited as an index for an (immutable) LSM disk segment, but
// pretty much useless for anything else
type DiskTree struct {
data []byte
}
type dtNode struct {
key []byte
startPos uint64
endPos uint64
leftChild int64
rightChild int64
}
func NewDiskTree(data []byte) *DiskTree {
return &DiskTree{
data: data,
}
}
func (t *DiskTree) Get(key []byte) (Node, error) {
if len(t.data) == 0 {
return Node{}, lsmkv.NotFound
}
var out Node
rw := byteops.NewReadWriter(t.data)
// jump to the buffer until the node with _key_ is found or return a NotFound error.
// This function avoids allocations by reusing the same buffer for all keys and avoids memory reads by only
// extracting the necessary pieces of information while skipping the rest
NodeKeyBuffer := make([]byte, len(key))
for {
// detect if there is no node with the wanted key.
if rw.Position+4 > uint64(len(t.data)) || rw.Position+4 < 4 {
return out, lsmkv.NotFound
}
keyLen := rw.ReadUint32()
if int(keyLen) > len(NodeKeyBuffer) {
NodeKeyBuffer = make([]byte, int(keyLen))
} else if int(keyLen) < len(NodeKeyBuffer) {
NodeKeyBuffer = NodeKeyBuffer[:keyLen]
}
_, err := rw.CopyBytesFromBuffer(uint64(keyLen), NodeKeyBuffer)
if err != nil {
return out, fmt.Errorf("copy node key: %w", err)
}
keyEqual := bytes.Compare(key, NodeKeyBuffer)
if keyEqual == 0 {
out.Key = NodeKeyBuffer
out.Start = rw.ReadUint64()
out.End = rw.ReadUint64()
return out, nil
} else if keyEqual < 0 {
rw.MoveBufferPositionForward(2 * 8) // jump over start+end position
rw.Position = rw.ReadUint64() // left child
} else {
rw.MoveBufferPositionForward(3 * 8) // jump over start+end position and left child
rw.Position = rw.ReadUint64() // right child
}
}
}
func (t *DiskTree) readNodeAt(offset int64) (dtNode, error) {
retNode, _, err := t.readNode(t.data[offset:])
return retNode, err
}
func (t *DiskTree) readNode(in []byte) (dtNode, int, error) {
var out dtNode
// in buffer needs at least 36 bytes of data:
// 4bytes for key length, 32bytes for position and children
if len(in) < 36 {
return out, 0, io.EOF
}
rw := byteops.NewReadWriter(in)
keyLen := uint64(rw.ReadUint32())
copiedBytes, err := rw.CopyBytesFromBuffer(keyLen, nil)
if err != nil {
return out, int(rw.Position), fmt.Errorf("copy node key: %w", err)
}
out.key = copiedBytes
out.startPos = rw.ReadUint64()
out.endPos = rw.ReadUint64()
out.leftChild = int64(rw.ReadUint64())
out.rightChild = int64(rw.ReadUint64())
return out, int(rw.Position), nil
}
func (t *DiskTree) Seek(key []byte) (Node, error) {
if len(t.data) == 0 {
return Node{}, lsmkv.NotFound
}
return t.seekAt(0, key)
}
func (t *DiskTree) seekAt(offset int64, key []byte) (Node, error) {
node, err := t.readNodeAt(offset)
if err != nil {
return Node{}, err
}
self := Node{
Key: node.key,
Start: node.startPos,
End: node.endPos,
}
if bytes.Equal(key, node.key) {
return self, nil
}
if bytes.Compare(key, node.key) < 0 {
if node.leftChild < 0 {
return self, nil
}
left, err := t.seekAt(node.leftChild, key)
if err == nil {
return left, nil
}
if errors.Is(err, lsmkv.NotFound) {
return self, nil
}
return Node{}, err
} else {
if node.rightChild < 0 {
return Node{}, lsmkv.NotFound
}
return t.seekAt(node.rightChild, key)
}
}
// AllKeys is a relatively expensive operation as it basically does a full disk
// read of the index. It is meant for one of operations, such as initializing a
// segment where we need access to all keys, e.g. to build a bloom filter. This
// should not run at query time.
//
// The binary tree is traversed in Level-Order so keys have no meaningful
// order. Do not use this method if an In-Order traversal is required, but only
// for use cases who don't require a specific order, such as building a
// bloom filter.
func (t *DiskTree) AllKeys() ([][]byte, error) {
var out [][]byte
bufferPos := 0
for {
node, readLength, err := t.readNode(t.data[bufferPos:])
bufferPos += readLength
if errors.Is(err, io.EOF) {
break
}
if err != nil {
return nil, err
}
out = append(out, node.key)
}
return out, nil
}
func (t *DiskTree) Size() int {
return len(t.data)
}