Repository: cockroachdb/pebble
Pull request #3827 (2024-08-06 00:16:15)

colblk: add PrefixBytes.AppendAt,SetNextAt
**colblk: add PrefixBytes.AppendAt**

Add a new PrefixBytes.AppendAt method that appends the slice at the i'th index
to the provided destination slice. This commit also adjusts the
BenchmarkPrefixBytes microbenchmark to access slices sequentially.

```
goos: darwin
goarch: arm64
pkg: github.com/cockroachdb/pebble/sstable/colblk
                                     │ prefixbytes-old.txt │
                                     │       sec/op        │
PrefixBytes/alphaLen=2/iteration-10            22.61n ± 0%
PrefixBytes/alphaLen=5/iteration-10            23.18n ± 2%
PrefixBytes/alphaLen=26/iteration-10           23.62n ± 0%
geomean                                        23.13n
```

**colblk: use pointer receiver for RawBytes, PrefixBytes**
```
goos: darwin
goarch: arm64
pkg: github.com/cockroachdb/pebble/sstable/colblk
                                     │ prefixbytes-old.txt │         prefixbytes-new.txt         │
                                     │       sec/op        │   sec/op     vs base                │
PrefixBytes/alphaLen=2/iteration-10            22.61n ± 0%   18.99n ± 1%  -16.03% (p=0.000 n=20)
PrefixBytes/alphaLen=5/iteration-10            23.18n ± 2%   19.68n ± 4%  -15.12% (p=0.000 n=20)
PrefixBytes/alphaLen=26/iteration-10           23.62n ± 0%   19.68n ± 0%  -16.68% (p=0.000 n=20)
geomean                                        23.13n        19.44n       -15.95%
```

**colblk: inline SharedPrefix in PrefixBytes.AppendAt**

```
goos: darwin
goarch: arm64
pkg: github.com/cockroachdb/pebble/sstable/colblk
                                     │ prefixbytes-old.txt │        prefixbytes-new.txt         │
                                     │       sec/op        │   sec/op     vs base               │
PrefixBytes/alphaLen=2/iteration-10            18.99n ± 1%   17.73n ± 5%  -6.61% (p=0.000 n=20)
PrefixBytes/alphaLen=5/iteration-10            19.68n ± 4%   18.65n ± 1%  -5.18% (p=0.000 n=20)
PrefixBytes/alphaLen=26/iteration-10           19.68n ± 0%   18.80n ± 2%  -4.42% (p=0.000 n=20)
geomean                                        19.44n        18.39n       -5.41%
```

**colblk: add PrefixBytes.SetNextAt**

Add a PrefixBytes.SetNextAt that may be used when materializing the next slice
in a PrefixBytes into a buffer known to contain the previous slice. SetNextAt
can take advantage of the previous slice to always avoid copying the shared
prefix, often avoid copying the bundle prefix and occasionally avoid copying
the row suffix.

```
goos: darwin
goarch: arm64
pkg: github.com/cockroachdb/pebble/sstable/colblk
                                     │ prefixbytes-old.txt │         prefixbytes-new.txt         │
                                     │       sec/op        │   sec/op     vs base                │
PrefixBytes/alphaLen=2/iteration-10           17.735n ± 5%   7.811n ± 1%  -55.96% (p=0.000 n=20)
PrefixBytes/alphaLen=5/iteration-10           18.655n ± 1%   7.717n ± 1%  -58.63% (p=0.000 n=20)
PrefixBytes/alphaLen=26/iteration-10          18.805n ± 2%   7.600n ± 1%  -59.58% (p=0.000 n=20)
geomean                                        18.39n        7.709n       -58.08%
```

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/3827)
<!-- Reviewable:end -->

jbowens: TFTR!

Patch:
diff --git a/sstable/colblk/prefix_bytes.go b/sstable/colblk/prefix_bytes.go
index 559ee96f..1d638196 100644
--- a/sstable/colblk/prefix_bytes.go
+++ b/sstable/colblk/prefix_bytes.go
@@ -238,10 +238,55 @@ func (b PrefixBytes) At(i int) []byte {
 	return slices.Concat(b.SharedPrefix(), b.RowBundlePrefix(i), b.RowSuffix(i))
 }
 
+// AppendAt appends the i'th []byte slice in the PrefixBytes onto the provided
+// bytes slice.
+func (b *PrefixBytes) AppendAt(dst []byte, i int) []byte {
+	// Inline dst = append(dst, b.SharedPrefix()...)
+	dst = append(dst, unsafe.Slice((*byte)(b.rawBytes.data), b.rawBytes.offsets.At(0))...)
+	dst = append(dst, b.RowBundlePrefix(i)...)
+	return append(dst, b.RowSuffix(i)...)
+}
+
+// SetNextAt updates the provided buf byte slice to hold the i'th []byte slice
+// in the PrefixBytes. SetNextAt requires the provided buf to currently hold the
+// i-1'th []byte slice in the PrefixBytes. If the i'th slice does not fit in
+// buf, SetNextAt allocates a new slice.
+func (b *PrefixBytes) SetNextAt(buf []byte, i int) []byte {
+	if invariants.Enabled {
+		if x := b.At(i - 1); !bytes.Equal(buf, x) {
+			panic(errors.AssertionFailedf("buf (%q) does not hold the previous slice (%q)", buf, x))
+		}
+	}
+
+	bundleOffsetIndex := b.bundleOffsetIndexForRow(i)
+	rowSuffixIndex := b.rowSuffixIndex(i)
+	if rowSuffixIndex == bundleOffsetIndex+1 {
+		// The key i is the first key in a new bundle. We need to replace the
+		// bundle prefix in buf. We can copy the entirety of the bundle prefix
+		// and the row suffix because they're stored contiguously.
+		s := b.rawBytes.offsets.At(bundleOffsetIndex)
+		e := b.rawBytes.offsets.At(rowSuffixIndex + 1)
+		return append(buf[:b.sharedPrefixLen], b.rawBytes.slice(s, e)...)
+	}
+
+	// The key i is a new key in the same bundle.
+	rowSuffixStart := b.rawBytes.offsets.At(rowSuffixIndex)
+	rowSuffixEnd := b.rawBytes.offsets.At(rowSuffixIndex + 1)
+	if rowSuffixStart == rowSuffixEnd {
+		// The start and end offsets are equal, indicating that the key is a
+		// duplicate. Since it's identical to the previous key, there's nothing
+		// left to do, we can return buf as-is.
+		return buf
+	}
+	bundlePrefixLen := b.rawBytes.offsets.At(bundleOffsetIndex+1) - b.rawBytes.offsets.At(bundleOffsetIndex)
+	return append(buf[:b.sharedPrefixLen+int(bundlePrefixLen)], b.rawBytes.slice(rowSuffixStart, rowSuffixEnd)...)
+
+}
+
 // SharedPrefix return a []byte of the shared prefix that was extracted from
 // all of the values in the Bytes vector. The returned slice should not be
 // mutated.
-func (b PrefixBytes) SharedPrefix() []byte {
+func (b *PrefixBytes) SharedPrefix() []byte {
 	// The very first slice is the prefix for the entire column.
 	return b.rawBytes.slice(0, b.rawBytes.offsets.At(0))
 }
@@ -249,7 +294,7 @@ func (b PrefixBytes) SharedPrefix() []byte {
 // RowBundlePrefix takes a row index and returns a []byte of the prefix shared
 // among all the keys in the row's bundle, but without the block-level shared
 // prefix for the column. The returned slice should not be mutated.
-func (b PrefixBytes) RowBundlePrefix(row int) []byte {
+func (b *PrefixBytes) RowBundlePrefix(row int) []byte {
 	i := b.bundleOffsetIndexForRow(row)
 	return b.rawBytes.slice(b.rawBytes.offsets.At(i), b.rawBytes.offsets.At(i+1))
 }
@@ -257,7 +302,7 @@ func (b PrefixBytes) RowBundlePrefix(row int) []byte {
 // BundlePrefix returns the prefix of the i-th bundle in the column. The
 // provided i must be in the range [0, BundleCount()). The returned slice should
 // not be mutated.
-func (b PrefixBytes) BundlePrefix(i int) []byte {
+func (b *PrefixBytes) BundlePrefix(i int) []byte {
 	j := b.offsetIndexByBundleIndex(i)
 	return b.rawBytes.slice(b.rawBytes.offsets.At(j), b.rawBytes.offsets.At(j+1))
 }
@@ -267,7 +312,7 @@ func (b PrefixBytes) BundlePrefix(i int) []byte {
 // RowSuffix().
 //
 // The returned slice should not be mutated.
-func (b PrefixBytes) RowSuffix(row int) []byte {
+func (b *PrefixBytes) RowSuffix(row int) []byte {
 	i := b.rowSuffixIndex(row)
 	// Retrieve the low and high offsets indicating the start and end of the
 	// row's suffix slice.
@@ -298,12 +343,12 @@ func (b PrefixBytes) RowSuffix(row int) []byte {
 }
 
 // Rows returns the count of rows whose keys are encoded within the PrefixBytes.
-func (b PrefixBytes) Rows() int {
+func (b *PrefixBytes) Rows() int {
 	return b.rows
 }
 
 // BundleCount returns the count of bundles within the PrefixBytes.
-func (b PrefixBytes) BundleCount() int {
+func (b *PrefixBytes) BundleCount() int {
 	return b.bundleCount(b.rows)
 }
 
@@ -311,7 +356,7 @@ func (b PrefixBytes) BundleCount() int {
 // equal to k, returning the index of the key and whether an equal key was
 // found. If multiple keys are equal, the index of the first such key is
 // returned. If all keys are < k, Search returns Rows() for the row index.
-func (b PrefixBytes) Search(k []byte) (rowIndex int, isEqual bool) {
+func (b *PrefixBytes) Search(k []byte) (rowIndex int, isEqual bool) {
 	// First compare to the block-level shared prefix.
 	n := min(len(k), b.sharedPrefixLen)
 	c := bytes.Compare(k[:n], unsafe.Slice((*byte)(b.rawBytes.data), b.sharedPrefixLen))
diff --git a/sstable/colblk/raw_bytes.go b/sstable/colblk/raw_bytes.go
index 8d72f231..43c9d474 100644
--- a/sstable/colblk/raw_bytes.go
+++ b/sstable/colblk/raw_bytes.go
@@ -95,11 +95,11 @@ func rawBytesToBinFormatter(f *binfmt.Formatter, count int, sliceFormatter func(
 	}
 }
 
-func (b RawBytes) ptr(offset uint32) unsafe.Pointer {
+func (b *RawBytes) ptr(offset uint32) unsafe.Pointer {
 	return unsafe.Pointer(uintptr(b.data) + uintptr(offset))
 }
 
-func (b RawBytes) slice(start, end uint32) []byte {
+func (b *RawBytes) slice(start, end uint32) []byte {
 	return unsafe.Slice((*byte)(b.ptr(start)), end-start)
 }
 
@@ -109,7 +109,7 @@ func (b RawBytes) At(i int) []byte {
 }
 
 // Slices returns the number of []byte slices encoded within the RawBytes.
-func (b RawBytes) Slices() int {
+func (b *RawBytes) Slices() int {
 	return b.slices
 }
 


Test patch:
diff --git a/sstable/colblk/prefix_bytes_test.go b/sstable/colblk/prefix_bytes_test.go
index d535070f..df762132 100644
--- a/sstable/colblk/prefix_bytes_test.go
+++ b/sstable/colblk/prefix_bytes_test.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/cockroachdb/datadriven"
 	"github.com/cockroachdb/pebble/internal/binfmt"
+	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/rand"
 )
@@ -273,16 +274,17 @@ func BenchmarkPrefixBytes(b *testing.B) {
 			buf = build(n)
 			pb, _ := DecodePrefixBytes(buf, 0, n)
 			b.ResetTimer()
+			var k []byte
 			for i := 0; i < b.N; i++ {
-				k := append([]byte(nil), pb.SharedPrefix()...)
-				l := len(k)
-				for i := 0; i < n; i++ {
-					j := rand.Intn(n)
-					k = append(append(k[:l], pb.RowBundlePrefix(j)...), pb.RowSuffix(j)...)
-					if !bytes.Equal(k, userKeys[j]) {
-						b.Fatalf("Constructed key %q (%q, %q, %q) for index %d; expected %q",
-							k, pb.SharedPrefix(), pb.RowBundlePrefix(j), pb.RowSuffix(j), j, userKeys[j])
-					}
+				j := i % n
+				if j == 0 {
+					k = pb.AppendAt(k[:0], j)
+				} else {
+					k = pb.SetNextAt(k, j)
+				}
+				if invariants.Enabled && !bytes.Equal(k, userKeys[j]) {
+					b.Fatalf("Constructed key %q (%q, %q, %q) for index %d; expected %q",
+						k, pb.SharedPrefix(), pb.RowBundlePrefix(j), pb.RowSuffix(j), j, userKeys[j])
 				}
 			}
 		})