chore(engine): add concurrency to Pipeline.Open (#20839)

pkg/dataobj/metastore/index_sections_reader.go

@@ -20,6 +20,7 @@ import (
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/user"
 	"github.com/prometheus/prometheus/model/labels"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/grafana/loki/v3/pkg/dataobj"
 	"github.com/grafana/loki/v3/pkg/dataobj/sections/pointers"
@@ -131,40 +132,106 @@ func (r *indexSectionsReader) init(ctx context.Context) error {
 		predicateKeys[predicate.Name] = struct{}{}
 	}
 
-	for _, section := range r.obj.Sections().Filter(streams.CheckSection) {
+	var (
+		unopenedStreams  []*dataobj.Section
+		unopenedPointers []*dataobj.Section
+	)
+
+	for _, section := range r.obj.Sections() {
 		if section.Tenant != targetTenant {
 			continue
 		}
 
-		sec, err := streams.Open(ctx, section)
-		if err != nil {
-			return fmt.Errorf("opening streams section: %w", err)
-		} else if err := r.openStreamsReader(ctx, sec, sStart, sEnd, predicateKeys); err != nil {
-			return err
+		switch {
+		case streams.CheckSection(section):
+			unopenedStreams = append(unopenedStreams, section)
+		case pointers.CheckSection(section):
+			unopenedPointers = append(unopenedPointers, section)
 		}
 	}
 
-	for _, section := range r.obj.Sections().Filter(pointers.CheckSection) {
-		if section.Tenant != targetTenant {
-			continue
-		}
+	r.streamsReaders = make([]*streams.Reader, len(unopenedStreams))
+	r.pointersReaders = make([]*pointers.Reader, len(unopenedPointers))
+	r.bloomReaders = make([]*pointers.Reader, len(unopenedPointers))
 
-		sec, err := pointers.Open(ctx, section)
-		if err != nil {
-			return fmt.Errorf("opening pointers section: %w", err)
-		}
+	g, ctx := errgroup.WithContext(ctx)
 
-		if err := r.openStreamPointersReader(ctx, sec, sStart, sEnd); err != nil {
-			return err
-		}
-		if err := r.openBloomPointersReader(ctx, sec); err != nil {
-			return err
-		}
+	for i, section := range unopenedStreams {
+		g.Go(func() error {
+			sec, err := streams.Open(ctx, section)
+			if err != nil {
+				return fmt.Errorf("opening streams section: %w", err)
+			}
+
+			reader, err := r.openStreamsReader(ctx, sec, sStart, sEnd, predicateKeys)
+			if err != nil {
+				return err
+			}
+
+			r.streamsReaders[i] = reader
+			return nil
+		})
+	}
+
+	for i, section := range unopenedPointers {
+		g.Go(func() error {
+			sec, err := pointers.Open(ctx, section)
+			if err != nil {
+				return fmt.Errorf("opening pointers section: %w", err)
+			}
+
+			sg, ctx := errgroup.WithContext(ctx)
+
+			sg.Go(func() error {
+				pointersReader, err := r.openStreamPointersReader(ctx, sec, sStart, sEnd)
+				if err != nil {
+					return err
+				}
+
+				r.pointersReaders[i] = pointersReader
+				return nil
+			})
+
+			sg.Go(func() error {
+				bloomReader, err := r.openBloomPointersReader(ctx, sec)
+				if err != nil {
+					return err
+				}
+
+				r.bloomReaders[i] = bloomReader
+				return nil
+			})
+
+			return sg.Wait()
+		})
+	}
+
+	if err := g.Wait(); err != nil {
+		closeAll(r.streamsReaders)
+		closeAll(r.pointersReaders)
+		closeAll(r.bloomReaders)
+		return err
 	}
 
 	return nil
 }
 
+type closable interface {
+	comparable
+	io.Closer
+}
+
+func closeAll[C closable](cs []C) {
+	var zero C
+
+	for _, c := range cs {
+		if c == zero {
+			continue
+		}
+		_ = c.Close()
+	}
+}
+
 func (r *indexSectionsReader) scalarTimestamps() (*scalar.Timestamp, *scalar.Timestamp) {
 	sStart := scalar.NewTimestampScalar(arrow.Timestamp(r.start.UnixNano()), arrow.FixedWidthTypes.Timestamp_ns)
 	sEnd := scalar.NewTimestampScalar(arrow.Timestamp(r.end.UnixNano()), arrow.FixedWidthTypes.Timestamp_ns)
@@ -176,10 +243,10 @@ func (r *indexSectionsReader) openStreamsReader(
 	sec *streams.Section,
 	sStart, sEnd *scalar.Timestamp,
 	predicateKeys map[string]struct{},
-) error {
+) (*streams.Reader, error) {
 	predicates, err := buildStreamReaderPredicate(sec, sStart, sEnd, r.matchers)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	targetColumns := make([]*streams.Column, 0, len(sec.Columns()))
@@ -198,18 +265,17 @@ func (r *indexSectionsReader) openStreamsReader(
 		Allocator:  memory.DefaultAllocator,
 	})
 	if err := reader.Open(ctx); err != nil {
-		return fmt.Errorf("opening streams reader: %w", err)
+		return nil, fmt.Errorf("opening streams reader: %w", err)
 	}
 
-	r.streamsReaders = append(r.streamsReaders, reader)
-	return nil
+	return reader, nil
 }
 
 func (r *indexSectionsReader) openStreamPointersReader(
 	ctx context.Context,
 	sec *pointers.Section,
 	sStart, sEnd *scalar.Timestamp,
-) error {
+) (*pointers.Reader, error) {
 	cols, err := findPointersColumnsByTypes(
 		sec.Columns(),
 		pointers.ColumnTypePath,
@@ -223,7 +289,7 @@ func (r *indexSectionsReader) openStreamPointersReader(
 		pointers.ColumnTypeUncompressedSize,
 	)
 	if err != nil {
-		return fmt.Errorf("finding pointers columns: %w", err)
+		return nil, fmt.Errorf("finding pointers columns: %w", err)
 	}
 
 	var (
@@ -247,7 +313,7 @@ func (r *indexSectionsReader) openStreamPointersReader(
 
 	if colStreamID == nil || colMinTimestamp == nil || colMaxTimestamp == nil || colPointerKind == nil {
 		// Section has no rows with stream-based indices, skip it
-		return nil
+		return nil, nil
 	}
 
 	reader := pointers.NewReader(pointers.ReaderOptions{
@@ -263,16 +329,15 @@ func (r *indexSectionsReader) openStreamPointersReader(
 		StreamIDToLabelNames: r.labelNamesByStream,
 	})
 	if err := reader.Open(ctx); err != nil {
-		return fmt.Errorf("opening pointers reader: %w", err)
+		return nil, fmt.Errorf("opening pointers reader: %w", err)
 	}
 
-	r.pointersReaders = append(r.pointersReaders, reader)
-	return nil
+	return reader, nil
 }
 
-func (r *indexSectionsReader) openBloomPointersReader(ctx context.Context, sec *pointers.Section) error {
+func (r *indexSectionsReader) openBloomPointersReader(ctx context.Context, sec *pointers.Section) (*pointers.Reader, error) {
 	if len(r.predicates) == 0 {
-		return nil
+		return nil, nil
 	}
 
 	pointerCols, err := findPointersColumnsByTypes(
@@ -284,7 +349,7 @@ func (r *indexSectionsReader) openBloomPointersReader(ctx context.Context, sec *
 		pointers.ColumnTypeValuesBloomFilter,
 	)
 	if err != nil {
-		return fmt.Errorf("finding bloom pointers columns: %w", err)
+		return nil, fmt.Errorf("finding bloom pointers columns: %w", err)
 	}
 
 	var (
@@ -312,7 +377,7 @@ func (r *indexSectionsReader) openBloomPointersReader(ctx context.Context, sec *
 
 	if colPath == nil || colSection == nil || colColumnName == nil || colBloom == nil || colPointerKind == nil {
 		// The section has no rows for blooms and can be ignored completely
-		return nil
+		return nil, nil
 	}
 
 	// Because a row only contains a bloom filter for one column, we must use OR
@@ -349,11 +414,10 @@ func (r *indexSectionsReader) openBloomPointersReader(ctx context.Context, sec *
 		Allocator: memory.DefaultAllocator,
 	})
 	if err := reader.Open(ctx); err != nil {
-		return fmt.Errorf("opening bloom pointers reader: %w", err)
+		return nil, fmt.Errorf("opening bloom pointers reader: %w", err)
 	}
 
-	r.bloomReaders = append(r.bloomReaders, reader)
-	return nil
+	return reader, nil
 }
 
 func (r *indexSectionsReader) Read(ctx context.Context) (arrow.RecordBatch, error) {
@@ -468,15 +532,9 @@ func (r *indexSectionsReader) allLabelNames() []string {
 }
 
 func (r *indexSectionsReader) Close() {
-	for _, sr := range r.streamsReaders {
-		_ = sr.Close()
-	}
-	for _, pr := range r.pointersReaders {
-		_ = pr.Close()
-	}
-	for _, br := range r.bloomReaders {
-		_ = br.Close()
-	}
+	closeAll(r.streamsReaders)
+	closeAll(r.pointersReaders)
+	closeAll(r.bloomReaders)
 }
 
 func (r *indexSectionsReader) addLabelNamesForStream(streamID int64, names []string) {
@@ -706,6 +764,12 @@ func (r *indexSectionsReader) readMatchedSectionKeys(ctx context.Context) (map[S
 	sectionMatches := make(map[SectionKey]map[int]struct{})
 
 	for _, br := range r.bloomReaders {
+		if br == nil {
+			// We can have nil readers when a section was skipped due to not
+			// having relevant data.
+			continue
+		}
+
 		var (
 			pathColumnIndex       = slices.IndexFunc(br.Columns(), func(c *pointers.Column) bool { return c.Type == pointers.ColumnTypePath })
 			sectionColumnIndex    = slices.IndexFunc(br.Columns(), func(c *pointers.Column) bool { return c.Type == pointers.ColumnTypeSection })

pkg/engine/internal/executor/dataobjscan.go

@@ -13,6 +13,7 @@ import (
 	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/apache/arrow-go/v18/arrow/scalar"
 	"github.com/go-kit/log"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/grafana/loki/v3/pkg/dataobj/sections/logs"
 	"github.com/grafana/loki/v3/pkg/dataobj/sections/streams"
@@ -73,13 +74,40 @@ func (s *dataobjScan) init(ctx context.Context) error {
 		return nil
 	}
 
-	// [dataobjScan.initLogs] depends on the result of [dataobjScan.initStreams]
-	// (to know whether label columns are needed), so we must initialize streams
-	// first.
-	if err := s.initStreams(ctx); err != nil {
-		return fmt.Errorf("initializing streams: %w", err)
-	} else if err := s.initLogs(ctx); err != nil {
-		return fmt.Errorf("initializing logs: %w", err)
+	columnsToRead := projectedLabelColumns(s.opts.StreamsSection, s.opts.Projections)
+	includeStreamID := len(columnsToRead) > 0 || len(s.opts.StreamIDs) > 0
+
+	g, ctx := errgroup.WithContext(ctx)
+
+	g.Go(func() error {
+		streams, streamsInjector, err := s.initStreams(ctx, columnsToRead)
+		if err != nil {
+			return fmt.Errorf("initializing streams: %w", err)
+		}
+
+		s.streams = streams
+		s.streamsInjector = streamsInjector
+		return nil
+	})
+
+	g.Go(func() error {
+		reader, desiredSchema, err := s.initLogs(ctx, includeStreamID)
+		if err != nil {
+			return fmt.Errorf("initializing logs: %w", err)
+		}
+		s.reader = reader
+		s.desiredSchema = desiredSchema
+		return nil
+	})
+
+	if err := g.Wait(); err != nil {
+		if s.streams != nil {
+			s.streams.Close()
+		}
+		if s.reader != nil {
+			_ = s.reader.Close()
+		}
+		return err
 	}
 
 	s.initialized = true
@@ -87,29 +115,25 @@ func (s *dataobjScan) init(ctx context.Context) error {
 	return nil
 }
 
-func (s *dataobjScan) initStreams(ctx context.Context) error {
+func (s *dataobjScan) initStreams(ctx context.Context, columnsToRead []*streams.Column) (*streamsView, *streamInjector, error) {
 	if s.opts.StreamsSection == nil {
-		return fmt.Errorf("no streams section provided")
+		return nil, nil, fmt.Errorf("no streams section provided")
 	}
 
-	columnsToRead := projectedLabelColumns(s.opts.StreamsSection, s.opts.Projections)
 	if len(columnsToRead) == 0 {
-		s.streams = nil
-		s.streamsInjector = nil
-		return nil
+		return nil, nil, nil
 	}
 
-	s.streams = newStreamsView(s.opts.StreamsSection, &streamsViewOptions{
+	streamsView := newStreamsView(s.opts.StreamsSection, &streamsViewOptions{
 		StreamIDs:    s.opts.StreamIDs,
 		LabelColumns: columnsToRead,
 		BatchSize:    int(s.opts.BatchSize),
 	})
-	if err := s.streams.Open(ctx); err != nil {
-		return fmt.Errorf("opening streams view: %w", err)
+	if err := streamsView.Open(ctx); err != nil {
+		return nil, nil, fmt.Errorf("opening streams view: %w", err)
 	}
 
-	s.streamsInjector = newStreamInjector(s.streams)
-	return nil
+	return streamsView, newStreamInjector(streamsView), nil
 }
 
 // projectedLabelColumns returns the label columns to read from the given
@@ -167,17 +191,17 @@ func projectedLabelColumns(sec *streams.Section, projections []physical.ColumnEx
 	return found
 }
 
-func (s *dataobjScan) initLogs(ctx context.Context) error {
+func (s *dataobjScan) initLogs(ctx context.Context, includeStreamID bool) (*logs.Reader, *arrow.Schema, error) {
 	if s.opts.LogsSection == nil {
-		return fmt.Errorf("no logs section provided")
+		return nil, nil, fmt.Errorf("no logs section provided")
 	}
 
 	predicates := s.opts.Predicates
 
 	var columnsToRead []*logs.Column
 
-	if s.streams != nil || len(s.opts.StreamIDs) > 0 {
-		// We're reading sreams, so we need to include the stream ID column.
+	if includeStreamID {
+		// We're reading streams, so we need to include the stream ID column.
 		var streamIDColumn *logs.Column
 
 		for _, col := range s.opts.LogsSection.Columns() {
@@ -189,7 +213,7 @@ func (s *dataobjScan) initLogs(ctx context.Context) error {
 		}
 
 		if streamIDColumn == nil {
-			return fmt.Errorf("logs section does not contain stream ID column")
+			return nil, nil, fmt.Errorf("logs section does not contain stream ID column")
 		}
 
 		if len(s.opts.StreamIDs) > 0 {
@@ -204,7 +228,7 @@ func (s *dataobjScan) initLogs(ctx context.Context) error {
 
 	columnsToRead = append(columnsToRead, projectedLogsColumns(s.opts.LogsSection, s.opts.Projections)...)
 
-	s.reader = logs.NewReader(logs.ReaderOptions{
+	reader := logs.NewReader(logs.ReaderOptions{
 		// TODO(rfratto): is it possible to hit an edge case where len(columnsToRead)
 		// == 0, indicating that we don't need to read any logs at all? How should we
 		// handle that?
@@ -213,14 +237,15 @@ func (s *dataobjScan) initLogs(ctx context.Context) error {
 		Predicates: predicates,
 		Allocator:  memory.DefaultAllocator,
 	})
-	if err := s.reader.Open(ctx); err != nil {
-		return fmt.Errorf("opening logs reader: %w", err)
+	if err := reader.Open(ctx); err != nil {
+		return nil, nil, fmt.Errorf("opening logs reader: %w", err)
 	}
 
 	// Create the engine-compatible expected schema for the logs section.
-	origSchema := s.reader.Schema()
+	origSchema := reader.Schema()
 	if got, want := origSchema.NumFields(), len(columnsToRead); got != want {
-		return fmt.Errorf("logs.Reader returned schema with %d fields, expected %d", got, want)
+		_ = reader.Close()
+		return nil, nil, fmt.Errorf("logs.Reader returned schema with %d fields, expected %d", got, want)
 	}
 
 	// Convert the logs columns to engine-compatible fields.
@@ -228,13 +253,13 @@ func (s *dataobjScan) initLogs(ctx context.Context) error {
 	for _, col := range columnsToRead {
 		field, err := logsColumnToEngineField(col)
 		if err != nil {
-			return err
+			_ = reader.Close()
+			return nil, nil, err
 		}
 		desiredFields = append(desiredFields, field)
 	}
 
-	s.desiredSchema = arrow.NewSchema(desiredFields, nil)
-	return nil
+	return reader, arrow.NewSchema(desiredFields, nil), nil
 }
 
 func makeScalars[S ~[]E, E any](s S) []scalar.Scalar {

pkg/engine/internal/executor/executor.go

@@ -14,6 +14,7 @@ import (
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/grafana/loki/v3/pkg/dataobj"
 	"github.com/grafana/loki/v3/pkg/dataobj/metastore"
@@ -155,6 +156,9 @@ func (c *Context) executeDataObjScan(ctx context.Context, node *physical.DataObj
 	span.AddEvent("opened dataobj")
 
 	var (
+		foundStreamsSection *dataobj.Section
+		foundLogsSection    *dataobj.Section
+
 		streamsSection *streams.Section
 		logsSection    *logs.Section
 	)
@@ -164,41 +168,60 @@ func (c *Context) executeDataObjScan(ctx context.Context, node *physical.DataObj
 		return errorPipeline(ctx, fmt.Errorf("missing org ID: %w", err))
 	}
 
-	for _, sec := range obj.Sections().Filter(streams.CheckSection) {
+	var logsSectionIndex int
+	for _, sec := range obj.Sections() {
 		if sec.Tenant != tenant {
+			if logs.CheckSection(sec) {
+				logsSectionIndex++
+			}
 			continue
 		}
 
-		if streamsSection != nil {
-			return errorPipeline(ctx, fmt.Errorf("multiple streams sections found in data object %q", node.Location))
-		}
+		switch {
+		case streams.CheckSection(sec):
+			if foundStreamsSection != nil {
+				return errorPipeline(ctx, fmt.Errorf("multiple streams sections found in data object %q", node.Location))
+			}
+			foundStreamsSection = sec
 
+		case logs.CheckSection(sec):
+			if logsSectionIndex == node.Section {
+				foundLogsSection = sec
+			}
+			logsSectionIndex++
+		}
+	}
+
+	if foundStreamsSection == nil {
+		return errorPipeline(ctx, fmt.Errorf("streams section not found in data object %q", node.Location))
+	} else if foundLogsSection == nil {
+		return errorPipeline(ctx, fmt.Errorf("logs section %d not found in data object %q", node.Section, node.Location))
+	}
+
+	g, ctx := errgroup.WithContext(ctx)
+
+	g.Go(func() error {
 		var err error
-		streamsSection, err = streams.Open(ctx, sec)
+		streamsSection, err = streams.Open(ctx, foundStreamsSection)
 		if err != nil {
-			return errorPipeline(ctx, fmt.Errorf("opening streams section %q: %w", sec.Type, err))
+			return fmt.Errorf("opening streams section %q: %w", foundStreamsSection.Type, err)
 		}
 		span.AddEvent("opened streams section")
-	}
-	if streamsSection == nil {
-		return errorPipeline(ctx, fmt.Errorf("streams section not found in data object %q", node.Location))
-	}
-
-	for i, sec := range obj.Sections().Filter(logs.CheckSection) {
-		if i != node.Section {
-			continue
-		}
+		return nil
+	})
 
+	g.Go(func() error {
 		var err error
-		logsSection, err = logs.Open(ctx, sec)
+		logsSection, err = logs.Open(ctx, foundLogsSection)
 		if err != nil {
-			return errorPipeline(ctx, fmt.Errorf("opening logs section %q: %w", sec.Type, err))
+			return fmt.Errorf("opening logs section %q: %w", foundLogsSection.Type, err)
 		}
 		span.AddEvent("opened logs section")
-		break
-	}
-	if logsSection == nil {
-		return errorPipeline(ctx, fmt.Errorf("logs section %d not found in data object %q", node.Section, node.Location))
+		return nil
+	})
+
+	if err := g.Wait(); err != nil {
+		return errorPipeline(ctx, err)
 	}
 
 	// Filter streams if a filterer is configured

pkg/engine/internal/executor/merge.go

@@ -50,12 +50,7 @@ func newMergePipeline(inputs []Pipeline, maxPrefetch int) (*Merge, error) {
 
 // Open opens all children pipelines.
 func (m *Merge) Open(ctx context.Context) error {
-	for _, p := range m.inputs {
-		if err := p.Open(ctx); err != nil {
-			return err
-		}
-	}
-	return nil
+	return openInputsConcurrently(ctx, m.inputs)
 }
 
 func (m *Merge) init(ctx context.Context) {

pkg/engine/internal/executor/pipeline.go

@@ -10,6 +10,7 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/grafana/loki/v3/pkg/xcap"
 )
@@ -88,14 +89,17 @@ func newGenericPipeline(read readFunc, inputs ...Pipeline) *GenericPipeline {
 
 var _ Pipeline = (*GenericPipeline)(nil)
 
+func openInputsConcurrently(ctx context.Context, inputs []Pipeline) error {
+	g, ctx := errgroup.WithContext(ctx)
+	for _, in := range inputs {
+		g.Go(func() error { return in.Open(ctx) })
+	}
+	return g.Wait()
+}
+
 // Open implements Pipeline.
 func (p *GenericPipeline) Open(ctx context.Context) error {
-	for _, inp := range p.inputs {
-		if err := inp.Open(ctx); err != nil {
-			return err
-		}
-	}
-	return nil
+	return openInputsConcurrently(ctx, p.inputs)
 }
 
 // Read implements Pipeline.

pkg/engine/internal/executor/range_aggregation.go

@@ -114,12 +114,7 @@ func (r *rangeAggregationPipeline) init() {
 
 // Open opens all input pipelines.
 func (r *rangeAggregationPipeline) Open(ctx context.Context) error {
-	for _, input := range r.inputs {
-		if err := input.Open(ctx); err != nil {
-			return err
-		}
-	}
-	return nil
+	return openInputsConcurrently(ctx, r.inputs)
 }
 
 // Read reads the next value into its state.

pkg/engine/internal/executor/topk.go

@@ -123,12 +123,7 @@ func guessLokiType(ref types.ColumnRef) (types.DataType, error) {
 
 // Open opens all input pipelines.
 func (p *topkPipeline) Open(ctx context.Context) error {
-	for _, in := range p.inputs {
-		if err := in.Open(ctx); err != nil {
-			return err
-		}
-	}
-	return nil
+	return openInputsConcurrently(ctx, p.inputs)
 }
 
 // Read computes the topk as the next record. Read blocks until all input

pkg/engine/internal/executor/vector_aggregate.go

@@ -88,12 +88,7 @@ func newVectorAggregationPipeline(inputs []Pipeline, evaluator *expressionEvalua
 
 // Open opens all input pipelines.
 func (v *vectorAggregationPipeline) Open(ctx context.Context) error {
-	for _, input := range v.inputs {
-		if err := input.Open(ctx); err != nil {
-			return err
-		}
-	}
-	return nil
+	return openInputsConcurrently(ctx, v.inputs)
 }
 
 // Read reads the next value into its state.