internal/base: introduce KVMeta, iterators

Informs: #5466

Author: annrpom

internal/base/internal.go

@@ -626,6 +626,29 @@ type InternalKV struct {
 	V InternalValue
 }
 
+type TieringAttribute uint64
+
+// KVMeta describes optional metadata associated with an `InternalKV`.
+// It's currently produced only by sstable-backed iterators and is not embedded
+// within `InternalKV` to avoid overhead on the common iteration path.
+// Instead, select iterators expose methods that return the metadata alongside
+// the key/value:
+//   - `levelIter.FirstWithMeta` / `levelIter.NextWithMeta`
+//   - sstable iterators' `FirstWithMeta` / `NextWithMeta`
+//
+// These methods exist to support compaction-only logic (eg, `compaction.Iter`).
+// Regular iteration should use the standard methods that do not surface metadata.
+type KVMeta struct {
+	TieringSpanID uint64
+	// TieringAttribute is a user-specified attribute for the key-value pair.
+	//
+	// TODO(sumeer): For CockroachDB decide on units for this attribute, which
+	// will be a timestamp, since unix nanos is unnecessarily large.
+	// log2(24*365*100) = 19.74, i.e., number of hours in 100 years fits in 3
+	// bytes.
+	TieringAttribute TieringAttribute
+}
+
 // Kind returns the KV's internal key kind.
 func (kv *InternalKV) Kind() InternalKeyKind {
 	return kv.K.Kind()

internal/base/iterator.go

@@ -517,6 +517,13 @@ func (s *InternalIteratorStats) SafeFormat(p redact.SafePrinter, verb rune) {
 	}
 }
 
+// MetaDecoder is an optional interface that can be implemented by iterators
+// to provide metadata about the current key-value pair.
+type MetaDecoder interface {
+	// DecodeMeta returns metadata for the current iterator position.
+	DecodeMeta() KVMeta
+}
+
 // IteratorDebug is an interface implemented by all internal iterators and
 // fragment iterators.
 type IteratorDebug interface {

internal/compact/iterator.go

@@ -18,6 +18,13 @@ import (
 	"github.com/cockroachdb/redact"
 )
 
+// MetaIterator is an interface for iterators that support metadata extraction.
+type MetaIterator interface {
+	base.InternalIterator
+	FirstWithMeta() (*base.InternalKV, base.KVMeta)
+	NextWithMeta() (*base.InternalKV, base.KVMeta)
+}
+
 // Iter provides a forward-only iterator that encapsulates the logic for
 // collapsing entries during compaction. It wraps an internal iterator and
 // collapses entries that are no longer necessary because they are shadowed by

level_iter.go

@@ -729,6 +729,37 @@ func (l *levelIter) First() *base.InternalKV {
 	return l.verify(l.skipEmptyFileForward())
 }
 
+// FirstWithMeta moves the iterator to the first key/value pair and returns
+// both the key/value and the associated metadata. This method is used by
+// compaction iterators that need access to tiering metadata without adding
+// overhead to the common iteration path.
+func (l *levelIter) FirstWithMeta() (*base.InternalKV, base.KVMeta) {
+	return l.First(), l.extractMetaFromCurrentPosition()
+}
+
+// NextWithMeta moves the iterator to the next key/value pair and returns
+// both the key/value and the associated metadata. This method is used by
+// compaction iterators that need access to tiering metadata without adding
+// overhead to the common iteration path.
+func (l *levelIter) NextWithMeta() (*base.InternalKV, base.KVMeta) {
+	return l.Next(), l.extractMetaFromCurrentPosition()
+}
+
+// extractMetaFromCurrentPosition extracts KVMeta from the current iterator
+// position. This method delegates to the underlying iterator if it supports the
+// specialized methods.
+func (l *levelIter) extractMetaFromCurrentPosition() base.KVMeta {
+	if l.iter == nil {
+		return base.KVMeta{}
+	}
+
+	if metaDecoder, ok := l.iter.(base.MetaDecoder); ok {
+		return metaDecoder.DecodeMeta()
+	}
+
+	return base.KVMeta{}
+}
+
 func (l *levelIter) Last() *base.InternalKV {
 	if invariants.Enabled && l.upper != nil {
 		panic(errors.AssertionFailedf("levelIter Last called while upper bound %q is set", l.upper))

sstable/colblk/data_block.go

@@ -993,6 +993,7 @@ func (v *DataBlockValidator) Validate(
 }
 
 var _ blockiter.Data = (*DataBlockIter)(nil)
+var _ base.MetaDecoder = (*DataBlockIter)(nil)
 
 // DataBlockIter iterates over a columnar data block.
 type DataBlockIter struct {
@@ -1267,6 +1268,33 @@ func (i *DataBlockIter) First() *base.InternalKV {
 	return i.decodeRow()
 }
 
+// FirstWithMeta moves the iterator to the first key/value pair and returns
+// both the key/value and the associated metadata. This method is used by
+// compaction iterators that need access to tiering metadata without adding
+// overhead to the common iteration path.
+func (i *DataBlockIter) FirstWithMeta() (*base.InternalKV, base.KVMeta) {
+	return i.First(), i.decodeMeta()
+}
+
+// NextWithMeta moves the iterator to the next key/value pair and returns
+// both the key/value and the associated metadata. This method is used by
+// compaction iterators that need access to tiering metadata without adding
+// overhead to the common iteration path.
+func (i *DataBlockIter) NextWithMeta() (*base.InternalKV, base.KVMeta) {
+	return i.Next(), i.decodeMeta()
+}
+
+// decodeMeta extracts the KVMeta for the current row.
+func (i *DataBlockIter) decodeMeta() base.KVMeta {
+	// TODO: Implement tiering metadata extraction when the fields are available
+	return base.KVMeta{}
+}
+
+// DecodeMeta implements the base.MetaDecoder interface.
+func (i *DataBlockIter) DecodeMeta() base.KVMeta {
+	return i.decodeMeta()
+}
+
 // Last implements the base.InternalIterator interface.
 func (i *DataBlockIter) Last() *base.InternalKV {
 	if i.d == nil {

sstable/reader_iter.go

@@ -25,6 +25,7 @@ import (
 // https://go.googlesource.com/proposal/+/refs/heads/master/design/43651-type-parameters.md#pointer-method-example
 type dataBlockIterator[D any] interface {
 	blockiter.Data
+	base.MetaDecoder
 
 	*D // non-interface type constraint element
 }

sstable/reader_iter_single_lvl.go

@@ -1226,6 +1226,33 @@ func (i *singleLevelIterator[I, PI, D, PD]) First() *base.InternalKV {
 	return i.firstInternal()
 }
 
+// FirstWithMeta moves the iterator to the first key/value pair and returns
+// both the key/value and the associated metadata. This method is used by
+// compaction iterators that need access to tiering metadata without adding
+// overhead to the common iteration path.
+func (i *singleLevelIterator[I, PI, D, PD]) FirstWithMeta() (*base.InternalKV, base.KVMeta) {
+	return i.First(), i.extractMetaFromCurrentPosition()
+}
+
+// NextWithMeta moves the iterator to the next key/value pair and returns
+// both the key/value and the associated metadata. This method is used by
+// compaction iterators that need access to tiering metadata without adding
+// overhead to the common iteration path.
+func (i *singleLevelIterator[I, PI, D, PD]) NextWithMeta() (*base.InternalKV, base.KVMeta) {
+	return i.Next(), i.extractMetaFromCurrentPosition()
+}
+
+// extractMetaFromCurrentPosition extracts KVMeta from the current iterator position.
+// This method delegates to the underlying data block iterator if it supports
+// the specialized methods.
+func (i *singleLevelIterator[I, PI, D, PD]) extractMetaFromCurrentPosition() base.KVMeta {
+	if PD(&i.data).IsDataInvalidated() {
+		return base.KVMeta{}
+	}
+
+	return PD(&i.data).DecodeMeta()
+}
+
 // firstInternal is a helper used for absolute positioning in a single-level
 // index file, or for positioning in the second-level index in a two-level
 // index file. For the latter, one cannot make any claims about absolute

sstable/reader_iter_two_lvl.go

@@ -856,6 +856,95 @@ func (i *twoLevelIterator[I, PI, D, PD]) First() *base.InternalKV {
 	return i.skipForward()
 }
 
+// FirstWithMeta moves the iterator to the first key/value pair and returns
+// both the key/value and the associated metadata. This method is used by
+// compaction iterators that need access to tiering metadata without adding
+// overhead to the common iteration path.
+func (i *twoLevelIterator[I, PI, D, PD]) FirstWithMeta() (*base.InternalKV, base.KVMeta) {
+	i.lastOpWasSeekPrefixGE.Set(false)
+	// The synthetic key is no longer relevant and must be cleared.
+	i.secondLevel.synthetic.atSyntheticKey = false
+
+	// If we have a lower bound, use SeekGE. Note that in general this is not
+	// supported usage, except when the lower bound is there because the table is
+	// virtual.
+	if i.secondLevel.lower != nil {
+		kv := i.SeekGE(i.secondLevel.lower, base.SeekGEFlagsNone)
+		if kv == nil {
+			return nil, base.KVMeta{}
+		}
+		meta := i.extractMetaFromCurrentPosition()
+		return kv, meta
+	}
+	i.secondLevel.exhaustedBounds = 0
+	i.secondLevel.err = nil // clear cached iteration error
+	// Seek optimization only applies until iterator is first positioned after SetBounds.
+	i.secondLevel.boundsCmp = 0
+
+	if !i.ensureTopLevelIndexLoaded() {
+		return nil, base.KVMeta{}
+	}
+
+	if !PI(&i.topLevelIndex).First() {
+		return nil, base.KVMeta{}
+	}
+	result := i.loadSecondLevelIndexBlock(+1)
+	if result == loadBlockFailed {
+		return nil, base.KVMeta{}
+	}
+	if result == loadBlockOK {
+		if ikv := i.secondLevel.First(); ikv != nil {
+			meta := i.extractMetaFromCurrentPosition()
+			return ikv, meta
+		}
+		// Else fall through to skipForward.
+	} else {
+		// result == loadBlockIrrelevant. Enforce the upper bound here since
+		// don't want to bother moving to the next entry in the top level index
+		// if upper bound is already exceeded. Note that the next entry starts
+		// with keys >= topLevelIndex.Separator() since even though this is the
+		// block separator, the same user key can span multiple index blocks.
+		// If upper is exclusive we pass orEqual=true below, else we require the
+		// separator to be strictly greater than upper.
+		if i.secondLevel.upper != nil && PI(&i.topLevelIndex).SeparatorGT(
+			i.secondLevel.upper, !i.secondLevel.endKeyInclusive) {
+			i.secondLevel.exhaustedBounds = +1
+		}
+	}
+	// NB: skipForward checks whether exhaustedBounds is already +1.
+	kv := i.skipForward()
+	if kv == nil {
+		return nil, base.KVMeta{}
+	}
+	meta := i.extractMetaFromCurrentPosition()
+	return kv, meta
+}
+
+// NextWithMeta moves the iterator to the next key/value pair and returns
+// both the key/value and the associated metadata. This method is used by
+// compaction iterators that need access to tiering metadata without adding
+// overhead to the common iteration path.
+func (i *twoLevelIterator[I, PI, D, PD]) NextWithMeta() (*base.InternalKV, base.KVMeta) {
+	kv := i.Next()
+	if kv == nil {
+		return nil, base.KVMeta{}
+	}
+	meta := i.extractMetaFromCurrentPosition()
+	return kv, meta
+}
+
+// extractMetaFromCurrentPosition extracts KVMeta from the current iterator position.
+// This method delegates to the underlying second level iterator if it supports
+// the specialized methods.
+func (i *twoLevelIterator[I, PI, D, PD]) extractMetaFromCurrentPosition() base.KVMeta {
+	if PD(&i.secondLevel.data).IsDataInvalidated() {
+		return base.KVMeta{}
+	}
+
+	// The dataBlockIterator constraint guarantees that PD(&i.secondLevel.data) implements MetaDecoder
+	return PD(&i.secondLevel.data).DecodeMeta()
+}
+
 // Last implements internalIterator.Last, as documented in the pebble
 // package. Note that Last only checks the lower bound. It is up to the caller
 // to ensure that key is less than the upper bound (e.g. via a call to

sstable/rowblk/rowblk_iter.go

@@ -1573,6 +1573,12 @@ func (i *Iter) Error() error {
 	return nil // infallible
 }
 
+// DecodeMeta implements the base.MetaDecoder interface.
+// Row-oriented blocks don't have tiering metadata, so this always returns empty metadata.
+func (i *Iter) DecodeMeta() base.KVMeta {
+	return base.KVMeta{}
+}
+
 // Close implements internalIterator.Close, as documented in the pebble
 // package.
 func (i *Iter) Close() error {