feat(collection): optional progress bar for Collection.upsert (#10)

setup.py

@@ -84,5 +84,6 @@ def read_package_variable(key, filename="__init__.py"):
             "mike",
         ],
         "text_embedding": ["sentence-transformers==2.*"],
+        "progress": ["tqdm"],
     },
 )

src/tests/test_upsert_progress.py

@@ -0,0 +1,77 @@
+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)

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]: