diff --git a/backend/go/qdrant-store/store.go b/backend/go/qdrant-store/store.go index 82f5742d1..8769fdd34 100644 --- a/backend/go/qdrant-store/store.go +++ b/backend/go/qdrant-store/store.go @@ -29,6 +29,9 @@ const ( // Results carry up to 50 MiB of base64-encoded values. maxRecvMsgBytes = 128 << 20 + + // upsertChunkSize bounds the points sent in one Upsert request. + upsertChunkSize = 500 ) var pointNamespace = uuid.NewSHA1(uuid.NameSpaceURL, []byte("https://localai.io/qdrant-store")) @@ -135,12 +138,12 @@ func (s *QdrantStore) StoresSet(opts *pb.StoresSetOptions) error { } } - err := s.upsert(points) + err := upsertChunked(points, upsertChunkSize, s.upsert) if isNotFound(err) { if err := s.createCollection(len(opts.Keys[0].Floats)); err != nil { return fmt.Errorf("qdrant-store: Set: %w", err) } - err = s.upsert(points) + err = upsertChunked(points, upsertChunkSize, s.upsert) } if err != nil { return fmt.Errorf("qdrant-store: Set: %w", err) @@ -148,6 +151,17 @@ func (s *QdrantStore) StoresSet(opts *pb.StoresSetOptions) error { return nil } +// upsertChunked calls send with batches of at most size points; send applies +// its own request timeout per call, so a large batch is not bounded by one request_timeout_ms. +func upsertChunked(points []*qdrant.PointStruct, size int, send func([]*qdrant.PointStruct) error) error { + for start := 0; start < len(points); start += size { + if err := send(points[start:min(start+size, len(points))]); err != nil { + return err + } + } + return nil +} + func (s *QdrantStore) upsert(points []*qdrant.PointStruct) error { ctx, cancel := s.ctx() defer cancel() diff --git a/backend/go/qdrant-store/store_test.go b/backend/go/qdrant-store/store_test.go index e91f627a7..6f78b5831 100644 --- a/backend/go/qdrant-store/store_test.go +++ b/backend/go/qdrant-store/store_test.go @@ -1,6 +1,7 @@ package main import ( + "errors" "math" "time" @@ -73,3 +74,36 @@ var _ = Describe("QdrantStore", func() { Expect(err).To(MatchError(ContainSubstring("topK"))) }) }) + +var _ = Describe("upsertChunked", func() { + points := func(n int) []*qdrant.PointStruct { + return make([]*qdrant.PointStruct, n) + } + + It("splits a batch into chunks of at most the given size", func() { + var sizes []int + err := upsertChunked(points(1201), 500, func(c []*qdrant.PointStruct) error { + sizes = append(sizes, len(c)) + return nil + }) + Expect(err).NotTo(HaveOccurred()) + Expect(sizes).To(Equal([]int{500, 500, 201})) + }) + + It("sends a small batch in one call", func() { + calls := 0 + Expect(upsertChunked(points(3), 500, func([]*qdrant.PointStruct) error { calls++; return nil })).To(Succeed()) + Expect(calls).To(Equal(1)) + }) + + It("stops at the first failing chunk", func() { + boom := errors.New("boom") + calls := 0 + err := upsertChunked(points(1200), 500, func([]*qdrant.PointStruct) error { + calls++ + return boom + }) + Expect(err).To(MatchError(boom)) + Expect(calls).To(Equal(1)) + }) +})