feat: Introduce FetchInto method for zero-allocation batch fetch operations

- Added FetchInto method to Segment interface for optimized data retrieval using a pre-allocated FetchArena.
- Implemented FetchInto in various segment types (diskann, flat, memtable) to reduce allocations during batch fetches.
- Enhanced metadata handling with UnmarshalBinaryInto for efficient deserialization into pre-allocated maps.
- Updated benchmarks to evaluate performance improvements with FetchArena pooling.
- Added new search options: WithMetadata and WithPayload for more granular control over returned data.
- Introduced FetchArena struct to manage pre-allocated scratch space for fetch operations, significantly reducing memory allocations.

Author: hupe1980

benchmark_test/baseline.txt

@@ -123,6 +123,95 @@ func BenchmarkFastSearch(b *testing.B) {
 	}
 }
 
+// BenchmarkFastSearchWithData benchmarks search WITH metadata/vectors returned.
+// This exercises the FetchInto/Clone path which is the allocation hotspot.
+func BenchmarkFastSearchWithData(b *testing.B) {
+	ctx := context.Background()
+
+	fixtureName := "uniform_128d_50k"
+	if testing.Short() {
+		fixtureName = "uniform_128d_10k"
+	}
+
+	if !FixtureExists(fixtureName) {
+		b.Skipf("fixture %q not found", fixtureName)
+	}
+
+	db, err := OpenFixture(ctx, fixtureName)
+	if err != nil {
+		b.Fatalf("open fixture: %v", err)
+	}
+	defer db.Close()
+
+	data, err := LoadFixtureData(fixtureName)
+	if err != nil {
+		b.Fatalf("load fixture data: %v", err)
+	}
+
+	queries := data.Queries
+	const k = 10
+
+	// Test different data inclusion patterns
+	// Default includes metadata+payload, WithoutData excludes all, WithVector adds vector
+	b.Run("id_only", func(b *testing.B) {
+		b.ReportAllocs()
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			q := queries[i%len(queries)]
+			results, err := db.Search(ctx, q, k, vecgo.WithoutData())
+			if err != nil {
+				b.Fatal(err)
+			}
+			if len(results) == 0 {
+				b.Fatal("no results")
+			}
+		}
+
+		b.StopTimer()
+		b.ReportMetric(float64(b.N)/b.Elapsed().Seconds(), "qps")
+	})
+
+	b.Run("with_metadata_payload", func(b *testing.B) {
+		// Default behavior - includes metadata and payload
+		b.ReportAllocs()
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			q := queries[i%len(queries)]
+			results, err := db.Search(ctx, q, k)
+			if err != nil {
+				b.Fatal(err)
+			}
+			if len(results) == 0 {
+				b.Fatal("no results")
+			}
+		}
+
+		b.StopTimer()
+		b.ReportMetric(float64(b.N)/b.Elapsed().Seconds(), "qps")
+	})
+
+	b.Run("with_vector", func(b *testing.B) {
+		b.ReportAllocs()
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			q := queries[i%len(queries)]
+			results, err := db.Search(ctx, q, k, vecgo.WithVector())
+			if err != nil {
+				b.Fatal(err)
+			}
+			if len(results) == 0 {
+				b.Fatal("no results")
+			}
+		}
+
+		b.StopTimer()
+		b.ReportMetric(float64(b.N)/b.Elapsed().Seconds(), "qps")
+	})
+}
+
 // BenchmarkFastBatchSearch benchmarks batch search performance.
 func BenchmarkFastBatchSearch(b *testing.B) {
 	ctx := context.Background()

benchmark_test/fast_bench_test.go

@@ -1035,25 +1035,47 @@ func (e *Engine) SearchIter(ctx context.Context, q []float32, k int, opts ...fun
 						s.Results[i+k].ID = s.ScratchForeignIDs[k]
 					}
 				} else {
-					batch, fErr := seg.Fetch(ctx, s.ScratchIDs, cols)
+					// Use FetchInto with pooled arena for efficient deserialization.
+					// Arena provides pre-sized maps and backing arrays to reduce allocations
+					// during unmarshaling. We then COPY to user-owned memory for safe ownership.
+					arena := segment.GetFetchArena()
+					arena.Reset(countBatch)
+					arena.EnsureCapacity(countBatch, e.dim)
+					arena.SetDimension(e.dim)
+
+					batch, fErr := seg.FetchInto(ctx, s.ScratchIDs, cols, arena)
 					if fErr != nil {
+						segment.PutFetchArena(arena)
 						err = fErr
 						yield(model.Candidate{}, err)
 						return
 					}
 
+					// Copy arena data to user-owned memory for safe ownership.
+					// Arena is returned to pool immediately after copy.
 					for k := 0; k < countBatch; k++ {
 						s.Results[i+k].ID = batch.ID(k)
 						if options.IncludeVector {
-							s.Results[i+k].Vector = batch.Vector(k)
+							// Clone vector - user needs to own this
+							if src := batch.Vector(k); src != nil {
+								s.Results[i+k].Vector = slices.Clone(src)
+							}
 						}
 						if options.IncludeMetadata {
-							s.Results[i+k].Metadata = batch.Metadata(k)
+							// Clone metadata map - user needs to own this
+							if src := batch.Metadata(k); src != nil {
+								s.Results[i+k].Metadata = src.Clone()
+							}
 						}
 						if options.IncludePayload {
-							s.Results[i+k].Payload = batch.Payload(k)
+							// Clone payload - user needs to own this
+							if src := batch.Payload(k); src != nil {
+								s.Results[i+k].Payload = slices.Clone(src)
+							}
 						}
 					}
+					// Return arena to pool immediately - all data has been copied
+					segment.PutFetchArena(arena)
 				}
 			}
 			i = j

internal/engine/search.go

@@ -284,6 +284,9 @@ func (s *Searcher) Reset() {
 		s.QueryBitmap.Clear()
 	}
 
+	// Note: FetchArena is managed by the engine via sync.Pool, not by Searcher.
+	// This avoids an import cycle (segment -> searcher -> segment).
+
 	s.OpsPerformed = 0
 	s.FilterGateStats.Reset()
 }

internal/searcher/searcher.go

@@ -856,6 +856,115 @@ func (s *Segment) Fetch(ctx context.Context, rows []uint32, cols []string) (segm
 	return batch, nil
 }
 
+// FetchInto resolves RowIDs to payload columns using a pre-allocated arena.
+// This is the zero-allocation hot path for batch fetch operations.
+func (s *Segment) FetchInto(ctx context.Context, rows []uint32, cols []string, arena *segment.FetchArena) (*segment.SimpleRecordBatch, error) {
+	fetchVectors := cols == nil
+	fetchMetadata := cols == nil
+	fetchPayloads := cols == nil
+
+	if cols != nil {
+		fetchVectors = false
+		for _, c := range cols {
+			switch c {
+			case "vector":
+				fetchVectors = true
+			case "metadata":
+				fetchMetadata = true
+			case "payload":
+				fetchPayloads = true
+			}
+		}
+	}
+
+	dim := int(s.header.Dim)
+	batchSize := len(rows)
+
+	// Ensure arena has capacity and reset for this batch
+	arena.EnsureCapacity(batchSize, dim)
+	arena.Reset(batchSize)
+	arena.SetDimension(dim)
+
+	// Pre-size slices
+	arena.IDs = arena.IDs[:batchSize]
+	if fetchVectors {
+		arena.Vectors = arena.Vectors[:batchSize]
+	}
+	if fetchMetadata {
+		arena.Metadatas = arena.Metadatas[:batchSize]
+	}
+	if fetchPayloads {
+		arena.Payloads = arena.Payloads[:batchSize]
+	}
+
+	for i, rowID := range rows {
+		// Periodic context check (every 64 rows)
+		if i&63 == 0 {
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			default:
+			}
+		}
+
+		if rowID >= s.header.RowCount {
+			return nil, fmt.Errorf("rowID %d out of bounds", rowID)
+		}
+
+		// Fetch ID
+		if id, ok := s.GetID(ctx, rowID); ok {
+			arena.IDs[i] = id
+		} else {
+			return nil, fmt.Errorf("failed to get ID for row %d", rowID)
+		}
+
+		// Fetch Vector
+		if fetchVectors {
+			vec, err := s.Get(ctx, rowID)
+			if err != nil {
+				return nil, err
+			}
+			// Use arena's backing array
+			dst := arena.AllocVectorSlice(i)
+			copy(dst, vec)
+			arena.Vectors[i] = dst
+		}
+
+		// Fetch Metadata (zero-alloc via pooled map)
+		if fetchMetadata {
+			md, err := s.readMetadata(ctx, rowID)
+			if err != nil {
+				return nil, err
+			}
+			// Copy into arena's pooled metadata map
+			if md != nil {
+				pooledMd := arena.AcquireMetadata(i)
+				for k, v := range md {
+					(*pooledMd)[k] = v
+				}
+				arena.Metadatas[i] = *pooledMd
+			}
+		}
+
+		// Fetch Payload
+		// TODO: Optimize with arena.PayloadBacking for batch reads
+		if fetchPayloads && s.payloadBlob != nil && rowID < s.payloadCount {
+			start := s.payloadOffsets[rowID]
+			end := s.payloadOffsets[rowID+1]
+			size := end - start
+			dataOffset := 4 + uint64(s.payloadCount+1)*8 + start
+
+			p := make([]byte, size)
+			if _, err := s.payloadBlob.ReadAt(ctx, p, int64(dataOffset)); err != nil {
+				return nil, err
+			}
+			arena.Payloads[i] = p
+		}
+	}
+
+	return arena.BuildRecordBatch(fetchVectors, fetchMetadata, fetchPayloads), nil
+}
+
 func (s *Segment) FetchIDs(ctx context.Context, rows []uint32, dst []model.ID) error {
 	if len(dst) != len(rows) {
 		return fmt.Errorf("dst length mismatch")