Make use of the BinaryJoinIterator and remove JoinIterator

Author: simonswine

pkg/phlaredb/block_querier.go

@@ -988,26 +988,26 @@ func (b *singleBlockQuerier) SelectMatchingProfiles(ctx context.Context, params
 	}
 
 	var (
-		buf       [][]parquet.Value
-		joinIters []query.Iterator
+		buf [][]parquet.Value
+	)
+
+	pIt := query.NewBinaryJoinIterator(
+		0,
+		b.profiles.columnIter(ctx, "SeriesIndex", newMapPredicate(lblsPerRef), "SeriesIndex"),
+		b.profiles.columnIter(ctx, "TimeNanos", query.NewIntBetweenPredicate(model.Time(params.Start).UnixNano(), model.Time(params.End).UnixNano()), "TimeNanos"),
 	)
 
 	if b.meta.Version >= 2 {
-		joinIters = []query.Iterator{
-			b.profiles.columnIter(ctx, "SeriesIndex", newMapPredicate(lblsPerRef), "SeriesIndex"),
-			b.profiles.columnIter(ctx, "TimeNanos", query.NewIntBetweenPredicate(model.Time(params.Start).UnixNano(), model.Time(params.End).UnixNano()), "TimeNanos"),
+		pIt = query.NewBinaryJoinIterator(
+			0,
+			pIt,
 			b.profiles.columnIter(ctx, "StacktracePartition", nil, "StacktracePartition"),
-		}
+		)
 		buf = make([][]parquet.Value, 3)
 	} else {
-		joinIters = []query.Iterator{
-			b.profiles.columnIter(ctx, "SeriesIndex", newMapPredicate(lblsPerRef), "SeriesIndex"),
-			b.profiles.columnIter(ctx, "TimeNanos", query.NewIntBetweenPredicate(model.Time(params.Start).UnixNano(), model.Time(params.End).UnixNano()), "TimeNanos"),
-		}
 		buf = make([][]parquet.Value, 2)
 	}
 
-	pIt := query.NewJoinIterator(0, joinIters, nil)
 	iters := make([]iter.Iterator[Profile], 0, len(lblsPerRef))
 	defer pIt.Close()
 

pkg/phlaredb/head_queriers.go

@@ -48,14 +48,13 @@ func (q *headOnDiskQuerier) SelectMatchingProfiles(ctx context.Context, params *
 		start = model.Time(params.Start)
 		end   = model.Time(params.End)
 	)
-	pIt := query.NewJoinIterator(
-		0,
-		[]query.Iterator{
+	pIt := query.NewBinaryJoinIterator(0,
+		query.NewBinaryJoinIterator(
+			0,
 			rowIter,
 			q.rowGroup().columnIter(ctx, "TimeNanos", query.NewIntBetweenPredicate(start.UnixNano(), end.UnixNano()), "TimeNanos"),
-			q.rowGroup().columnIter(ctx, "StacktracePartition", nil, "StacktracePartition"),
-		},
-		nil,
+		),
+		q.rowGroup().columnIter(ctx, "StacktracePartition", nil, "StacktracePartition"),
 	)
 	defer pIt.Close()
 

pkg/phlaredb/query/iters.go

@@ -487,151 +487,6 @@ type JoinIterator struct {
 	result *IteratorResult
 }
 
-var _ Iterator = (*JoinIterator)(nil)
-
-func NewJoinIterator(definitionLevel int, iters []Iterator, pred GroupPredicate) *JoinIterator {
-	j := JoinIterator{
-		definitionLevel: definitionLevel,
-		iters:           iters,
-		peeks:           make([]*IteratorResult, len(iters)),
-		pred:            pred,
-	}
-	return &j
-}
-
-func (j *JoinIterator) At() *IteratorResult {
-	return j.result
-}
-
-func (j *JoinIterator) Next() bool {
-	// Here is the algorithm for joins:  On each pass of the iterators
-	// we remember which ones are pointing at the earliest rows. If all
-	// are the lowest (and therefore pointing at the same thing) then
-	// there is a successful join and return the result.
-	// Else we progress the iterators and try again.
-	// There is an optimization here in that we can seek to the highest
-	// row seen. It's impossible to have joins before that row.
-	for {
-		lowestRowNumber := MaxRowNumber()
-		highestRowNumber := EmptyRowNumber()
-		lowestIters := make([]int, 0, len(j.iters))
-
-		for iterNum := range j.iters {
-			res := j.peek(iterNum)
-
-			if res == nil {
-				// Iterator exhausted, no more joins possible
-				j.result = nil
-				return false
-			}
-
-			c := CompareRowNumbers(j.definitionLevel, res.RowNumber, lowestRowNumber)
-			switch c {
-			case -1:
-				// New lowest, reset
-				lowestIters = lowestIters[:0]
-				lowestRowNumber = res.RowNumber
-				fallthrough
-
-			case 0:
-				// Same, append
-				lowestIters = append(lowestIters, iterNum)
-			}
-
-			if CompareRowNumbers(j.definitionLevel, res.RowNumber, highestRowNumber) == 1 {
-				// New high water mark
-				highestRowNumber = res.RowNumber
-			}
-		}
-
-		// All iterators pointing at same row?
-		if len(lowestIters) == len(j.iters) {
-			// Get the data
-			result := j.collect(lowestRowNumber)
-
-			// Keep group?
-			if j.pred == nil || j.pred.KeepGroup(result) {
-				// Yes
-				j.result = result
-				return true
-			}
-		}
-
-		// Skip all iterators to the highest row seen, it's impossible
-		// to find matches before that.
-		j.seekAll(RowNumberWithDefinitionLevel{RowNumber: highestRowNumber, DefinitionLevel: j.definitionLevel})
-	}
-}
-
-func (j *JoinIterator) Seek(to RowNumberWithDefinitionLevel) bool {
-	j.seekAll(to)
-	return j.Next()
-}
-
-func (j *JoinIterator) seekAll(to RowNumberWithDefinitionLevel) {
-	to.RowNumber = TruncateRowNumber(to)
-	for iterNum, iter := range j.iters {
-		if j.peeks[iterNum] == nil || CompareRowNumbers(to.DefinitionLevel, j.peeks[iterNum].RowNumber, to.RowNumber) == -1 {
-			iteratorResultPoolPut(j.peeks[iterNum])
-			if iter.Seek(to) {
-				j.peeks[iterNum] = iter.At()
-			} else {
-				j.peeks[iterNum] = nil
-			}
-		}
-	}
-}
-
-func (j *JoinIterator) peek(iterNum int) *IteratorResult {
-	if j.peeks[iterNum] == nil {
-		if j.iters[iterNum].Next() {
-			j.peeks[iterNum] = j.iters[iterNum].At()
-		}
-	}
-	return j.peeks[iterNum]
-}
-
-// Collect data from the given iterators until they point at
-// the next row (according to the configured definition level)
-// or are exhausted.
-func (j *JoinIterator) collect(rowNumber RowNumber) *IteratorResult {
-	result := iteratorResultPoolGet()
-	result.RowNumber = rowNumber
-
-	for i := range j.iters {
-		for j.peeks[i] != nil && CompareRowNumbers(j.definitionLevel, j.peeks[i].RowNumber, rowNumber) == 0 {
-
-			result.Append(j.peeks[i])
-
-			iteratorResultPoolPut(j.peeks[i])
-
-			if j.iters[i].Next() {
-				j.peeks[i] = j.iters[i].At()
-			} else {
-				j.peeks[i] = nil
-			}
-		}
-	}
-	return result
-}
-
-func (j *JoinIterator) Close() error {
-	var merr multierror.MultiError
-	for _, i := range j.iters {
-		merr.Add(i.Close())
-	}
-	return merr.Err()
-}
-
-func (j *JoinIterator) Err() error {
-	for _, i := range j.iters {
-		if err := i.Err(); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 // UnionIterator produces all results for all given iterators.  When iterators
 // align to the same row, based on the configured definition level, then the results
 // are returned together. Else the next matching iterator is returned.