chore(metastore): use indexpointers columnar reader

Author: ivkalita

pkg/dataobj/metastore/iter.go

@@ -12,6 +12,7 @@ import (
 	"github.com/grafana/dskit/user"
 
 	"github.com/grafana/loki/v3/pkg/dataobj"
+	"github.com/grafana/loki/v3/pkg/dataobj/sections/indexpointers"
 	"github.com/grafana/loki/v3/pkg/dataobj/sections/pointers"
 )
 
@@ -332,6 +333,123 @@ func forEachMatchedPointerSectionKey(
 	return nil
 }
 
+func forEachIndexPointer(
+	ctx context.Context,
+	object *dataobj.Object,
+	sStart, sEnd *scalar.Timestamp,
+	f func(pointer indexpointers.IndexPointer),
+) error {
+	targetTenant, err := user.ExtractOrgID(ctx)
+	if err != nil {
+		return fmt.Errorf("extracting org ID: %w", err)
+	}
+	var reader indexpointers.Reader
+	defer reader.Close()
+
+	const batchSize = 1024
+	buf := make([]indexpointers.IndexPointer, 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 object.Sections().Filter(indexpointers.CheckSection) {
+		if section.Tenant != targetTenant {
+			continue
+		}
+
+		sec, err := indexpointers.Open(ctx, section)
+		if err != nil {
+			return fmt.Errorf("opening section: %w", err)
+		}
+
+		var (
+			colPath         *indexpointers.Column
+			colMinTimestamp *indexpointers.Column
+			colMaxTimestamp *indexpointers.Column
+		)
+
+		for _, c := range sec.Columns() {
+			if c.Type == indexpointers.ColumnTypePath {
+				colPath = c
+			}
+			if c.Type == indexpointers.ColumnTypeMinTimestamp {
+				colMinTimestamp = c
+			}
+			if c.Type == indexpointers.ColumnTypeMaxTimestamp {
+				colMaxTimestamp = c
+			}
+			if colPath != nil && colMinTimestamp != nil && colMaxTimestamp != nil {
+				break
+			}
+		}
+
+		if colPath == nil || colMinTimestamp == nil || colMaxTimestamp == nil {
+			return fmt.Errorf("one of the mandatory columns is missing: (path=%t, minTimestamp=%t, maxTimestamp=%t)", colPath == nil, colMinTimestamp == nil, colMaxTimestamp == nil)
+		}
+
+		reader.Reset(indexpointers.ReaderOptions{
+			Columns: sec.Columns(),
+			Predicates: []indexpointers.Predicate{
+				indexpointers.WhereTimeRangeOverlapsWith(colMinTimestamp, colMaxTimestamp, sStart, sEnd),
+			},
+		})
+
+		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 := sec.Columns()[colIdx]
+
+				switch pointerCol.Type {
+				case indexpointers.ColumnTypePath:
+					values := col.(*array.String)
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].Path = values.Value(rIdx)
+					}
+				case indexpointers.ColumnTypeMinTimestamp:
+					values := col.(*array.Timestamp)
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].StartTs = time.Unix(0, int64(values.Value(rIdx)))
+					}
+				case indexpointers.ColumnTypeMaxTimestamp:
+					values := col.(*array.Timestamp)
+					for rIdx := range numRows {
+						if col.IsNull(rIdx) {
+							continue
+						}
+						buf[rIdx].EndTs = time.Unix(0, int64(values.Value(rIdx)))
+					}
+				default:
+					continue
+				}
+			}
+
+			for rowIdx := range numRows {
+				f(buf[rowIdx])
+			}
+
+			if errors.Is(readErr, io.EOF) {
+				break
+			}
+		}
+	}
+
+	return nil
+}
+
 func findPointersColumnsByTypes(allColumns []*pointers.Column, columnTypes ...pointers.ColumnType) ([]*pointers.Column, error) {
 	result := make([]*pointers.Column, 0, len(columnTypes))
 

pkg/dataobj/metastore/object.go

@@ -401,10 +401,13 @@ func (m *ObjectMetastore) listObjectsFromTables(ctx context.Context, tablePaths
 	objects := make([][]string, len(tablePaths))
 	g, ctx := errgroup.WithContext(ctx)
 
+	sStart := scalar.NewTimestampScalar(arrow.Timestamp(start.UnixNano()), arrow.FixedWidthTypes.Timestamp_ns)
+	sEnd := scalar.NewTimestampScalar(arrow.Timestamp(end.UnixNano()), arrow.FixedWidthTypes.Timestamp_ns)
+
 	for i, path := range tablePaths {
 		g.Go(func() error {
 			var err error
-			objects[i], err = m.listObjects(ctx, path, start, end)
+			objects[i], err = m.listObjects(ctx, path, sStart, sEnd)
 			// If the metastore object is not found, it means it's outside of any existing window
 			// and we can safely ignore it.
 			if err != nil && !m.bucket.IsObjNotFoundErr(err) {
@@ -627,7 +630,7 @@ func addLabels(mtx *sync.Mutex, streams map[uint64][]*labels.Labels, newLabels *
 	streams[key] = append(streams[key], newLabels)
 }
 
-func (m *ObjectMetastore) listObjects(ctx context.Context, path string, start, end time.Time) ([]string, error) {
+func (m *ObjectMetastore) listObjects(ctx context.Context, path string, sStart, sEnd *scalar.Timestamp) ([]string, error) {
 	var buf bytes.Buffer
 	objectReader, err := m.bucket.Get(ctx, path)
 	if err != nil {
@@ -645,12 +648,7 @@ func (m *ObjectMetastore) listObjects(ctx context.Context, path string, start, e
 	}
 	var objectPaths []string
 
-	// Read all relevant entries from the table of contents
-	predicate := indexpointers.TimeRangeRowPredicate{
-		Start: start.UTC(),
-		End:   end.UTC(),
-	}
-	err = forEachIndexPointer(ctx, object, predicate, func(indexPointer indexpointers.IndexPointer) {
+	err = forEachIndexPointer(ctx, object, sStart, sEnd, func(indexPointer indexpointers.IndexPointer) {
 		objectPaths = append(objectPaths, indexPointer.Path)
 	})
 	if err != nil {
@@ -660,50 +658,6 @@ func (m *ObjectMetastore) listObjects(ctx context.Context, path string, start, e
 	return objectPaths, nil
 }
 
-func forEachIndexPointer(ctx context.Context, object *dataobj.Object, predicate indexpointers.RowPredicate, f func(indexpointers.IndexPointer)) error {
-	targetTenant, err := user.ExtractOrgID(ctx)
-	if err != nil {
-		return fmt.Errorf("extracting org ID: %w", err)
-	}
-	var reader indexpointers.RowReader
-	defer reader.Close()
-
-	buf := make([]indexpointers.IndexPointer, 1024)
-
-	for _, section := range object.Sections().Filter(indexpointers.CheckSection) {
-		if section.Tenant != targetTenant {
-			continue
-		}
-		sec, err := indexpointers.Open(ctx, section)
-		if err != nil {
-			return fmt.Errorf("opening section: %w", err)
-		}
-
-		reader.Reset(sec)
-		if predicate != nil {
-			err := reader.SetPredicate(predicate)
-			if err != nil {
-				return err
-			}
-		}
-
-		for {
-			num, err := reader.Read(ctx, buf)
-			if err != nil && !errors.Is(err, io.EOF) {
-				return err
-			}
-			if num == 0 && errors.Is(err, io.EOF) {
-				break
-			}
-			for _, indexPointer := range buf[:num] {
-				f(indexPointer)
-			}
-		}
-	}
-
-	return nil
-}
-
 func forEachStream(ctx context.Context, object *dataobj.Object, predicate streams.RowPredicate, f func(streams.Stream)) error {
 	targetTenant, err := user.ExtractOrgID(ctx)
 	if err != nil {