chore(metastore): stream pointers columnar reader

Use columnar reader to read the stream pointers

Author: ivkalita

pkg/dataobj/consumer/logsobj/builder.go

@@ -127,10 +127,10 @@ type BuilderConfig struct {
 // RegisterFlagsWithPrefix registers flags with the given prefix.
 func (cfg *BuilderConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
 	// Set defaults for base builder configuration
-	_ = cfg.BuilderBaseConfig.TargetPageSize.Set("2MB")
-	_ = cfg.BuilderBaseConfig.TargetObjectSize.Set("1GB")
-	_ = cfg.BuilderBaseConfig.BufferSize.Set("16MB")
-	_ = cfg.BuilderBaseConfig.TargetSectionSize.Set("128MB")
+	_ = cfg.TargetPageSize.Set("2MB")
+	_ = cfg.TargetObjectSize.Set("1GB")
+	_ = cfg.BufferSize.Set("16MB")
+	_ = cfg.TargetSectionSize.Set("128MB")
 	cfg.BuilderBaseConfig.RegisterFlagsWithPrefix(prefix, f)
 
 	f.StringVar(&cfg.DataobjSortOrder, prefix+"dataobj-sort-order", sortStreamASC, "The desired sort order of the logs section. Can either be `stream-asc` (order by streamID ascending and timestamp descending) or `timestamp-desc` (order by timestamp descending and streamID ascending).")

pkg/dataobj/index/config.go

@@ -21,10 +21,10 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 
 func (cfg *Config) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
 	// Set defaults for base builder configuration
-	_ = cfg.BuilderBaseConfig.TargetPageSize.Set("128KB")
-	_ = cfg.BuilderBaseConfig.TargetObjectSize.Set("64MB")
-	_ = cfg.BuilderBaseConfig.BufferSize.Set("2MB")
-	_ = cfg.BuilderBaseConfig.TargetSectionSize.Set("16MB")
+	_ = cfg.TargetPageSize.Set("128KB")
+	_ = cfg.TargetObjectSize.Set("64MB")
+	_ = cfg.BufferSize.Set("2MB")
+	_ = cfg.TargetSectionSize.Set("16MB")
 	cfg.BuilderBaseConfig.RegisterFlagsWithPrefix(prefix, f)
 
 	f.IntVar(&cfg.EventsPerIndex, prefix+"events-per-index", 32, "Experimental: The number of events to batch before building an index")

pkg/dataobj/metastore/iter.go

@@ -0,0 +1,260 @@
+package metastore
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"time"
+
+	"github.com/apache/arrow-go/v18/arrow/array"
+	"github.com/apache/arrow-go/v18/arrow/scalar"
+	"github.com/grafana/dskit/user"
+
+	"github.com/grafana/loki/v3/pkg/dataobj"
+	"github.com/grafana/loki/v3/pkg/dataobj/sections/pointers"
+)
+
+// forEachStreamSectionPointer iterates over all the section pointers that point to one of the
+// [streamIDs] in a given [indexObj] that overlap [sStart, sEnd] (inclusive) time range and
+// calls [f] on every found SectionPointer.
+func forEachStreamSectionPointer(
+	ctx context.Context,
+	indexObj *dataobj.Object,
+	sStart, sEnd *scalar.Timestamp,
+	streamIDs []int64,
+	f func(pointers.SectionPointer),
+) error {
+	targetTenant, err := user.ExtractOrgID(ctx)
+	if err != nil {
+		return fmt.Errorf("extracting org ID: %w", err)
+	}
+	var reader pointers.Reader
+	defer reader.Close()
+
+	// prepare streamIDs scalars only once, they are used in the predicate later
+	var sStreamIDs []scalar.Scalar
+	for _, streamID := range streamIDs {
+		sStreamIDs = append(sStreamIDs, scalar.NewInt64Scalar(streamID))
+	}
+
+	const batchSize = 128
+	buf := make([]streamSectionPointerBuilder, batchSize)
+
+	// iterate over the sections and fill buf column by column
+	// once the read operation is over invoke client's [f] on every read row (numRows not always the same as len(buf))
+	for _, section := range indexObj.Sections().Filter(pointers.CheckSection) {
+		if section.Tenant != targetTenant {
+			continue
+		}
+
+		sec, err := pointers.Open(ctx, section)
+		if err != nil {
+			return fmt.Errorf("opening section: %w", err)
+		}
+
+		pointerCols, err := findPointersColumnsByTypes(
+			sec.Columns(),
+			pointers.ColumnTypePath,
+			pointers.ColumnTypeSection,
+			pointers.ColumnTypeStreamID,
+			pointers.ColumnTypeStreamIDRef,
+			pointers.ColumnTypeMinTimestamp,
+			pointers.ColumnTypeMaxTimestamp,
+			pointers.ColumnTypeRowCount,
+			pointers.ColumnTypeUncompressedSize,
+		)
+		if err != nil {
+			return fmt.Errorf("finding pointers columns: %w", err)
+		}
+
+		var (
+			colStreamID     *pointers.Column
+			colMinTimestamp *pointers.Column
+			colMaxTimestamp *pointers.Column
+		)
+
+		for _, c := range pointerCols {
+			if c.Type == pointers.ColumnTypeStreamID {
+				colStreamID = c
+			}
+			if c.Type == pointers.ColumnTypeMinTimestamp {
+				colMinTimestamp = c
+			}
+			if c.Type == pointers.ColumnTypeMaxTimestamp {
+				colMaxTimestamp = c
+			}
+			if colStreamID != nil && colMinTimestamp != nil && colMaxTimestamp != nil {
+				break
+			}
+		}
+
+		if colStreamID == nil || colMinTimestamp == nil || colMaxTimestamp == nil {
+			return fmt.Errorf(
+				"one of mandatory columns is missing: (streamID=%t, minTimestamp=%t, maxTimestamp=%t)",
+				colStreamID == nil, colMinTimestamp == nil, colMaxTimestamp == nil,
+			)
+		}
+
+		reader.Reset(pointers.ReaderOptions{
+			Columns: pointerCols,
+			Predicates: []pointers.Predicate{
+				buildPointersTimeRangePredicate(colMinTimestamp, colMaxTimestamp, sStart, sEnd),
+				pointers.InPredicate{
+					Column: colStreamID,
+					Values: sStreamIDs,
+				},
+			},
+		})
+
+		for {
+			rec, readErr := reader.Read(ctx, batchSize)
+			if readErr != nil && !errors.Is(readErr, io.EOF) {
+				return fmt.Errorf("reading recordBatch: %w", readErr)
+			}
+			numRows := int(rec.NumRows())
+			if numRows == 0 && errors.Is(readErr, io.EOF) {
+				break
+			}
+
+			for colIdx := range int(rec.NumCols()) {
+				col := rec.Column(colIdx)
+				pointerCol := pointerCols[colIdx]
+
+				switch pointerCol.Type {
+				case pointers.ColumnTypePath:
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].path = col.(*array.String).Value(rIdx)
+					}
+				case pointers.ColumnTypeSection:
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].section = col.(*array.Int64).Value(rIdx)
+					}
+				case pointers.ColumnTypeStreamID:
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].streamID = col.(*array.Int64).Value(rIdx)
+					}
+				case pointers.ColumnTypeStreamIDRef:
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].streamIDRef = col.(*array.Int64).Value(rIdx)
+					}
+				case pointers.ColumnTypeMinTimestamp:
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].start = time.Unix(0, int64(col.(*array.Timestamp).Value(rIdx)))
+					}
+				case pointers.ColumnTypeMaxTimestamp:
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].end = time.Unix(0, int64(col.(*array.Timestamp).Value(rIdx)))
+					}
+				case pointers.ColumnTypeRowCount:
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].lineCount = col.(*array.Int64).Value(rIdx)
+					}
+				case pointers.ColumnTypeUncompressedSize:
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].uncompressedSize = col.(*array.Int64).Value(rIdx)
+					}
+				default:
+					continue
+				}
+			}
+
+			for rowIdx := range numRows {
+				b := buf[rowIdx]
+				f(pointers.SectionPointer{
+					Path:             b.path,
+					Section:          b.section,
+					PointerKind:      pointers.PointerKindStreamIndex,
+					StreamID:         b.streamID,
+					StreamIDRef:      b.streamIDRef,
+					StartTs:          b.start,
+					EndTs:            b.end,
+					LineCount:        b.lineCount,
+					UncompressedSize: b.uncompressedSize,
+				})
+			}
+
+			if errors.Is(readErr, io.EOF) {
+				break
+			}
+		}
+	}
+
+	return nil
+}
+
+type streamSectionPointerBuilder struct {
+	path             string
+	section          int64
+	streamID         int64
+	streamIDRef      int64
+	start            time.Time
+	end              time.Time
+	lineCount        int64
+	uncompressedSize int64
+}
+
+func (b streamSectionPointerBuilder) build() streamSectionPointerBuilder {
+	return streamSectionPointerBuilder{
+		path:             b.path,
+		section:          b.section,
+		streamID:         b.streamID,
+		streamIDRef:      b.streamIDRef,
+		start:            b.start,
+		end:              b.end,
+		lineCount:        b.lineCount,
+		uncompressedSize: b.uncompressedSize,
+	}
+}
+
+func buildPointersTimeRangePredicate(
+	colMinTimestamp, colMaxTimestamp *pointers.Column,
+	sStart scalar.Scalar, sEnd scalar.Scalar,
+) pointers.Predicate {
+	return pointers.AndPredicate{
+		Left: pointers.OrPredicate{
+			Left: pointers.EqualPredicate{
+				Column: colMaxTimestamp,
+				Value:  sStart,
+			},
+			Right: pointers.GreaterThanPredicate{
+				Column: colMaxTimestamp,
+				Value:  sStart,
+			},
+		},
+		Right: pointers.OrPredicate{
+			Left: pointers.EqualPredicate{
+				Column: colMinTimestamp,
+				Value:  sEnd,
+			},
+			Right: pointers.LessThanPredicate{
+				Column: colMinTimestamp,
+				Value:  sEnd,
+			},
+		},
+	}
+}

pkg/dataobj/metastore/object.go

@@ -14,6 +14,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/apache/arrow-go/v18/arrow/scalar"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/tenant"
@@ -206,11 +208,7 @@ func (m *ObjectMetastore) Sections(ctx context.Context, start, end time.Time, ma
 
 	// Search the stream sections of the matching objects to find matching streams
 	streamMatchers := streamPredicateFromMatchers(start, end, matchers...)
-	pointerPredicate := pointers.TimeRangeRowPredicate{
-		Start: start,
-		End:   end,
-	}
-	streamSectionPointers, err := m.getSectionsForStreams(ctx, indexObjects, streamMatchers, pointerPredicate)
+	streamSectionPointers, err := m.getSectionsForStreams(ctx, indexObjects, streamMatchers, start, end)
 	if err != nil {
 		return nil, err
 	}
@@ -504,12 +502,15 @@ func (m *ObjectMetastore) listStreamIDsFromLogObjects(ctx context.Context, objec
 
 // getSectionsForStreams reads the section data from matching streams and aggregates them into section descriptors.
 // This is an exact lookup and includes metadata from the streams in each section: the stream IDs, the min-max timestamps, the number of bytes & number of lines.
-func (m *ObjectMetastore) getSectionsForStreams(ctx context.Context, indexObjects []*dataobj.Object, streamPredicate streams.RowPredicate, timeRangePredicate pointers.TimeRangeRowPredicate) ([]*DataobjSectionDescriptor, error) {
+func (m *ObjectMetastore) getSectionsForStreams(ctx context.Context, indexObjects []*dataobj.Object, streamPredicate streams.RowPredicate, start, end time.Time) ([]*DataobjSectionDescriptor, error) {
 	if streamPredicate == nil {
 		// At least one stream matcher is required, currently.
 		return nil, nil
 	}
 
+	sStart := scalar.NewTimestampScalar(arrow.Timestamp(start.UnixNano()), arrow.FixedWidthTypes.Timestamp_ns)
+	sEnd := scalar.NewTimestampScalar(arrow.Timestamp(end.UnixNano()), arrow.FixedWidthTypes.Timestamp_ns)
+
 	timer := prometheus.NewTimer(m.metrics.streamFilterTotalDuration)
 	defer timer.ObserveDuration()
 
@@ -540,7 +541,8 @@ func (m *ObjectMetastore) getSectionsForStreams(ctx context.Context, indexObject
 
 			objectSectionDescriptors := make(map[SectionKey]*DataobjSectionDescriptor)
 			sectionPointerReadTimer := prometheus.NewTimer(m.metrics.streamFilterPointersReadDuration)
-			err = forEachObjPointer(ctx, indexObject, timeRangePredicate, matchingStreamIDs, func(pointer pointers.SectionPointer) {
+
+			err = forEachStreamSectionPointer(ctx, indexObject, sStart, sEnd, matchingStreamIDs, func(pointer pointers.SectionPointer) {
 				key.ObjectPath = pointer.Path
 				key.SectionIdx = pointer.Section
 
@@ -551,6 +553,7 @@ func (m *ObjectMetastore) getSectionsForStreams(ctx context.Context, indexObject
 				}
 				sectionDescriptor.Merge(pointer)
 			})
+
 			if err != nil {
 				return fmt.Errorf("reading section pointers from index: %w", err)
 			}
@@ -827,3 +830,19 @@ func dedupeAndSort(objects [][]string) []string {
 	sort.Strings(paths)
 	return paths
 }
+
+func findPointersColumnsByTypes(allColumns []*pointers.Column, columnTypes ...pointers.ColumnType) ([]*pointers.Column, error) {
+	result := make([]*pointers.Column, 0, len(columnTypes))
+
+	for _, c := range allColumns {
+		for _, neededType := range columnTypes {
+			if neededType != c.Type {
+				continue
+			}
+
+			result = append(result, c)
+		}
+	}
+
+	return result, nil
+}