Line data Source code
1 : // Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use
2 : // of this source code is governed by a BSD-style license that can be found in
3 : // the LICENSE file.
4 :
5 : package colblk
6 :
7 : import (
8 : "bytes"
9 :
10 : "github.com/cockroachdb/pebble/internal/base"
11 : "github.com/cockroachdb/pebble/internal/binfmt"
12 : "github.com/cockroachdb/pebble/sstable/block"
13 : )
14 :
15 : const indexBlockCustomHeaderSize = 0
16 :
17 : // IndexBlockWriter writes columnar index blocks. The writer is used for both
18 : // first-level and second-level index blocks. The index block schema consists of
19 : // three primary columns:
20 : // - Separators: a user key that is ≥ the largest user key in the
21 : // corresponding entry, and ≤ the smallest user key in the next entry.
22 : // Note that this allows consecutive separators to be equal. This is
23 : // possible when snapshots required we preserve duplicate user keys at
24 : // different sequence numbers.
25 : // - Offsets: the offset of the end of the corresponding block.
26 : // - Lengths: the length in bytes of the corresponding block.
27 : // - Block properties: a slice encoding arbitrary user-defined block
28 : // properties.
29 : //
30 : // TODO(jackson): Consider splitting separators into prefixes and suffixes (even
31 : // without user-defined columns). This would allow us to use prefix compression
32 : // for the prefix. Separators should typically be suffixless unless two KVs with
33 : // the same prefix straddle a block boundary. We would need to use a buffer to
34 : // materialize the separator key when we need to use it outside the context of
35 : // seeking within the block.
36 : type IndexBlockWriter struct {
37 : separators RawBytesBuilder
38 : offsets UintBuilder
39 : lengths UintBuilder
40 : blockProperties RawBytesBuilder
41 : rows int
42 : enc blockEncoder
43 : }
44 :
45 : const (
46 : indexBlockColumnSeparator = iota
47 : indexBlockColumnOffsets
48 : indexBlockColumnLengths
49 : indexBlockColumnBlockProperties
50 : indexBlockColumnCount
51 : )
52 :
53 : // Init initializes the index block writer.
54 1 : func (w *IndexBlockWriter) Init(blockProperties ...ColumnWriter) {
55 1 : w.separators.Init()
56 1 : w.offsets.Init()
57 1 : w.lengths.InitWithDefault()
58 1 : w.blockProperties.Init()
59 1 : w.rows = 0
60 1 : }
61 :
62 : // Reset resets the index block writer to its initial state, retaining buffers.
63 0 : func (w *IndexBlockWriter) Reset() {
64 0 : w.separators.Reset()
65 0 : w.offsets.Reset()
66 0 : w.lengths.Reset()
67 0 : w.blockProperties.Reset()
68 0 : w.rows = 0
69 0 : w.enc.reset()
70 0 : }
71 :
72 : // Rows returns the number of entries in the index block so far.
73 1 : func (w *IndexBlockWriter) Rows() int {
74 1 : return w.rows
75 1 : }
76 :
77 : // AddBlockHandle adds a new separator and end offset of a data block to the
78 : // index block. Add returns the index of the row.
79 : //
80 : // AddBlockHandle should only be used for first-level index blocks.
81 : func (w *IndexBlockWriter) AddBlockHandle(
82 : separator []byte, handle block.Handle, blockProperties []byte,
83 1 : ) int {
84 1 : idx := w.rows
85 1 : w.separators.Put(separator)
86 1 : w.offsets.Set(w.rows, handle.Offset)
87 1 : w.lengths.Set(w.rows, handle.Length)
88 1 : w.blockProperties.Put(blockProperties)
89 1 : w.rows++
90 1 : return idx
91 1 : }
92 :
93 : // UnsafeSeparator returns the separator of the i'th entry.
94 1 : func (w *IndexBlockWriter) UnsafeSeparator(i int) []byte {
95 1 : return w.separators.UnsafeGet(i)
96 1 : }
97 :
98 : // Size returns the size of the pending index block.
99 1 : func (w *IndexBlockWriter) Size() int {
100 1 : off := blockHeaderSize(indexBlockColumnCount, indexBlockCustomHeaderSize)
101 1 : off = w.separators.Size(w.rows, off)
102 1 : off = w.offsets.Size(w.rows, off)
103 1 : off = w.lengths.Size(w.rows, off)
104 1 : off = w.blockProperties.Size(w.rows, off)
105 1 : off++
106 1 : return int(off)
107 1 : }
108 :
109 : // Finish serializes the pending index block.
110 1 : func (w *IndexBlockWriter) Finish() []byte {
111 1 : w.enc.init(w.Size(), Header{
112 1 : Version: Version1,
113 1 : Columns: indexBlockColumnCount,
114 1 : Rows: uint32(w.rows),
115 1 : }, indexBlockCustomHeaderSize)
116 1 : w.enc.encode(w.rows, &w.separators)
117 1 : w.enc.encode(w.rows, &w.offsets)
118 1 : w.enc.encode(w.rows, &w.lengths)
119 1 : w.enc.encode(w.rows, &w.blockProperties)
120 1 : return w.enc.finish()
121 1 : }
122 :
123 : // An IndexReader reads columnar index blocks.
124 : type IndexReader struct {
125 : separators RawBytes
126 : offsets UnsafeUints
127 : lengths UnsafeUints // only used for second-level index blocks
128 : blockProps RawBytes
129 : br BlockReader
130 : }
131 :
132 : // Init initializes the index reader with the given serialized index block.
133 1 : func (r *IndexReader) Init(data []byte) {
134 1 : r.br.Init(data, indexBlockCustomHeaderSize)
135 1 : r.separators = r.br.RawBytes(indexBlockColumnSeparator)
136 1 : r.offsets = r.br.Uints(indexBlockColumnOffsets)
137 1 : r.lengths = r.br.Uints(indexBlockColumnLengths)
138 1 : r.blockProps = r.br.RawBytes(indexBlockColumnBlockProperties)
139 1 : }
140 :
141 : // DebugString prints a human-readable explanation of the keyspan block's binary
142 : // representation.
143 1 : func (r *IndexReader) DebugString() string {
144 1 : f := binfmt.New(r.br.data).LineWidth(20)
145 1 : f.CommentLine("index block header")
146 1 : r.br.headerToBinFormatter(f)
147 1 : for i := 0; i < indexBlockColumnCount; i++ {
148 1 : r.br.columnToBinFormatter(f, i, int(r.br.header.Rows))
149 1 : }
150 1 : return f.String()
151 : }
152 :
153 : // IndexIter is an iterator over the block entries in an index block.
154 : type IndexIter struct {
155 : r *IndexReader
156 : n int
157 : row int
158 :
159 : h block.BufferHandle
160 : allocReader IndexReader
161 : }
162 :
163 : // Assert that IndexIter satisfies the block.IndexBlockIterator interface.
164 : var _ block.IndexBlockIterator = (*IndexIter)(nil)
165 :
166 : // Init initializes an index iterator from the provided reader.
167 1 : func (i *IndexIter) Init(r *IndexReader) {
168 1 : *i = IndexIter{r: r, n: int(r.br.header.Rows)}
169 1 : }
170 :
171 : // InitHandle initializes an iterator from the provided block handle.
172 : func (i *IndexIter) InitHandle(
173 : cmp base.Compare, split base.Split, block block.BufferHandle, transforms block.IterTransforms,
174 0 : ) error {
175 0 : // TODO(jackson): Handle the transforms.
176 0 :
177 0 : // TODO(jackson): If block.h != nil, use a *IndexReader that's allocated
178 0 : // when the block is loaded into the block cache. On cache hits, this will
179 0 : // reduce the amount of setup necessary to use an iterator. (It's relatively
180 0 : // common to open an iterator and perform just a few seeks, so avoiding the
181 0 : // overhead can be material.)
182 0 : i.h.Release()
183 0 : i.h = block
184 0 : i.allocReader.Init(i.h.Get())
185 0 : i.r = &i.allocReader
186 0 : return nil
187 0 : }
188 :
189 : // RowIndex returns the index of the block entry at the iterator's current
190 : // position.
191 0 : func (i *IndexIter) RowIndex() int {
192 0 : return i.row
193 0 : }
194 :
195 : // ResetForReuse resets the IndexIter for reuse, retaining buffers to avoid
196 : // future allocations.
197 0 : func (i *IndexIter) ResetForReuse() IndexIter {
198 0 : return IndexIter{ /* nothing to retain */ }
199 0 : }
200 :
201 : // Valid returns true if the iterator is currently positioned at a valid block
202 : // handle.
203 0 : func (i *IndexIter) Valid() bool {
204 0 : return 0 <= i.row && i.row < int(i.r.br.header.Rows)
205 0 : }
206 :
207 : // Invalidate invalidates the block iterator, removing references to the block
208 : // it was initialized with.
209 0 : func (i *IndexIter) Invalidate() {
210 0 : i.r = nil
211 0 : i.h = block.BufferHandle{}
212 0 : }
213 :
214 : // IsDataInvalidated returns true when the iterator has been invalidated
215 : // using an Invalidate call. NB: this is different from Valid.
216 0 : func (i *IndexIter) IsDataInvalidated() bool {
217 0 : return i.r == nil
218 0 : }
219 :
220 : // Handle returns the underlying block buffer handle, if the iterator was
221 : // initialized with one.
222 0 : func (i *IndexIter) Handle() block.BufferHandle {
223 0 : return i.h
224 0 : }
225 :
226 : // Separator returns the separator at the iterator's current position. The
227 : // iterator must be positioned at a valid row.
228 0 : func (i *IndexIter) Separator() []byte {
229 0 : return i.r.separators.At(i.row)
230 0 : }
231 :
232 : // BlockHandleWithProperties decodes the block handle with any encoded
233 : // properties at the iterator's current position.
234 1 : func (i *IndexIter) BlockHandleWithProperties() (block.HandleWithProperties, error) {
235 1 : return block.HandleWithProperties{
236 1 : Handle: block.Handle{
237 1 : Offset: i.r.offsets.At(i.row),
238 1 : Length: i.r.lengths.At(i.row),
239 1 : },
240 1 : Props: i.r.blockProps.At(i.row),
241 1 : }, nil
242 1 : }
243 :
244 : // SeekGE seeks the index iterator to the first block entry with a separator key
245 : // greater or equal to the given key. It returns false if the seek key is
246 : // greater than all index block separators.
247 1 : func (i *IndexIter) SeekGE(key []byte) bool {
248 1 : // Define f(-1) == false and f(upper) == true.
249 1 : // Invariant: f(index-1) == false, f(upper) == true.
250 1 : index, upper := 0, i.n
251 1 : for index < upper {
252 1 : h := int(uint(index+upper) >> 1) // avoid overflow when computing h
253 1 : // index ≤ h < upper
254 1 :
255 1 : // TODO(jackson): Is Bytes.At or Bytes.Slice(Bytes.Offset(h),
256 1 : // Bytes.Offset(h+1)) faster in this code?
257 1 : c := bytes.Compare(key, i.r.separators.At(h))
258 1 : if c > 0 {
259 1 : index = h + 1 // preserves f(index-1) == false
260 1 : } else {
261 1 : upper = h // preserves f(upper) == true
262 1 : }
263 : }
264 : // index == upper, f(index-1) == false, and f(upper) (= f(index)) == true => answer is index.
265 1 : i.row = index
266 1 : return index < i.n
267 : }
268 :
269 : // First seeks index iterator to the first block entry. It returns false if the
270 : // index block is empty.
271 1 : func (i *IndexIter) First() bool {
272 1 : i.row = 0
273 1 : return i.n > 0
274 1 : }
275 :
276 : // Last seeks index iterator to the last block entry. It returns false if the
277 : // index block is empty.
278 1 : func (i *IndexIter) Last() bool {
279 1 : i.row = i.n - 1
280 1 : return i.n > 0
281 1 : }
282 :
283 : // Next steps the index iterator to the next block entry. It returns false if
284 : // the index block is exhausted.
285 1 : func (i *IndexIter) Next() bool {
286 1 : i.row++
287 1 : return i.row < i.n
288 1 : }
289 :
290 : // Prev steps the index iterator to the previous block entry. It returns false
291 : // if the index block is exhausted.
292 1 : func (i *IndexIter) Prev() bool {
293 1 : i.row--
294 1 : return i.row >= 0
295 1 : }
296 :
297 : // Close closes the iterator, releasing any resources it holds.
298 0 : func (i *IndexIter) Close() error {
299 0 : i.h.Release()
300 0 : *i = IndexIter{}
301 0 : return nil
302 0 : }
|