From e8ff36f088bf027e5afbd07a26afbb5dee48ecb2 Mon Sep 17 00:00:00 2001 From: Zohaib Hassnain <109234410+ZohaibHassan16@users.noreply.github.com> Date: Mon, 31 Aug 2026 03:01:37 +0500 Subject: [PATCH] feat(vector_store): add weaviate iter_all --- semantica/vector_store/weaviate_store.py | 94 +++++++++++ tests/vector_store/test_weaviate_store.py | 192 ++++++++++++++++++++++ 2 files changed, 286 insertions(+) create mode 100644 tests/vector_store/test_weaviate_store.py diff --git a/semantica/vector_store/weaviate_store.py b/semantica/vector_store/weaviate_store.py index 63f30475..55c8544b 100644 --- a/semantica/vector_store/weaviate_store.py +++ b/semantica/vector_store/weaviate_store.py @@ -611,6 +611,100 @@ class WeaviateStore: self.logger.warning(f"Failed to fetch Weaviate objects by metadata filter: {e}") return results if results else [] + def iter_all(self, batch_size: int = 500): + """ + Iterate over every stored object using Weaviate's native cursor. + + Weaviate paginates by the UUID of the last object seen (``after``), + not by row offset, so this is exposed instead of + scan_vectors(offset, limit): a row number such as 100000 cannot be + translated into the correct UUID without walking there. + VectorStore.iter_vectors() prefers this method when it is present. + + Assumes a single unnamed vector per object, matching how get_vector() + and filter_by_metadata() already read them back. Collections + configured with named vectors return a mapping here and are not + handled. + + Args: + batch_size: Objects to request per fetch_objects() call + + Yields: + Result dicts with 'id', 'metadata', and 'vector', in cursor order + + Raises: + ProcessingError: If the collection is not initialized. Errors are + raised rather than swallowed because a scan that silently + yields nothing is indistinguishable from an empty source, + which would let a caller such as `store migrate` report + success having copied nothing (issue #1083). + """ + if self.collection is None or not WEAVIATE_AVAILABLE: + raise ProcessingError( + "Collection not initialized. Call get_collection() first." + ) + + after_cursor = None + scanned_count = 0 + use_after_cursor = True + + while True: + kwargs = {"limit": batch_size, "include_vector": True} + if use_after_cursor and after_cursor is not None: + kwargs["after"] = after_cursor + + try: + objs = self.collection.query.fetch_objects(**kwargs) + except TypeError: + # Client versions without `after` fall back to numeric offset, + # then to no pagination argument at all. Mirrors the + # degradation chain in filter_by_metadata(). + if "after" not in kwargs: + raise + use_after_cursor = False + kwargs.pop("after", None) + try: + objs = self.collection.query.fetch_objects( + offset=scanned_count, **kwargs + ) + except TypeError: + objs = self.collection.query.fetch_objects(**kwargs) + + batch_objects = getattr(objs, "objects", None) if objs else None + if not batch_objects: + return + + for obj in batch_objects: + obj_uuid = getattr(obj, "uuid", None) + raw_vector = getattr(obj, "vector", None) + yield { + "id": str(obj_uuid) if obj_uuid is not None else None, + "metadata": getattr(obj, "properties", None) or {}, + "vector": ( + np.array(raw_vector) + if raw_vector is not None and len(raw_vector) > 0 + else None + ), + } + + scanned_count += len(batch_objects) + + # A short page means the collection is exhausted. + if len(batch_objects) < batch_size: + return + + last_uuid = getattr(batch_objects[-1], "uuid", None) + if last_uuid is None: + return + + next_cursor = str(last_uuid) + # There is no result cap on a full scan, so a cursor that stops + # advancing would loop forever. Stop instead of re-reading the + # same page. + if use_after_cursor and next_cursor == after_cursor: + return + after_cursor = next_cursor + def query_vectors( self, diff --git a/tests/vector_store/test_weaviate_store.py b/tests/vector_store/test_weaviate_store.py new file mode 100644 index 00000000..7ad7f7db --- /dev/null +++ b/tests/vector_store/test_weaviate_store.py @@ -0,0 +1,192 @@ +"""Tests for WeaviateStore.iter_all() cursor enumeration. + +weaviate-client is not installed in this environment, so these drive the real +WeaviateStore against MagicMocks, following the pattern already used for +weaviate in test_backend_metadata_filtering.py. +""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from semantica.utils.exceptions import ProcessingError +from semantica.vector_store.weaviate_store import WeaviateStore + + +def _obj(uuid, properties=None, vector=None): + """Stand-in for a weaviate v4 returned object.""" + obj = MagicMock() + obj.uuid = uuid + obj.properties = properties + obj.vector = vector + return obj + + +def _page(objects): + """Stand-in for a fetch_objects() response.""" + response = MagicMock() + response.objects = objects + return response + + +def _store_with_pages(*pages): + store = WeaviateStore() + store.collection = MagicMock() + store.collection.query.fetch_objects.side_effect = list(pages) + return store + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_threads_uuid_cursor_across_pages(): + """The next page must continue after the last object's UUID.""" + store = _store_with_pages( + _page([_obj("uuid-1"), _obj("uuid-2")]), + _page([_obj("uuid-3")]), + ) + + result = list(store.iter_all(batch_size=2)) + + assert [item["id"] for item in result] == ["uuid-1", "uuid-2", "uuid-3"] + calls = store.collection.query.fetch_objects.call_args_list + assert "after" not in calls[0][1] + assert calls[1][1]["after"] == "uuid-2" + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_stops_on_short_page(): + """A page smaller than batch_size means the collection is exhausted.""" + store = _store_with_pages(_page([_obj("uuid-1")])) + + result = list(store.iter_all(batch_size=5)) + + assert [item["id"] for item in result] == ["uuid-1"] + assert store.collection.query.fetch_objects.call_count == 1 + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_stops_when_cursor_stops_advancing(): + """A full scan has no result cap, so a stalled cursor would loop forever. + + Pathological rather than expected from a real server, but termination is + the property that matters: this must not hang. + """ + store = WeaviateStore() + store.collection = MagicMock() + store.collection.query.fetch_objects.return_value = _page( + [_obj("same-uuid"), _obj("same-uuid")] + ) + + result = list(store.iter_all(batch_size=2)) + + assert store.collection.query.fetch_objects.call_count == 2 + assert len(result) == 4 + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_empty_collection_yields_nothing(): + store = _store_with_pages(_page([])) + + assert list(store.iter_all()) == [] + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_converts_objects_to_the_shared_result_shape(): + store = _store_with_pages( + _page([_obj("uuid-7", properties={"tag": "x"}, vector=[0.1, 0.2, 0.3])]), + ) + + item = list(store.iter_all())[0] + + assert item["id"] == "uuid-7" + assert item["metadata"] == {"tag": "x"} + np.testing.assert_allclose(item["vector"], np.array([0.1, 0.2, 0.3])) + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_handles_missing_properties_and_vector(): + store = _store_with_pages(_page([_obj("uuid-1", properties=None, vector=None)])) + + item = list(store.iter_all())[0] + + assert item["metadata"] == {} + assert item["vector"] is None + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_treats_empty_vector_as_none(): + store = _store_with_pages(_page([_obj("uuid-1", vector=[])])) + + assert list(store.iter_all())[0]["vector"] is None + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_requests_vectors(): + """Weaviate omits vectors unless include_vector is set.""" + store = _store_with_pages(_page([])) + + list(store.iter_all(batch_size=64)) + + kwargs = store.collection.query.fetch_objects.call_args[1] + assert kwargs["include_vector"] is True + assert kwargs["limit"] == 64 + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_falls_back_to_offset_when_after_unsupported(): + """Older clients reject `after`; the scan degrades to numeric offset.""" + store = WeaviateStore() + store.collection = MagicMock() + seen = {"calls": 0} + + def _fetch(**kwargs): + if "after" in kwargs: + raise TypeError("unexpected keyword argument 'after'") + seen["calls"] += 1 + if seen["calls"] == 1: + return _page([_obj("uuid-1"), _obj("uuid-2")]) + return _page([_obj("uuid-3")]) + + store.collection.query.fetch_objects.side_effect = _fetch + + result = list(store.iter_all(batch_size=2)) + + assert [item["id"] for item in result] == ["uuid-1", "uuid-2", "uuid-3"] + offsets = [ + c[1]["offset"] + for c in store.collection.query.fetch_objects.call_args_list + if "offset" in c[1] + ] + assert offsets == [2] + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_raises_when_collection_not_initialized(): + """Must fail loudly, not yield nothing. + + An empty scan is indistinguishable from an empty source, which would let + `store migrate` report success having copied nothing (issue #1083). + """ + store = WeaviateStore() + + with pytest.raises(ProcessingError, match="Collection not initialized"): + list(store.iter_all()) + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", False) +def test_iter_all_raises_when_weaviate_unavailable(): + store = WeaviateStore() + store.collection = MagicMock() + + with pytest.raises(ProcessingError): + list(store.iter_all()) + + +@patch("semantica.vector_store.weaviate_store.WEAVIATE_AVAILABLE", True) +def test_iter_all_propagates_fetch_errors(): + store = WeaviateStore() + store.collection = MagicMock() + store.collection.query.fetch_objects.side_effect = RuntimeError("connection reset") + + with pytest.raises(RuntimeError, match="connection reset"): + list(store.iter_all())