From 7fa996261d0a18d292a010557450a2146d6895c6 Mon Sep 17 00:00:00 2001 From: Jiangtian Feng Date: Tue, 23 Jun 2026 18:59:16 +0800 Subject: [PATCH] perf: batch BM25 sparse encoding in Qdrant insert (#5592) --- mem0/vector_stores/qdrant.py | 46 +++++++++++++-- tests/vector_stores/test_qdrant.py | 93 ++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 6 deletions(-) diff --git a/mem0/vector_stores/qdrant.py b/mem0/vector_stores/qdrant.py index 25a688058..3f29edc4b 100644 --- a/mem0/vector_stores/qdrant.py +++ b/mem0/vector_stores/qdrant.py @@ -191,6 +191,43 @@ class Qdrant(VectorStoreBase): ids (list, optional): List of IDs corresponding to vectors. Defaults to None. """ logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}") + + # Pre-compute BM25 sparse vectors in a single batch call. fastembed's + # embed() accepts a list of texts, so batching avoids per-row encoder + # overhead (model dispatch, tokenizer setup, etc.). + bm25_sparse_vectors: list[Optional[SparseVector]] = [None] * len(vectors) + if self._has_bm25_slot and payloads: + texts_for_bm25: list[str] = [] + indices_for_bm25: list[int] = [] + for idx, payload in enumerate(payloads): + text = payload.get("text_lemmatized") or payload.get("data", "") + if text: + texts_for_bm25.append(text) + indices_for_bm25.append(idx) + + if texts_for_bm25: + encoder = self._get_bm25_encoder() + if encoder is not None: + try: + sparse_results = list(encoder.embed(texts_for_bm25)) + if len(sparse_results) != len(texts_for_bm25): + logger.warning( + f"BM25 batch returned {len(sparse_results)} results for " + f"{len(texts_for_bm25)} texts; falling back to per-row encoding" + ) + raise ValueError("count mismatch") + for i, sparse in enumerate(sparse_results): + bm25_sparse_vectors[indices_for_bm25[i]] = SparseVector( + indices=sparse.indices.tolist(), + values=sparse.values.tolist(), + ) + except Exception as e: + # Fall back to per-row encoding so a single bad input + # doesn't drop BM25 for the whole batch. + logger.debug(f"Batch BM25 encoding failed, falling back to per-row: {e}") + for i, text in enumerate(texts_for_bm25): + bm25_sparse_vectors[indices_for_bm25[i]] = self._encode_bm25(text) + points = [] for idx, vector in enumerate(vectors): payload = payloads[idx] if payloads else {} @@ -198,12 +235,8 @@ class Qdrant(VectorStoreBase): # Build named vectors: dense + optional BM25 sparse (only if collection has the slot). named_vectors = {"": vector} - if self._has_bm25_slot: - text_for_bm25 = payload.get("text_lemmatized") or payload.get("data", "") - if text_for_bm25: - sparse = self._encode_bm25(text_for_bm25) - if sparse is not None: - named_vectors["bm25"] = sparse + if self._has_bm25_slot and bm25_sparse_vectors[idx] is not None: + named_vectors["bm25"] = bm25_sparse_vectors[idx] points.append(PointStruct(id=point_id, vector=named_vectors, payload=payload)) @@ -476,6 +509,7 @@ class Qdrant(VectorStoreBase): if self._has_bm25_slot: text_for_bm25 = payload.get("text_lemmatized") or payload.get("data", "") if text_for_bm25: + # Single-item update: per-row encoding is correct here; see insert() for the batch path. sparse = self._encode_bm25(text_for_bm25) if sparse is not None: named_vectors["bm25"] = sparse diff --git a/tests/vector_stores/test_qdrant.py b/tests/vector_stores/test_qdrant.py index 83de59cd3..3bee29052 100644 --- a/tests/vector_stores/test_qdrant.py +++ b/tests/vector_stores/test_qdrant.py @@ -18,6 +18,7 @@ from qdrant_client.models import ( PointStruct, PointVectors, Range, + SparseVector, SparseVectorParams, VectorParams, ) @@ -86,6 +87,98 @@ class TestQdrant(unittest.TestCase): self.assertEqual(points[0].payload, payloads[0]) + def test_insert_batches_bm25_encoding(self): + """BM25 encoding must be done in a single batch call, not per-row. + + Regression test for the optimization that switched from looping over + `_encode_bm25(text)` to a single `encoder.embed(all_texts)` call. + """ + # Pretend the collection has the bm25 sparse slot (set up by create_col + # on a v3 collection) and stub the encoder to return predictable sparse + # vectors mimicking fastembed's output (objects with .indices/.values). + self.qdrant._has_bm25_slot = True + + encoder_mock = MagicMock() + encoder_mock.embed.return_value = iter( + [ + MagicMock(indices=MagicMock(tolist=lambda: [1, 2]), values=MagicMock(tolist=lambda: [0.5, 0.3])), + MagicMock(indices=MagicMock(tolist=lambda: [3]), values=MagicMock(tolist=lambda: [0.8])), + ] + ) + self.qdrant._bm25_encoder = encoder_mock + + vectors = [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]] + payloads = [ + {"data": "hello world"}, + {"text_lemmatized": "foo bar"}, + {"key": "no text — should skip BM25"}, + ] + ids = [str(uuid.uuid4()) for _ in range(3)] + + self.qdrant.insert(vectors=vectors, payloads=payloads, ids=ids) + + # 1. encoder.embed must be called exactly once (batched), not per-row. + encoder_mock.embed.assert_called_once() + called_texts = encoder_mock.embed.call_args[0][0] + self.assertEqual(called_texts, ["hello world", "foo bar"]) + + # 2. Resulting points carry the right named vectors: + # - Rows with text get both "" (dense) and "bm25" sparse. + # - Row without text gets dense only. + points = self.client_mock.upsert.call_args[1]["points"] + self.assertEqual(len(points), 3) + self.assertIn("bm25", points[0].vector) + self.assertEqual(points[0].vector["bm25"].indices, [1, 2]) + self.assertEqual(points[0].vector["bm25"].values, [0.5, 0.3]) + self.assertIn("bm25", points[1].vector) + self.assertEqual(points[1].vector["bm25"].indices, [3]) + self.assertNotIn("bm25", points[2].vector) + + def test_insert_skips_bm25_when_slot_missing(self): + """Pre-v3 collections (no bm25 slot) must not invoke the encoder at all.""" + self.qdrant._has_bm25_slot = False + encoder_mock = MagicMock() + self.qdrant._bm25_encoder = encoder_mock + + self.qdrant.insert( + vectors=[[0.1, 0.2]], + payloads=[{"data": "hello"}], + ids=[str(uuid.uuid4())], + ) + + encoder_mock.embed.assert_not_called() + points = self.client_mock.upsert.call_args[1]["points"] + self.assertNotIn("bm25", points[0].vector) + + def test_insert_falls_back_to_per_row_on_batch_failure(self): + """If batch encoding raises, each row should be re-encoded individually + so a single bad input doesn't drop BM25 for the whole batch.""" + self.qdrant._has_bm25_slot = True + + # Batch call raises; per-row _encode_bm25 should be used as fallback. + encoder_mock = MagicMock() + encoder_mock.embed.side_effect = RuntimeError("batch failed") + self.qdrant._bm25_encoder = encoder_mock + + # Use a real SparseVector — PointStruct validates the vector dict via + # Pydantic, which would coerce a bare MagicMock into [] (MagicMock is + # iterable). _encode_bm25's real return type is SparseVector. + fallback_sparse = SparseVector(indices=[7], values=[0.9]) + with patch.object(self.qdrant, "_encode_bm25", return_value=fallback_sparse) as fallback: + self.qdrant.insert( + vectors=[[0.1, 0.2], [0.3, 0.4]], + payloads=[{"data": "a"}, {"data": "b"}], + ids=[str(uuid.uuid4()), str(uuid.uuid4())], + ) + + encoder_mock.embed.assert_called_once() + self.assertEqual(fallback.call_count, 2) + self.assertEqual([c.args[0] for c in fallback.call_args_list], ["a", "b"]) + points = self.client_mock.upsert.call_args[1]["points"] + for point in points: + self.assertEqual(point.vector["bm25"].indices, [7]) + self.assertEqual(point.vector["bm25"].values, [0.9]) + def test_search(self): vectors = [[0.1, 0.2]] mock_point = MagicMock(id=str(uuid.uuid4()), score=0.95, payload={"key": "value"})