Skip to content

Commit

Permalink
colblk: make KeySeeker.SeekGE return whether prefix is equal
Browse files Browse the repository at this point in the history
This is useful for synthetic suffix replacement and it will also be
useful to implement `SeekPrefixGE` optimizations.
  • Loading branch information
RaduBerinde committed Oct 3, 2024
1 parent 41145ee commit 3b9551c
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 11 deletions.
8 changes: 5 additions & 3 deletions sstable/colblk/cockroach_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -213,14 +213,16 @@ func (ks *cockroachKeySeeker) IsLowerBound(k []byte) bool {
}

// SeekGE is part of the KeySeeker interface.
func (ks *cockroachKeySeeker) SeekGE(key []byte, currRow int, dir int8) (row int) {
func (ks *cockroachKeySeeker) SeekGE(
key []byte, boundRow int, searchDir int8,
) (row int, equalPrefix bool) {
// TODO(jackson): Inline crdbtest.Split.
si := crdbtest.Split(key)
row, eq := ks.prefixes.Search(key[:si-1])
if eq {
return ks.seekGEOnSuffix(row, key[si:])
return ks.seekGEOnSuffix(row, key[si:]), true
}
return row
return row, false
}

// seekGEOnSuffix is a helper function for SeekGE when a seek key's prefix
Expand Down
19 changes: 11 additions & 8 deletions sstable/colblk/data_block.go
Original file line number Diff line number Diff line change
Expand Up @@ -103,8 +103,8 @@ type KeySeeker interface {
// IsLowerBound returns true if all keys in the data block are >= the given
// key. If the data block contains no keys, returns true.
IsLowerBound(k []byte) bool
// SeekGE returns the index of the first row with a key greater than or
// equal to [key].
// SeekGE returns the index of the first row with a key greater than or equal
// to [key], and whether that row has the same prefix as [key].
//
// If the caller externally knows a bound on where the key is located, it
// may indicate it through [boundRow] and [searchDir]. A [searchDir] value
Expand All @@ -113,7 +113,7 @@ type KeySeeker interface {
// ≥ [boundRow]. Implementations may use this information to constrain the
// search. See (base.SeekGEFlags).TrySeekUsingNext for context on when this
// may be set in practice.
SeekGE(key []byte, boundRow int, searchDir int8) (row int)
SeekGE(key []byte, boundRow int, searchDir int8) (row int, equalPrefix bool)
// MaterializeUserKey materializes the user key of the specified row,
// returning a slice of the materialized user key.
//
Expand Down Expand Up @@ -306,13 +306,15 @@ func (ks *defaultKeySeeker) IsLowerBound(k []byte) bool {
}

// SeekGE is part of the KeySeeker interface.
func (ks *defaultKeySeeker) SeekGE(key []byte, currRow int, dir int8) (row int) {
func (ks *defaultKeySeeker) SeekGE(
key []byte, boundRow int, searchDir int8,
) (row int, equalPrefix bool) {
si := ks.comparer.Split(key)
row, eq := ks.prefixes.Search(key[:si])
if eq {
return ks.seekGEOnSuffix(row, key[si:])
return ks.seekGEOnSuffix(row, key[si:]), true
}
return row
return row, false
}

// seekGEOnSuffix is a helper function for SeekGE when a seek key's prefix
Expand Down Expand Up @@ -909,7 +911,7 @@ func (i *DataBlockIter) SeekGE(key []byte, flags base.SeekGEFlags) *base.Interna
if flags.TrySeekUsingNext() {
searchDir = +1
}
i.row = i.keySeeker.SeekGE(key, i.row, searchDir)
i.row, _ = i.keySeeker.SeekGE(key, i.row, searchDir)
if i.transforms.HideObsoletePoints {
i.nextObsoletePoint = i.r.isObsolete.SeekSetBitGE(i.row)
if i.atObsoletePointForward() {
Expand Down Expand Up @@ -941,7 +943,8 @@ func (i *DataBlockIter) SeekLT(key []byte, _ base.SeekLTFlags) *base.InternalKV
if i.r == nil {
return nil
}
i.row = i.keySeeker.SeekGE(key, i.row, 0 /* searchDir */) - 1
geRow, _ := i.keySeeker.SeekGE(key, i.row, 0 /* searchDir */)
i.row = geRow - 1
if i.transforms.HideObsoletePoints {
i.nextObsoletePoint = i.r.isObsolete.SeekSetBitGE(max(i.row, 0))
if i.atObsoletePointBackward() {
Expand Down

0 comments on commit 3b9551c

Please sign in to comment.