Skip to content

Commit 1e3ee7b

Browse files
committed
sstable: refactor prop block load
We now pass in an iterator to `load` in order to range all over the keys and values.
1 parent aa61b6f commit 1e3ee7b

File tree

6 files changed

+63
-58
lines changed

6 files changed

+63
-58
lines changed

sstable/colblk/key_value_block.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
package colblk
66

77
import (
8+
"iter"
9+
810
"github.com/cockroachdb/pebble/internal/binfmt"
911
"github.com/cockroachdb/pebble/internal/treeprinter"
1012
)
@@ -124,3 +126,14 @@ func (r *KeyValueBlockDecoder) KeyAt(i int) []byte {
124126
func (r *KeyValueBlockDecoder) ValueAt(i int) []byte {
125127
return r.values.At(i)
126128
}
129+
130+
// All returns an iterator that ranges over all key-value pairs in the block.
131+
func (r *KeyValueBlockDecoder) All() iter.Seq2[[]byte, []byte] {
132+
return func(yield func([]byte, []byte) bool) {
133+
for i := 0; i < r.BlockDecoder().Rows(); i++ {
134+
if !yield(r.KeyAt(i), r.ValueAt(i)) {
135+
return
136+
}
137+
}
138+
}
139+
}

sstable/layout.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -574,11 +574,17 @@ func decodeLayout(comparer *base.Comparer, data []byte, tableFormat TableFormat)
574574
return Layout{}, errors.Wrap(err, "decompressing properties")
575575
}
576576
if tableFormat >= TableFormatPebblev6 {
577-
if err = props.load(decompressedProps, map[string]struct{}{}); err != nil {
577+
var decoder colblk.KeyValueBlockDecoder
578+
decoder.Init(decompressedProps)
579+
if err = props.load(decoder.All(), map[string]struct{}{}); err != nil {
578580
return Layout{}, err
579581
}
580582
} else {
581-
if err = props.loadFromRowBlock(decompressedProps, map[string]struct{}{}); err != nil {
583+
i, err := rowblk.NewRawIter(bytes.Compare, decompressedProps)
584+
if err != nil {
585+
return Layout{}, err
586+
}
587+
if err = props.load(i.All(), map[string]struct{}{}); err != nil {
582588
return Layout{}, err
583589
}
584590
}

sstable/properties.go

Lines changed: 10 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"bytes"
99
"encoding/binary"
1010
"fmt"
11+
"iter"
1112
"maps"
1213
"math"
1314
"reflect"
@@ -274,66 +275,26 @@ func (p *Properties) String() string {
274275
return buf.String()
275276
}
276277

277-
func (p *Properties) loadFromRowBlock(b []byte, deniedUserProperties map[string]struct{}) error {
278-
i, err := rowblk.NewRawIter(bytes.Compare, b)
279-
if err != nil {
280-
return err
281-
}
282-
p.Loaded = make(map[uintptr]struct{})
283-
v := reflect.ValueOf(p).Elem()
284-
285-
for valid := i.First(); valid; valid = i.Next() {
286-
if f, ok := propTagMap[string(i.Key().UserKey)]; ok {
287-
p.Loaded[f.Offset] = struct{}{}
288-
field := v.FieldByIndex(f.Index)
289-
switch f.Type.Kind() {
290-
case reflect.Bool:
291-
field.SetBool(bytes.Equal(i.Value(), propBoolTrue))
292-
case reflect.Uint32:
293-
field.SetUint(uint64(binary.LittleEndian.Uint32(i.Value())))
294-
case reflect.Uint64:
295-
n, _ := binary.Uvarint(i.Value())
296-
field.SetUint(n)
297-
case reflect.String:
298-
field.SetString(intern.Bytes(i.Value()))
299-
default:
300-
panic("not reached")
301-
}
302-
continue
303-
}
304-
if p.UserProperties == nil {
305-
p.UserProperties = make(map[string]string)
306-
}
307-
308-
if _, denied := deniedUserProperties[string(i.Key().UserKey)]; !denied {
309-
p.UserProperties[intern.Bytes(i.Key().UserKey)] = string(i.Value())
310-
}
311-
}
312-
return nil
313-
}
314-
315-
func (p *Properties) load(b []byte, deniedUserProperties map[string]struct{}) error {
316-
var decoder colblk.KeyValueBlockDecoder
317-
decoder.Init(b)
278+
func (p *Properties) load(
279+
i iter.Seq2[[]byte, []byte], deniedUserProperties map[string]struct{},
280+
) error {
318281
p.Loaded = make(map[uintptr]struct{})
319282
v := reflect.ValueOf(p).Elem()
320283

321-
for i := 0; i < decoder.BlockDecoder().Rows(); i++ {
322-
key := decoder.KeyAt(i)
323-
value := decoder.ValueAt(i)
284+
for key, val := range i {
324285
if f, ok := propTagMap[string(key)]; ok {
325286
p.Loaded[f.Offset] = struct{}{}
326287
field := v.FieldByIndex(f.Index)
327288
switch f.Type.Kind() {
328289
case reflect.Bool:
329-
field.SetBool(bytes.Equal(value, propBoolTrue))
290+
field.SetBool(bytes.Equal(val, propBoolTrue))
330291
case reflect.Uint32:
331-
field.SetUint(uint64(binary.LittleEndian.Uint32(value)))
292+
field.SetUint(uint64(binary.LittleEndian.Uint32(val)))
332293
case reflect.Uint64:
333-
n, _ := binary.Uvarint(value)
294+
n, _ := binary.Uvarint(val)
334295
field.SetUint(n)
335296
case reflect.String:
336-
field.SetString(intern.Bytes(value))
297+
field.SetString(intern.Bytes(val))
337298
default:
338299
panic("not reached")
339300
}
@@ -344,7 +305,7 @@ func (p *Properties) load(b []byte, deniedUserProperties map[string]struct{}) er
344305
}
345306

346307
if _, denied := deniedUserProperties[string(key)]; !denied {
347-
p.UserProperties[intern.Bytes(key)] = string(value)
308+
p.UserProperties[intern.Bytes(key)] = string(val)
348309
}
349310
}
350311
return nil

sstable/properties_test.go

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
package sstable
66

77
import (
8+
"bytes"
89
randv1 "math/rand"
910
"path/filepath"
1011
"reflect"
@@ -106,7 +107,9 @@ func TestPropertiesSave(t *testing.T) {
106107
require.NoError(t, e.saveToRowWriter(TableFormatPebblev2, &w))
107108
var props Properties
108109

109-
require.NoError(t, props.loadFromRowBlock(w.Finish(), make(map[string]struct{})))
110+
i, err := rowblk.NewRawIter(bytes.Compare, w.Finish())
111+
require.NoError(t, err)
112+
require.NoError(t, props.load(i.All(), make(map[string]struct{})))
110113
props.Loaded = nil
111114
if diff := pretty.Diff(*e, props); diff != nil {
112115
t.Fatalf("%s", strings.Join(diff, "\n"))
@@ -137,6 +140,8 @@ func BenchmarkPropertiesLoad(b *testing.B) {
137140
p := &Properties{}
138141
for i := 0; i < b.N; i++ {
139142
*p = Properties{}
140-
require.NoError(b, p.loadFromRowBlock(block, nil))
143+
it, err := rowblk.NewRawIter(bytes.Compare, block)
144+
require.NoError(b, err)
145+
require.NoError(b, p.load(it.All(), nil))
141146
}
142147
}

sstable/reader.go

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
package sstable
66

77
import (
8+
"bytes"
89
"cmp"
910
"context"
1011
"fmt"
@@ -446,14 +447,21 @@ func (r *Reader) readMetaindex(
446447
}
447448
r.propertiesBH = bh
448449
if r.tableFormat >= TableFormatPebblev6 {
449-
err = r.Properties.load(b.BlockData(), deniedUserProperties)
450+
var decoder colblk.KeyValueBlockDecoder
451+
decoder.Init(b.BlockData())
452+
if err := r.Properties.load(decoder.All(), deniedUserProperties); err != nil {
453+
return err
454+
}
450455
} else {
451-
err = r.Properties.loadFromRowBlock(b.BlockData(), deniedUserProperties)
456+
i, err := rowblk.NewRawIter(bytes.Compare, b.BlockData())
457+
if err != nil {
458+
return err
459+
}
460+
if err := r.Properties.load(i.All(), deniedUserProperties); err != nil {
461+
return err
462+
}
452463
}
453464
b.Release()
454-
if err != nil {
455-
return err
456-
}
457465
} else {
458466
return errors.New("did not read any value for the properties block in the meta index")
459467
}

sstable/rowblk/rowblk_iter.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"context"
1010
"encoding/binary"
1111
"io"
12+
"iter"
1213
"slices"
1314
"sort"
1415
"unsafe"
@@ -1943,6 +1944,17 @@ func (i *RawIter) Describe(tp treeprinter.Node, fmtKV DescribeKV) {
19431944
}
19441945
}
19451946

1947+
// All returns an iterator that ranges over all key-value pairs in the block.
1948+
func (i *RawIter) All() iter.Seq2[[]byte, []byte] {
1949+
return func(yield func([]byte, []byte) bool) {
1950+
for valid := i.First(); valid; valid = i.Next() {
1951+
if !yield(i.Key().UserKey, i.Value()) {
1952+
return
1953+
}
1954+
}
1955+
}
1956+
}
1957+
19461958
func decodeVarint(ptr unsafe.Pointer) (uint32, unsafe.Pointer) {
19471959
if a := *((*uint8)(ptr)); a < 128 {
19481960
return uint32(a),

0 commit comments

Comments
 (0)