manifest: add range annotations

This change adds a "range annotation" feature to Annotators , which
are computations that aggregate some value over a specific key range within a level. Range annotations use the same B-tree caching behavior as regular annotations, so queries remain fast even with thousands of tables because they avoid a sequential iteration over a level's files.

This PR only sets up range annotations without changing any existing
behavior. See cockroachdb#3793 for some potential use cases.

`BenchmarkNumFilesRangeAnnotation` shows that range annotations are
significantly faster than using `version.Overlaps` to aggregate over
a key range:
```
pkg: github.com/cockroachdb/pebble/internal/manifest
BenchmarkNumFilesRangeAnnotation/annotator-10         	  306010	      4015 ns/op	      48 B/op	       6 allocs/op
BenchmarkNumFilesRangeAnnotation/overlaps-10          	    2223	    513519 ns/op	     336 B/op	       8 allocs/op
```

internal/manifest/annotator.go

@@ -4,6 +4,12 @@
 
 package manifest
 
+import (
+ "sort"
+
+ "github.com/cockroachdb/pebble/internal/base"
+)
+
 // The Annotator type defined below is used by other packages to lazily
 // compute a value over a B-Tree. Each node of the B-Tree stores one
 // `annotation` per annotator, containing the result of the computation over
@@ -24,6 +30,10 @@ package manifest
 // computed incrementally as edits are applied to a level.
 type Annotator[T any] struct {
  Aggregator AnnotationAggregator[T]
+
+ // scratch is used to hold the aggregated annotation value when computing
+ // range annotations in order to avoid additional allocations.
+ scratch *T
 }
 
 // An AnnotationAggregator defines how an annotation should be accumulated
@@ -116,6 +126,80 @@ func (a *Annotator[T]) nodeAnnotation(n *node) (_ *T, cacheOK bool) {
  return t, annot.valid
 }
 
+// accumulateRangeAnnotation computes this annotator's annotation across all
+// files in the node's subtree which overlap with the range defined by bounds.
+// The computed annotation is accumulated into a.scratch.
+func (a *Annotator[T]) accumulateRangeAnnotation(
+ n *node,
+ cmp base.Compare,
+ bounds base.UserKeyBounds,
+ // fullyWithinLowerBound and fullyWithinUpperBound indicate whether this
+ // node's subtree is already known to be within each bound.
+ fullyWithinLowerBound bool,
+ fullyWithinUpperBound bool,
+) {
+ // If this node's subtree is fully within the bounds, compute a regular
+ // annotation.
+ if fullyWithinLowerBound && fullyWithinUpperBound {
+ v, _ := a.nodeAnnotation(n)
+ a.scratch = a.Aggregator.Merge(v, a.scratch)
+ return
+ }
+
+ // We will accumulate annotations from each item in the end-exclusive
+ // range [leftItem, rightItem).
+ leftItem, rightItem := 0, int(n.count)
+ if !fullyWithinLowerBound {
+ // leftItem is the index of the first item that overlaps the lower bound.
+ leftItem = sort.Search(int(n.count), func(i int) bool {
+ return cmp(bounds.Start, n.items[i].Largest.UserKey) <= 0
+ })
+ }
+ if !fullyWithinUpperBound {
+ // rightItem is the index of the first item that does not overlap the
+ // upper bound.
+ rightItem = sort.Search(int(n.count), func(i int) bool {
+ return !bounds.End.IsUpperBoundFor(cmp, n.items[i].Smallest.UserKey)
+ })
+ }
+
+ // Accumulate annotations from every item that overlaps the bounds.
+ for i := leftItem; i < rightItem; i++ {
+ v, _ := a.Aggregator.Accumulate(n.items[i], a.scratch)
+ a.scratch = v
+ }
+
+ if !n.leaf {
+ // We will accumulate annotations from each child in the end-inclusive
+ // range [leftChild, rightChild].
+ leftChild, rightChild := leftItem, rightItem
+ // If the lower bound overlaps with the child at leftItem, there is no
+ // need to accumulate annotations from the child to its left.
+ if leftItem < int(n.count) && cmp(bounds.Start, n.items[leftItem].Smallest.UserKey) >= 0 {
+ leftChild++
+ }
+ // If the upper bound spans beyond the child at rightItem, we must also
+ // accumulate annotations from the child to its right.
+ if rightItem < int(n.count) && bounds.End.IsUpperBoundFor(cmp, n.items[rightItem].Largest.UserKey) {
+ rightChild++
+ }
+
+ for i := leftChild; i <= rightChild; i++ {
+ a.accumulateRangeAnnotation(
+ n.children[i],
+ cmp,
+ bounds,
+ // If this child is to the right of leftItem, then its entire
+ // subtree is within the lower bound.
+ fullyWithinLowerBound || i > leftItem,
+ // If this child is to the left of rightItem, then its entire
+ // subtree is within the upper bound.
+ fullyWithinUpperBound || i < rightItem,
+ )
+ }
+ }
+}
+
 // InvalidateAnnotation removes any existing cached annotations from this
 // annotator from a node's subtree.
 func (a *Annotator[T]) invalidateNodeAnnotation(n *node) {
@@ -142,8 +226,8 @@ func (a *Annotator[T]) LevelAnnotation(lm LevelMetadata) *T {
  return v
 }
 
-// LevelAnnotation calculates the annotation defined by this Annotator for all
-// files across the given levels. A pointer to the Annotator is used as the
+// MultiLevelAnnotation calculates the annotation defined by this Annotator for
+// all files across the given levels. A pointer to the Annotator is used as the
 // key for pre-calculated values, so the same Annotator must be used to avoid
 // duplicate computation. Annotation must not be called concurrently, and in
 // practice this is achieved by requiring callers to hold DB.mu.
@@ -158,6 +242,22 @@ func (a *Annotator[T]) MultiLevelAnnotation(lms []LevelMetadata) *T {
  return aggregated
 }
 
+// LevelRangeAnnotation calculates the annotation defined by this Annotator for
+// the files within LevelMetadata which are within the range
+// [lowerBound, upperBound). A pointer to the Annotator is used as the key for
+// pre-calculated values, so the same Annotator must be used to avoid duplicate
+// computation. Annotation must not be called concurrently, and in practice this
+// is achieved by requiring callers to hold DB.mu.
+func (a *Annotator[T]) LevelRangeAnnotation(lm LevelMetadata, bounds base.UserKeyBounds) *T {
+ if lm.Empty() {
+ return a.Aggregator.Zero(nil)
+ }
+
+ a.scratch = a.Aggregator.Zero(a.scratch)
+ a.accumulateRangeAnnotation(lm.tree.root, lm.tree.cmp, bounds, false, false)
+ return a.scratch
+}
+
 // InvalidateAnnotation clears any cached annotations defined by Annotator. A
 // pointer to the Annotator is used as the key for pre-calculated values, so
 // the same Annotator must be used to clear the appropriate cached annotation.
@@ -206,6 +306,14 @@ func SumAnnotator(accumulate func(f *FileMetadata) (v uint64, cacheOK bool)) *An
  }
 }
 
+// NumFilesAnnotator is an Annotator which computes an annotation value
+// equal to the number of files included in the annotation. Particularly, it
+// can be used to efficiently calculate the number of files in a given key
+// range using range annotations.
+var NumFilesAnnotator = SumAnnotator(func(f *FileMetadata) (uint64, bool) {
+ return 1, true
+})
+
 // PickFileAggregator implements the AnnotationAggregator interface. It defines
 // an aggregator that picks a single file from a set of eligible files.
 type PickFileAggregator struct {

internal/manifest/annotator_test.go

@@ -5,54 +5,47 @@
 package manifest
 
 import (
+ "math/rand"
  "testing"
 
  "github.com/cockroachdb/pebble/internal/base"
  "github.com/stretchr/testify/require"
 )
 
-func makeTestLevelMetadata(count int) (LevelMetadata, []*FileMetadata) {
- files := make([]*FileMetadata, count)
- for i := 0; i < count; i++ {
- files[i] = newItem(key(i))
+// Creates a version with numFiles files in level 6.
+func makeTestVersion(numFiles int) (*Version, []*FileMetadata) {
+ files := make([]*FileMetadata, numFiles)
+ for i := 0; i < numFiles; i++ {
+ // Each file spans 10 keys, e.g. [0->9], [10->19], etc.
+ files[i] = (&FileMetadata{}).ExtendPointKeyBounds(
+ base.DefaultComparer.Compare, key(i*10), key(i*10+9),
+ )
+ files[i].InitPhysicalBacking()
  }
 
- lm := MakeLevelMetadata(base.DefaultComparer.Compare, 6, files)
- return lm, files
-}
+ var levelFiles [7][]*FileMetadata
+ levelFiles[6] = files
 
-// NumFilesAnnotator is an Annotator which computes an annotation value
-// equal to the number of files included in the annotation.
-var NumFilesAnnotator = SumAnnotator(func(f *FileMetadata) (uint64, bool) {
- return 1, true
-})
+ v := NewVersion(base.DefaultComparer, 0, levelFiles)
+ return v, files
+}
 
 func TestNumFilesAnnotator(t *testing.T) {
  const count = 1000
- lm, _ := makeTestLevelMetadata(0)
+ v, _ := makeTestVersion(0)
 
  for i := 1; i <= count; i++ {
- lm.tree.Insert(newItem(key(i)))
- numFiles := *NumFilesAnnotator.LevelAnnotation(lm)
+ v.Levels[6].tree.Insert(newItem(key(i)))
+ numFiles := *NumFilesAnnotator.LevelAnnotation(v.Levels[6])
  require.EqualValues(t, i, numFiles)
  }
-
- numFiles := *NumFilesAnnotator.LevelAnnotation(lm)
- require.EqualValues(t, count, numFiles)
-
- numFiles = *NumFilesAnnotator.LevelAnnotation(lm)
- require.EqualValues(t, count, numFiles)
-
- lm.tree.Delete(newItem(key(count / 2)))
- numFiles = *NumFilesAnnotator.LevelAnnotation(lm)
- require.EqualValues(t, count-1, numFiles)
 }
 
 func BenchmarkNumFilesAnnotator(b *testing.B) {
- lm, _ := makeTestLevelMetadata(0)
+ v, _ := makeTestVersion(0)
  for i := 1; i <= b.N; i++ {
- lm.tree.Insert(newItem(key(i)))
- numFiles := *NumFilesAnnotator.LevelAnnotation(lm)
+ v.Levels[6].tree.Insert(newItem(key(i)))
+ numFiles := *NumFilesAnnotator.LevelAnnotation(v.Levels[6])
  require.EqualValues(b, uint64(i), numFiles)
  }
 }
@@ -70,12 +63,115 @@ func TestPickFileAggregator(t *testing.T) {
  },
  }
 
- lm, files := makeTestLevelMetadata(1)
+ v, files := makeTestVersion(1)
 
  for i := 1; i <= count; i++ {
- lm.tree.Insert(newItem(key(i)))
- pickedFile := a.LevelAnnotation(lm)
+ v.Levels[6].tree.Insert(newItem(key(i)))
+ pickedFile := a.LevelAnnotation(v.Levels[6])
  // The picked file should always be the one with the smallest key.
  require.Same(t, files[0], pickedFile)
  }
 }
+
+func bounds(i int, j int, exclusive bool) base.UserKeyBounds {
+ b := base.UserKeyBoundsEndExclusiveIf(key(i).UserKey, key(j).UserKey, exclusive)
+ return b
+}
+
+func randomBounds(rng *rand.Rand, count int) base.UserKeyBounds {
+ first := rng.Intn(count)
+ second := rng.Intn(count)
+ exclusive := rng.Intn(2) == 0
+ return bounds(min(first, second), max(first, second), exclusive)
+}
+
+func requireMatchOverlaps(t *testing.T, v *Version, bounds base.UserKeyBounds) {
+ overlaps := v.Overlaps(6, bounds)
+ numFiles := *NumFilesAnnotator.LevelRangeAnnotation(v.Levels[6], bounds)
+ require.EqualValues(t, overlaps.length, numFiles)
+}
+
+func TestNumFilesRangeAnnotationEmptyRanges(t *testing.T) {
+ const count = 5_000
+ v, files := makeTestVersion(count)
+
+ // Delete files containing key ranges [0, 999] and [24_000, 25_999].
+ for i := 0; i < 100; i++ {
+ v.Levels[6].tree.Delete(files[i])
+ }
+ for i := 2400; i < 2600; i++ {
+ v.Levels[6].tree.Delete(files[i])
+ }
+
+ // Ranges that are completely empty.
+ requireMatchOverlaps(t, v, bounds(1, 999, false))
+ requireMatchOverlaps(t, v, bounds(0, 1000, true))
+ requireMatchOverlaps(t, v, bounds(50_000, 60_000, false))
+ requireMatchOverlaps(t, v, bounds(24_500, 25_500, false))
+ requireMatchOverlaps(t, v, bounds(24_000, 26_000, true))
+
+ // Partial overlaps with empty ranges.
+ requireMatchOverlaps(t, v, bounds(0, 1000, false))
+ requireMatchOverlaps(t, v, bounds(20, 1001, true))
+ requireMatchOverlaps(t, v, bounds(20, 1010, true))
+ requireMatchOverlaps(t, v, bounds(23_000, 27_000, true))
+ requireMatchOverlaps(t, v, bounds(25_000, 40_000, false))
+ requireMatchOverlaps(t, v, bounds(25_500, 26_001, true))
+
+ // Ranges which only spans a single table.
+ requireMatchOverlaps(t, v, bounds(45_000, 45_000, true))
+ requireMatchOverlaps(t, v, bounds(30_000, 30_001, true))
+ requireMatchOverlaps(t, v, bounds(23_000, 23_000, false))
+}
+
+func TestNumFilesRangeAnnotationRandomized(t *testing.T) {
+ const count = 10_000
+ const numIterations = 10_000
+
+ v, _ := makeTestVersion(count)
+
+ rng := rand.New(rand.NewSource(int64(0)))
+ for i := 0; i < numIterations; i++ {
+ requireMatchOverlaps(t, v, randomBounds(rng, count*11))
+ }
+}
+
+func BenchmarkNumFilesRangeAnnotation(b *testing.B) {
+ const count = 100_000
+ v, files := makeTestVersion(count)
+
+ rng := rand.New(rand.NewSource(int64(0)))
+ b.Run("annotator", func(b *testing.B) {
+ for i := 0; i < b.N; i++ {
+ b := randomBounds(rng, count*11)
+ // Randomly delete and reinsert a file to verify that range
+ // annotations are still fast despite small mutations.
+ toDelete := rng.Intn(count)
+ v.Levels[6].tree.Delete(files[toDelete])
+
+ NumFilesAnnotator.LevelRangeAnnotation(v.Levels[6], b)
+
+ v.Levels[6].tree.Insert(files[toDelete])
+ }
+ })
+
+ // Also benchmark an equivalent aggregation using version.Overlaps to show
+ // the difference in performance.
+ b.Run("overlaps", func(b *testing.B) {
+ for i := 0; i < b.N; i++ {
+ b := randomBounds(rng, count*11)
+ toDelete := rng.Intn(count)
+ v.Levels[6].tree.Delete(files[toDelete])
+
+ overlaps := v.Overlaps(6, b)
+ iter := overlaps.Iter()
+ numFiles := 0
+ for f := iter.First(); f != nil; f = iter.Next() {
+ numFiles++
+ }
+
+ v.Levels[6].tree.Insert(files[toDelete])
+ }
+ })
+
+}