Compare commits

...
Author SHA1 Message Date
Zohaib Hassnain af829f5f20 fix(vector_store): don't let iterator close() mask the real scan error, dedupe milvus result shaping, split unavailable/uninitialized messages 2026-09-01 23:09:53 +05:00
Zohaib Hassnain 930e7f9b71 docs(vector_store): note the milvus schema assumption 2026-09-01 23:09:53 +05:00
Zohaib Hassnain e335971dcd feat(vector_store): add milvus iter_all 2026-09-01 23:09:53 +05:00
Zohaib Hassnain 1227947be5 fix(vector_store): dedupe qdrant record conversion, don't abort iter_all on a live cursor with an empty page 2026-09-01 22:51:32 +05:00
Zohaib Hassnain b4a14d87f5 making it clean 2026-09-01 22:51:32 +05:00
Zohaib Hassnain 3bf89e523f fix(vector_store): raise instead of truncating when a qdrant scan cannot advance 2026-09-01 22:51:32 +05:00
Zohaib Hassnain 2b5b62bb8d fix(vector_store): drop qdrant migrate wiring, keep iter_all only
VectorStore cannot actually migrate to or from qdrant yet. _init_backend_store constructs QdrantStore without connecting or selecting a collection, so reads raise a Collection not initialized error, and the facade store_vectors dispatches only to add/add_vectors while QdrantStore exposes insert_vectors, so writes raise NotImplementedError.

Both are pre-existing facade gaps that nothing had exposed, since migrate previously only allowed faiss/sqlite/pgvector. Adding qdrant to the allowlist claimed support that does not work end to end, so it is removed along with the dimension inference that only fires for backends missing a .dimension attribute. Tracked separately; this PR keeps just the iter_all primitive.
2026-09-01 22:51:32 +05:00
Zohaib Hassnain 3a0f3f672a feat(vector_store): add iter_all enumeration and wire up qdrant migration 2026-09-01 22:51:32 +05:00
6 changed files with 591 additions and 22 deletions
+82 -11
View File
@@ -621,6 +621,15 @@ class MilvusStore:
except Exception:
return None
@staticmethod
def _record_to_result(item: Dict[str, Any]) -> Dict[str, Any]:
vec = item.get("vector")
return {
"id": str(item.get("id")),
"metadata": item.get("metadata") or {},
"vector": np.array(vec) if vec is not None else None,
}
def filter_by_metadata(
self, filters: Dict[str, Any], limit: int = 10
) -> List[Dict[str, Any]]:
@@ -665,21 +674,83 @@ class MilvusStore:
limit=limit,
output_fields=["id", "vector", "metadata"],
)
results = []
for item in query_results:
vec = item.get("vector")
results.append(
{
"id": str(item.get("id")),
"metadata": item.get("metadata") or {},
"vector": np.array(vec) if vec is not None else None,
}
)
return results
return [self._record_to_result(item) for item in query_results]
except Exception as e:
self.logger.warning(f"Failed to query Milvus vectors by metadata expression: {e}")
return []
def iter_all(self, batch_size: int = 500):
"""
Iterate over every stored entity using Milvus's query iterator.
Paginates by primary-key cursor rather than row offset, which is why
this exists instead of scan_vectors(offset, limit). query(offset=...)
is capped by the 16384 result window and would truncate anything
larger.
Assumes the schema create_collection() builds: a VARCHAR `id` primary
key plus vector and metadata fields, as get_vector() and
filter_by_metadata() already do. get_collection() does not validate
schema, so a collection with an integer key or no metadata field fails
here.
Args:
batch_size: Entities to request per iterator batch
Yields:
Result dicts with 'id', 'metadata', and 'vector', in cursor order
Raises:
ProcessingError: If the collection is not initialized, or the
installed pymilvus does not expose query_iterator().
"""
if self.collection is None:
raise ProcessingError(
"Collection not initialized. Call create_collection() or get_collection() first."
)
if not MILVUS_AVAILABLE:
raise ProcessingError("Milvus not available")
query_iterator = getattr(self.collection.collection, "query_iterator", None)
if not callable(query_iterator):
raise ProcessingError(
"This pymilvus version does not expose Collection.query_iterator(), "
"which full enumeration requires. Falling back to query(offset=...) "
"is not safe here: it is capped by the 16384 result window and would "
"silently truncate a larger collection."
)
# Query operations need a loaded collection. Idempotent, and once per
# scan rather than per batch.
self.collection.load()
# Milvus rejects an empty expression; this match-all form is what
# filter_by_metadata() already uses.
iterator = query_iterator(
batch_size=batch_size,
expr="id != ''",
output_fields=["id", "vector", "metadata"],
)
try:
while True:
batch = iterator.next()
if not batch:
return
for item in batch:
yield self._record_to_result(item)
finally:
# Release the server-side iterator even if the consumer stops early.
# Swallowed so a broken connection at cleanup time doesn't replace
# whatever real exception was already propagating out of the try.
close = getattr(iterator, "close", None)
if callable(close):
try:
close()
except Exception as e:
self.logger.warning(f"Failed to close Milvus query iterator: {e}")
def get_stats(self, collection_name: Optional[str] = None) -> Dict[str, Any]:
"""Get collection statistics."""
if self.collection is None and collection_name:
+67 -10
View File
@@ -590,20 +590,77 @@ class QdrantStore:
with_payload=True,
with_vectors=True,
)
results = []
for rec in records:
results.append(
{
"id": str(rec.id),
"metadata": rec.payload or {},
"vector": np.array(rec.vector) if rec.vector is not None else None,
}
)
return results
return [self._record_to_result(rec) for rec in records]
except Exception as e:
self.logger.warning(f"Failed to scroll Qdrant points by metadata filter: {e}")
return []
@staticmethod
def _record_to_result(rec: Any) -> Dict[str, Any]:
return {
"id": str(rec.id),
"metadata": rec.payload or {},
"vector": np.array(rec.vector) if rec.vector is not None else None,
}
def iter_all(self, batch_size: int = 500):
"""
Iterate over every stored point using Qdrant's scroll cursor.
Paginates by point-ID cursor rather than row offset, which is why this
exists instead of scan_vectors(offset, limit). An integer offset is a
point ID, not a rank.
Assumes a single unnamed vector per point, as insert_vectors() and
get_vector() already do. Named and multi-vector collections are not
handled.
Args:
batch_size: Points to request per scroll call
Yields:
Result dicts with 'id', 'metadata', and 'vector', in scroll order
Raises:
ProcessingError: If the collection or client is not initialized, or
if the cursor stops advancing before the scan completes.
"""
if self.collection is None or self.client is None or not QDRANT_AVAILABLE:
raise ProcessingError(
"Collection not initialized. Call create_collection() or get_collection() first."
)
next_offset = None
last_offset = object()
while True:
records, next_offset = self.client.scroll(
collection_name=self.collection.collection_name,
limit=batch_size,
offset=next_offset,
with_payload=True,
with_vectors=True,
)
for rec in records:
yield self._record_to_result(rec)
# A final page can carry records alongside a null cursor, so they
# are yielded above before stopping. Passing offset=None back to
# scroll() would restart from the beginning, not continue.
if next_offset is None:
return
# An empty page with a live cursor isn't necessarily truncation —
# a batch window that lands entirely on deleted points comes back
# this way too, and there's more to scan past it. Only treat it as
# stuck if the cursor itself stops moving.
if not records and next_offset == last_offset:
raise ProcessingError(
"Qdrant scroll cursor stopped advancing without reaching "
"the end of the collection, so the scan cannot complete."
)
last_offset = next_offset
def delete_vectors(
self, point_ids: List[Union[str, int]], **options
) -> Dict[str, Any]:
+11 -1
View File
@@ -867,12 +867,22 @@ class VectorStore:
"""
Iterate over every stored vector, one page at a time.
Cursor-based backends expose iter_all() because they cannot support a
positional offset; it takes precedence when present. Everything else
falls through to the scan_vectors() offset loop.
Args:
batch_size: Number of vectors to fetch per underlying scan_vectors() call
batch_size: Number of vectors to fetch per underlying call
Yields:
Result dicts with 'id', 'metadata', and 'vector', in scan order
"""
if self.backend != "inmemory" and self._backend_store is not None:
iter_all = getattr(self._backend_store, "iter_all", None)
if callable(iter_all):
yield from iter_all(batch_size=batch_size)
return
offset = 0
while True:
page = self.scan_vectors(offset=offset, limit=batch_size)
+176
View File
@@ -0,0 +1,176 @@
"""Tests for MilvusStore.iter_all() query-iterator enumeration.
pymilvus is not installed in this environment, so these drive the real
MilvusStore against MagicMocks, following the pattern already used for milvus
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.milvus_store import MilvusStore
def _store_with_batches(*batches):
"""MilvusStore whose query_iterator yields the given batches then stops.
The attribute path is doubled here: the pymilvus Collection sits at
wrapper.collection.
"""
store = MilvusStore()
wrapper = MagicMock()
inner = MagicMock()
iterator = MagicMock()
iterator.next.side_effect = list(batches)
inner.query_iterator.return_value = iterator
wrapper.collection = inner
store.collection = wrapper
return store, wrapper, inner, iterator
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_yields_batches_until_exhausted():
"""Exhaustion is an empty list, not StopIteration."""
store, _, _, iterator = _store_with_batches(
[{"id": 1, "vector": [0.1], "metadata": {}}],
[{"id": 2, "vector": [0.2], "metadata": {}}],
[],
)
result = list(store.iter_all(batch_size=1))
assert [item["id"] for item in result] == ["1", "2"]
assert iterator.next.call_count == 3
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_requests_the_fields_needed_for_the_result_shape():
store, _, inner, _ = _store_with_batches([])
list(store.iter_all(batch_size=64))
kwargs = inner.query_iterator.call_args[1]
assert kwargs["batch_size"] == 64
assert kwargs["output_fields"] == ["id", "vector", "metadata"]
# Milvus rejects an empty expression, so a match-all form is required.
assert kwargs["expr"] == "id != ''"
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_loads_the_collection_before_querying():
"""Milvus requires a loaded collection for query operations."""
store, wrapper, _, _ = _store_with_batches([])
list(store.iter_all())
assert wrapper.load.called
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_closes_the_iterator_on_exhaustion():
store, _, _, iterator = _store_with_batches([])
list(store.iter_all())
assert iterator.close.called
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_closes_the_iterator_when_consumer_stops_early():
"""Abandoning the generator early must still release the iterator."""
store, _, _, iterator = _store_with_batches(
[{"id": 1, "vector": [0.1], "metadata": {}}],
[{"id": 2, "vector": [0.2], "metadata": {}}],
[],
)
generator = store.iter_all(batch_size=1)
next(generator)
assert not iterator.close.called
generator.close()
assert iterator.close.called
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_converts_entities_to_the_shared_result_shape():
store, _, _, _ = _store_with_batches(
[{"id": 7, "vector": [0.1, 0.2, 0.3], "metadata": {"tag": "x"}}], []
)
item = list(store.iter_all())[0]
assert item["id"] == "7"
assert item["metadata"] == {"tag": "x"}
np.testing.assert_allclose(item["vector"], np.array([0.1, 0.2, 0.3]))
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_handles_missing_vector_and_metadata():
store, _, _, _ = _store_with_batches([{"id": 1, "vector": None, "metadata": None}], [])
item = list(store.iter_all())[0]
assert item["metadata"] == {}
assert item["vector"] is None
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_empty_collection_yields_nothing():
store, _, _, _ = _store_with_batches([])
assert list(store.iter_all()) == []
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_raises_when_query_iterator_is_unavailable():
"""Older pymilvus lacks query_iterator; falling back to query(offset=...)
would truncate at the 16384 window."""
store = MilvusStore()
wrapper = MagicMock()
wrapper.collection = MagicMock(spec=["query"])
store.collection = wrapper
with pytest.raises(ProcessingError, match="query_iterator"):
list(store.iter_all())
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_raises_when_collection_not_initialized():
"""Must fail loudly: an empty scan reads the same as an empty source."""
store = MilvusStore()
with pytest.raises(ProcessingError, match="Collection not initialized"):
list(store.iter_all())
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", False)
def test_iter_all_raises_when_milvus_unavailable():
store = MilvusStore()
store.collection = MagicMock()
with pytest.raises(ProcessingError):
list(store.iter_all())
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_propagates_iterator_errors():
store, _, _, iterator = _store_with_batches()
iterator.next.side_effect = RuntimeError("connection reset")
with pytest.raises(RuntimeError, match="connection reset"):
list(store.iter_all())
@patch("semantica.vector_store.milvus_store.MILVUS_AVAILABLE", True)
def test_iter_all_closes_the_iterator_when_a_batch_fails():
store, _, _, iterator = _store_with_batches()
iterator.next.side_effect = RuntimeError("connection reset")
with pytest.raises(RuntimeError):
list(store.iter_all())
assert iterator.close.called
+160
View File
@@ -0,0 +1,160 @@
"""Tests for QdrantStore.iter_all() cursor enumeration.
Qdrant is not installed in this environment, so these drive the real
QdrantStore against a MagicMock standing in for the qdrant_client, following
the pattern already used for qdrant 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.qdrant_store import QdrantStore
def _record(point_id, payload=None, vector=None):
"""Build a stand-in for a qdrant_client Record."""
rec = MagicMock()
rec.id = point_id
rec.payload = payload
rec.vector = vector
return rec
def _store_with_scroll(*pages):
"""QdrantStore whose client.scroll() returns the given (records, cursor) pages."""
store = QdrantStore()
store.client = MagicMock()
store.client.scroll.side_effect = list(pages)
store.collection = MagicMock()
store.collection.collection_name = "test_collection"
return store
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_threads_cursor_across_pages():
"""The next call continues from the previous page's cursor."""
store = _store_with_scroll(
([_record(1), _record(2)], "cursor-1"),
([_record(3)], None),
)
result = list(store.iter_all(batch_size=2))
assert [item["id"] for item in result] == ["1", "2", "3"]
calls = store.client.scroll.call_args_list
assert len(calls) == 2
assert calls[0][1]["offset"] is None
assert calls[0][1]["limit"] == 2
assert calls[1][1]["offset"] == "cursor-1"
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_yields_final_page_that_reports_no_next_cursor():
"""Records and a null cursor can arrive together; those records must still
be yielded or every scan loses its tail."""
store = _store_with_scroll(([_record(1), _record(2)], None))
result = list(store.iter_all(batch_size=10))
assert [item["id"] for item in result] == ["1", "2"]
assert store.client.scroll.call_count == 1
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_converts_records_to_the_shared_result_shape():
store = _store_with_scroll(
([_record(7, payload={"tag": "x"}, vector=[0.1, 0.2, 0.3])], None),
)
item = list(store.iter_all())[0]
assert item["id"] == "7"
assert item["metadata"] == {"tag": "x"}
np.testing.assert_allclose(item["vector"], np.array([0.1, 0.2, 0.3]))
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_handles_missing_payload_and_vector():
store = _store_with_scroll(([_record(1, payload=None, vector=None)], None))
item = list(store.iter_all())[0]
assert item["metadata"] == {}
assert item["vector"] is None
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_empty_collection_yields_nothing():
store = _store_with_scroll(([], None))
assert list(store.iter_all()) == []
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_continues_past_empty_page_with_advancing_cursor():
store = _store_with_scroll(
([], "cursor-1"),
([_record(1)], None),
)
result = list(store.iter_all())
assert [item["id"] for item in result] == ["1"]
assert store.client.scroll.call_count == 2
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_raises_when_cursor_stops_advancing():
store = _store_with_scroll(
([], "stuck-cursor"),
([], "stuck-cursor"),
)
with pytest.raises(ProcessingError, match="stopped advancing"):
list(store.iter_all())
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_raises_when_collection_not_initialized():
"""Must fail loudly: an empty scan reads the same as an empty source."""
store = QdrantStore()
with pytest.raises(ProcessingError, match="Collection not initialized"):
list(store.iter_all())
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", False)
def test_iter_all_raises_when_qdrant_unavailable():
store = QdrantStore()
store.client = MagicMock()
store.collection = MagicMock()
with pytest.raises(ProcessingError):
list(store.iter_all())
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_propagates_scroll_errors():
store = QdrantStore()
store.client = MagicMock()
store.client.scroll.side_effect = RuntimeError("connection reset")
store.collection = MagicMock()
store.collection.collection_name = "test_collection"
with pytest.raises(RuntimeError, match="connection reset"):
list(store.iter_all())
@patch("semantica.vector_store.qdrant_store.QDRANT_AVAILABLE", True)
def test_iter_all_requests_payload_and_vectors():
store = _store_with_scroll(([], None))
list(store.iter_all())
kwargs = store.client.scroll.call_args[1]
assert kwargs["with_payload"] is True
assert kwargs["with_vectors"] is True
assert kwargs["collection_name"] == "test_collection"
@@ -26,6 +26,7 @@ from unittest.mock import MagicMock, patch
import numpy as np
from semantica.utils.exceptions import ProcessingError
from semantica.vector_store.vector_store import VectorStore, VectorManager
@@ -138,6 +139,34 @@ class _NonScanningBackendStore:
"""Fake persistent backend store without any scan capability."""
class _IterAllBackendStore:
"""Fake cursor-based store: iter_all() only, no usable scan_vectors()."""
def __init__(self, items):
self._items = items
self.batch_sizes = []
def iter_all(self, batch_size=500):
self.batch_sizes.append(batch_size)
for item in self._items:
yield item
def scan_vectors(self, offset=0, limit=100):
raise AssertionError("scan_vectors() must not be called when iter_all() exists")
class _MisShapedIterAllBackendStore:
"""Backend store whose ``iter_all`` attribute is not callable."""
iter_all = 42 # plain attribute, not a method
def __init__(self, items):
self._items = items
def scan_vectors(self, offset=0, limit=100):
return self._items[offset:offset + limit]
class VectorStoreScanVectorsTests(unittest.TestCase):
"""VectorStore.scan_vectors() / iter_vectors() backend-agnostic accessors."""
@@ -192,6 +221,72 @@ class VectorStoreScanVectorsTests(unittest.TestCase):
self.assertEqual(list(store.iter_vectors(batch_size=2)), [])
# ---------------------------------------------------------------------------
# VectorStore.iter_vectors() preference for a native iter_all()
# ---------------------------------------------------------------------------
class VectorStoreIterAllDispatchTests(unittest.TestCase):
"""iter_vectors() prefers a backend's native iter_all() when present."""
def _persistent_store(self, backend_store, backend_name="qdrant"):
store = VectorStore(backend="inmemory", dimension=2)
store.backend = backend_name
store._backend_store = backend_store
return store
def test_iter_vectors_uses_iter_all_when_available(self):
items = [
{"id": "a", "vector": None, "metadata": {"n": 1}},
{"id": "b", "vector": None, "metadata": {"n": 2}},
]
backend = _IterAllBackendStore(items)
store = self._persistent_store(backend)
self.assertEqual(list(store.iter_vectors(batch_size=7)), items)
def test_iter_vectors_forwards_batch_size_to_iter_all(self):
backend = _IterAllBackendStore([])
store = self._persistent_store(backend)
list(store.iter_vectors(batch_size=32))
self.assertEqual(backend.batch_sizes, [32])
def test_iter_vectors_falls_back_to_scan_vectors_without_iter_all(self):
items = [{"id": "a", "vector": None, "metadata": {}}]
store = self._persistent_store(_ScanningBackendStore(items))
self.assertEqual(list(store.iter_vectors(batch_size=2)), items)
def test_iter_vectors_falls_back_when_iter_all_not_callable(self):
# Mirrors the count() precedent in _MisShapedBackendStore.
items = [{"id": "a", "vector": None, "metadata": {}}]
store = self._persistent_store(_MisShapedIterAllBackendStore(items))
self.assertEqual(list(store.iter_vectors(batch_size=2)), items)
def test_iter_vectors_inmemory_ignores_iter_all(self):
store = VectorStore(backend="inmemory", dimension=2)
store.store_vectors([np.array([1.0, 0.0])], [{"type": "a"}])
store._backend_store = _IterAllBackendStore([{"id": "wrong"}])
collected = list(store.iter_vectors(batch_size=2))
self.assertEqual([item["metadata"] for item in collected], [{"type": "a"}])
def test_iter_vectors_propagates_iter_all_errors(self):
# Silently yielding nothing would read as an empty source (#1083).
class _FailingIterAll:
def iter_all(self, batch_size=500):
raise ProcessingError("backend unreachable")
yield # pragma: no cover - makes this a generator
store = self._persistent_store(_FailingIterAll())
with self.assertRaises(ProcessingError):
list(store.iter_vectors(batch_size=2))
# ---------------------------------------------------------------------------
# VectorManager tests — inmemory backend
# ---------------------------------------------------------------------------