enhance: pipeline sealed BM25 stats materialization

internal/streamingnode/server/wal/vchannel/idf/oracle.go

@@ -593,15 +593,23 @@ func (r *oracleRuntime) materialize(call *materializationCall) {
 		resultErr = merr.Wrapf(err, "get sealed BM25 resources for data version %s", call.target.String())
 		return
 	}
-	sealed, err = r.provider.acquireSealedContributions(call.ctx, resources)
+	var stats bm25Stats
+	sealed, err = r.provider.acquireSealedContributionsWithConsumer(
+		call.ctx,
+		resources,
+		func(contribution sealedContribution) {
+			if stats == nil {
+				stats = newBM25StatsFromSchema(r.schema)
+			}
+			stats.merge(contribution.stats)
+		},
+	)
 	if err != nil {
 		resultErr = merr.Wrapf(err, "load sealed BM25 stats for data version %s", call.target.String())
 		return
 	}
-
-	stats := newBM25StatsFromSchema(r.schema)
-	for _, contribution := range sealed {
-		stats.merge(contribution.stats)
+	if stats == nil {
+		stats = newBM25StatsFromSchema(r.schema)
 	}
 
 	var oldSealed map[int64]sealedContribution
@@ -1081,6 +1089,14 @@ func (p *Provider) getSealedBM25Resources(
 func (p *Provider) acquireSealedContributions(
 	ctx context.Context,
 	resources []*datapb.StreamingNodeBM25Resource,
+) (map[int64]sealedContribution, error) {
+	return p.acquireSealedContributionsWithConsumer(ctx, resources, nil)
+}
+
+func (p *Provider) acquireSealedContributionsWithConsumer(
+	ctx context.Context,
+	resources []*datapb.StreamingNodeBM25Resource,
+	consume func(sealedContribution),
 ) (map[int64]sealedContribution, error) {
 	loaded := make([]sealedContribution, len(resources))
 	keepLeases := false
@@ -1099,13 +1115,25 @@ func (p *Provider) acquireSealedContributions(
 	if limiter == nil {
 		limiter = getGlobalSealedStatsLoadLimiter()
 	}
+	results := make(chan sealedContribution, len(resources))
+	contributions := make(map[int64]sealedContribution, len(resources))
+	collectorDone := make(chan struct{})
+	go func() {
+		defer close(collectorDone)
+		for contribution := range results {
+			contributions[contribution.segmentID] = contribution
+			if consume != nil {
+				consume(contribution)
+			}
+		}
+	}()
+
 	group, groupCtx := errgroup.WithContext(ctx)
+	var acquireErr error
 	for i, resource := range resources {
 		if err := limiter.Acquire(groupCtx, 1); err != nil {
-			if groupErr := group.Wait(); groupErr != nil {
-				return nil, groupErr
-			}
-			return nil, err
+			acquireErr = err
+			break
 		}
 		i, resource := i, resource
 		group.Go(func() error {
@@ -1120,16 +1148,18 @@ func (p *Provider) acquireSealedContributions(
 				stats:       stats,
 				lease:       lease,
 			}
+			results <- loaded[i]
 			return nil
 		})
 	}
-	if err := group.Wait(); err != nil {
-		return nil, err
+	groupErr := group.Wait()
+	close(results)
+	<-collectorDone
+	if groupErr != nil {
+		return nil, groupErr
 	}
-
-	contributions := make(map[int64]sealedContribution, len(loaded))
-	for _, contribution := range loaded {
-		contributions[contribution.segmentID] = contribution
+	if acquireErr != nil {
+		return nil, acquireErr
 	}
 	keepLeases = true
 	return contributions, nil

internal/streamingnode/server/wal/vchannel/idf/oracle_test.go

@@ -190,6 +190,101 @@ func TestAcquireSealedContributionsUsesSharedLimit(t *testing.T) {
 	require.Equal(t, int32(2), maxActive.Load())
 }
 
+func TestAcquireSealedContributionsConsumesCompletedLoadsImmediately(t *testing.T) {
+	stats := storage.NewBM25Stats()
+	stats.Append(map[uint32]float32{1: 1})
+	statsBytes, err := stats.Serialize()
+	require.NoError(t, err)
+
+	secondReadStarted := make(chan struct{})
+	releaseSecondRead := make(chan struct{})
+	chunkManager := mocks.NewChunkManager(t)
+	chunkManager.EXPECT().Read(mock.Anything, mock.Anything).
+		RunAndReturn(func(_ context.Context, path string) ([]byte, error) {
+			if path == "stats-2" {
+				close(secondReadStarted)
+				<-releaseSecondRead
+			}
+			return statsBytes, nil
+		}).Times(2)
+
+	provider := &Provider{
+		chunkManager:           chunkManager,
+		sealedCache:            newSegmentCache(),
+		sealedStatsLoadLimiter: semaphore.NewWeighted(2),
+	}
+	consumed := make(chan int64, 2)
+	var aggregate bm25Stats
+	type result struct {
+		contributions map[int64]sealedContribution
+		err           error
+	}
+	resultCh := make(chan result, 1)
+	go func() {
+		contributions, err := provider.acquireSealedContributionsWithConsumer(
+			context.Background(),
+			testSealedBM25Resources(1, 2),
+			func(contribution sealedContribution) {
+				if aggregate == nil {
+					aggregate = newBM25StatsFromSchema(testBM25WALView(qviews.DataVersion{}).Schema)
+				}
+				aggregate.merge(contribution.stats)
+				consumed <- contribution.segmentID
+			},
+		)
+		resultCh <- result{contributions: contributions, err: err}
+	}()
+
+	<-secondReadStarted
+	select {
+	case segmentID := <-consumed:
+		require.Equal(t, int64(1), segmentID)
+	case <-time.After(time.Second):
+		t.Fatal("completed sealed stats were not consumed while another load was pending")
+	}
+	close(releaseSecondRead)
+	acquired := <-resultCh
+	require.NoError(t, acquired.err)
+	require.Len(t, acquired.contributions, 2)
+	require.Equal(t, int64(2), aggregate[testBM25OutputFieldID].NumRow())
+	for _, contribution := range acquired.contributions {
+		contribution.lease.Close()
+	}
+}
+
+func TestAcquireSealedContributionsCancellationWhileWaitingForLimit(t *testing.T) {
+	readStarted := make(chan struct{})
+	chunkManager := mocks.NewChunkManager(t)
+	chunkManager.EXPECT().Read(mock.Anything, mock.Anything).
+		RunAndReturn(func(ctx context.Context, _ string) ([]byte, error) {
+			close(readStarted)
+			<-ctx.Done()
+			return nil, ctx.Err()
+		}).Once()
+
+	provider := &Provider{
+		chunkManager:           chunkManager,
+		sealedCache:            newSegmentCache(),
+		sealedStatsLoadLimiter: semaphore.NewWeighted(1),
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	resultCh := make(chan error, 1)
+	go func() {
+		_, err := provider.acquireSealedContributions(ctx, testSealedBM25Resources(1, 2))
+		resultCh <- err
+	}()
+
+	<-readStarted
+	cancel()
+	select {
+	case err := <-resultCh:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("sealed stats acquisition did not stop after cancellation")
+	}
+	require.Empty(t, provider.sealedCache.entries)
+}
+
 func TestAcquireSealedContributionsReleasesLeasesAfterError(t *testing.T) {
 	stats := storage.NewBM25Stats()
 	stats.Append(map[uint32]float32{1: 1})

internal/streamingnode/server/wal/vchannel/idf/segment_cache.go

@@ -36,9 +36,8 @@ func (c *segmentCache) acquire(
 	chunkManager storage.ChunkManager,
 	resource *datapb.StreamingNodeBM25Resource,
 ) (bm25Stats, *segmentCacheLease, error) {
-	aggregate := make(bm25Stats)
 	if chunkManager == nil || resource == nil {
-		return aggregate, nil, nil
+		return make(bm25Stats), nil, nil
 	}
 	key, err := buildSealedCacheKey(resource)
 	if err != nil {
@@ -48,8 +47,7 @@ func (c *segmentCache) acquire(
 	if err != nil {
 		return nil, nil, err
 	}
-	aggregate.merge(stats)
-	return aggregate, &segmentCacheLease{cache: c, keys: []sealedCacheKey{key}}, nil
+	return stats, &segmentCacheLease{cache: c, keys: []sealedCacheKey{key}}, nil
 }
 
 func (c *segmentCache) retain(
@@ -146,7 +144,6 @@ func loadSealedSegmentStats(
 	}
 	stats := make(bm25Stats)
 	for fieldID, paths := range pathsByField {
-		fieldStats := stats.getOrCreate(fieldID)
 		for _, path := range paths {
 			bytes, err := chunkManager.Read(ctx, path)
 			if err != nil {
@@ -156,7 +153,12 @@ func loadSealedSegmentStats(
 			if err != nil {
 				return nil, err
 			}
-			fieldStats.Merge(loaded)
+			fieldStats := stats[fieldID]
+			if fieldStats == nil {
+				stats[fieldID] = loaded
+			} else {
+				fieldStats.Merge(loaded)
+			}
 		}
 	}
 	return stats, nil

internal/streamingnode/server/wal/vchannel/idf/segment_cache_test.go

@@ -87,3 +87,38 @@ func TestLoadSealedSegmentStatsUsesLegacyBinlogsForStorageV2(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, int64(1), loaded[102].NumRow())
 }
+
+func TestSegmentCacheAcquireReturnsCachedStats(t *testing.T) {
+	paramtable.Init()
+	ctx := context.Background()
+	chunkManager := storage.NewLocalChunkManager()
+	stats := storage.NewBM25Stats()
+	stats.Append(map[uint32]float32{1: 1})
+	bytes, err := stats.Serialize()
+	require.NoError(t, err)
+	statsPath := t.TempDir() + "/bm25-stats"
+	require.NoError(t, chunkManager.Write(ctx, statsPath, bytes))
+	resource := &datapb.StreamingNodeBM25Resource{
+		SegmentId:      3,
+		StorageVersion: storage.StorageV2,
+		Bm25Binlogs: []*datapb.FieldBinlog{{
+			FieldID: 102,
+			Binlogs: []*datapb.Binlog{{LogPath: statsPath}},
+		}},
+	}
+
+	cache := newSegmentCache()
+	first, firstLease, err := cache.acquire(ctx, chunkManager, resource)
+	require.NoError(t, err)
+	require.NotNil(t, firstLease)
+	second, secondLease, err := cache.acquire(ctx, chunkManager, resource)
+	require.NoError(t, err)
+	require.NotNil(t, secondLease)
+	key, err := buildSealedCacheKey(resource)
+	require.NoError(t, err)
+	require.Same(t, cache.entries[key].stats[102], first[102])
+	require.Same(t, first[102], second[102])
+
+	firstLease.Close()
+	secondLease.Close()
+}