Repository: cockroachdb/pebble
Pull request #3924 (2024-09-06 18:44:58)

colblk: bug fixes and improvements to description of binary format

cockroach-teamcity: 

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/cockroachdb/pebble/3924)
<!-- Reviewable:end -->


Patch:
diff --git a/internal/aligned/aligned.go b/internal/aligned/aligned.go
index eaf4f99d..ca166a2a 100644
--- a/internal/aligned/aligned.go
+++ b/internal/aligned/aligned.go
@@ -14,6 +14,9 @@ import (
 // make([]byte, n) is aligned. In practice it often is, especially for larger n,
 // but small n can often be misaligned.
 func ByteSlice(n int) []byte {
+	if n == 0 {
+		return []byte{}
+	}
 	a := make([]uint64, (n+7)/8)
 	b := unsafe.Slice((*byte)(unsafe.Pointer(&a[0])), n)
 
diff --git a/internal/binfmt/binfmt.go b/internal/binfmt/binfmt.go
index 2eded917..3478eb0b 100644
--- a/internal/binfmt/binfmt.go
+++ b/internal/binfmt/binfmt.go
@@ -28,16 +28,38 @@ func New(data []byte) *Formatter {
 
 // Formatter is a utility for formatting binary data with descriptive comments.
 type Formatter struct {
-	buf   bytes.Buffer
-	lines [][2]string // (binary data, comment) tuples
-	data  []byte
-	off   int
+	buf       bytes.Buffer
+	lines     [][2]string // (binary data, comment) tuples
+	data      []byte
+	off       int
+	anchorOff int
 
 	// config
 	lineWidth       int
 	offsetFormatStr string
 }
 
+// SetAnchorOffset sets the reference point for relative offset calculations to
+// the current offset. Future calls to RelativeOffset() will return an offset
+// relative to the current offset.
+func (f *Formatter) SetAnchorOffset() {
+	f.anchorOff = f.off
+}
+
+// RelativeOffset retrieves the current offset relative to the offset at the
+// last time SetAnchorOffset was called. If SetAnchorOffset was never called,
+// RelativeOffset is equivalent to Offset.
+func (f *Formatter) RelativeOffset() int {
+	return f.off - f.anchorOff
+}
+
+// RelativeData returns the subslice of the original data slice beginning at the
+// offset at which SetAnchorOffset was last called. If SetAnchorOffset was never
+// called, RelativeData is equivalent to Data.
+func (f *Formatter) RelativeData() []byte {
+	return f.data[f.anchorOff:]
+}
+
 // LineWidth sets the Formatter's maximum line width for binary data.
 func (f *Formatter) LineWidth(width int) *Formatter {
 	f.lineWidth = width
diff --git a/sstable/colblk/bitmap.go b/sstable/colblk/bitmap.go
index 3b5aacd6..33dc2230 100644
--- a/sstable/colblk/bitmap.go
+++ b/sstable/colblk/bitmap.go
@@ -384,8 +384,8 @@ func bitmapToBinFormatter(f *binfmt.Formatter, rows int) {
 	if encoding != defaultBitmapEncoding {
 		panic(fmt.Sprintf("unknown bitmap encoding %d", encoding))
 	}
-	if aligned := align(f.Offset(), 8); aligned-f.Offset() != 0 {
-		f.HexBytesln(aligned-f.Offset(), "padding to align to 64-bit boundary")
+	if aligned := align(f.RelativeOffset(), 8); aligned-f.RelativeOffset() != 0 {
+		f.HexBytesln(aligned-f.RelativeOffset(), "padding to align to 64-bit boundary")
 	}
 	bitmapWords := (rows + 63) / 64
 	for i := 0; i < bitmapWords; i++ {
diff --git a/sstable/colblk/data_block.go b/sstable/colblk/data_block.go
index 40c9d610..77ab72ca 100644
--- a/sstable/colblk/data_block.go
+++ b/sstable/colblk/data_block.go
@@ -562,13 +562,24 @@ func (r *DataBlockReader) Init(schema KeySchema, data []byte) {
 	r.maximumKeyLength = binary.LittleEndian.Uint32(data[:dataBlockCustomHeaderSize])
 }
 
-func (r *DataBlockReader) toFormatter(f *binfmt.Formatter) {
+// Describe descirbes the binary format of the data block, assuming f.Offset()
+// is positioned at the beginning of the same data block described by r.
+func (r *DataBlockReader) Describe(f *binfmt.Formatter) {
+	// Set the relative offset. When loaded into memory, the beginning of blocks
+	// are aligned. Padding that ensures alignment is done relative to the
+	// current offset. Setting the relative offset ensures that if we're
+	// describing this block within a larger structure (eg, f.Offset()>0), we
+	// compute padding appropriately assuming the current byte f.Offset() is
+	// aligned.
+	f.SetAnchorOffset()
+
 	f.CommentLine("data block header")
 	f.HexBytesln(4, "maximum key length: %d", r.maximumKeyLength)
 	r.r.headerToBinFormatter(f)
 	for i := 0; i < int(r.r.header.Columns); i++ {
 		r.r.columnToBinFormatter(f, i, int(r.r.header.Rows))
 	}
+	f.HexBytesln(1, "block padding byte")
 }
 
 // DataBlockIter iterates over a columnar data block.
diff --git a/sstable/colblk/index_block.go b/sstable/colblk/index_block.go
index 3f815d4c..29a20eb3 100644
--- a/sstable/colblk/index_block.go
+++ b/sstable/colblk/index_block.go
@@ -142,12 +142,27 @@ func (r *IndexReader) Init(data []byte) {
 // representation.
 func (r *IndexReader) DebugString() string {
 	f := binfmt.New(r.br.data).LineWidth(20)
+	r.Describe(f)
+	return f.String()
+}
+
+// Describe describes the binary format of the index block, assuming f.Offset()
+// is positioned at the beginning of the same index block described by r.
+func (r *IndexReader) Describe(f *binfmt.Formatter) {
+	// Set the relative offset. When loaded into memory, the beginning of blocks
+	// are aligned. Padding that ensures alignment is done relative to the
+	// current offset. Setting the relative offset ensures that if we're
+	// describing this block within a larger structure (eg, f.Offset()>0), we
+	// compute padding appropriately assuming the current byte f.Offset() is
+	// aligned.
+	f.SetAnchorOffset()
+
 	f.CommentLine("index block header")
 	r.br.headerToBinFormatter(f)
 	for i := 0; i < indexBlockColumnCount; i++ {
 		r.br.columnToBinFormatter(f, i, int(r.br.header.Rows))
 	}
-	return f.String()
+	f.HexBytesln(1, "block padding byte")
 }
 
 // IndexIter is an iterator over the block entries in an index block.
diff --git a/sstable/colblk/keyspan.go b/sstable/colblk/keyspan.go
index 2462e863..e9d84585 100644
--- a/sstable/colblk/keyspan.go
+++ b/sstable/colblk/keyspan.go
@@ -231,6 +231,22 @@ func (r *KeyspanReader) Init(data []byte) {
 // representation.
 func (r *KeyspanReader) DebugString() string {
 	f := binfmt.New(r.blockReader.data).LineWidth(20)
+	r.Describe(f)
+	return f.String()
+}
+
+// Describe describes the binary format of the keyspan block, assuming
+// f.Offset() is positioned at the beginning of the same keyspan block described
+// by r.
+func (r *KeyspanReader) Describe(f *binfmt.Formatter) {
+	// Set the relative offset. When loaded into memory, the beginning of blocks
+	// are aligned. Padding that ensures alignment is done relative to the
+	// current offset. Setting the relative offset ensures that if we're
+	// describing this block within a larger structure (eg, f.Offset()>0), we
+	// compute padding appropriately assuming the current byte f.Offset() is
+	// aligned.
+	f.SetAnchorOffset()
+
 	f.CommentLine("keyspan block header")
 	f.HexBytesln(4, "user key count: %d", r.boundaryKeysCount)
 	r.blockReader.headerToBinFormatter(f)
@@ -246,7 +262,7 @@ func (r *KeyspanReader) DebugString() string {
 		}
 		r.blockReader.columnToBinFormatter(f, i, rows)
 	}
-	return f.String()
+	f.HexBytesln(1, "block padding byte")
 }
 
 // searchBoundaryKeys returns the index of the first boundary key greater than
diff --git a/sstable/colblk/prefix_bytes.go b/sstable/colblk/prefix_bytes.go
index de44d8db..03280403 100644
--- a/sstable/colblk/prefix_bytes.go
+++ b/sstable/colblk/prefix_bytes.go
@@ -570,11 +570,11 @@ func prefixBytesToBinFormatter(f *binfmt.Formatter, count int, sliceFormatter fu
 	if sliceFormatter == nil {
 		sliceFormatter = defaultSliceFormatter
 	}
-	pb, _ := DecodePrefixBytes(f.Data(), uint32(f.Offset()), count)
+	pb, _ := DecodePrefixBytes(f.RelativeData(), uint32(f.RelativeOffset()), count)
 	f.CommentLine("PrefixBytes")
 	f.HexBytesln(1, "bundleSize: %d", 1<<pb.bundleShift)
 	f.CommentLine("Offsets table")
-	dataOffset := uint64(f.Offset()) + uint64(uintptr(pb.rawBytes.data)-uintptr(pb.rawBytes.start))
+	dataOffset := uint64(f.RelativeOffset()) + uint64(uintptr(pb.rawBytes.data)-uintptr(pb.rawBytes.start))
 	uintsToBinFormatter(f, pb.rawBytes.slices+1,
 		func(offsetDelta, offsetBase uint64) string {
 			// NB: offsetBase will always be zero for PrefixBytes columns.
diff --git a/sstable/colblk/raw_bytes.go b/sstable/colblk/raw_bytes.go
index 75b1c2d2..38c98937 100644
--- a/sstable/colblk/raw_bytes.go
+++ b/sstable/colblk/raw_bytes.go
@@ -81,12 +81,15 @@ func defaultSliceFormatter(x []byte) string {
 }
 
 func rawBytesToBinFormatter(f *binfmt.Formatter, count int, sliceFormatter func([]byte) string) {
+	if count == 0 {
+		return
+	}
 	if sliceFormatter == nil {
 		sliceFormatter = defaultSliceFormatter
 	}
 
-	rb, _ := DecodeRawBytes(f.Data(), uint32(f.Offset()), count)
-	dataOffset := uint64(f.Offset()) + uint64(uintptr(rb.data)-uintptr(rb.start))
+	rb, _ := DecodeRawBytes(f.RelativeData(), uint32(f.RelativeOffset()), count)
+	dataOffset := uint64(f.RelativeOffset()) + uint64(uintptr(rb.data)-uintptr(rb.start))
 	f.CommentLine("rawbytes")
 	f.CommentLine("offsets table")
 	uintsToBinFormatter(f, count+1, func(offset, base uint64) string {
diff --git a/sstable/colblk/uints.go b/sstable/colblk/uints.go
index 5efb6cb6..e5dcc505 100644
--- a/sstable/colblk/uints.go
+++ b/sstable/colblk/uints.go
@@ -12,6 +12,7 @@ import (
 	"math/bits"
 	"unsafe"
 
+	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/binfmt"
 	"github.com/cockroachdb/pebble/internal/invariants"
 	"golang.org/x/exp/constraints"
@@ -206,8 +207,16 @@ func (b *UintBuilder) Reset() {
 }
 
 // Get gets the value of the provided row index. The provided row must have been
-// Set.
+// Set or the builder must have been initialized with InitWithDefault.
 func (b *UintBuilder) Get(row int) uint64 {
+	// If the UintBuilder is configured to use a zero value for unset rows, it's
+	// possible that the array has not been grown to a size that includes [row].
+	if b.array.n <= row {
+		if invariants.Enabled && !b.useDefault {
+			panic(errors.AssertionFailedf("Get(%d) on UintBuilder with array of size %d", row, b.array.n))
+		}
+		return 0
+	}
 	return b.array.elems.At(row)
 }
 
@@ -391,6 +400,9 @@ func computeMinMax[I constraints.Unsigned](values []I) (I, I) {
 func uintsToBinFormatter(
 	f *binfmt.Formatter, rows int, uintFormatter func(el, base uint64) string,
 ) {
+	if rows == 0 {
+		return
+	}
 	if uintFormatter == nil {
 		uintFormatter = func(v, base uint64) string {
 			if base == 0 {
@@ -417,8 +429,8 @@ func uintsToBinFormatter(
 		return
 	}
 
-	if off := align(f.Offset(), width); off != f.Offset() {
-		f.HexBytesln(off-f.Offset(), "padding (aligning to %d-bit boundary)", width*8)
+	if off := align(f.RelativeOffset(), width); off != f.RelativeOffset() {
+		f.HexBytesln(off-f.RelativeOffset(), "padding (aligning to %d-bit boundary)", width*8)
 	}
 	for i := 0; i < rows; i++ {
 		f.HexBytesln(width, "data[%d] = %s", i, uintFormatter(f.PeekUint(width), base))
diff --git a/sstable/colblk/unsafe_slice.go b/sstable/colblk/unsafe_slice.go
index 956020ad..3f3ab9b0 100644
--- a/sstable/colblk/unsafe_slice.go
+++ b/sstable/colblk/unsafe_slice.go
@@ -57,6 +57,14 @@ var _ Array[uint64] = UnsafeUints{}
 // DecodeUnsafeUints decodes the structure of a slice of uints from a
 // byte slice.
 func DecodeUnsafeUints(b []byte, off uint32, rows int) (_ UnsafeUints, endOffset uint32) {
+	if rows == 0 {
+		// NB: &b[off] is actually pointing beyond the uints serialization.  We
+		// ensure this is always valid at the block-level by appending a
+		// trailing 0x00 block padding byte to all serialized columnar blocks.
+		// This means &b[off] will always point to a valid, allocated byte even
+		// if this is the last column of the block.
+		return makeUnsafeUints(0, unsafe.Pointer(&b[off]), 0), off
+	}
 	encoding := UintEncoding(b[off])
 	if !encoding.IsValid() {
 		panic(errors.AssertionFailedf("invalid encoding 0x%x", b[off:off+1]))


Test patch:
diff --git a/sstable/colblk/data_block_test.go b/sstable/colblk/data_block_test.go
index ec3e1b06..18c2e305 100644
--- a/sstable/colblk/data_block_test.go
+++ b/sstable/colblk/data_block_test.go
@@ -61,7 +61,7 @@ func TestDataBlock(t *testing.T) {
 				block := w.Finish()
 				r.Init(testKeysSchema, block)
 				f := binfmt.New(r.r.data).LineWidth(20)
-				r.toFormatter(f)
+				r.Describe(f)
 
 				return f.String()
 			case "iter":
diff --git a/sstable/colblk/raw_bytes_test.go b/sstable/colblk/raw_bytes_test.go
index b18c5af1..48e61046 100644
--- a/sstable/colblk/raw_bytes_test.go
+++ b/sstable/colblk/raw_bytes_test.go
@@ -8,9 +8,9 @@ import (
 	"bytes"
 	"fmt"
 	"io"
-	"strings"
 	"testing"
 
+	"github.com/cockroachdb/crlib/crstrings"
 	"github.com/cockroachdb/datadriven"
 	"github.com/cockroachdb/pebble/internal/aligned"
 	"github.com/cockroachdb/pebble/internal/binfmt"
@@ -32,7 +32,7 @@ func TestRawBytes(t *testing.T) {
 			td.ScanArgs(t, "offset", &startOffset)
 
 			var count int
-			for _, k := range strings.Split(strings.TrimSpace(td.Input), "\n") {
+			for _, k := range crstrings.Lines(td.Input) {
 				builder.Put([]byte(k))
 				count++
 			}
diff --git a/sstable/colblk/testdata/data_block/bundle_search b/sstable/colblk/testdata/data_block/bundle_search
index feb4fbd9..cddba3d1 100644
--- a/sstable/colblk/testdata/data_block/bundle_search
+++ b/sstable/colblk/testdata/data_block/bundle_search
@@ -509,6 +509,7 @@ finish
 717-718: x 76                                                               # data[65]: v
 # data for column 5
 718-719: x 01                                                               # bitmap encoding
+719-720: x 00                                                               # block padding byte
 
 iter
 seek-ge backache
diff --git a/sstable/colblk/testdata/data_block/external_value b/sstable/colblk/testdata/data_block/external_value
index 861dfbb1..7b1052ac 100644
--- a/sstable/colblk/testdata/data_block/external_value
+++ b/sstable/colblk/testdata/data_block/external_value
@@ -258,6 +258,7 @@ finish
 623-624: x 00                                                               # bitmap encoding
 624-632: b 1100111011111011000001110000000000000000000000000000000000000000 # bitmap word 0
 632-640: b 0000000100000000000000000000000000000000000000000000000000000000 # bitmap summary word 0-63
+640-641: x 00                                                               # block padding byte
 
 # Scan across the block using next.
 iter
diff --git a/sstable/colblk/testdata/data_block/next_prefix b/sstable/colblk/testdata/data_block/next_prefix
index 4e2a7ba7..84e18799 100644
--- a/sstable/colblk/testdata/data_block/next_prefix
+++ b/sstable/colblk/testdata/data_block/next_prefix
@@ -233,6 +233,7 @@ finish
 406-406: x                                                                  # data[19]:
 # data for column 5
 406-407: x 01                                                               # bitmap encoding
+407-408: x 00                                                               # block padding byte
 
 # Scan across the block using next prefix.
 
diff --git a/sstable/colblk/testdata/data_block/simple b/sstable/colblk/testdata/data_block/simple
index 7b35202c..9a24252f 100644
--- a/sstable/colblk/testdata/data_block/simple
+++ b/sstable/colblk/testdata/data_block/simple
@@ -135,6 +135,7 @@ finish
 167-167: x                                                                  # data[6]:
 # data for column 5
 167-168: x 01                                                               # bitmap encoding
+168-169: x 00                                                               # block padding byte
 
 iter
 first
@@ -484,6 +485,7 @@ finish
 331-332: x 61                                                               # data[16]: a
 # data for column 5
 332-333: x 01                                                               # bitmap encoding
+333-334: x 00                                                               # block padding byte
 
 iter
 seek-ge aaa
diff --git a/sstable/colblk/testdata/index_block b/sstable/colblk/testdata/index_block
index 0f1a3ed6..c409b2ff 100644
--- a/sstable/colblk/testdata/index_block
+++ b/sstable/colblk/testdata/index_block
@@ -74,6 +74,7 @@ UnsafeSeparator(Rows()-1) = "bonifide"
 121-121: x                      # data[3]:
 121-124: x 627035               # data[4]: bp5
 124-127: x 627036               # data[5]: bp6
+127-128: x 00                   # block padding byte
 
 iter
 seek-ge zoo
@@ -193,6 +194,7 @@ UnsafeSeparator(Rows()-1) = "coat"
 114-114: x                      # data[2]:
 114-117: x 627034               # data[3]: bp4
 117-117: x                      # data[4]:
+117-118: x 00                   # block padding byte
 
 iter
 first
diff --git a/sstable/colblk/testdata/keyspan_block b/sstable/colblk/testdata/keyspan_block
index 83cf7281..03295cb9 100644
--- a/sstable/colblk/testdata/keyspan_block
+++ b/sstable/colblk/testdata/keyspan_block
@@ -123,6 +123,7 @@ Boundaries: a#0,RANGEDEL — e#inf,RANGEDEL
 72-72: x          # data[4]:
 72-72: x          # data[5]:
 72-72: x          # data[6]:
+72-73: x 00       # block padding byte
 
 # Test iterating over the block's spans.
 
@@ -277,6 +278,7 @@ Boundaries: b#4,RANGEKEYSET — d#inf,RANGEKEYSET
 61-62: x 07               # data[1] = 7 [69 overall]
 # data
 62-69: x 636f636f6e7574   # data[0]: coconut
+69-70: x 00               # block padding byte
 
 iter
 seek-ge a
@@ -374,6 +376,7 @@ Boundaries: b#4,RANGEKEYSET — g#inf,RANGEKEYSET
 # data
 68-75: x 636f636f6e7574 # data[0]: coconut
 75-79: x 74726565       # data[1]: tree
+79-80: x 00             # block padding byte
 
 iter
 seek-ge dog
diff --git a/sstable/colblk/testdata/raw_bytes b/sstable/colblk/testdata/raw_bytes
index 22f70ea2..ce6f9cc7 100644
--- a/sstable/colblk/testdata/raw_bytes
+++ b/sstable/colblk/testdata/raw_bytes
@@ -208,3 +208,9 @@ cat
 orange
 zucchini
 lemon
+
+# Test serializing a RawBytes structure with no rows. The serialization should
+# require zero bytes to represent.
+
+build offset=0
+----
diff --git a/sstable/colblk/testdata/uints b/sstable/colblk/testdata/uints
index 84ab8de9..e75cf320 100644
--- a/sstable/colblk/testdata/uints
+++ b/sstable/colblk/testdata/uints
@@ -13,6 +13,12 @@ size rows=(100)
 ----
 Size(100, 0) = 1
 
+get indices=(0, 10, 99)
+----
+b.Get(0) = 0
+b.Get(10) = 0
+b.Get(99) = 0
+
 finish rows=100
 ----
 0-1: x 00 # encoding: zero
@@ -68,6 +74,13 @@ Size(2, 0) = 1
 Size(1, 0) = 1
 Size(0, 0) = 0
 
+get indices=(0, 5, 6, 7)
+----
+b.Get(0) = 0
+b.Get(5) = 10
+b.Get(6) = 0
+b.Get(7) = 10
+
 # Check width=8 encoding.
 finish rows=8
 ----
@@ -102,7 +115,6 @@ size rows=(9, 8) offset=2
 Size(9, 2) = 22 [20 w/o offset]
 Size(8, 2) = 11 [9 w/o offset]
 
-
 size rows=(9, 8) offset=5
 ----
 Size(9, 5) = 24 [19 w/o offset]
@@ -391,6 +403,18 @@ size rows=5
 ----
 Size(5, 0) = 6
 
+get indices=(0, 1, 2, 3, 4, 5, 6, 7, 20)
+----
+b.Get(0) = 1
+b.Get(1) = 0
+b.Get(2) = 92
+b.Get(3) = 1
+b.Get(4) = 0
+b.Get(5) = 0
+b.Get(6) = 0
+b.Get(7) = 86
+b.Get(20) = 221
+
 finish rows=5
 ----
 0-1: x 01 # encoding: 1b
@@ -486,3 +510,15 @@ finish rows=3 offset=1
 12-16: x 00000000         # data[0] = 0 + 1 = 1
 16-20: x bff90000         # data[1] = 63935 + 1 = 63936
 20-24: x ffffffff         # data[2] = 4294967295 + 1 = 4294967296
+
+# Test serializing with zero rows. The representation should require zero bytes.
+
+init
+----
+
+size rows=(0) offset=0
+----
+Size(0, 0) = 0
+
+finish rows=0 offset=0
+----
diff --git a/sstable/colblk/uints_test.go b/sstable/colblk/uints_test.go
index 8b359175..9ef2ba84 100644
--- a/sstable/colblk/uints_test.go
+++ b/sstable/colblk/uints_test.go
@@ -27,8 +27,10 @@ func TestUintEncoding(t *testing.T) {
 
 func TestUints(t *testing.T) {
 	var b UintBuilder
+	var out bytes.Buffer
 
 	datadriven.RunTest(t, "testdata/uints", func(t *testing.T, td *datadriven.TestData) string {
+		out.Reset()
 		switch td.Cmd {
 		case "init":
 			defaultZero := td.HasArg("default-zero")
@@ -48,12 +50,18 @@ func TestUints(t *testing.T) {
 				b.Set(i, v)
 			}
 			return ""
+		case "get":
+			var indices []int
+			td.ScanArgs(t, "indices", &indices)
+			for _, i := range indices {
+				fmt.Fprintf(&out, "b.Get(%d) = %d\n", i, b.Get(i))
+			}
+			return out.String()
 		case "size":
 			var offset uint32
 			var rowCounts []int
 			td.ScanArgs(t, "rows", &rowCounts)
 			td.MaybeScanArgs(t, "offset", &offset)
-			var out bytes.Buffer
 			for _, rows := range rowCounts {
 				sz := b.Size(rows, offset)
 				if offset > 0 {