Compare commits

...
Author SHA1 Message Date
KaifAhmad1 4dd88375e1 fix(vector_store): don't short-circuit pinecone iter_all() on an empty page
if not vector_ids: return fired before the continuation token was ever
checked. Pinecone's actual pagination contract is that a scan is only
exhausted when the response carries no pagination token -- a page can
legitimately list zero ids while pagination.next is still set (sparse
or filtered namespaces, eventual-consistency windows on serverless
indexes). This was flagged in review but the fix commit that followed
only addressed the separate repeated-token stall case, not this one.

Reproduced concretely against the unfixed code: a page with data,
followed by an empty page with a live token, followed by a page with
more data -- the last page was silently dropped with no error raised,
exactly the #1083 failure mode (store migrate reporting success after
copying only part of a collection).

Now the empty-page case skips the pointless fetch() call but still
falls through to the same next_token check every other path already
goes through, so a live token continues the scan and only a genuinely
absent token (or one that's stopped advancing) ends it.

Added test_continues_past_an_empty_page_with_a_live_token, the
"empty page + non-None next token" case the original review asked for
and that wasn't otherwise covered.
2026-09-02 19:54:24 +05:30
Zohaib Hassnain b8299b1427 chore: clean it 2026-09-02 17:19:19 +05:30
Zohaib Hassnain bbd423c50a fix(vector_store): raise on stalled pinecone pagination instead of truncating 2026-09-02 17:19:19 +05:30
Zohaib Hassnain 6b36379f15 feat(vector_store): add pinecone iter_all 2026-09-02 17:19:19 +05:30
2 changed files with 308 additions and 0 deletions
+123
View File
@@ -299,6 +299,44 @@ class PineconeSearch:
)
def _pinecone_listed_ids(response: Any) -> List[str]:
"""Extract vector IDs from a list_paginated() response.
Accepts record objects, bare id strings and dicts, since what listing
returns has changed across pinecone SDK major versions.
"""
records = getattr(response, "vectors", None)
if records is None and isinstance(response, dict):
records = response.get("vectors")
ids: List[str] = []
for record in records or []:
if isinstance(record, str):
ids.append(record)
elif isinstance(record, dict):
if record.get("id") is not None:
ids.append(record["id"])
else:
record_id = getattr(record, "id", None)
if record_id is not None:
ids.append(record_id)
return ids
def _pinecone_next_token(response: Any) -> Optional[str]:
"""Return the continuation token, or None when the listing is exhausted."""
pagination = getattr(response, "pagination", None)
if pagination is None and isinstance(response, dict):
pagination = response.get("pagination")
if pagination is None:
return None
token = getattr(pagination, "next", None)
if token is None and isinstance(pagination, dict):
token = pagination.get("next")
return token or None
class PineconeStore:
"""
Pinecone store for vector storage and similarity search.
@@ -735,6 +773,91 @@ class PineconeStore:
self.logger.warning(f"Failed to filter Pinecone vectors by metadata: {e}")
return []
def iter_all(self, batch_size: int = 500, namespace: str = ""):
"""
Iterate over every stored vector by listing IDs then fetching them.
Paginates with an opaque continuation token, which is why this exists
instead of scan_vectors(offset, limit): the token for page N cannot be
constructed without walking there.
Needs two calls per page, unlike the other backends, because listing
returns IDs only. Both calls are namespace scoped and must agree, and
listing covers one namespace rather than the whole index.
Args:
batch_size: IDs to request per list_paginated() call
namespace: Namespace to enumerate (default: the default namespace)
Yields:
Result dicts with 'id', 'metadata', and 'vector', in listing order
Raises:
ProcessingError: If the index is not initialized, if the installed
SDK does not expose list_paginated(), or if the listing stops
advancing.
"""
if self.index is None or not PINECONE_AVAILABLE:
raise ProcessingError(
"Index not initialized. Call create_index() or get_index() first."
)
# list_paginated() rather than list(): list() is an auto-paging
# iterator in current SDKs but reads as plain id lists in older
# examples. Threading the token explicitly is version-agnostic.
list_paginated = getattr(self.index.index, "list_paginated", None)
if not callable(list_paginated):
raise ProcessingError(
"This pinecone SDK version does not expose Index.list_paginated(), "
"which full enumeration requires."
)
token = None
while True:
kwargs: Dict[str, Any] = {"limit": batch_size, "namespace": namespace}
if token is not None:
kwargs["pagination_token"] = token
response = list_paginated(**kwargs)
vector_ids = _pinecone_listed_ids(response)
# A page listing zero ids is not necessarily exhaustion: Pinecone's
# contract is that a scan ends only when there's no pagination
# token, and a page can legitimately come back empty while
# pagination.next is still set (sparse/filtered namespaces,
# eventual-consistency windows on serverless indexes). Skip the
# fetch (nothing to hydrate) but still fall through to the token
# check below instead of returning early, or a gap like that
# silently truncates the scan with no error.
if vector_ids:
fetched = self.index.fetch_vectors(vector_ids, namespace=namespace)
vectors = fetched.get("vectors") or {}
for vector_id in vector_ids:
entry = vectors.get(vector_id)
if entry is None:
# fetch() omits ids it cannot find: deleted since listing.
continue
values = entry.get("values")
yield {
"id": vector_id,
"metadata": entry.get("metadata") or {},
"vector": np.array(values) if values is not None else None,
}
next_token = _pinecone_next_token(response)
if not next_token:
return
if next_token == token:
# Distinct from exhaustion above: a partial scan here would be
# indistinguishable from a complete one.
raise ProcessingError(
"Pinecone returned the same pagination token twice, so the "
"listing is not advancing. Refusing to return a truncated "
"scan."
)
token = next_token
def fetch_vectors(
self, vector_ids: List[str], namespace: str = "", **options
) -> Dict[str, Any]:
+185
View File
@@ -249,6 +249,191 @@ class TestPineconeIndex(unittest.TestCase):
mock_index.query.assert_called_once()
class TestPineconeIterAll(unittest.TestCase):
"""PineconeStore.iter_all() list-then-fetch enumeration."""
def _page(self, ids, next_token):
"""Stand-in for a list_paginated() response."""
response = MagicMock()
response.vectors = [MagicMock(id=vector_id) for vector_id in ids]
response.pagination = MagicMock(next=next_token)
return response
def _store(self, pages, fetch_results):
store = PineconeStore()
wrapper = MagicMock()
raw_index = MagicMock()
raw_index.list_paginated.side_effect = list(pages)
wrapper.index = raw_index
wrapper.fetch_vectors.side_effect = list(fetch_results)
store.index = wrapper
return store, wrapper, raw_index
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_threads_pagination_token_across_pages(self):
store, _, raw_index = self._store(
[self._page(["a", "b"], "token-1"), self._page(["c"], None)],
[
{"vectors": {"a": {"values": [0.1], "metadata": {}},
"b": {"values": [0.2], "metadata": {}}}},
{"vectors": {"c": {"values": [0.3], "metadata": {}}}},
],
)
result = list(store.iter_all(batch_size=2))
self.assertEqual([item["id"] for item in result], ["a", "b", "c"])
calls = raw_index.list_paginated.call_args_list
self.assertNotIn("pagination_token", calls[0][1])
self.assertEqual(calls[1][1]["pagination_token"], "token-1")
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_hydrates_listed_ids_with_a_fetch(self):
"""Listing returns ids only, so each page needs a fetch()."""
store, wrapper, _ = self._store(
[self._page(["a"], None)],
[{"vectors": {"a": {"values": [0.1, 0.2], "metadata": {"tag": "x"}}}}],
)
item = list(store.iter_all())[0]
self.assertEqual(item["id"], "a")
self.assertEqual(item["metadata"], {"tag": "x"})
np.testing.assert_allclose(item["vector"], np.array([0.1, 0.2]))
wrapper.fetch_vectors.assert_called_once_with(["a"], namespace="")
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_list_and_fetch_use_the_same_namespace(self):
store, wrapper, raw_index = self._store(
[self._page(["a"], None)],
[{"vectors": {"a": {"values": [0.1], "metadata": {}}}}],
)
list(store.iter_all(namespace="prod"))
self.assertEqual(raw_index.list_paginated.call_args[1]["namespace"], "prod")
wrapper.fetch_vectors.assert_called_once_with(["a"], namespace="prod")
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_skips_ids_deleted_between_list_and_fetch(self):
"""fetch() omits ids it cannot find rather than returning blanks."""
store, _, _ = self._store(
[self._page(["a", "gone"], None)],
[{"vectors": {"a": {"values": [0.1], "metadata": {}}}}],
)
result = list(store.iter_all())
self.assertEqual([item["id"] for item in result], ["a"])
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_raises_when_pagination_token_repeats(self):
"""A stalled token must not loop forever, nor quietly return a partial
scan that reads as a complete one."""
store, _, raw_index = self._store(
[self._page(["a"], "same"), self._page(["b"], "same")],
[
{"vectors": {"a": {"values": [0.1], "metadata": {}}}},
{"vectors": {"b": {"values": [0.2], "metadata": {}}}},
],
)
with self.assertRaises(ProcessingError):
list(store.iter_all())
self.assertEqual(raw_index.list_paginated.call_count, 2)
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_empty_listing_yields_nothing_without_fetching(self):
store, wrapper, _ = self._store([self._page([], None)], [])
self.assertEqual(list(store.iter_all()), [])
wrapper.fetch_vectors.assert_not_called()
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_continues_past_an_empty_page_with_a_live_token(self):
"""An empty page is not necessarily the end: Pinecone can legitimately
list zero ids for a page while pagination.next is still set (sparse
or filtered namespaces, eventual-consistency windows on serverless
indexes). Only the absence of a next token means exhaustion."""
store, wrapper, raw_index = self._store(
[
self._page(["a"], "token-1"),
self._page([], "token-2"), # empty page, but the token still advances
self._page(["b"], None),
],
[
{"vectors": {"a": {"values": [0.1], "metadata": {}}}},
{"vectors": {"b": {"values": [0.2], "metadata": {}}}},
],
)
result = list(store.iter_all(batch_size=1))
self.assertEqual([item["id"] for item in result], ["a", "b"])
self.assertEqual(raw_index.list_paginated.call_count, 3)
# Nothing to hydrate on the empty page, so only two fetches happen.
self.assertEqual(wrapper.fetch_vectors.call_count, 2)
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_accepts_plain_string_ids_from_listing(self):
"""SDK generations differ on what listing yields."""
store, _, _ = self._store(
[self._page([], None)],
[{"vectors": {"a": {"values": [0.1], "metadata": {}}}}],
)
response = MagicMock()
response.vectors = ["a"]
response.pagination = MagicMock(next=None)
store.index.index.list_paginated.side_effect = [response]
self.assertEqual([item["id"] for item in store.iter_all()], ["a"])
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_handles_missing_values_and_metadata(self):
store, _, _ = self._store(
[self._page(["a"], None)],
[{"vectors": {"a": {"values": None, "metadata": None}}}],
)
item = list(store.iter_all())[0]
self.assertIsNone(item["vector"])
self.assertEqual(item["metadata"], {})
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_raises_when_list_paginated_unavailable(self):
store = PineconeStore()
wrapper = MagicMock()
wrapper.index = MagicMock(spec=["query", "fetch"])
store.index = wrapper
with self.assertRaises(ProcessingError):
list(store.iter_all())
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_raises_when_index_not_initialized(self):
"""Must fail loudly: an empty scan reads the same as an empty source."""
with self.assertRaises(ProcessingError):
list(PineconeStore().iter_all())
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', False)
def test_raises_when_pinecone_unavailable(self):
store = PineconeStore()
store.index = MagicMock()
with self.assertRaises(ProcessingError):
list(store.iter_all())
@patch('semantica.vector_store.pinecone_store.PINECONE_AVAILABLE', True)
def test_propagates_listing_errors(self):
store, _, raw_index = self._store([], [])
raw_index.list_paginated.side_effect = RuntimeError("connection reset")
with self.assertRaises(RuntimeError):
list(store.iter_all())
if __name__ == '__main__':
print("DEBUG: Starting unittest.main()")
unittest.main()