colblk: add Bitmap, BitmapBuilder

Introduce the sstable/colblk package for columnar block primitives, and add a
first primitive for storing bitmaps.

The Bitmap uses 1+1/64 physical bits per logical bit. Bits are organized into
64-bit words, allowing constant-time access of an individual bit. Additionally,
after the bitmap a summary bitmap consisting of 1 bit per 64-bit word of the
bitmap provides faster lookup of preceding or successive set bits. This will be
used by a bitmap indicating when a key prefix changes to quickly find the next
key within a block with a new prefix.

sstable/colblk/base.go

@@ -0,0 +1,45 @@
+// Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+// Package colblk implements a columnar block format.
+package colblk
+
+import "golang.org/x/exp/constraints"
+
+// align returns the next value greater than or equal to offset that's divisible
+// by val.
+func align[T constraints.Integer](offset, val T) T {
+	return (offset + val - 1) & ^(val - 1)
+}
+
+// alignWithZeroes aligns the provided offset to val, and writing zeroes to any
+// bytes in buf between the old offset and new aligned offset. This provides
+// determinism when reusing memory that has not been zeroed.
+func alignWithZeroes[T constraints.Integer](buf []byte, offset, val T) T {
+	aligned := align[T](offset, val)
+	for i := offset; i < aligned; i++ {
+		buf[i] = 0
+	}
+	return aligned
+}
+
+const (
+	align16 = 2
+	align32 = 4
+	align64 = 8
+)
+
+// When multiplying or dividing by align{16,32,64} using signed integers, it's
+// faster to shift to the left to multiply or shift to the right to divide. (The
+// compiler optimization is limited to unsigned integers.) The below constants
+// define the shift amounts corresponding to the above align constants.  (eg,
+// alignNShift = log(alignN)).
+//
+// TODO(jackson): Consider updating usages to use uints? They can be less
+// ergonomic.
+const (
+	align16Shift = 1
+	align32Shift = 2
+	align64Shift = 3
+)

sstable/colblk/bitmap.go

@@ -0,0 +1,317 @@
+// Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package colblk
+
+import (
+	"fmt"
+	"io"
+	"math/bits"
+	"slices"
+	"strings"
+	"unsafe"
+
+	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble/internal/binfmt"
+)
+
+// Bitmap is a bitmap structure built on a []uint64. A bitmap utilizes ~1
+// physical bit/logical bit (~0.125 bytes/row). The bitmap is encoded into an
+// 8-byte aligned array of 64-bit words which is (nRows+63)/64 words in length.
+//
+// A summary bitmap is stored after the primary bitmap in which each bit in the
+// summary bitmap corresponds to 1 word in the primary bitmap. A bit is set in
+// the summary bitmap if the corresponding word in the primary bitmap is
+// non-zero. The summary bitmap accelerates predecessor and successor
+// operations.
+type Bitmap struct {
+	data     UnsafeRawSlice[uint64]
+	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, bitCount int) Bitmap {
+	if len(b) != bitmapRequiredSize(bitCount) {
+		panic(errors.AssertionFailedf("bitmap of %d bits requires at %d bytes; provided with %d-byte slice",
+			bitCount, bitmapRequiredSize(bitCount), len(b)))
+	}
+	return Bitmap{
+		data:     makeUnsafeRawSlice[uint64](unsafe.Pointer(&b[0])),
+		bitCount: bitCount,
+	}
+}
+
+// Get returns true if the bit at position i is set and false otherwise.
+func (b Bitmap) Get(i int) bool {
+	return (b.data.At(i>>6 /* i/64 */) & (1 << uint(i%64))) != 0
+}
+
+// Successor returns the next bit greater than or equal to i set in the bitmap.
+// The i parameter must be in [0, bitCount). Returns the number of bits
+// represented by the bitmap if no next bit is set.
+func (b Bitmap) Successor(i int) int {
+	// nextInWord returns the index of the smallest set bit with an index >= bit
+	// within the provided word.  The returned index is an index local to the
+	// word.
+	nextInWord := func(word uint64, bit uint) int {
+		// We want to find the index of the next set bit. We can accomplish this
+		// by clearing the trailing `bit` bits from the word and counting the
+		// number of trailing zeros. For example, consider the word and bit=37:
+		//
+		//           word: 1010101010111111111110000001110101010101011111111111000000111011
+		//
+		//         1<<bit: 0000000000000000000000000010000000000000000000000000000000000000
+		//       1<<bit-1: 0000000000000000000000000001111111111111111111111111111111111111
+		//      ^1<<bit-1: 1111111111111111111111111110000000000000000000000000000000000000
+		// word&^1<<bit-1: 1010101010111111111110000000000000000000000000000000000000000000
+		//
+		// Counting the trailing zeroes of this last value gives us 43. For
+		// visualizing, 1<<43 is:
+		//
+		//                 0000000000000000000010000000000000000000000000000000000000000000
+		//
+		return bits.TrailingZeros64(word &^ ((1 << bit) - 1))
+	}
+
+	wordIdx := i >> 6 // i/64
+	// Fast path for common case of reasonably dense bitmaps; if the there's a
+	// bit > i set in the same word, return it.
+	if next := nextInWord(b.data.At(wordIdx), uint(i%64)); next < 64 {
+		return wordIdx<<6 + next
+	}
+
+	// Consult summary structure to find the next word with a set bit. The word
+	// we just checked (wordIdx) is represented by the wordIdx%64'th bit in the
+	// wordIdx/64'th summary word. We want to know if any of the other later
+	// words that are summarized together have a set bit. We call [nextInWord]
+	// on the summary word to get the index of which word has a set bit, if any.
+	summaryTableOffset, summaryTableEnd := b.summaryTableBounds()
+	summaryWordIdx := summaryTableOffset + wordIdx>>6
+	summaryNext := nextInWord(b.data.At(summaryWordIdx), uint(wordIdx%64)+1)
+	// If [summaryNext] == 64, then there are no set bits in any of the earlier
+	// words represented by the summary word at [summaryWordIdx]. In that case,
+	// we need to keep scanning the summary table forwards.
+	if summaryNext == 64 {
+		for summaryWordIdx++; ; summaryWordIdx++ {
+			// When we fall off the end of the summary table, we've determined
+			// there are no set bits after i across the entirety of the bitmap.
+			if summaryWordIdx >= summaryTableEnd {
+				return b.bitCount
+			}
+			if summaryWord := b.data.At(summaryWordIdx); summaryWord != 0 {
+				summaryNext = bits.TrailingZeros64(summaryWord)
+				break
+			}
+		}
+	}
+	// The summary word index and the summaryNext together tell us which word
+	// has a set bit. The number of leading zeros in the word itself tell us
+	// which bit is set.
+	wordIdx = ((summaryWordIdx - summaryTableOffset) << 6) + summaryNext
+	return (wordIdx << 6) + bits.TrailingZeros64(b.data.At(wordIdx))
+}
+
+// Predecessor returns the previous bit less than or equal to i set in the
+// bitmap. The i parameter must be in [0, bitCount). Returns -1 if no previous
+// bit is set.
+func (b Bitmap) Predecessor(i int) int {
+	// prevInWord returns the index of the largest set bit ≤ bit within the
+	// provided word. The returned index is an index local to the word. Returns
+	// -1 if no set bit is found.
+	prevInWord := func(word uint64, bit uint) int {
+		// We want to find the index of the previous set bit. We can accomplish
+		// this by clearing the leading `bit` bits from the word and counting
+		// the number of leading zeros. For example, consider the word and
+		// bit=42:
+		//
+		//              word: 1010101010111111111110000001110101010101011111111111000000111011
+		//
+		//        1<<(bit+1): 0000000000000000000010000000000000000000000000000000000000000000
+		//      1<<(bit+1)-1: 0000000000000000000001111111111111111111111111111111111111111111
+		// word&1<<(bit+1)-1: 0000000000000000000000000001110101010101011111111111000000111011
+		//
+		// Counting the leading zeroes of this last value gives us 27 leading
+		// zeros. 63-27 gives index 36. For visualizing, 1<<36 is:
+		//
+		//                    0000000000000000000000000001000000000000000000000000000000000000
+		//
+		return 63 - bits.LeadingZeros64(word&((1<<(bit+1))-1))
+	}
+
+	wordIdx := i >> 6 // i/64
+	// Fast path for common case of reasonably dense bitmaps; if the there's a
+	// bit < i set in the same word, return it.
+	if prev := prevInWord(b.data.At(wordIdx), uint(i%64)); prev >= 0 {
+		return (wordIdx << 6) + prev
+	}
+
+	// Consult summary structure to find the next word with a set bit. The word
+	// we just checked (wordIdx) is represented by the wordIdx%64'th bit in the
+	// wordIdx/64'th summary word. We want to know if any of other earlier words
+	// that are summarized together have a set bit. We call [prevInWord] on the
+	// summary word to get the index of which word has a set bit, if any.
+	summaryTableOffset, _ := b.summaryTableBounds()
+	summaryWordIdx := summaryTableOffset + wordIdx>>6
+	summaryPrev := prevInWord(b.data.At(summaryWordIdx), uint(wordIdx%64)-1)
+	// If [summaryPrev] is negative, then there are no set bits in any of the
+	// earlier words represented by the summary word at [summaryWordIdx]. In
+	// that case, we need to keep scanning the summary table backwards.
+	if summaryPrev < 0 {
+		for summaryWordIdx--; ; summaryWordIdx-- {
+			// When we fall below the beginning of the summary table, we've
+			// determined there are no set bits before i across the entirety of
+			// the bitmap.
+			if summaryWordIdx < summaryTableOffset {
+				return -1
+			}
+			if summaryWord := b.data.At(summaryWordIdx); summaryWord != 0 {
+				summaryPrev = 63 - bits.LeadingZeros64(summaryWord)
+				break
+			}
+		}
+	}
+	// The summary word index and the summary prev together tell us which word
+	// has a set bit. The number of trailing zeros in the word itself tell us
+	// which bit is set.
+	wordIdx = ((summaryWordIdx - summaryTableOffset) << 6) + summaryPrev
+	return (wordIdx << 6) + 63 - bits.LeadingZeros64(b.data.At(wordIdx))
+}
+
+func (b Bitmap) summaryTableBounds() (startOffset, endOffset int) {
+	startOffset = (b.bitCount + 63) >> 6
+	endOffset = startOffset + startOffset>>6
+	return startOffset, endOffset
+}
+
+// String returns a string representation of the entire bitmap.
+func (b Bitmap) String() string {
+	var sb strings.Builder
+	for w := 0; w < (b.bitCount+63)/64; w++ {
+		fmt.Fprintf(&sb, "%064b", b.data.At(w))
+	}
+	return sb.String()
+}
+
+// BitmapBuilder constructs a Bitmap. Bits are default false.
+type BitmapBuilder struct {
+	words []uint64
+}
+
+func bitmapRequiredSize(total int) int {
+	nWords := (total + 63) >> 6          // divide by 64
+	nSummaryWords := (nWords + 63) >> 6  // divide by 64
+	return (nWords + nSummaryWords) << 3 // multiply by 8
+}
+
+// Set sets the bit at position i if v is true and clears the bit at position i
+// otherwise. Callers need not call Set if v is false and Set(i, true) has not
+// been called yet.
+func (b *BitmapBuilder) Set(i int, v bool) {
+	w := i >> 6 // divide by 64
+	for len(b.words) <= w {
+		b.words = append(b.words, 0)
+	}
+	if v {
+		b.words[w] |= 1 << uint(i%64)
+	} else {
+		b.words[w] &^= 1 << uint(i%64)
+	}
+}
+
+// Reset resets the bitmap to the empty state.
+func (b *BitmapBuilder) Reset() {
+	clear(b.words)
+	b.words = b.words[:0]
+}
+
+// NumColumns implements the ColumnWriter interface.
+func (b *BitmapBuilder) NumColumns() int { return 1 }
+
+// Size implements the ColumnWriter interface.
+func (b *BitmapBuilder) Size(rows int, offset uint32) uint32 {
+	offset = align(offset, align64)
+	return offset + uint32(bitmapRequiredSize(rows))
+}
+
+// Invert inverts the bitmap, setting all bits that are not set and clearing all
+// bits that are set. If the bitmap's tail is sparse and is not large enough to
+// represent nRows rows, it's first materialized.
+func (b *BitmapBuilder) Invert(nRows int) {
+	// If the tail of b is sparse, fill in zeroes before inverting.
+	nBitmapWords := (nRows + 63) >> 6
+	b.words = slices.Grow(b.words, nBitmapWords-len(b.words))[:nBitmapWords]
+	for i := range b.words {
+		b.words[i] = ^b.words[i]
+	}
+}
+
+// Finish finalizes the bitmap, computing the per-word summary bitmap and
+// writing the resulting data to buf at offset.
+func (b *BitmapBuilder) Finish(col, nRows int, offset uint32, buf []byte) (uint32, ColumnDesc) {
+	offset = alignWithZeroes(buf, offset, align64)
+	dest := makeUnsafeRawSlice[uint64](unsafe.Pointer(&buf[offset]))
+
+	nBitmapWords := (nRows + 63) >> 6
+	// Truncate the bitmap to the number of words required to represent nRows.
+	// The caller may have written more bits than nRows and no longer cares to
+	// write them out.
+	if len(b.words) > nBitmapWords {
+		b.words = b.words[:nBitmapWords]
+	}
+	// Ensure the last word of the bitmap does not contain any set bits beyond
+	// the last row. This is not just for determinism but also to ensure that
+	// the summary bitmap is correct (which is necessary for Bitmap.Successor
+	// correctness).
+	if i := nRows % 64; len(b.words) >= nBitmapWords && i != 0 {
+		b.words[nBitmapWords-1] &= (1 << i) - 1
+	}
+
+	// Copy all the words of the bitmap into the destination buffer.
+	offset += uint32(copy(dest.Slice(len(b.words)), b.words)) << align64Shift
+
+	// The caller may have written fewer than nRows rows if the tail is all
+	// zeroes, relying on these bits being implicitly zero. If the tail of b is
+	// sparse, fill in zeroes.
+	for i := len(b.words); i < nBitmapWords; i++ {
+		dest.set(i, 0)
+		offset += align64
+	}
+
+	// Add the summary bitmap.
+	nSummaryWords := (nBitmapWords + 63) >> 6
+	for i := 0; i < nSummaryWords; i++ {
+		wordsOff := (i << 6) // i*64
+		nWords := min(64, len(b.words)-wordsOff)
+		var summaryWord uint64
+		for j := 0; j < nWords; j++ {
+			if (b.words)[wordsOff+j] != 0 {
+				summaryWord |= 1 << j
+			}
+		}
+		dest.set(nBitmapWords+i, summaryWord)
+	}
+	offset += uint32(nSummaryWords) << align64Shift
+	return offset, ColumnDesc{DataType: DataTypeBool, Encoding: EncodingDefault}
+}
+
+// WriteDebug implements the ColumnWriter interface.
+func (b *BitmapBuilder) WriteDebug(w io.Writer, rows int) {
+	// TODO(jackson): Add more detailed debugging information.
+	fmt.Fprint(w, "bitmap")
+}
+
+func bitmapToBinFormatter(f *binfmt.Formatter, rows int) int {
+	bitmapWords := (rows + 63) / 64
+	for i := 0; i < bitmapWords; i++ {
+		f.Line(8).Append("b ").Binary(8).Done("bitmap word %d", i)
+	}
+	summaryWords := (bitmapWords + 63) / 64
+	for i := 0; i < summaryWords; i++ {
+		f.Line(8).Append("b ").Binary(8).Done("bitmap summary word %d-%d", i*64, i*64+63)
+	}
+	return (bitmapWords + summaryWords) * align64
+}