From db8682bc7dc34612205b628ee61625ea17ea7257 Mon Sep 17 00:00:00 2001 From: Srijan Arya Date: Sat, 5 Sep 2026 20:26:40 +0530 Subject: [PATCH 1/2] feat(collection): optional progress bar for Collection.upsert (#10) Adds opt-in show_progress to Collection.upsert. Uses tqdm.auto so it picks a notebook widget or terminal bar; reports a known total when records is Sized, otherwise an indeterminate bar. tqdm is an optional extra (vecs[progress]); default behaviour is unchanged. Scoped to upsert; create_index progress is left for a follow-up. --- setup.py | 1 + src/tests/test_upsert_progress.py | 76 +++++++++++++++++++++++++++++++ src/vecs/collection.py | 39 +++++++++++++++- 3 files changed, 114 insertions(+), 2 deletions(-) create mode 100644 src/tests/test_upsert_progress.py diff --git a/setup.py b/setup.py index 7bce72c..aec6936 100644 --- a/setup.py +++ b/setup.py @@ -84,5 +84,6 @@ def read_package_variable(key, filename="__init__.py"): "mike", ], "text_embedding": ["sentence-transformers==2.*"], + "progress": ["tqdm"], }, ) diff --git a/src/tests/test_upsert_progress.py b/src/tests/test_upsert_progress.py new file mode 100644 index 0000000..23c2428 --- /dev/null +++ b/src/tests/test_upsert_progress.py @@ -0,0 +1,76 @@ +import builtins +import sys +import types + +import numpy as np +import pytest + +import vecs + + +def _records(n: int, dim: int = 4): + return [(f"vec{ix}", vec, {}) for ix, vec in enumerate(np.random.random((n, dim)))] + + +class _FakeBar: + """Records tqdm calls so tests don't need a TTY or notebook.""" + + instances = [] + + def __init__(self, total=None, **_): + self.total = total + self.updates = [] + self.closed = False + _FakeBar.instances.append(self) + + def update(self, n): + self.updates.append(n) + + def close(self): + self.closed = True + + +@pytest.fixture +def fake_tqdm(monkeypatch): + _FakeBar.instances.clear() + auto = types.ModuleType("tqdm.auto") + auto.tqdm = _FakeBar + pkg = types.ModuleType("tqdm") + pkg.auto = auto + monkeypatch.setitem(sys.modules, "tqdm", pkg) + monkeypatch.setitem(sys.modules, "tqdm.auto", auto) + return _FakeBar + + +def test_upsert_default_has_no_progress(client: vecs.Client, fake_tqdm) -> None: + coll = client.get_or_create_collection(name="prog_default", dimension=4) + coll.upsert(_records(10)) + assert fake_tqdm.instances == [] + assert len(coll) == 10 + + +def test_upsert_progress_sized_reports_total(client: vecs.Client, fake_tqdm) -> None: + coll = client.get_or_create_collection(name="prog_sized", dimension=4) + coll.upsert(_records(1200), show_progress=True) + (bar,) = fake_tqdm.instances + assert bar.total == 1200 + assert bar.updates == [500, 500, 200] + assert bar.closed + assert len(coll) == 1200 + + +def test_upsert_progress_generator_is_indeterminate(client: vecs.Client, fake_tqdm) -> None: + coll = client.get_or_create_collection(name="prog_gen", dimension=4) + coll.upsert((r for r in _records(7)), show_progress=True) + (bar,) = fake_tqdm.instances + assert bar.total is None + assert sum(bar.updates) == 7 + assert bar.closed + + +def test_upsert_progress_without_tqdm_raises(client: vecs.Client, monkeypatch) -> None: + monkeypatch.setitem(sys.modules, "tqdm", None) + monkeypatch.setitem(sys.modules, "tqdm.auto", None) + coll = client.get_or_create_collection(name="prog_missing", dimension=4) + with pytest.raises(ImportError, match="vecs\\[progress\\]"): + coll.upsert(_records(3), show_progress=True) diff --git a/src/vecs/collection.py b/src/vecs/collection.py index eab8f93..8ad36ae 100644 --- a/src/vecs/collection.py +++ b/src/vecs/collection.py @@ -12,7 +12,17 @@ import warnings from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + List, + Optional, + Sized, + Tuple, + Union, +) from flupy import flu from pgvector.sqlalchemy import Vector @@ -319,7 +329,10 @@ def _drop(self): return self def upsert( - self, records: Iterable[Tuple[str, Any, Metadata]], skip_adapter: bool = False + self, + records: Iterable[Tuple[str, Any, Metadata]], + skip_adapter: bool = False, + show_progress: bool = False, ) -> None: """ Inserts or updates *vectors* records in the collection. @@ -334,10 +347,27 @@ def upsert( skip_adapter (bool): Should the adapter be skipped while upserting. i.e. if vectors are being provided, rather than a media type that needs to be transformed + + show_progress (bool): Display a progress bar while upserting. Requires the + optional `tqdm` dependency (`pip install vecs[progress]`). The total is + shown when *records* is sized (e.g. a list); otherwise the bar is + indeterminate. """ chunk_size = 500 + progress = None + if show_progress: + try: + from tqdm.auto import tqdm # picks notebook widget or terminal bar + except ImportError as exc: + raise ImportError( + "show_progress=True requires the optional 'tqdm' dependency. " + "Install it with `pip install vecs[progress]`." + ) from exc + total = len(records) if isinstance(records, Sized) else None + progress = tqdm(total=total, desc="upsert", unit="records") + if skip_adapter: pipeline = flu(records).chunk(chunk_size) else: @@ -357,6 +387,11 @@ def upsert( ), ) sess.execute(stmt) + if progress is not None: + progress.update(len(chunk)) + + if progress is not None: + progress.close() return None def fetch(self, ids: Iterable[str]) -> List[Record]: From ef2ba233157316e905022f37e0768d052a8b19d2 Mon Sep 17 00:00:00 2001 From: Srijan Arya Date: Sat, 5 Sep 2026 20:57:50 +0530 Subject: [PATCH 2/2] test: black formatting, drop unused import --- src/tests/test_upsert_progress.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/tests/test_upsert_progress.py b/src/tests/test_upsert_progress.py index 23c2428..b6f91e3 100644 --- a/src/tests/test_upsert_progress.py +++ b/src/tests/test_upsert_progress.py @@ -1,4 +1,3 @@ -import builtins import sys import types @@ -59,7 +58,9 @@ def test_upsert_progress_sized_reports_total(client: vecs.Client, fake_tqdm) -> assert len(coll) == 1200 -def test_upsert_progress_generator_is_indeterminate(client: vecs.Client, fake_tqdm) -> None: +def test_upsert_progress_generator_is_indeterminate( + client: vecs.Client, fake_tqdm +) -> None: coll = client.get_or_create_collection(name="prog_gen", dimension=4) coll.upsert((r for r in _records(7)), show_progress=True) (bar,) = fake_tqdm.instances