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%
```

sstable/colblk/prefix_bytes.go

@@ -247,6 +247,42 @@ func (b *PrefixBytes) AppendAt(dst []byte, i int) []byte {
 	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.

sstable/colblk/prefix_bytes_test.go

@@ -277,7 +277,11 @@ func BenchmarkPrefixBytes(b *testing.B) {
 			var k []byte
 			for i := 0; i < b.N; i++ {
 				j := i % n
-				k = pb.AppendAt(k[:0], j)
+				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])