Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

mongodb: [performance] Increase DEFAULT_INSERT_BATCH_SIZE to 100,000 and introduce sizing constraints #19608

Merged
merged 7 commits into from
May 14, 2024
Merged

mongodb: [performance] Increase DEFAULT_INSERT_BATCH_SIZE to 100,000 and introduce sizing constraints #19608

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
0acb929
Remove DEFAULT_INSERT_BATCH_SIZE
Jibola Mar 26, 2024
cf77ba5
Merge branch 'master' into unbounded-batch-insert
Jibola Apr 23, 2024
44589ed
updated BATCH_SIZE calculations
Jibola Apr 30, 2024
20581bc
Merge branch 'master' into unbounded-batch-insert
Jibola Apr 30, 2024
561d726
Merge branch 'master' into unbounded-batch-insert
Jibola May 13, 2024
4d92b52
Merge branch 'master' into unbounded-batch-insert
Jibola May 14, 2024
d36dbff
Merge branch 'master' into unbounded-batch-insert
efriis May 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 17 additions & 11 deletions libs/partners/mongodb/langchain_mongodb/vectorstores.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@

logger = logging.getLogger(__name__)

DEFAULT_INSERT_BATCH_SIZE = 100
Jibola marked this conversation as resolved.
Show resolved Hide resolved
DEFAULT_INSERT_BATCH_SIZE = 100_000


class MongoDBAtlasVectorSearch(VectorStore):
Expand Down Expand Up @@ -151,18 +151,24 @@ def add_texts(
"""
batch_size = kwargs.get("batch_size", DEFAULT_INSERT_BATCH_SIZE)
_metadatas: Union[List, Generator] = metadatas or ({} for _ in texts)
texts_batch = []
metadatas_batch = []
texts_batch = texts
metadatas_batch = _metadatas
result_ids = []
for i, (text, metadata) in enumerate(zip(texts, _metadatas)):
texts_batch.append(text)
metadatas_batch.append(metadata)
if (i + 1) % batch_size == 0:
result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))
texts_batch = []
metadatas_batch = []
if batch_size:
texts_batch = []
metadatas_batch = []
size = 0
for i, (text, metadata) in enumerate(zip(texts, _metadatas)):
size += len(text) + len(metadata)
texts_batch.append(text)
metadatas_batch.append(metadata)
if (i + 1) % batch_size == 0 or size >= 47_000_000:
result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))
texts_batch = []
metadatas_batch = []
size = 0
if texts_batch:
result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))
result_ids.extend(self._insert_texts(texts_batch, metadatas_batch)) # type: ignore
Copy link

@ShaneHarvey ShaneHarvey Mar 28, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The one concern here is what happens when someone calls add_texts() with a generator that yields a large amount of data. With this new code we'll inflate the entire thing into memory which can degrade performance. For example:

add_texts((random_large_string() for i in range(1_000_000))

In this case it would be best to limit the batches roughly based on the string size in addition to batch_size. We could do that by tracking the sum total of the batch's text:

                size += len(text)
                texts_batch.append(text)
                metadatas_batch.append(metadata)
                if (i + 1) % batch_size == 0 or size >= 47_000_000:

47MB is the maxMessageSizeBytes that pymongo can batch in one message minus 1MB of overhead to account for other bson data.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, that's fine with me. We'll re-introduce BATCH_SIZE as 100_000 and also have an additional bounding of 47_000_000 on the text aggregate text length. In this case, then, should we also include the length of the accompanying metadata?

return result_ids

def _insert_texts(self, texts: List[str], metadatas: List[Dict[str, Any]]) -> List:
Expand Down
Loading