Repository: cockroachdb/pebble
Pull request #3777 (2024-07-18 20:41:40)

colblk: impose consistent interfaces
**colblk: make decoder function signatures consistent**

Unify the various functions that decode data structures from a columnar block
so that they all have consistent function signatures, and define a generic
DecodeFunc that captures this consistency. In future work, the DecodeFunc will
be used to aid in composing decoders.

**colblk: implement a common interface across column data accessors**

Introduce a new generic `Array[T]` type that allows indexed access to data and
implement it across the various accessors.

**colblk: randomize BundleSizes in block writer randomized test**

Randomize the BundleSize used for PrefixBytes columns in the
TestBlockWriterRandomized test.

**colblk: add facility for cloning an Array to a []T**

Add a `Clone[T](Array[T], int)` func that clones an array of n elements into a
heap-allocated Go slice. It's primarily intended for use in tests.

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


Patch:
diff --git a/sstable/colblk/bitmap.go b/sstable/colblk/bitmap.go
index 2fd1429d..a89d256d 100644
--- a/sstable/colblk/bitmap.go
+++ b/sstable/colblk/bitmap.go
@@ -30,24 +30,31 @@ type Bitmap struct {
 	bitCount int
 }
 
-// MakeBitmap returns a Bitmap that reads from b supporting bitCount logical
-// bits. No bounds checking is performed, so the caller must guarantee the
-// bitmap is appropriately sized and the provided bitCount correctly identifies
-// the number of bits in the bitmap.
-func MakeBitmap(b []byte, off uint32, bitCount int) Bitmap {
-	if len(b) < int(off)+bitmapRequiredSize(bitCount) {
+// Assert that Bitmap implements Array[bool].
+var _ Array[bool] = Bitmap{}
+
+// DecodeBitmap decodes the structure of a Bitmap and returns a Bitmap that
+// reads from b supporting bitCount logical bits. No bounds checking is
+// performed, so the caller must guarantee the bitmap is appropriately sized and
+// the provided bitCount correctly identifies the number of bits in the bitmap.
+func DecodeBitmap(b []byte, off uint32, bitCount int) (bitmap Bitmap, endOffset uint32) {
+	sz := bitmapRequiredSize(bitCount)
+	off = align(off, align64)
+	if len(b) < int(off)+sz {
 		panic(errors.AssertionFailedf("bitmap of %d bits requires at least %d bytes; provided with %d-byte slice",
 			bitCount, bitmapRequiredSize(bitCount), len(b[off:])))
 	}
-	off = align(off, align64)
 	return Bitmap{
 		data:     makeUnsafeRawSlice[uint64](unsafe.Pointer(&b[off])),
 		bitCount: bitCount,
-	}
+	}, off + uint32(sz)
 }
 
-// Get returns true if the bit at position i is set and false otherwise.
-func (b Bitmap) Get(i int) bool {
+// Assert that DecodeBitmap implements DecodeFunc.
+var _ DecodeFunc[Bitmap] = DecodeBitmap
+
+// At returns true if the bit at position i is set and false otherwise.
+func (b Bitmap) At(i int) bool {
 	return (b.data.At(i>>6 /* i/64 */) & (1 << uint(i%64))) != 0
 }
 
diff --git a/sstable/colblk/block.go b/sstable/colblk/block.go
index 96526a6c..fb40480a 100644
--- a/sstable/colblk/block.go
+++ b/sstable/colblk/block.go
@@ -165,6 +165,22 @@ func FinishBlock(rows int, writers []ColumnWriter) []byte {
 	return buf
 }
 
+// DecodeColumn decodes the col'th column of the provided reader's block as a
+// column of dataType using decodeFunc.
+func DecodeColumn[V any](r *BlockReader, col int, dataType DataType, decodeFunc DecodeFunc[V]) V {
+	if uint16(col) >= r.header.Columns {
+		panic(errors.AssertionFailedf("column %d is out of range [0, %d)", col, r.header.Columns))
+	}
+	if dt := r.dataType(col); dt != dataType {
+		panic(errors.AssertionFailedf("column %d is type %s; not %s", col, dt, dataType))
+	}
+	v, endOffset := decodeFunc(r.data, r.pageStart(col), int(r.header.Rows))
+	if nextColumnOff := r.pageStart(col + 1); endOffset != nextColumnOff {
+		panic(errors.AssertionFailedf("column %d decoded to offset %d; expected %d", col, endOffset, nextColumnOff))
+	}
+	return v
+}
+
 // A BlockReader holds metadata for accessing the columns of a columnar block.
 type BlockReader struct {
 	data             []byte
@@ -196,74 +212,53 @@ func (r *BlockReader) DataType(col int) DataType {
 	if uint16(col) >= r.header.Columns {
 		panic(errors.AssertionFailedf("column %d is out of range [0, %d)", col, r.header.Columns))
 	}
-	return DataType(*(*uint8)(r.pointer(r.customHeaderSize + 7 + 5*uint32(col))))
+	return r.dataType(col)
+}
+
+func (r *BlockReader) dataType(col int) DataType {
+	return DataType(*(*uint8)(r.pointer(r.customHeaderSize + blockHeaderBaseSize + columnHeaderSize*uint32(col))))
 }
 
 // Bitmap retrieves the col'th column as a bitmap. The column must be of type
 // DataTypeBool.
 func (r *BlockReader) Bitmap(col int) Bitmap {
-	if dt := r.DataType(col); dt != DataTypeBool {
-		panic(errors.AssertionFailedf("column %d is not a Bitmap; holds data type %s", dt))
-	}
-	return MakeBitmap(r.data, r.pageStart(col), int(r.header.Rows))
+	return DecodeColumn(r, col, DataTypeBool, DecodeBitmap)
 }
 
 // RawBytes retrieves the col'th column as a column of byte slices. The column
 // must be of type DataTypeBytes.
 func (r *BlockReader) RawBytes(col int) RawBytes {
-	if dt := r.DataType(col); dt != DataTypeBytes {
-		panic(errors.AssertionFailedf("column %d is not a RawBytes column; holds data type %s", dt))
-	}
-	return MakeRawBytes(int(r.header.Rows), r.data, r.pageStart(col))
+	return DecodeColumn(r, col, DataTypeBytes, DecodeRawBytes)
 }
 
 // PrefixBytes retrieves the col'th column as a prefix-compressed byte slice column. The column
 // must be of type DataTypePrefixBytes.
 func (r *BlockReader) PrefixBytes(col int) PrefixBytes {
-	if dt := r.DataType(col); dt != DataTypePrefixBytes {
-		panic(errors.AssertionFailedf("column %d is not a PrefixBytes column; holds data type %s", dt))
-	}
-	return MakePrefixBytes(int(r.header.Rows), r.data, r.pageStart(col))
+	return DecodeColumn(r, col, DataTypePrefixBytes, DecodePrefixBytes)
 }
 
 // Uint8s retrieves the col'th column as a column of uint8s. The column must be
 // of type DataTypeUint8.
 func (r *BlockReader) Uint8s(col int) UnsafeUint8s {
-	if dt := r.DataType(col); dt != DataTypeUint8 {
-		panic(errors.AssertionFailedf("column %d is not a Uint8 column; holds data type %s", col, dt))
-	}
-	_, s := readUnsafeIntegerSlice[uint8](int(r.header.Rows), r.data, r.pageStart(col))
-	return s
+	return DecodeColumn(r, col, DataTypeUint8, DecodeUnsafeIntegerSlice[uint8])
 }
 
 // Uint16s retrieves the col'th column as a column of uint8s. The column must be
 // of type DataTypeUint16.
 func (r *BlockReader) Uint16s(col int) UnsafeUint16s {
-	if dt := r.DataType(col); dt != DataTypeUint16 {
-		panic(errors.AssertionFailedf("column %d is not a Uint16 column; holds data type %s", col, dt))
-	}
-	_, s := readUnsafeIntegerSlice[uint16](int(r.header.Rows), r.data, r.pageStart(col))
-	return s
+	return DecodeColumn(r, col, DataTypeUint16, DecodeUnsafeIntegerSlice[uint16])
 }
 
 // Uint32s retrieves the col'th column as a column of uint32s. The column must be
 // of type DataTypeUint32.
 func (r *BlockReader) Uint32s(col int) UnsafeUint32s {
-	if dt := r.DataType(col); dt != DataTypeUint32 {
-		panic(errors.AssertionFailedf("column %d is not a Uint32 column; holds data type %s", col, dt))
-	}
-	_, s := readUnsafeIntegerSlice[uint32](int(r.header.Rows), r.data, r.pageStart(col))
-	return s
+	return DecodeColumn(r, col, DataTypeUint32, DecodeUnsafeIntegerSlice[uint32])
 }
 
 // Uint64s retrieves the col'th column as a column of uint64s. The column must be
 // of type DataTypeUint64.
 func (r *BlockReader) Uint64s(col int) UnsafeUint64s {
-	if dt := r.DataType(col); dt != DataTypeUint64 {
-		panic(errors.AssertionFailedf("column %d is not a Uint64 column; holds data type %s", col, dt))
-	}
-	_, s := readUnsafeIntegerSlice[uint64](int(r.header.Rows), r.data, r.pageStart(col))
-	return s
+	return DecodeColumn(r, col, DataTypeUint64, DecodeUnsafeIntegerSlice[uint64])
 }
 
 func (r *BlockReader) pageStart(col int) uint32 {
diff --git a/sstable/colblk/column.go b/sstable/colblk/column.go
index 9b1d6733..d69db595 100644
--- a/sstable/colblk/column.go
+++ b/sstable/colblk/column.go
@@ -108,3 +108,24 @@ type Encoder interface {
 	// state to the provided writer.
 	WriteDebug(w io.Writer, rows int)
 }
+
+// A DecodeFunc decodes a data structure from a byte slice, returning an
+// accessor for the data and the offset of the first byte after the structure.
+// The rows argument must be number of logical rows encoded within the data
+// structure.
+type DecodeFunc[T any] func(buf []byte, offset uint32, rows int) (decoded T, nextOffset uint32)
+
+// An Array provides indexed access to an array of values.
+type Array[V any] interface {
+	// At returns the i'th value in the array.
+	At(i int) V
+}
+
+// Clone clones the first n elements of the array a.
+func Clone[V any](a Array[V], n int) []V {
+	c := make([]V, n)
+	for i := 0; i < n; i++ {
+		c[i] = a.At(i)
+	}
+	return c
+}
diff --git a/sstable/colblk/prefix_bytes.go b/sstable/colblk/prefix_bytes.go
index fbb5c961..7469cafa 100644
--- a/sstable/colblk/prefix_bytes.go
+++ b/sstable/colblk/prefix_bytes.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"io"
 	"math/bits"
+	"slices"
 	"strings"
 	"unsafe"
 
@@ -171,10 +172,16 @@ type PrefixBytes struct {
 	rawBytes        RawBytes
 }
 
-// MakePrefixBytes constructs an accessor for an array of lexicographically
-// sorted byte slices constructed by PrefixBytesBuilder. Count must be the
-// number of logical slices within the array.
-func MakePrefixBytes(count int, b []byte, offset uint32) PrefixBytes {
+// Assert that PrefixBytes implements Array[[]byte].
+var _ Array[[]byte] = PrefixBytes{}
+
+// DecodePrefixBytes decodes the structure of a PrefixBytes, constructing an
+// accessor for an array of lexicographically sorted byte slices constructed by
+// PrefixBytesBuilder. Count must be the number of logical slices within the
+// array.
+func DecodePrefixBytes(
+	b []byte, offset uint32, count int,
+) (prefixBytes PrefixBytes, endOffset uint32) {
 	if count == 0 {
 		panic(errors.AssertionFailedf("empty PrefixBytes"))
 	}
@@ -185,17 +192,28 @@ func MakePrefixBytes(count int, b []byte, offset uint32) PrefixBytes {
 	calc := makeBundleCalc(bundleShift)
 	nBundles := calc.bundleCount(count)
 
+	rb, endOffset := DecodeRawBytes(b, offset+1, count+nBundles)
 	pb := PrefixBytes{
 		bundleCalc: calc,
 		rows:       count,
-		rawBytes:   MakeRawBytes(count+nBundles, b, offset+1),
+		rawBytes:   rb,
 	}
 	// We always set the base to zero.
 	if pb.rawBytes.offsets.base != 0 {
 		panic(errors.AssertionFailedf("unexpected non-zero base in offsets"))
 	}
 	pb.sharedPrefixLen = int(pb.rawBytes.offsets.At(0))
-	return pb
+	return pb, endOffset
+}
+
+// Assert that DecodePrefixBytes implements DecodeFunc.
+var _ DecodeFunc[PrefixBytes] = DecodePrefixBytes
+
+// At returns the i'th []byte slice in the PrefixBytes. At must allocate, so
+// callers should prefer accessing a slice's constituent components through
+// SharedPrefix, BundlePrefix and RowSuffix.
+func (b PrefixBytes) At(i int) []byte {
+	return slices.Concat(b.SharedPrefix(), b.RowBundlePrefix(i), b.RowSuffix(i))
 }
 
 // SharedPrefix return a []byte of the shared prefix that was extracted from
@@ -419,7 +437,7 @@ func prefixBytesToBinFormatter(f *binfmt.Formatter, count int, sliceFormatter fu
 	if sliceFormatter == nil {
 		sliceFormatter = defaultSliceFormatter
 	}
-	pb := MakePrefixBytes(count, f.Data(), uint32(f.Offset()))
+	pb, _ := DecodePrefixBytes(f.Data(), uint32(f.Offset()), count)
 	f.CommentLine("PrefixBytes")
 	f.HexBytesln(1, "bundleSize: %d", 1<<pb.bundleShift)
 	f.CommentLine("Offsets table")
diff --git a/sstable/colblk/raw_bytes.go b/sstable/colblk/raw_bytes.go
index ee4ac621..c7eb8e67 100644
--- a/sstable/colblk/raw_bytes.go
+++ b/sstable/colblk/raw_bytes.go
@@ -49,21 +49,28 @@ type RawBytes struct {
 	data    unsafe.Pointer
 }
 
-// MakeRawBytes constructs an accessor for an array of byte slices constructed
-// by RawBytesBuilder. Count must be the number of byte slices within the array.
-func MakeRawBytes(count int, b []byte, offset uint32) RawBytes {
+// Assert that RawBytes implements Array[[]byte].
+var _ Array[[]byte] = RawBytes{}
+
+// DecodeRawBytes decodes the structure of a RawBytes, constructing an accessor
+// for an array of byte slices constructed by RawBytesBuilder. Count must be the
+// number of byte slices within the array.
+func DecodeRawBytes(b []byte, offset uint32, count int) (rawBytes RawBytes, endOffset uint32) {
 	if count == 0 {
-		return RawBytes{}
+		return RawBytes{}, 0
 	}
-	dataOff, offsets := readUnsafeIntegerSlice[uint32](count+1 /* +1 offset */, b, offset)
+	offsets, dataOff := DecodeUnsafeIntegerSlice[uint32](b, offset, count+1 /* +1 offset */)
 	return RawBytes{
 		slices:  count,
 		offsets: offsets,
 		start:   unsafe.Pointer(&b[offset]),
 		data:    unsafe.Pointer(&b[dataOff]),
-	}
+	}, dataOff + offsets.At(count)
 }
 
+// Assert that DecodeRawBytes implements DecodeFunc.
+var _ DecodeFunc[RawBytes] = DecodeRawBytes
+
 func defaultSliceFormatter(x []byte) string {
 	return string(x)
 }
@@ -73,7 +80,7 @@ func rawBytesToBinFormatter(f *binfmt.Formatter, count int, sliceFormatter func(
 		sliceFormatter = defaultSliceFormatter
 	}
 
-	rb := MakeRawBytes(count, f.Data(), uint32(f.Offset()))
+	rb, _ := DecodeRawBytes(f.Data(), uint32(f.Offset()), count)
 	dataOffset := uint64(f.Offset()) + uint64(uintptr(rb.data)-uintptr(rb.start))
 	f.CommentLine("rawbytes")
 	f.CommentLine("offsets table")
diff --git a/sstable/colblk/unsafe_slice.go b/sstable/colblk/unsafe_slice.go
index 3ab6851c..cb17dc6f 100644
--- a/sstable/colblk/unsafe_slice.go
+++ b/sstable/colblk/unsafe_slice.go
@@ -67,9 +67,14 @@ type UnsafeIntegerSlice[T constraints.Integer] struct {
 	deltaWidth uintptr
 }
 
-func readUnsafeIntegerSlice[T constraints.Integer](
-	rows int, b []byte, off uint32,
-) (endOffset uint32, slice UnsafeIntegerSlice[T]) {
+// Assert that UnsafeIntegerSlice implements Array.
+var _ Array[uint8] = UnsafeIntegerSlice[uint8]{}
+
+// DecodeUnsafeIntegerSlice decodes the structure of a slice of uints from a
+// byte slice.
+func DecodeUnsafeIntegerSlice[T constraints.Integer](
+	b []byte, off uint32, rows int,
+) (slice UnsafeIntegerSlice[T], endOffset uint32) {
 	delta := UintDeltaEncoding(b[off])
 	off++
 	switch delta {
@@ -91,9 +96,12 @@ func readUnsafeIntegerSlice[T constraints.Integer](
 	default:
 		panic("unreachable")
 	}
-	return off, slice
+	return slice, off
 }
 
+// Assert that DecodeUnsafeIntegerSlice implements DecodeFunc.
+var _ DecodeFunc[UnsafeUint8s] = DecodeUnsafeIntegerSlice[uint8]
+
 func makeUnsafeIntegerSlice[T constraints.Integer](
 	base T, deltaPtr unsafe.Pointer, deltaWidth int,
 ) UnsafeIntegerSlice[T] {
@@ -104,9 +112,6 @@ func makeUnsafeIntegerSlice[T constraints.Integer](
 	}
 }
 
-// TODO(jackson): Remove when more of the read path is hooked up.
-var _ = makeUnsafeIntegerSlice[uint64]
-
 // At returns the `i`-th element of the slice.
 func (s UnsafeIntegerSlice[T]) At(i int) T {
 	// TODO(jackson): Experiment with other alternatives that might be faster
@@ -129,12 +134,3 @@ func (s UnsafeIntegerSlice[T]) At(i int) T {
 		panic("unreachable")
 	}
 }
-
-// Clone allocates a new slice and copies the first `rows` elements.
-func (s UnsafeIntegerSlice[T]) Clone(rows int) []T {
-	result := make([]T, rows)
-	for i := 0; i < rows; i++ {
-		result[i] = s.At(i)
-	}
-	return result
-}


Test patch:
diff --git a/sstable/colblk/bitmap_test.go b/sstable/colblk/bitmap_test.go
index f3da1c3f..58c5e320 100644
--- a/sstable/colblk/bitmap_test.go
+++ b/sstable/colblk/bitmap_test.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/cockroachdb/datadriven"
 	"github.com/cockroachdb/pebble/internal/binfmt"
+	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/rand"
 )
 
@@ -43,7 +44,7 @@ func TestBitmapFixed(t *testing.T) {
 			}
 
 			_ = builder.Finish(0, n, 0, data)
-			bitmap = MakeBitmap(data, 0, n)
+			bitmap, _ = DecodeBitmap(data, 0, n)
 			dumpBitmap(&buf, bitmap)
 			fmt.Fprint(&buf, "\nBinary representation:\n")
 			f := binfmt.New(data)
@@ -75,7 +76,7 @@ func dumpBitmap(w io.Writer, b Bitmap) {
 		if i > 0 && i%64 == 0 {
 			w.Write([]byte{'\n'})
 		}
-		if b.Get(i) {
+		if b.At(i) {
 			w.Write([]byte{'1'})
 		} else {
 			w.Write([]byte{'0'})
@@ -100,33 +101,34 @@ func TestBitmapRandom(t *testing.T) {
 		}
 		data := make([]byte, builder.Size(size, 0))
 		_ = builder.Finish(0, size, 0, data)
-		bitmap := MakeBitmap(data, 0, size)
+		bitmap, endOffset := DecodeBitmap(data, 0, size)
+		require.Equal(t, uint32(len(data)), endOffset)
 		for i := 0; i < size; i++ {
-			if got := bitmap.Get(i); got != v[i] {
+			if got := bitmap.At(i); got != v[i] {
 				t.Fatalf("b.Get(%d) = %t; want %t", i, got, v[i])
 			}
 		}
 		for i := 0; i < size; i++ {
 			succ := bitmap.Successor(i)
 			// Ensure that Successor always returns the index of a set bit.
-			if succ != size && !bitmap.Get(succ) {
+			if succ != size && !bitmap.At(succ) {
 				t.Fatalf("b.Successor(%d) = %d; bit at index %d is not set", i, succ, succ)
 			}
 			pred := bitmap.Predecessor(i)
 			// Ensure that Predecessor always returns the index of a set bit.
-			if pred >= 0 && !bitmap.Get(pred) {
+			if pred >= 0 && !bitmap.At(pred) {
 				t.Fatalf("b.Predecessor(%d) = %d; bit at index %d is not set", i, pred, pred)
 			}
 
 			// Ensure there are no set bits between i and succ.
 			for j := i; j < succ; j++ {
-				if bitmap.Get(j) {
+				if bitmap.At(j) {
 					t.Fatalf("b.Successor(%d) = %d; bit at index %d is set", i, succ, j)
 				}
 			}
 			// Ensure there are no set bits between pred and i.
 			for j := pred + 1; j < i; j++ {
-				if bitmap.Get(j) {
+				if bitmap.At(j) {
 					t.Fatalf("b.Predecessor(%d) = %d; bit at index %d is set", i, pred, j)
 				}
 			}
diff --git a/sstable/colblk/block_test.go b/sstable/colblk/block_test.go
index eef8e92c..24fd8b0b 100644
--- a/sstable/colblk/block_test.go
+++ b/sstable/colblk/block_test.go
@@ -18,6 +18,11 @@ import (
 	"golang.org/x/exp/rand"
 )
 
+type ColumnSpec struct {
+	DataType
+	BundleSize int // Only used for DataTypePrefixBytes
+}
+
 func TestBlockWriter(t *testing.T) {
 	panicIfErr := func(dataType DataType, stringValue string, err error) {
 		if err != nil {
@@ -151,10 +156,10 @@ func dataTypeFromName(name string) DataType {
 // returns the serialized raw block and a []interface{} slice containing the
 // generated data. The type of each element of the slice is dependent on the
 // corresponding column's type.
-func randBlock(rng *rand.Rand, rows int, schema []DataType) ([]byte, []interface{}) {
+func randBlock(rng *rand.Rand, rows int, schema []ColumnSpec) ([]byte, []interface{}) {
 	data := make([]interface{}, len(schema))
 	for col := range data {
-		switch schema[col] {
+		switch schema[col].DataType {
 		case DataTypeBool:
 			v := make([]bool, rows)
 			for row := 0; row < rows; row++ {
@@ -208,10 +213,10 @@ func randBlock(rng *rand.Rand, rows int, schema []DataType) ([]byte, []interface
 	return buf, data
 }
 
-func buildBlock(schema []DataType, rows int, data []interface{}) []byte {
+func buildBlock(schema []ColumnSpec, rows int, data []interface{}) []byte {
 	cw := make([]ColumnWriter, len(schema))
 	for col := range schema {
-		switch schema[col] {
+		switch schema[col].DataType {
 		case DataTypeBool:
 			var bb BitmapBuilder
 			bb.Reset()
@@ -256,8 +261,7 @@ func buildBlock(schema []DataType, rows int, data []interface{}) []byte {
 			cw[col] = &b
 		case DataTypePrefixBytes:
 			var pbb PrefixBytesBuilder
-			pbb.Init(8)
-			// TODO(jackson): Randomize bundle size.
+			pbb.Init(schema[col].BundleSize)
 			colData := data[col].([][]byte)
 			for r, v := range colData {
 				sharedPrefix := 0
@@ -272,7 +276,7 @@ func buildBlock(schema []DataType, rows int, data []interface{}) []byte {
 	return FinishBlock(rows, cw)
 }
 
-func testRandomBlock(t *testing.T, rng *rand.Rand, rows int, schema []DataType) {
+func testRandomBlock(t *testing.T, rng *rand.Rand, rows int, schema []ColumnSpec) {
 	var sb strings.Builder
 	for i := range schema {
 		if i > 0 {
@@ -291,43 +295,29 @@ func testRandomBlock(t *testing.T, rng *rand.Rand, rows int, schema []DataType)
 			t.Fatalf("expected %d rows, but found %d\n", rows, r.header.Rows)
 		}
 		for col := range schema {
-			if schema[col] != r.DataType(col) {
+			if schema[col].DataType != r.DataType(col) {
 				t.Fatalf("schema mismatch: %s != %s\n", schema[col], r.DataType(col))
 			}
 		}
 
 		for col := range data {
+			spec := schema[col]
 			var got interface{}
-			switch schema[col] {
+			switch spec.DataType {
 			case DataTypeBool:
-				b := r.Bitmap(col)
-				vals := make([]bool, r.header.Rows)
-				for i := range vals {
-					vals[i] = b.Get(i)
-				}
-				got = vals
+				got = Clone(r.Bitmap(col), rows)
 			case DataTypeUint8:
-				got = r.Uint8s(col).Clone(rows)
+				got = Clone(r.Uint8s(col), rows)
 			case DataTypeUint16:
-				got = r.Uint16s(col).Clone(rows)
+				got = Clone(r.Uint16s(col), rows)
 			case DataTypeUint32:
-				got = r.Uint32s(col).Clone(rows)
+				got = Clone(r.Uint32s(col), rows)
 			case DataTypeUint64:
-				got = r.Uint64s(col).Clone(rows)
+				got = Clone(r.Uint64s(col), rows)
 			case DataTypeBytes:
-				vals2 := make([][]byte, rows)
-				vals := r.RawBytes(col)
-				for i := range vals2 {
-					vals2[i] = vals.At(i)
-				}
-				got = vals2
+				got = Clone(r.RawBytes(col), rows)
 			case DataTypePrefixBytes:
-				vals2 := make([][]byte, rows)
-				vals := r.PrefixBytes(col)
-				for i := range vals2 {
-					vals2[i] = slices.Concat(vals.SharedPrefix(), vals.RowBundlePrefix(i), vals.RowSuffix(i))
-				}
-				got = vals2
+				got = Clone(r.PrefixBytes(col), rows)
 			}
 			if !reflect.DeepEqual(data[col], got) {
 				t.Fatalf("%d: %s: expected\n%+v\ngot\n%+v\n% x",
@@ -344,20 +334,21 @@ func TestBlockWriterRandomized(t *testing.T) {
 	randInt := func(lo, hi int) int {
 		return lo + rng.Intn(hi-lo)
 	}
-	testRandomBlock(t, rng, randInt(1, 100), []DataType{DataTypeBool})
-	testRandomBlock(t, rng, randInt(1, 100), []DataType{DataTypeUint8})
-	testRandomBlock(t, rng, randInt(1, 100), []DataType{DataTypeUint16})
-	testRandomBlock(t, rng, randInt(1, 100), []DataType{DataTypeUint32})
-	testRandomBlock(t, rng, randInt(1, 100), []DataType{DataTypeUint64})
-	testRandomBlock(t, rng, randInt(1, 100), []DataType{DataTypeBytes})
-	testRandomBlock(t, rng, randInt(1, 100), []DataType{DataTypePrefixBytes})
+	testRandomBlock(t, rng, randInt(1, 100), []ColumnSpec{{DataType: DataTypeBool}})
+	testRandomBlock(t, rng, randInt(1, 100), []ColumnSpec{{DataType: DataTypeUint8}})
+	testRandomBlock(t, rng, randInt(1, 100), []ColumnSpec{{DataType: DataTypeUint16}})
+	testRandomBlock(t, rng, randInt(1, 100), []ColumnSpec{{DataType: DataTypeUint32}})
+	testRandomBlock(t, rng, randInt(1, 100), []ColumnSpec{{DataType: DataTypeUint64}})
+	testRandomBlock(t, rng, randInt(1, 100), []ColumnSpec{{DataType: DataTypeBytes}})
+	testRandomBlock(t, rng, randInt(1, 100), []ColumnSpec{{DataType: DataTypePrefixBytes, BundleSize: 1 << randInt(0, 6)}})
 
 	for i := 0; i < 100; i++ {
-		schema := make([]DataType, 2+rng.Intn(8))
+		schema := make([]ColumnSpec, 2+rng.Intn(8))
 		for j := range schema {
-			// TODO(jackson): Adjust this to generate DataTypePrefixBytes
-			// columns too once they're supported.
-			schema[j] = DataType(randInt(1, int(dataTypesCount)))
+			schema[j].DataType = DataType(randInt(1, int(dataTypesCount)))
+			if schema[j].DataType == DataTypePrefixBytes {
+				schema[j].BundleSize = 1 << randInt(0, 6)
+			}
 		}
 		testRandomBlock(t, rng, randInt(1, 100), schema)
 	}
diff --git a/sstable/colblk/prefix_bytes_test.go b/sstable/colblk/prefix_bytes_test.go
index 29a7cf0b..d535070f 100644
--- a/sstable/colblk/prefix_bytes_test.go
+++ b/sstable/colblk/prefix_bytes_test.go
@@ -69,7 +69,9 @@ func TestPrefixBytes(t *testing.T) {
 
 			f := binfmt.New(buf)
 			prefixBytesToBinFormatter(f, rows, nil)
-			pb = MakePrefixBytes(rows, buf, 0)
+			var endOffset uint32
+			pb, endOffset = DecodePrefixBytes(buf, 0, rows)
+			require.Equal(t, offset, endOffset)
 			require.Equal(t, rows, pb.Rows())
 			return f.String()
 		case "get":
@@ -169,7 +171,8 @@ func TestPrefixBytesRandomized(t *testing.T) {
 		t.Logf("Size: %d; NumUserKeys: %d (%d); Aggregate pre-compressed string data: %d",
 			size, n, len(userKeys), totalSize)
 
-		pb := MakePrefixBytes(n, buf, 0)
+		pb, endOffset := DecodePrefixBytes(buf, 0, n)
+		require.Equal(t, offset, endOffset)
 		f := binfmt.New(buf)
 		prefixBytesToBinFormatter(f, n, nil)
 		t.Logf("PrefixBytes:\n%s", f.String())
@@ -268,7 +271,7 @@ func BenchmarkPrefixBytes(b *testing.B) {
 		b.Run("iteration", func(b *testing.B) {
 			n := len(userKeys)
 			buf = build(n)
-			pb := MakePrefixBytes(n, buf, 0)
+			pb, _ := DecodePrefixBytes(buf, 0, n)
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
 				k := append([]byte(nil), pb.SharedPrefix()...)
diff --git a/sstable/colblk/raw_bytes_test.go b/sstable/colblk/raw_bytes_test.go
index 79732b94..b18c5af1 100644
--- a/sstable/colblk/raw_bytes_test.go
+++ b/sstable/colblk/raw_bytes_test.go
@@ -51,7 +51,9 @@ func TestRawBytes(t *testing.T) {
 				f.HexBytesln(startOffset, "start offset")
 			}
 			rawBytesToBinFormatter(f, count, nil)
-			rawBytes = MakeRawBytes(count, buf[startOffset:], 0)
+			var decodedEndOffset uint32
+			rawBytes, decodedEndOffset = DecodeRawBytes(buf[startOffset:], 0, count)
+			require.Equal(t, endOffset, decodedEndOffset+uint32(startOffset))
 			return f.String()
 		case "at":
 			var indices []int
@@ -126,7 +128,7 @@ func BenchmarkRawBytes(b *testing.B) {
 			b.Run(fmt.Sprintf("sliceLen=%d", sz), func(b *testing.B) {
 				slices := generateRandomSlices(sz, 32<<10 /* 32 KiB */)
 				data := buildRawBytes(slices)
-				rb := MakeRawBytes(len(slices), data, 0)
+				rb, _ := DecodeRawBytes(data, 0, len(slices))
 				b.ResetTimer()
 				b.SetBytes(32 << 10)
 				for i := 0; i < b.N; i++ {