Skip to content

Commit 4f542e0

Browse files
committed
feat(iterator): add DocIterator for full collection traversal
Add streaming full-collection traversal across C++/C/Python: - C++: Collection::CreateIterator + DocIterator (isolated Flush+snapshot scan, ConcatenatingReader across segments, FilteringReader for deletes, batch-prefetched vectors) - C API: zvec_collection_create_iterator/next/close + iterator options - Python: collection.iter_docs() generator (constant memory) - Tests: C++ (unit/integration/concurrency/perf), C API, Python Relates to #380
1 parent ec8a78e commit 4f542e0

16 files changed

Lines changed: 2145 additions & 0 deletions

File tree

python/tests/test_iter_docs.py

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
# Copyright 2025-present the zvec project
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""Tests for Collection.iter_docs (document iterator)."""
15+
16+
from __future__ import annotations
17+
18+
import pytest
19+
import zvec
20+
from zvec import (
21+
CollectionOption,
22+
DataType,
23+
Doc,
24+
FieldSchema,
25+
HnswIndexParam,
26+
VectorSchema,
27+
)
28+
29+
30+
@pytest.fixture(scope="session")
31+
def iter_schema():
32+
return zvec.CollectionSchema(
33+
name="iter_test_collection",
34+
fields=[
35+
FieldSchema("id", DataType.INT64, nullable=False),
36+
FieldSchema("name", DataType.STRING, nullable=False),
37+
FieldSchema("weight", DataType.FLOAT, nullable=True),
38+
],
39+
vectors=[
40+
VectorSchema(
41+
"dense",
42+
DataType.VECTOR_FP32,
43+
dimension=8,
44+
index_param=HnswIndexParam(),
45+
),
46+
],
47+
)
48+
49+
50+
@pytest.fixture(scope="function")
51+
def iter_collection(tmp_path_factory, iter_schema):
52+
temp_dir = tmp_path_factory.mktemp("zvec_iter")
53+
path = temp_dir / "iter_collection"
54+
coll = zvec.create_and_open(
55+
path=str(path),
56+
schema=iter_schema,
57+
option=CollectionOption(read_only=False, enable_mmap=True),
58+
)
59+
assert coll is not None
60+
try:
61+
yield coll
62+
finally:
63+
try:
64+
coll.destroy()
65+
except Exception as e:
66+
print(f"Warning: failed to destroy collection: {e}")
67+
68+
69+
def _make_docs(n: int) -> list[Doc]:
70+
return [
71+
Doc(
72+
id=f"{i}",
73+
fields={"id": i, "name": f"name_{i}", "weight": float(i)},
74+
vectors={"dense": [float(i)] * 8},
75+
)
76+
for i in range(n)
77+
]
78+
79+
80+
def test_iter_docs_basic(iter_collection):
81+
"""Insert N docs, iterate, verify count + PK + scalar fields."""
82+
n = 50
83+
result = iter_collection.insert(_make_docs(n))
84+
assert bool(result)
85+
iter_collection.flush()
86+
87+
seen_ids = set()
88+
count = 0
89+
for doc in iter_collection.iter_docs():
90+
assert isinstance(doc, Doc)
91+
assert doc.id != ""
92+
# scalar fields present
93+
assert doc.field("id") is not None
94+
assert doc.field("name") is not None
95+
seen_ids.add(doc.id)
96+
count += 1
97+
98+
assert count == n
99+
assert len(seen_ids) == n
100+
assert seen_ids == {f"{i}" for i in range(n)}
101+
102+
103+
def test_iter_docs_include_vector(iter_collection):
104+
"""include_vector=True (default) returns vectors of correct dimension."""
105+
iter_collection.insert(_make_docs(10))
106+
iter_collection.flush()
107+
108+
count = 0
109+
for doc in iter_collection.iter_docs(include_vector=True):
110+
vec = doc.vector("dense")
111+
assert vec is not None, f"dense vector missing for {doc.id}"
112+
assert len(vec) == 8
113+
count += 1
114+
assert count == 10
115+
116+
117+
def test_iter_docs_exclude_vector(iter_collection):
118+
"""include_vector=False omits vector fields."""
119+
iter_collection.insert(_make_docs(10))
120+
iter_collection.flush()
121+
122+
count = 0
123+
for doc in iter_collection.iter_docs(include_vector=False):
124+
# scalar present, vector absent
125+
assert doc.field("id") is not None
126+
# Doc.vector() returns {} (falsy) when no vectors are present.
127+
assert not doc.vector("dense")
128+
assert "dense" not in doc.vector_names()
129+
count += 1
130+
assert count == 10
131+
132+
133+
def test_iter_docs_output_fields(iter_collection):
134+
"""output_fields limits returned scalar fields."""
135+
iter_collection.insert(_make_docs(10))
136+
iter_collection.flush()
137+
138+
for doc in iter_collection.iter_docs(output_fields=["id"], include_vector=False):
139+
assert doc.field("id") is not None
140+
assert not doc.has_field("name")
141+
assert not doc.has_field("weight")
142+
143+
144+
def test_iter_docs_empty(iter_collection):
145+
"""Empty collection yields nothing."""
146+
docs = list(iter_collection.iter_docs())
147+
assert docs == []
148+
149+
150+
def test_iter_docs_after_delete(iter_collection):
151+
"""Deleted docs must not appear in iteration."""
152+
iter_collection.insert(_make_docs(20))
153+
# delete even ids
154+
to_delete = [f"{i}" for i in range(0, 20, 2)]
155+
iter_collection.delete(to_delete)
156+
iter_collection.flush()
157+
158+
deleted = set(to_delete)
159+
ids = []
160+
for doc in iter_collection.iter_docs():
161+
assert doc.id not in deleted
162+
ids.append(doc.id)
163+
assert len(ids) == 10
164+
165+
166+
def test_iter_docs_isolation(iter_collection):
167+
"""Docs written after the iterator is created are not visible."""
168+
iter_collection.insert(_make_docs(10))
169+
iter_collection.flush()
170+
171+
it = iter_collection.iter_docs()
172+
# Consume the first doc so the snapshot is established.
173+
first = next(it)
174+
assert first is not None
175+
176+
# Insert more docs after the iterator started.
177+
iter_collection.insert(
178+
[
179+
Doc(
180+
id=f"new_{i}",
181+
fields={"id": 1000 + i, "name": "new", "weight": 1.0},
182+
vectors={"dense": [1.0] * 8},
183+
)
184+
for i in range(5)
185+
]
186+
)
187+
iter_collection.flush()
188+
189+
# Count remaining from the original snapshot (should be 9, total 10).
190+
remaining = sum(1 for _ in it)
191+
assert remaining == 9
192+
193+
# A fresh iterator sees all 15.
194+
assert sum(1 for _ in iter_collection.iter_docs()) == 15
195+
196+
197+
def test_iter_docs_is_generator(iter_collection):
198+
"""iter_docs returns a lazy generator (constant memory)."""
199+
iter_collection.insert(_make_docs(5))
200+
iter_collection.flush()
201+
202+
gen = iter_collection.iter_docs()
203+
# It is an iterator: next() works and StopIteration terminates it.
204+
got = [next(gen) for _ in range(5)]
205+
assert len(got) == 5
206+
with pytest.raises(StopIteration):
207+
next(gen)

python/zvec/model/collection.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from __future__ import annotations
1515

1616
import warnings
17+
from collections.abc import Iterator
1718
from typing import Optional, Union, overload
1819

1920
from zvec._zvec import _Collection
@@ -378,6 +379,42 @@ def fetch(
378379
if (py_doc := convert_to_py_doc(core_doc, self.schema)) is not None
379380
}
380381

382+
def iter_docs(
383+
self,
384+
*,
385+
output_fields: Optional[list[str]] = None,
386+
include_vector: bool = True,
387+
) -> Iterator[Doc]:
388+
"""Iterate over all documents in the collection.
389+
390+
Streams documents one by one using an isolated snapshot taken at call
391+
time: memory usage stays constant regardless of collection size, and
392+
data written after the iterator is created is not visible.
393+
394+
Args:
395+
output_fields (Optional[list[str]], optional): Scalar fields to
396+
include. If None, all fields are returned. Defaults to None.
397+
include_vector (bool, optional): Whether to include vector data in
398+
each document. Defaults to True.
399+
400+
Yields:
401+
Doc: Each document in the collection.
402+
403+
Examples:
404+
>>> for doc in collection.iter_docs(include_vector=False):
405+
... print(doc.id, doc.field("title"))
406+
"""
407+
iterator = self._obj.CreateIterator(output_fields, include_vector)
408+
try:
409+
for core_doc in iterator:
410+
py_doc = convert_to_py_doc(core_doc, self.schema)
411+
if py_doc is not None:
412+
yield py_doc
413+
finally:
414+
# Release native resources (segments, file handles) even if the
415+
# caller stops early (e.g. breaks out of the loop).
416+
iterator.close()
417+
381418
# ========== Collection DQL-Query Methods ==========
382419

383420
def query(

0 commit comments

Comments
 (0)