From 4f542e0c955868ffc140504e455aadbe1c07c637 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=90=AA=E9=9C=84?= Date: Thu, 16 Jul 2026 21:00:47 +0800 Subject: [PATCH] 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 --- python/tests/test_iter_docs.py | 207 ++++++ python/zvec/model/collection.py | 37 ++ src/binding/c/c_api.cc | 132 ++++ .../python/include/python_collection.h | 1 + src/binding/python/model/python_collection.cc | 44 ++ src/db/collection.cc | 132 ++++ src/db/doc_iterator.cc | 387 +++++++++++ src/db/doc_iterator_internal.h | 56 ++ src/db/index/segment/concatenating_reader.h | 70 ++ src/db/index/segment/filtering_reader.h | 114 ++++ src/include/zvec/c_api.h | 87 +++ src/include/zvec/db/collection.h | 4 + src/include/zvec/db/doc_iterator.h | 48 ++ src/include/zvec/db/options.h | 8 + tests/c/c_api_test.c | 211 ++++++ tests/db/iterator_test.cc | 607 ++++++++++++++++++ 16 files changed, 2145 insertions(+) create mode 100644 python/tests/test_iter_docs.py create mode 100644 src/db/doc_iterator.cc create mode 100644 src/db/doc_iterator_internal.h create mode 100644 src/db/index/segment/concatenating_reader.h create mode 100644 src/db/index/segment/filtering_reader.h create mode 100644 src/include/zvec/db/doc_iterator.h create mode 100644 tests/db/iterator_test.cc diff --git a/python/tests/test_iter_docs.py b/python/tests/test_iter_docs.py new file mode 100644 index 000000000..2aa5c0124 --- /dev/null +++ b/python/tests/test_iter_docs.py @@ -0,0 +1,207 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for Collection.iter_docs (document iterator).""" + +from __future__ import annotations + +import pytest +import zvec +from zvec import ( + CollectionOption, + DataType, + Doc, + FieldSchema, + HnswIndexParam, + VectorSchema, +) + + +@pytest.fixture(scope="session") +def iter_schema(): + return zvec.CollectionSchema( + name="iter_test_collection", + fields=[ + FieldSchema("id", DataType.INT64, nullable=False), + FieldSchema("name", DataType.STRING, nullable=False), + FieldSchema("weight", DataType.FLOAT, nullable=True), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=8, + index_param=HnswIndexParam(), + ), + ], + ) + + +@pytest.fixture(scope="function") +def iter_collection(tmp_path_factory, iter_schema): + temp_dir = tmp_path_factory.mktemp("zvec_iter") + path = temp_dir / "iter_collection" + coll = zvec.create_and_open( + path=str(path), + schema=iter_schema, + option=CollectionOption(read_only=False, enable_mmap=True), + ) + assert coll is not None + try: + yield coll + finally: + try: + coll.destroy() + except Exception as e: + print(f"Warning: failed to destroy collection: {e}") + + +def _make_docs(n: int) -> list[Doc]: + return [ + Doc( + id=f"{i}", + fields={"id": i, "name": f"name_{i}", "weight": float(i)}, + vectors={"dense": [float(i)] * 8}, + ) + for i in range(n) + ] + + +def test_iter_docs_basic(iter_collection): + """Insert N docs, iterate, verify count + PK + scalar fields.""" + n = 50 + result = iter_collection.insert(_make_docs(n)) + assert bool(result) + iter_collection.flush() + + seen_ids = set() + count = 0 + for doc in iter_collection.iter_docs(): + assert isinstance(doc, Doc) + assert doc.id != "" + # scalar fields present + assert doc.field("id") is not None + assert doc.field("name") is not None + seen_ids.add(doc.id) + count += 1 + + assert count == n + assert len(seen_ids) == n + assert seen_ids == {f"{i}" for i in range(n)} + + +def test_iter_docs_include_vector(iter_collection): + """include_vector=True (default) returns vectors of correct dimension.""" + iter_collection.insert(_make_docs(10)) + iter_collection.flush() + + count = 0 + for doc in iter_collection.iter_docs(include_vector=True): + vec = doc.vector("dense") + assert vec is not None, f"dense vector missing for {doc.id}" + assert len(vec) == 8 + count += 1 + assert count == 10 + + +def test_iter_docs_exclude_vector(iter_collection): + """include_vector=False omits vector fields.""" + iter_collection.insert(_make_docs(10)) + iter_collection.flush() + + count = 0 + for doc in iter_collection.iter_docs(include_vector=False): + # scalar present, vector absent + assert doc.field("id") is not None + # Doc.vector() returns {} (falsy) when no vectors are present. + assert not doc.vector("dense") + assert "dense" not in doc.vector_names() + count += 1 + assert count == 10 + + +def test_iter_docs_output_fields(iter_collection): + """output_fields limits returned scalar fields.""" + iter_collection.insert(_make_docs(10)) + iter_collection.flush() + + for doc in iter_collection.iter_docs(output_fields=["id"], include_vector=False): + assert doc.field("id") is not None + assert not doc.has_field("name") + assert not doc.has_field("weight") + + +def test_iter_docs_empty(iter_collection): + """Empty collection yields nothing.""" + docs = list(iter_collection.iter_docs()) + assert docs == [] + + +def test_iter_docs_after_delete(iter_collection): + """Deleted docs must not appear in iteration.""" + iter_collection.insert(_make_docs(20)) + # delete even ids + to_delete = [f"{i}" for i in range(0, 20, 2)] + iter_collection.delete(to_delete) + iter_collection.flush() + + deleted = set(to_delete) + ids = [] + for doc in iter_collection.iter_docs(): + assert doc.id not in deleted + ids.append(doc.id) + assert len(ids) == 10 + + +def test_iter_docs_isolation(iter_collection): + """Docs written after the iterator is created are not visible.""" + iter_collection.insert(_make_docs(10)) + iter_collection.flush() + + it = iter_collection.iter_docs() + # Consume the first doc so the snapshot is established. + first = next(it) + assert first is not None + + # Insert more docs after the iterator started. + iter_collection.insert( + [ + Doc( + id=f"new_{i}", + fields={"id": 1000 + i, "name": "new", "weight": 1.0}, + vectors={"dense": [1.0] * 8}, + ) + for i in range(5) + ] + ) + iter_collection.flush() + + # Count remaining from the original snapshot (should be 9, total 10). + remaining = sum(1 for _ in it) + assert remaining == 9 + + # A fresh iterator sees all 15. + assert sum(1 for _ in iter_collection.iter_docs()) == 15 + + +def test_iter_docs_is_generator(iter_collection): + """iter_docs returns a lazy generator (constant memory).""" + iter_collection.insert(_make_docs(5)) + iter_collection.flush() + + gen = iter_collection.iter_docs() + # It is an iterator: next() works and StopIteration terminates it. + got = [next(gen) for _ in range(5)] + assert len(got) == 5 + with pytest.raises(StopIteration): + next(gen) diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py index 6663a1b49..820739ed1 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -14,6 +14,7 @@ from __future__ import annotations import warnings +from collections.abc import Iterator from typing import Optional, Union, overload from zvec._zvec import _Collection @@ -378,6 +379,42 @@ def fetch( if (py_doc := convert_to_py_doc(core_doc, self.schema)) is not None } + def iter_docs( + self, + *, + output_fields: Optional[list[str]] = None, + include_vector: bool = True, + ) -> Iterator[Doc]: + """Iterate over all documents in the collection. + + Streams documents one by one using an isolated snapshot taken at call + time: memory usage stays constant regardless of collection size, and + data written after the iterator is created is not visible. + + Args: + output_fields (Optional[list[str]], optional): Scalar fields to + include. If None, all fields are returned. Defaults to None. + include_vector (bool, optional): Whether to include vector data in + each document. Defaults to True. + + Yields: + Doc: Each document in the collection. + + Examples: + >>> for doc in collection.iter_docs(include_vector=False): + ... print(doc.id, doc.field("title")) + """ + iterator = self._obj.CreateIterator(output_fields, include_vector) + try: + for core_doc in iterator: + py_doc = convert_to_py_doc(core_doc, self.schema) + if py_doc is not None: + yield py_doc + finally: + # Release native resources (segments, file handles) even if the + # caller stops early (e.g. breaks out of the loop). + iterator.close() + # ========== Collection DQL-Query Methods ========== def query( diff --git a/src/binding/c/c_api.cc b/src/binding/c/c_api.cc index 1dfce56f2..2050feb5a 100644 --- a/src/binding/c/c_api.cc +++ b/src/binding/c/c_api.cc @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -7246,3 +7247,134 @@ zvec_error_code_t zvec_collection_fetch(zvec_collection_t *collection, } return convert_fetched_document_results(doc_map, results, doc_count);) } + +// ============================================================================= +// Document Iterator Interface implementation +// ============================================================================= + +zvec_iterator_options_t *zvec_iterator_options_create(void) { + ZVEC_TRY_RETURN_NULL( + "Failed to create zvec_iterator_options_t", + auto *options = new zvec::IteratorOptions(); + return reinterpret_cast(options);) + return nullptr; +} + +void zvec_iterator_options_destroy(zvec_iterator_options_t *options) { + if (options) { + delete reinterpret_cast(options); + } +} + +zvec_error_code_t zvec_iterator_options_set_output_fields( + zvec_iterator_options_t *options, const char *const *output_fields, + size_t count) { + ZVEC_CHECK_NOTNULL_ERRCODE(options, ZVEC_ERROR_INVALID_ARGUMENT, + "Iterator options pointer is null"); + + ZVEC_TRY_RETURN_ERROR( + "Failed to set output_fields", + auto *ptr = reinterpret_cast(options); + // NULL output_fields means "return all fields" (nullopt). A non-NULL + // array with count == 0 yields an empty selection (no scalar fields). + if (output_fields == nullptr) { + ptr->output_fields_ = std::nullopt; + return ZVEC_OK; + } + std::vector fields; + fields.reserve(count); + for (size_t i = 0; i < count; ++i) { + if (output_fields[i]) { + fields.emplace_back(output_fields[i]); + } else { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Null output_field at index " + std::to_string(i)); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + } + ptr->output_fields_ = std::move(fields); + return ZVEC_OK;) +} + +zvec_error_code_t zvec_iterator_options_set_include_vector( + zvec_iterator_options_t *options, bool include) { + ZVEC_CHECK_NOTNULL_ERRCODE(options, ZVEC_ERROR_INVALID_ARGUMENT, + "Iterator options pointer is null"); + auto *ptr = reinterpret_cast(options); + ptr->include_vector_ = include; + return ZVEC_OK; +} + +zvec_error_code_t zvec_collection_create_iterator( + zvec_collection_t *collection, const zvec_iterator_options_t *options, + zvec_doc_iterator_t **out_iter) { + ZVEC_CHECK_NOTNULL_ERRCODE(collection, ZVEC_ERROR_INVALID_ARGUMENT, + "Collection handle is null"); + ZVEC_CHECK_NOTNULL_ERRCODE(out_iter, ZVEC_ERROR_INVALID_ARGUMENT, + "out_iter pointer is null"); + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_collection_create_iterator", + // CreateIterator is non-const (it flushes), so use a non-const handle. + auto coll_ptr = + reinterpret_cast *>(collection); + + zvec::IteratorOptions iter_options; + if (options) { + iter_options = + *reinterpret_cast(options); + } + + auto result = (*coll_ptr)->CreateIterator(iter_options); + if (!result.has_value()) { + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, + "Failed to create iterator: " + + result.error().message()); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // Wrap the shared_ptr like other handles. + *out_iter = reinterpret_cast( + new std::shared_ptr(std::move(result.value()))); + return ZVEC_OK;) +} + +zvec_error_code_t zvec_doc_iterator_next(zvec_doc_iterator_t *iter, + zvec_doc_t **out_doc) { + ZVEC_CHECK_NOTNULL_ERRCODE(iter, ZVEC_ERROR_INVALID_ARGUMENT, + "Iterator handle is null"); + ZVEC_CHECK_NOTNULL_ERRCODE(out_doc, ZVEC_ERROR_INVALID_ARGUMENT, + "out_doc pointer is null"); + + *out_doc = nullptr; + + ZVEC_TRY_RETURN_ERROR( + "Exception in zvec_doc_iterator_next", + auto iter_ptr = + reinterpret_cast *>(iter); + + auto result = (*iter_ptr)->Next(); + if (!result.has_value()) { + SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR, + "Iterator next failed: " + result.error().message()); + return ZVEC_ERROR_INTERNAL_ERROR; + } + + // EOF: value() is nullptr → leave *out_doc = NULL, return OK. + auto doc = result.value(); + if (doc == nullptr) { + return ZVEC_OK; + } + + // Copy into a heap Doc owned by the caller (freed via zvec_doc_destroy). + *out_doc = reinterpret_cast(new zvec::Doc(*doc)); + return ZVEC_OK;) +} + +void zvec_doc_iterator_close(zvec_doc_iterator_t *iter) { + if (iter) { + // Deleting the shared_ptr wrapper releases the DocIterator (whose + // destructor calls Close() and releases segments/delete_store). + delete reinterpret_cast *>(iter); + } +} diff --git a/src/binding/python/include/python_collection.h b/src/binding/python/include/python_collection.h index 7c4177565..4b9653a14 100644 --- a/src/binding/python/include/python_collection.h +++ b/src/binding/python/include/python_collection.h @@ -27,6 +27,7 @@ class ZVecPyCollection { static void Initialize(py::module_ &m); private: + static void bind_iterator(py::module_ &m); static void bind_db_methods(py::class_ &col); static void bind_ddl_methods(py::class_ &col); static void bind_dml_methods(py::class_ &col); diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 4468a53ff..e15b57481 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -15,6 +15,7 @@ #include "python_collection.h" #include #include +#include namespace zvec { @@ -52,6 +53,8 @@ void ZVecPyCollection::Initialize(pybind11::module_ &m) { .def_readonly("group_by_value", &GroupResult::group_by_value_) .def_readonly("docs", &GroupResult::docs_); + bind_iterator(m); + py::class_ collection(m, "_Collection"); bind_db_methods(collection); bind_ddl_methods(collection); @@ -302,6 +305,25 @@ void ZVecPyCollection::bind_dql_methods( }, py::arg("pks"), py::arg("output_fields") = py::none(), py::arg("include_vector") = true) + .def( + "CreateIterator", + [](Collection &self, + const std::optional> &output_fields, + bool include_vector) { + IteratorOptions options; + options.output_fields_ = output_fields; + options.include_vector_ = include_vector; + Result result; + { + py::gil_scoped_release release; + result = self.CreateIterator(options); + } + // return DocIterator::Ptr -> _DocIterator + return unwrap_expected(result); + }, + py::arg("output_fields") = py::none(), + py::arg("include_vector") = true, + "Create a document iterator to traverse all documents.") .def( "_debug_hnsw_storage_mode", [](const Collection &self, const std::string &column_name) { @@ -316,4 +338,26 @@ void ZVecPyCollection::bind_dql_methods( "for introspection and testing only; not part of the stable API."); } +void ZVecPyCollection::bind_iterator(py::module_ &m) { + // Document iterator: Python iterator protocol (__iter__ / __next__). + // Constructed only via Collection.CreateIterator (no py::init). + py::class_(m, "_DocIterator") + .def("__iter__", [](py::object self) { return self; }) + .def("__next__", + [](DocIterator &self) { + Result result; + { + py::gil_scoped_release release; + result = self.Next(); + } + // !has_value() -> error (raises); value()==nullptr -> EOF + auto doc = unwrap_expected(result); + if (doc == nullptr) { + throw py::stop_iteration(); + } + return doc; + }) + .def("close", [](DocIterator &self) { self.Close(); }); +} + } // namespace zvec diff --git a/src/db/collection.cc b/src/db/collection.cc index 33750f3b1..70ea701af 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -28,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -37,11 +39,14 @@ #include "db/common/global_resource.h" #include "db/common/profiler.h" #include "db/common/typedef.h" +#include "db/doc_iterator_internal.h" #include "db/index/common/delete_store.h" #include "db/index/common/id_map.h" #include "db/index/common/index_filter.h" #include "db/index/common/type_helper.h" #include "db/index/common/version_manager.h" +#include "db/index/segment/concatenating_reader.h" +#include "db/index/segment/filtering_reader.h" #include "db/index/segment/segment.h" #include "db/index/segment/segment_helper.h" #include "db/index/segment/segment_manager.h" @@ -131,6 +136,9 @@ class CollectionImpl : public Collection { &output_fields = std::nullopt, bool include_vector = true) const override; + Result CreateIterator( + const IteratorOptions &options = {}) override; + Result DebugGetHnswStorageMode( const std::string &column_name) const override; @@ -167,6 +175,18 @@ class CollectionImpl : public Collection { Result write_impl(std::vector &docs, WriteMode mode); + // Build scan columns: prepend GLOBAL_DOC_ID + USER_ID to user-specified + // fields If output_fields is nullopt, include all forward_fields + std::vector build_scan_columns( + const std::optional> &output_fields) const; + + // Isolated scan: Flush + snapshot + clone + reader chain + // segments/delete_store filled by Scan for caller to keep alive + Result Scan(const IteratorOptions &options, + std::vector &segments, + DeleteStore::Ptr &delete_store, + CollectionSchema::Ptr &schema); + std::vector get_all_segments() const; std::vector get_all_persist_segments() const; @@ -2085,4 +2105,116 @@ std::vector CollectionImpl::get_all_persist_segments() const { return segment_manager_->get_segments(); } +std::vector CollectionImpl::build_scan_columns( + const std::optional> &output_fields) const { + // Same logic as Segment::Fetch() (segment.cc:1121-1139) + std::vector columns; + columns.push_back(GLOBAL_DOC_ID); // _zvec_g_doc_id_ + columns.push_back(USER_ID); // _zvec_uid_ + + if (!output_fields.has_value()) { + // nullopt = all forward fields + for (const auto &field : schema_->forward_fields()) { + columns.push_back(field->name()); + } + } else { + // Only requested fields + const auto &requested = *output_fields; + std::unordered_set requested_set(requested.begin(), + requested.end()); + for (const auto &field : schema_->forward_fields()) { + if (requested_set.count(field->name())) { + columns.push_back(field->name()); + } + } + } + + return columns; +} + +Result CollectionImpl::Scan( + const IteratorOptions &options, std::vector &segments, + DeleteStore::Ptr &delete_store, CollectionSchema::Ptr &schema) { + // shared_lock blocks Optimize (exclusive) and schema changes (exclusive) + // but allows concurrent Query/Fetch (shared) + std::shared_lock schema_lock(schema_handle_mtx_); + CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); + + // Capture the schema under schema_handle_mtx_ so the iterator interprets + // batches with the same schema used to build the reader (avoids a race with + // a concurrent schema change between Scan() and CreateIterator()). + schema = schema_; + + { + // write_mtx_ blocks concurrent writes (Insert/Upsert/Update/Delete) + // auto-released when scope exits + std::lock_guard write_lock(write_mtx_); + + // Flush writing segment if it has data: + // 1. dump() → sealed_=true (segment.cc:2211) + // 2. add_segment() → persist segment (segment_manager.cc:31) + // 3. create new writing segment (collection.cc:1596-1608) + // Skip for read-only collections: they cannot flush, and a read-only + // collection has no unpersisted writes (the writing segment is empty). + if (!options_.read_only_ && writing_segment_->doc_count() != 0) { + auto s = switch_to_new_segment_for_writing(); + CHECK_RETURN_STATUS_EXPECTED(s); + } + + // Snapshot: shared_ptr copies (no data copy) + segments = get_all_persist_segments(); + + // Deep clone DeleteStore bitmap (delete_store.h:130) + delete_store = delete_store_->clone(); + } // ← write_mtx_ released, concurrent writes resume + + // Build reader chain (shared_lock still held, blocks schema changes) + auto scan_columns = build_scan_columns(options.output_fields_); + + std::vector> readers; + for (const auto &seg : segments) { + auto scalar_reader = seg->scan(scan_columns); + if (scalar_reader) { + readers.push_back(scalar_reader); + } + // TODO(Day 2): VectorScanReader for vector columns + } + + // Concatenate across segments + auto concat_reader = ConcatenatingReader::Make(std::move(readers)); + + // Wrap with delete filter + auto filter = delete_store->make_filter(); + auto filtering_reader = FilteringReader::Make(concat_reader, filter); + + return filtering_reader; +} + +Result CollectionImpl::CreateIterator( + const IteratorOptions &options) { + // Note: iteration is a read-only operation and is intentionally allowed on + // read-only collections (a key export use-case). Scan() skips the writing + // segment flush when read-only. + + // Scan() takes shared_lock(schema_handle_mtx_) internally and captures + // the schema under that lock, consistent with the reader it builds. + std::vector segments; + DeleteStore::Ptr delete_store; + CollectionSchema::Ptr schema; + auto reader_result = Scan(options, segments, delete_store, schema); + if (!reader_result) { + return tl::make_unexpected(reader_result.error()); + } + + // Create DocIterator, keeping segments + delete_store alive + auto impl = std::make_unique(); + impl->reader = reader_result.value(); + impl->schema = std::move(schema); + impl->segments = std::move(segments); + impl->delete_store = std::move(delete_store); + impl->include_vector = options.include_vector_; + + return DocIterator::Ptr(new DocIterator(std::move(impl))); +} + } // namespace zvec diff --git a/src/db/doc_iterator.cc b/src/db/doc_iterator.cc new file mode 100644 index 000000000..4cbff8a07 --- /dev/null +++ b/src/db/doc_iterator.cc @@ -0,0 +1,387 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "db/common/constants.h" +#include "db/doc_iterator_internal.h" + +namespace zvec { + +// ── VectorDataBuffer → Doc field conversion ── +// Same logic as SegmentImpl::ConvertVectorDataBufferToDocField +// (segment.cc:1024-1093) but as a free function to avoid modifying Segment's +// interface. +static Status SetVectorFieldFromBuffer( + const FieldSchema::Ptr &field, + const vector_column_params::VectorDataBuffer &buf, Doc *doc) { + if (std::holds_alternative( + buf.vector_buffer)) { + const auto &dense = + std::get(buf.vector_buffer); + switch (field->data_type()) { + case DataType::VECTOR_FP32: { + const auto *p = reinterpret_cast(dense.data.data()); + size_t n = dense.data.size() / sizeof(float); + doc->set(field->name(), std::vector(p, p + n)); + break; + } + case DataType::VECTOR_FP16: { + const auto *p = reinterpret_cast(dense.data.data()); + size_t n = dense.data.size() / sizeof(float16_t); + doc->set(field->name(), std::vector(p, p + n)); + break; + } + case DataType::VECTOR_INT8: { + const auto *p = reinterpret_cast(dense.data.data()); + size_t n = dense.data.size() / sizeof(int8_t); + doc->set(field->name(), std::vector(p, p + n)); + break; + } + case DataType::VECTOR_FP64: { + const auto *p = reinterpret_cast(dense.data.data()); + size_t n = dense.data.size() / sizeof(double); + doc->set(field->name(), std::vector(p, p + n)); + break; + } + case DataType::VECTOR_INT16: { + const auto *p = reinterpret_cast(dense.data.data()); + size_t n = dense.data.size() / sizeof(int16_t); + doc->set(field->name(), std::vector(p, p + n)); + break; + } + default: + return Status::InvalidArgument("Unsupported dense vector type: ", + field->data_type()); + } + } else if (std::holds_alternative( + buf.vector_buffer)) { + const auto &sparse = + std::get(buf.vector_buffer); + switch (field->data_type()) { + case DataType::SPARSE_VECTOR_FP32: { + const auto *idx = + reinterpret_cast(sparse.indices.data()); + size_t idx_n = sparse.indices.size() / sizeof(uint32_t); + const auto *val = reinterpret_cast(sparse.values.data()); + size_t val_n = sparse.values.size() / sizeof(float); + doc->set(field->name(), + std::make_pair(std::vector(idx, idx + idx_n), + std::vector(val, val + val_n))); + break; + } + case DataType::SPARSE_VECTOR_FP16: { + const auto *idx = + reinterpret_cast(sparse.indices.data()); + size_t idx_n = sparse.indices.size() / sizeof(uint32_t); + const auto *val = + reinterpret_cast(sparse.values.data()); + size_t val_n = sparse.values.size() / sizeof(float16_t); + doc->set(field->name(), + std::make_pair(std::vector(idx, idx + idx_n), + std::vector(val, val + val_n))); + break; + } + default: + return Status::InvalidArgument("Unsupported sparse vector type: ", + field->data_type()); + } + } + return Status::OK(); +} + +// ── Extract a numeric list field +// (ARRAY_INT32/INT64/UINT32/UINT64/FLOAT/DOUBLE) from row `row` of a ListArray. +// Numeric arrays are contiguous, so raw_values() (already offset-adjusted for +// the sliced sub-array) can be copied directly. +template +static void SetNumericListField( + const std::shared_ptr &list_array, int64_t row, + const std::string &name, Doc *doc) { + auto values = + std::dynamic_pointer_cast(list_array->value_slice(row)); + if (values) { + doc->set(name, std::vector(values->raw_values(), + values->raw_values() + values->length())); + } +} + +// ── DocIterator implementation ── + +DocIterator::DocIterator(std::unique_ptr impl) : impl_(std::move(impl)) {} + +DocIterator::~DocIterator() { + Close(); +} + +void DocIterator::Close() { + if (impl_) { + impl_->closed = true; + // Release the Arrow reader/batch first (they reference segment files), + // then release the kept-alive snapshot resources so a closed iterator + // retains nothing. + impl_->reader.reset(); + impl_->current_batch.reset(); + impl_->vector_cache_.clear(); + impl_->segments.clear(); + impl_->delete_store.reset(); + impl_->schema.reset(); + } +} + +Result DocIterator::Next() { + if (!impl_ || impl_->closed) { + return tl::make_unexpected(Status::InternalError("Iterator is closed")); + } + + // If current batch exhausted or not loaded, read next batch + if (!impl_->current_batch || + impl_->current_row >= impl_->current_batch->num_rows()) { + auto status = impl_->reader->ReadNext(&impl_->current_batch); + if (!status.ok()) { + return tl::make_unexpected( + Status::InternalError("ReadNext failed: ", status.ToString())); + } + impl_->current_row = 0; + if (!impl_->current_batch) { + return Doc::Ptr(nullptr); // EOF + } + // Pre-fetch vectors for the new batch (batch-level, not per-doc) + impl_->vector_cache_.clear(); + if (impl_->include_vector && impl_->schema) { + auto &batch = *impl_->current_batch; + int gc = batch.schema()->GetFieldIndex(GLOBAL_DOC_ID); + if (gc >= 0) { + auto ga = + std::dynamic_pointer_cast(batch.column(gc)); + if (ga && batch.num_rows() > 0) { + uint64_t first_gdoc = ga->Value(0); + for (const auto &seg : impl_->segments) { + if (first_gdoc >= seg->meta()->min_doc_id() && + first_gdoc <= seg->meta()->max_doc_id()) { + uint64_t min_doc_id = seg->meta()->min_doc_id(); + for (const auto &field : impl_->schema->vector_fields()) { + auto indexer = seg->get_combined_vector_indexer(field->name()); + if (!indexer) continue; + std::vector< + std::optional> + bufs; + bufs.reserve(batch.num_rows()); + for (int64_t i = 0; i < batch.num_rows(); i++) { + uint32_t seg_doc_id = + static_cast(ga->Value(i) - min_doc_id); + auto fr = indexer->Fetch(seg_doc_id); + if (fr.has_value()) { + bufs.push_back(std::move(fr.value())); + } else { + // Match Segment::Fetch(): log and skip the field instead of + // fabricating an empty vector that would hide the error. + LOG_ERROR( + "vector prefetch failed, field: %s, g_doc_id: %llu: %s", + field->name().c_str(), + static_cast(ga->Value(i)), + fr.error().message().c_str()); + bufs.emplace_back(std::nullopt); + } + } + impl_->vector_cache_[field->name()] = std::move(bufs); + } + break; + } + } + } + } + } + } + + auto &batch = *impl_->current_batch; + int64_t row = impl_->current_row; + auto doc = std::make_shared(); + + // 1. Extract PK from _zvec_uid_ column + int uid_col = batch.schema()->GetFieldIndex(USER_ID); + if (uid_col >= 0) { + auto uid_array = + std::dynamic_pointer_cast(batch.column(uid_col)); + if (uid_array) { + // GetView avoids the per-row Scalar allocation of GetScalar()->ToString() + doc->set_pk(std::string(uid_array->GetView(row))); + } + } + + // 2. Extract doc_id from _zvec_g_doc_id_ column + int gdoc_col = batch.schema()->GetFieldIndex(GLOBAL_DOC_ID); + if (gdoc_col >= 0) { + auto gdoc_array = + std::dynamic_pointer_cast(batch.column(gdoc_col)); + if (gdoc_array) { + doc->set_doc_id(gdoc_array->Value(row)); + } + } + + // 3. Extract scalar fields from Arrow batch + if (impl_->schema) { + for (const auto &field : impl_->schema->forward_fields()) { + int col = batch.schema()->GetFieldIndex(field->name()); + if (col < 0) continue; + + auto array = batch.column(col); + if (array->IsNull(row)) continue; + + switch (field->data_type()) { + case DataType::STRING: { + auto str_array = std::dynamic_pointer_cast(array); + if (str_array) { + // GetView avoids a per-row Scalar allocation in the hot path + doc->set(field->name(), std::string(str_array->GetView(row))); + } + break; + } + case DataType::INT32: { + auto a = std::dynamic_pointer_cast(array); + if (a) doc->set(field->name(), a->Value(row)); + break; + } + case DataType::INT64: { + auto a = std::dynamic_pointer_cast(array); + if (a) doc->set(field->name(), a->Value(row)); + break; + } + case DataType::UINT32: { + auto a = std::dynamic_pointer_cast(array); + if (a) doc->set(field->name(), a->Value(row)); + break; + } + case DataType::UINT64: { + auto a = std::dynamic_pointer_cast(array); + if (a) doc->set(field->name(), a->Value(row)); + break; + } + case DataType::FLOAT: { + auto a = std::dynamic_pointer_cast(array); + if (a) doc->set(field->name(), a->Value(row)); + break; + } + case DataType::DOUBLE: { + auto a = std::dynamic_pointer_cast(array); + if (a) doc->set(field->name(), a->Value(row)); + break; + } + case DataType::BOOL: { + auto a = std::dynamic_pointer_cast(array); + if (a) doc->set(field->name(), a->Value(row)); + break; + } + case DataType::ARRAY_INT32: { + auto la = std::dynamic_pointer_cast(array); + if (la) + SetNumericListField( + la, row, field->name(), doc.get()); + break; + } + case DataType::ARRAY_INT64: { + auto la = std::dynamic_pointer_cast(array); + if (la) + SetNumericListField( + la, row, field->name(), doc.get()); + break; + } + case DataType::ARRAY_UINT32: { + auto la = std::dynamic_pointer_cast(array); + if (la) + SetNumericListField( + la, row, field->name(), doc.get()); + break; + } + case DataType::ARRAY_UINT64: { + auto la = std::dynamic_pointer_cast(array); + if (la) + SetNumericListField( + la, row, field->name(), doc.get()); + break; + } + case DataType::ARRAY_FLOAT: { + auto la = std::dynamic_pointer_cast(array); + if (la) + SetNumericListField( + la, row, field->name(), doc.get()); + break; + } + case DataType::ARRAY_DOUBLE: { + auto la = std::dynamic_pointer_cast(array); + if (la) + SetNumericListField( + la, row, field->name(), doc.get()); + break; + } + case DataType::ARRAY_BOOL: { + auto la = std::dynamic_pointer_cast(array); + if (la) { + auto vals = std::dynamic_pointer_cast( + la->value_slice(row)); + if (vals) { + std::vector vec; + vec.reserve(vals->length()); + for (int64_t i = 0; i < vals->length(); ++i) { + vec.push_back(vals->Value(i)); + } + doc->set(field->name(), vec); + } + } + break; + } + case DataType::ARRAY_STRING: { + auto la = std::dynamic_pointer_cast(array); + if (la) { + auto vals = std::dynamic_pointer_cast( + la->value_slice(row)); + if (vals) { + std::vector vec; + vec.reserve(vals->length()); + for (int64_t i = 0; i < vals->length(); ++i) { + vec.push_back(vals->GetString(i)); + } + doc->set(field->name(), vec); + } + } + break; + } + default: + break; + } + } + } + + // 4. Extract vector fields from pre-fetched cache + if (impl_->include_vector && impl_->schema) { + for (const auto &field : impl_->schema->vector_fields()) { + auto it = impl_->vector_cache_.find(field->name()); + if (it == impl_->vector_cache_.end()) continue; + if (row >= static_cast(it->second.size())) continue; + if (!it->second[row].has_value()) continue; // fetch failed: skip field + + auto s = SetVectorFieldFromBuffer(field, *(it->second[row]), doc.get()); + if (!s.ok()) { + LOG_ERROR("SetVectorFieldFromBuffer failed for %s: %s", + field->name().c_str(), s.message().c_str()); + } + } + } + + impl_->current_row++; + return doc; +} + +} // namespace zvec diff --git a/src/db/doc_iterator_internal.h b/src/db/doc_iterator_internal.h new file mode 100644 index 000000000..b9b387fda --- /dev/null +++ b/src/db/doc_iterator_internal.h @@ -0,0 +1,56 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Internal header — NOT in public includes (src/include/) +// Shared by collection.cc and doc_iterator.cc +#pragma once + +#include +#include +#include +#include +#include +#include "db/index/column/vector_column/combined_vector_column_indexer.h" +#include "db/index/column/vector_column/vector_column_params.h" +#include "db/index/common/delete_store.h" +#include "db/index/segment/segment.h" +#include "db/index/storage/base_forward_store.h" + +namespace zvec { + +struct DocIterator::Impl { + // Declaration order controls destruction order (reverse of declaration). + // segments must be declared FIRST → destroyed LAST. + // reader must be declared LAST → destroyed FIRST. + // This ensures Arrow file handles are released before Segment::cleanup() + // deletes files from disk (important on Windows). + std::vector segments; // keep Segment alive + DeleteStore::Ptr delete_store; // keep delete bitmap alive + CollectionSchema::Ptr schema; + int64_t current_row{0}; + bool closed{false}; + bool include_vector{false}; // whether to fetch vector fields + std::shared_ptr current_batch; + // Pre-fetched vector data for current batch. + // Key: field_name, Value: one optional per row in + // current_batch; nullopt marks a row whose vector fetch failed (field is + // skipped, matching Segment::Fetch, instead of fabricating an empty vector). + std::unordered_map< + std::string, + std::vector>> + vector_cache_; + RecordBatchReaderPtr reader; // destroyed before segments +}; + +} // namespace zvec diff --git a/src/db/index/segment/concatenating_reader.h b/src/db/index/segment/concatenating_reader.h new file mode 100644 index 000000000..708c3e211 --- /dev/null +++ b/src/db/index/segment/concatenating_reader.h @@ -0,0 +1,70 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include + +namespace zvec { + +// Concatenates multiple RecordBatchReaders in sequence. +// Each reader is consumed fully before moving to the next. +// Modeled after CombinedRecordBatchReader (segment.cc:2875). +class ConcatenatingReader : public arrow::RecordBatchReader { + public: + static std::shared_ptr Make( + std::vector> &&readers) { + return std::make_shared(std::move(readers)); + } + + explicit ConcatenatingReader( + std::vector> &&readers) + : readers_(std::move(readers)), current_index_(0) { + if (!readers_.empty()) { + schema_ = readers_[0]->schema(); + } + } + + ~ConcatenatingReader() override = default; + + std::shared_ptr schema() const override { + return schema_; + } + + arrow::Status ReadNext(std::shared_ptr *batch) override { + *batch = nullptr; + while (current_index_ < readers_.size()) { + auto status = readers_[current_index_]->ReadNext(batch); + if (!status.ok()) { + return status; + } + if (*batch) { + return arrow::Status::OK(); + } + // Current reader exhausted, move to next + current_index_++; + } + *batch = nullptr; + return arrow::Status::OK(); + } + + private: + std::vector> readers_; + size_t current_index_; + std::shared_ptr schema_; +}; + +} // namespace zvec diff --git a/src/db/index/segment/filtering_reader.h b/src/db/index/segment/filtering_reader.h new file mode 100644 index 000000000..616b65d55 --- /dev/null +++ b/src/db/index/segment/filtering_reader.h @@ -0,0 +1,114 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include "db/common/constants.h" +#include "db/index/common/index_filter.h" + +namespace zvec { + +// Wraps a RecordBatchReader and filters out deleted rows. +// Uses IndexFilter (from DeleteStore::make_filter()) to check each row's +// _zvec_g_doc_id_ against the delete bitmap. +// System columns are preserved (DocIterator extracts PK from them). +class FilteringReader : public arrow::RecordBatchReader { + public: + static std::shared_ptr Make( + std::shared_ptr inner_reader, + const IndexFilter::Ptr &filter) { + return std::make_shared(std::move(inner_reader), filter); + } + + FilteringReader(std::shared_ptr inner_reader, + const IndexFilter::Ptr &filter) + : inner_reader_(std::move(inner_reader)), filter_(filter) {} + + ~FilteringReader() override = default; + + std::shared_ptr schema() const override { + return inner_reader_->schema(); + } + + arrow::Status ReadNext(std::shared_ptr *batch) override { + while (true) { + ARROW_RETURN_NOT_OK(inner_reader_->ReadNext(batch)); + if (!*batch) { + return arrow::Status::OK(); + } + + // No filter → return as-is + if (!filter_) { + return arrow::Status::OK(); + } + + // Find the _zvec_g_doc_id_ column + int gdoc_col = (*batch)->schema()->GetFieldIndex(GLOBAL_DOC_ID); + if (gdoc_col < 0) { + // No global doc id column — can't filter, return as-is + return arrow::Status::OK(); + } + + auto gdoc_array = std::dynamic_pointer_cast( + (*batch)->column(gdoc_col)); + if (!gdoc_array) { + return arrow::Status::OK(); + } + + // Build filter mask: true = keep, false = skip (deleted) + arrow::BooleanBuilder mask_builder; + int64_t num_rows = (*batch)->num_rows(); + ARROW_RETURN_NOT_OK(mask_builder.Reserve(num_rows)); + + bool has_filtered = false; + for (int64_t i = 0; i < num_rows; ++i) { + uint64_t g_doc_id = gdoc_array->Value(i); + bool is_deleted = filter_->is_filtered(g_doc_id); + if (is_deleted) has_filtered = true; + mask_builder.UnsafeAppend(!is_deleted); + } + + // No rows filtered → return batch as-is + if (!has_filtered) { + return arrow::Status::OK(); + } + + // Apply filter + std::shared_ptr mask_array; + ARROW_RETURN_NOT_OK(mask_builder.Finish(&mask_array)); + + arrow::Datum result; + ARROW_ASSIGN_OR_RAISE(result, + arrow::compute::Filter(arrow::Datum(*batch), + arrow::Datum(mask_array))); + + *batch = result.record_batch(); + + // If all rows filtered out, continue to next batch + if ((*batch)->num_rows() == 0) { + continue; + } + + return arrow::Status::OK(); + } + } + + private: + std::shared_ptr inner_reader_; + IndexFilter::Ptr filter_; +}; + +} // namespace zvec diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 2d048689b..065fab3b5 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -3590,6 +3590,93 @@ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_collection_fetch( size_t count, const char *const *output_fields, size_t output_field_count, bool include_vector, zvec_doc_t ***documents, size_t *found_count); +// ============================================================================= +// Document Iterator Interface (full traversal) +// ============================================================================= + +/** + * @brief Opaque handle for a document iterator. + * + * Created by zvec_collection_create_iterator, released by + * zvec_doc_iterator_close. Iterates over all documents in a collection using + * an isolated snapshot (data written after creation is not visible). + */ +typedef struct zvec_doc_iterator_t zvec_doc_iterator_t; + +/** + * @brief Opaque handle for iterator options. + * + * Follows the same pattern as zvec_collection_options_t (opaque type + setters) + * for ABI stability: adding new options later does not break the ABI. + */ +typedef struct zvec_iterator_options_t zvec_iterator_options_t; + +/** + * @brief Create an iterator options object with default values + * (output_fields = all, include_vector = true). + * @return Options handle, or NULL on allocation failure. + */ +ZVEC_EXPORT zvec_iterator_options_t *ZVEC_CALL +zvec_iterator_options_create(void); + +/** + * @brief Destroy an iterator options object. + * @param options Options handle (may be NULL). + */ +ZVEC_EXPORT void ZVEC_CALL +zvec_iterator_options_destroy(zvec_iterator_options_t *options); + +/** + * @brief Set the scalar fields to return. + * @param options Options handle + * @param output_fields Array of field names; NULL means return all fields + * @param count Number of entries in output_fields. If output_fields is + * non-NULL and count is 0, no scalar fields are returned (only + * the primary key / system columns). + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_iterator_options_set_output_fields( + zvec_iterator_options_t *options, const char *const *output_fields, + size_t count); + +/** + * @brief Set whether to include vector fields in the returned documents. + * @param options Options handle + * @param include true to include vectors, false to skip them + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_iterator_options_set_include_vector(zvec_iterator_options_t *options, + bool include); + +/** + * @brief Create a document iterator over the collection. + * @param collection Collection handle + * @param options Iterator options (may be NULL to use defaults) + * @param[out] out_iter Returned iterator handle (release with + * zvec_doc_iterator_close) + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_collection_create_iterator( + zvec_collection_t *collection, const zvec_iterator_options_t *options, + zvec_doc_iterator_t **out_iter); + +/** + * @brief Advance the iterator and return the next document. + * @param iter Iterator handle + * @param[out] out_doc Returned document (release with zvec_doc_destroy). + * Set to NULL when iteration reaches the end (EOF). + * @return zvec_error_code_t ZVEC_OK on success or EOF; error code otherwise. + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_doc_iterator_next(zvec_doc_iterator_t *iter, zvec_doc_t **out_doc); + +/** + * @brief Close the iterator and release all its resources. + * @param iter Iterator handle (may be NULL). + */ +ZVEC_EXPORT void ZVEC_CALL zvec_doc_iterator_close(zvec_doc_iterator_t *iter); + // ============================================================================= // Document Related Structures // ============================================================================= diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index 83a289c52..5c3448a22 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -109,6 +110,9 @@ class Collection { &output_fields = std::nullopt, bool include_vector = true) const = 0; + virtual Result CreateIterator( + const IteratorOptions &options = {}) = 0; + public: //! Debug-only: retrieve the storage mode string of an HNSW index on the //! given vector column. Returns one of {"mmap", "buffer_pool", diff --git a/src/include/zvec/db/doc_iterator.h b/src/include/zvec/db/doc_iterator.h new file mode 100644 index 000000000..91ec48c37 --- /dev/null +++ b/src/include/zvec/db/doc_iterator.h @@ -0,0 +1,48 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include + +namespace zvec { + +class CollectionImpl; + +class DocIterator { + public: + using Ptr = std::shared_ptr; + + // !has_value() → error + // has_value() && value() == nullptr → EOF + // has_value() && value() != nullptr → success + Result Next(); + + void Close(); + + ~DocIterator(); + + private: + friend class CollectionImpl; + + // Pimpl: implementation holds Arrow types (RecordBatchReader, RecordBatch) + struct Impl; + std::unique_ptr impl_; + + // Called by CollectionImpl::CreateIterator + explicit DocIterator(std::unique_ptr impl); +}; + +} // namespace zvec diff --git a/src/include/zvec/db/options.h b/src/include/zvec/db/options.h index 8c3e8c594..5e57ea8ec 100644 --- a/src/include/zvec/db/options.h +++ b/src/include/zvec/db/options.h @@ -14,6 +14,9 @@ #pragma once #include +#include +#include +#include namespace zvec { @@ -66,4 +69,9 @@ struct AlterColumnOptions { int concurrency_{0}; }; +struct IteratorOptions { + std::optional> output_fields_; + bool include_vector_{true}; +}; + } // namespace zvec \ No newline at end of file diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index 7365dcf1a..6a8fc0454 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -6353,6 +6353,210 @@ void test_diskann_wiring_on_vector_query(void) { // Main function // ============================================================================= +// ============================================================================= +// Document iterator tests (zvec_collection_create_iterator / next / close) +// ============================================================================= + +// Insert `n` docs into the collection and flush. +static void iter_insert_docs(zvec_collection_t *collection, + const zvec_collection_schema_t *schema, int n) { + for (int i = 0; i < n; ++i) { + zvec_doc_t *doc = zvec_test_create_doc((uint64_t)i, schema, NULL); + zvec_doc_t *docs[] = {doc}; + size_t ok_count = 0, err_count = 0; + zvec_error_code_t err = zvec_collection_insert( + collection, (const zvec_doc_t **)docs, 1, &ok_count, &err_count); + TEST_ASSERT(err == ZVEC_OK && ok_count == 1 && err_count == 0); + zvec_doc_destroy(doc); + } + TEST_ASSERT(zvec_collection_flush(collection) == ZVEC_OK); +} + +// Basic iteration: count all docs, PK non-null. +void test_iterator_basic(void) { + TEST_START(); + char dir[] = "./zvec_test_c_iter_basic"; + zvec_test_delete_dir(dir); + + zvec_collection_schema_t *schema = zvec_test_create_temp_schema(); + TEST_ASSERT(schema != NULL); + + zvec_collection_t *collection = NULL; + zvec_error_code_t err = + zvec_collection_create_and_open(dir, schema, NULL, &collection); + TEST_ASSERT(err == ZVEC_OK && collection != NULL); + + const int N = 20; + iter_insert_docs(collection, schema, N); + + zvec_doc_iterator_t *iter = NULL; + err = zvec_collection_create_iterator(collection, NULL, &iter); + TEST_ASSERT(err == ZVEC_OK && iter != NULL); + + int count = 0; + while (1) { + zvec_doc_t *doc = NULL; + err = zvec_doc_iterator_next(iter, &doc); + TEST_ASSERT(err == ZVEC_OK); + if (err != ZVEC_OK) break; + if (doc == NULL) break; // EOF + + const char *pk = zvec_doc_get_pk_pointer(doc); + TEST_ASSERT(pk != NULL && strlen(pk) > 0); + zvec_doc_destroy(doc); + count++; + } + TEST_ASSERT(count == N); + + zvec_doc_iterator_close(iter); + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + zvec_test_delete_dir(dir); + TEST_END(); +} + +// include_vector=false: dense vector field should be absent. +void test_iterator_exclude_vector(void) { + TEST_START(); + char dir[] = "./zvec_test_c_iter_novec"; + zvec_test_delete_dir(dir); + + zvec_collection_schema_t *schema = zvec_test_create_temp_schema(); + zvec_collection_t *collection = NULL; + zvec_error_code_t err = + zvec_collection_create_and_open(dir, schema, NULL, &collection); + TEST_ASSERT(err == ZVEC_OK && collection != NULL); + + iter_insert_docs(collection, schema, 5); + + zvec_iterator_options_t *opts = zvec_iterator_options_create(); + TEST_ASSERT(opts != NULL); + err = zvec_iterator_options_set_include_vector(opts, false); + TEST_ASSERT(err == ZVEC_OK); + + zvec_doc_iterator_t *iter = NULL; + err = zvec_collection_create_iterator(collection, opts, &iter); + TEST_ASSERT(err == ZVEC_OK && iter != NULL); + + int count = 0; + while (1) { + zvec_doc_t *doc = NULL; + err = zvec_doc_iterator_next(iter, &doc); + TEST_ASSERT(err == ZVEC_OK); + if (err != ZVEC_OK || doc == NULL) break; + + TEST_ASSERT(zvec_doc_has_field(doc, "id")); + TEST_ASSERT(!zvec_doc_has_field(doc, "dense")); + zvec_doc_destroy(doc); + count++; + } + TEST_ASSERT(count == 5); + + zvec_doc_iterator_close(iter); + zvec_iterator_options_destroy(opts); + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + zvec_test_delete_dir(dir); + TEST_END(); +} + +// output_fields={"id"}: only "id" scalar returned, "name" absent. +void test_iterator_output_fields(void) { + TEST_START(); + char dir[] = "./zvec_test_c_iter_fields"; + zvec_test_delete_dir(dir); + + zvec_collection_schema_t *schema = zvec_test_create_temp_schema(); + zvec_collection_t *collection = NULL; + zvec_error_code_t err = + zvec_collection_create_and_open(dir, schema, NULL, &collection); + TEST_ASSERT(err == ZVEC_OK && collection != NULL); + + iter_insert_docs(collection, schema, 5); + + zvec_iterator_options_t *opts = zvec_iterator_options_create(); + const char *fields[] = {"id"}; + err = zvec_iterator_options_set_output_fields(opts, fields, 1); + TEST_ASSERT(err == ZVEC_OK); + err = zvec_iterator_options_set_include_vector(opts, false); + TEST_ASSERT(err == ZVEC_OK); + + zvec_doc_iterator_t *iter = NULL; + err = zvec_collection_create_iterator(collection, opts, &iter); + TEST_ASSERT(err == ZVEC_OK && iter != NULL); + + int count = 0; + while (1) { + zvec_doc_t *doc = NULL; + err = zvec_doc_iterator_next(iter, &doc); + if (err != ZVEC_OK || doc == NULL) break; + + TEST_ASSERT(zvec_doc_has_field(doc, "id")); + TEST_ASSERT(!zvec_doc_has_field(doc, "name")); + zvec_doc_destroy(doc); + count++; + } + TEST_ASSERT(count == 5); + + zvec_doc_iterator_close(iter); + zvec_iterator_options_destroy(opts); + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + zvec_test_delete_dir(dir); + TEST_END(); +} + +// Empty collection yields immediate EOF. +void test_iterator_empty(void) { + TEST_START(); + char dir[] = "./zvec_test_c_iter_empty"; + zvec_test_delete_dir(dir); + + zvec_collection_schema_t *schema = zvec_test_create_temp_schema(); + zvec_collection_t *collection = NULL; + zvec_error_code_t err = + zvec_collection_create_and_open(dir, schema, NULL, &collection); + TEST_ASSERT(err == ZVEC_OK && collection != NULL); + + zvec_doc_iterator_t *iter = NULL; + err = zvec_collection_create_iterator(collection, NULL, &iter); + TEST_ASSERT(err == ZVEC_OK && iter != NULL); + + zvec_doc_t *doc = NULL; + err = zvec_doc_iterator_next(iter, &doc); + TEST_ASSERT(err == ZVEC_OK); + TEST_ASSERT(doc == NULL); // EOF on empty collection + + zvec_doc_iterator_close(iter); + zvec_collection_destroy(collection); + zvec_collection_schema_destroy(schema); + zvec_test_delete_dir(dir); + TEST_END(); +} + +// Null-argument handling. +void test_iterator_null_args(void) { + TEST_START(); + + zvec_doc_iterator_t *iter = NULL; + zvec_error_code_t err = zvec_collection_create_iterator(NULL, NULL, &iter); + TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + + zvec_doc_t *doc = NULL; + err = zvec_doc_iterator_next(NULL, &doc); + TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + + err = zvec_iterator_options_set_include_vector(NULL, true); + TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + err = zvec_iterator_options_set_output_fields(NULL, NULL, 0); + TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + + // close / destroy with NULL must be safe (no crash). + zvec_doc_iterator_close(NULL); + zvec_iterator_options_destroy(NULL); + TEST_END(); +} + int main(void) { printf("Starting comprehensive C API tests...\n\n"); @@ -6431,6 +6635,13 @@ int main(void) { test_query_params_functions(); test_actual_vector_queries(); + // Document iterator tests + test_iterator_basic(); + test_iterator_exclude_vector(); + test_iterator_output_fields(); + test_iterator_empty(); + test_iterator_null_args(); + // FTS tests test_fts_index_params_functions(); test_fts_query_params_functions(); diff --git a/tests/db/iterator_test.cc b/tests/db/iterator_test.cc new file mode 100644 index 000000000..416ca8ea2 --- /dev/null +++ b/tests/db/iterator_test.cc @@ -0,0 +1,607 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "db/common/file_helper.h" +#include "index/utils/utils.h" + +using namespace zvec; +using namespace zvec::test; + +static std::string iter_test_path = "test_iterator_collection"; + +class IteratorTest : public ::testing::Test { + protected: + void SetUp() override { + zvec::ailego::MemoryLimitPool::get_instance().init(2 * 1024ll * 1024ll * + 1024ll); + FileHelper::RemoveDirectory(iter_test_path); + } + + void TearDown() override { + FileHelper::RemoveDirectory(iter_test_path); + } +}; + +// Test 1: Basic iteration — insert 100 docs, iterate, verify count + PK +TEST_F(IteratorTest, BasicIteration) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + options.enable_mmap_ = true; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()) << result.error().message(); + auto collection = std::move(result.value()); + + // Insert 100 docs + const int N = 100; + std::vector docs; + for (int i = 0; i < N; i++) { + docs.push_back(TestHelper::CreateDoc(i, *schema)); + } + auto insert_result = collection->Insert(docs); + ASSERT_TRUE(insert_result.has_value()); + + // Flush to ensure all docs are in persist segments + collection->Flush(); + + // Create iterator + auto iter_result = collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()) << iter_result.error().message(); + auto iter = iter_result.value(); + + // Iterate and collect PKs + verify scalar fields + std::set pks; + int count = 0; + while (true) { + auto r = iter->Next(); + if (!r.has_value()) { + FAIL() << "Iterator error: " << r.error().message(); + } + if (r.value() == nullptr) { + break; // EOF + } + auto doc = r.value(); + pks.insert(doc->pk()); + // Verify scalar field extraction (int32 field exists in + // TestHelper::CreateNormalSchema) + auto int32_val = doc->get("int32"); + EXPECT_TRUE(int32_val.has_value()) + << "int32 field missing for doc " << doc->pk(); + count++; + } + + EXPECT_EQ(count, N) << "Expected " << N << " docs, got " << count; + EXPECT_EQ(pks.size(), N) << "Expected " << N << " unique PKs"; + + iter->Close(); + collection->Destroy(); +} + +// Test 2: Empty collection — iterator should immediately return EOF +TEST_F(IteratorTest, EmptyCollection) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + // No docs inserted, just create iterator + auto iter_result = collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + // Next should return EOF immediately + auto r = iter->Next(); + ASSERT_TRUE(r.has_value()); + EXPECT_EQ(r.value(), nullptr) << "Expected EOF on empty collection"; + + collection->Destroy(); +} + +// Test 3: Deleted docs are filtered out +TEST_F(IteratorTest, DeletedDocsFiltered) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + // Insert 50 docs + const int N = 50; + std::vector docs; + std::vector pks_to_delete; + for (int i = 0; i < N; i++) { + auto doc = TestHelper::CreateDoc(i, *schema); + docs.push_back(doc); + if (i % 2 == 0) { + pks_to_delete.push_back(doc.pk()); + } + } + auto insert_result = collection->Insert(docs); + ASSERT_TRUE(insert_result.has_value()); + + // Delete every other doc + auto delete_result = collection->Delete(pks_to_delete); + ASSERT_TRUE(delete_result.has_value()); + + collection->Flush(); + + // Create iterator + auto iter_result = collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + // Iterate + std::set deleted_set(pks_to_delete.begin(), pks_to_delete.end()); + int count = 0; + while (true) { + auto r = iter->Next(); + if (!r.has_value()) { + FAIL() << "Iterator error: " << r.error().message(); + } + if (r.value() == nullptr) break; + + auto pk = r.value()->pk(); + EXPECT_EQ(deleted_set.count(pk), 0) + << "Deleted doc " << pk << " should not appear in iteration"; + count++; + } + + // Should have N - deleted_count docs + EXPECT_EQ(count, N - static_cast(pks_to_delete.size())); + + collection->Destroy(); +} + +// Test 4: Iterator after Close() returns error +TEST_F(IteratorTest, CloseThenNext) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + // Insert a few docs + std::vector docs; + for (int i = 0; i < 5; i++) { + docs.push_back(TestHelper::CreateDoc(i, *schema)); + } + auto insert_result = collection->Insert(docs); + ASSERT_TRUE(insert_result.has_value()); + + collection->Flush(); + + auto iter_result = collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + // Close iterator + iter->Close(); + + // Next() after Close should return error + auto r = iter->Next(); + EXPECT_FALSE(r.has_value()) << "Expected error after Close()"; + + collection->Destroy(); +} + +// Test 5: Iterator with include_vector=true — verify vector fields are present +TEST_F(IteratorTest, IncludeVector) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + options.enable_mmap_ = true; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + // Insert 10 docs + const int N = 10; + std::vector docs; + for (int i = 0; i < N; i++) { + docs.push_back(TestHelper::CreateDoc(i, *schema)); + } + auto insert_result = collection->Insert(docs); + ASSERT_TRUE(insert_result.has_value()); + collection->Flush(); + + // Create iterator with include_vector=true (default) + IteratorOptions iter_opts; + iter_opts.include_vector_ = true; + auto iter_result = collection->CreateIterator(iter_opts); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + int count = 0; + while (true) { + auto r = iter->Next(); + if (!r.has_value()) FAIL() << r.error().message(); + if (r.value() == nullptr) break; + + auto doc = r.value(); + // Verify PK + EXPECT_FALSE(doc->pk().empty()); + + // Verify vector field exists (dense_fp32 is in + // TestHelper::CreateNormalSchema) + auto vec = doc->get>("dense_fp32"); + EXPECT_TRUE(vec.has_value()) + << "dense_fp32 vector missing for doc " << doc->pk(); + if (vec.has_value()) { + EXPECT_EQ(vec->size(), 128) << "dense_fp32 dimension should be 128"; + } + + count++; + } + + EXPECT_EQ(count, N); + collection->Destroy(); +} + +// Test 6: Iterator with include_vector=false — verify no vector fields +TEST_F(IteratorTest, ExcludeVector) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + std::vector docs; + for (int i = 0; i < 5; i++) { + docs.push_back(TestHelper::CreateDoc(i, *schema)); + } + collection->Insert(docs); + collection->Flush(); + + IteratorOptions iter_opts; + iter_opts.include_vector_ = false; + auto iter_result = collection->CreateIterator(iter_opts); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + int count = 0; + while (true) { + auto r = iter->Next(); + if (!r.has_value()) FAIL() << r.error().message(); + if (r.value() == nullptr) break; + + auto doc = r.value(); + // Scalar field should be present + auto int32_val = doc->get("int32"); + EXPECT_TRUE(int32_val.has_value()); + + // Vector field should NOT be present (include_vector=false) + auto vec = doc->get>("dense_fp32"); + EXPECT_FALSE(vec.has_value()) + << "Vector should not be present with include_vector=false"; + + count++; + } + + EXPECT_EQ(count, 5); + collection->Destroy(); +} + +// Test 7: Scalar type mapping — every scalar/array Arrow type extracted +// correctly. CreateNormalSchema covers 8 base types + 8 array types. +TEST_F(IteratorTest, ScalarTypeMapping) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + // doc_id = 7 → deterministic values (see TestHelper::CreateDoc). + const uint64_t kId = 7; + std::vector docs{TestHelper::CreateDoc(kId, *schema)}; + ASSERT_TRUE(collection->Insert(docs).has_value()); + collection->Flush(); + + auto iter_result = collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + auto r = iter->Next(); + ASSERT_TRUE(r.has_value()); + ASSERT_NE(r.value(), nullptr); + auto doc = r.value(); + + // ── base scalar types ── + EXPECT_EQ(doc->get("int32").value_or(-1), (int32_t)kId); + EXPECT_EQ(doc->get("int64").value_or(-1), (int64_t)kId); + EXPECT_EQ(doc->get("uint32").value_or(0), (uint32_t)kId); + EXPECT_EQ(doc->get("uint64").value_or(0), (uint64_t)kId); + EXPECT_FLOAT_EQ(doc->get("float").value_or(-1), (float)kId); + EXPECT_DOUBLE_EQ(doc->get("double").value_or(-1), (double)kId); + EXPECT_EQ(doc->get("string").value_or(""), + "value_" + std::to_string(kId)); + EXPECT_EQ(doc->get("bool").value_or(true), kId % 10 == 0); + + // ── array types (each element == kId, length 10) ── + auto a_i32 = doc->get>("array_int32"); + ASSERT_TRUE(a_i32.has_value()); + EXPECT_EQ(a_i32->size(), 10u); + EXPECT_EQ((*a_i32)[0], (int32_t)kId); + + auto a_i64 = doc->get>("array_int64"); + ASSERT_TRUE(a_i64.has_value()); + EXPECT_EQ((*a_i64)[0], (int64_t)kId); + + auto a_u32 = doc->get>("array_uint32"); + ASSERT_TRUE(a_u32.has_value()); + EXPECT_EQ((*a_u32)[0], (uint32_t)kId); + + auto a_u64 = doc->get>("array_uint64"); + ASSERT_TRUE(a_u64.has_value()); + EXPECT_EQ((*a_u64)[0], (uint64_t)kId); + + auto a_f = doc->get>("array_float"); + ASSERT_TRUE(a_f.has_value()); + EXPECT_FLOAT_EQ((*a_f)[0], (float)kId); + + auto a_d = doc->get>("array_double"); + ASSERT_TRUE(a_d.has_value()); + EXPECT_DOUBLE_EQ((*a_d)[0], (double)kId); + + auto a_b = doc->get>("array_bool"); + ASSERT_TRUE(a_b.has_value()); + EXPECT_EQ(a_b->size(), 10u); + + auto a_s = doc->get>("array_string"); + ASSERT_TRUE(a_s.has_value()); + EXPECT_EQ((*a_s)[0], "value_" + std::to_string(kId)); + + collection->Destroy(); +} + +// Test 8: Integration — 1000 docs, verify count + PK + scalar + vector values. +TEST_F(IteratorTest, Integration1000Docs) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + const int N = 1000; + std::vector docs; + docs.reserve(N); + for (int i = 0; i < N; i++) { + docs.push_back(TestHelper::CreateDoc(i, *schema)); + } + ASSERT_TRUE(collection->Insert(docs).has_value()); + collection->Flush(); + + auto iter_result = collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + int count = 0; + std::set seen_pks; + while (true) { + auto r = iter->Next(); + ASSERT_TRUE(r.has_value()) << r.error().message(); + if (r.value() == nullptr) break; + auto doc = r.value(); + + // PK format is "pk_" (TestHelper::MakePK); derive id back. + std::string pk = doc->pk(); + seen_pks.insert(pk); + + // int32 field == the doc's id; verify vector value matches id + 0.1. + auto id32 = doc->get("int32"); + ASSERT_TRUE(id32.has_value()); + uint64_t id = static_cast(*id32); + + auto vec = doc->get>("dense_fp32"); + ASSERT_TRUE(vec.has_value()) << "vector missing for " << pk; + EXPECT_EQ(vec->size(), 128u); + EXPECT_FLOAT_EQ((*vec)[0], float(id + 0.1)); + + // scalar string value matches id. + EXPECT_EQ(doc->get("string").value_or(""), + "value_" + std::to_string(id)); + count++; + } + + EXPECT_EQ(count, N); + EXPECT_EQ(seen_pks.size(), (size_t)N); + collection->Destroy(); +} + +// Test 9: Concurrency — docs inserted after iterator creation are not visible. +TEST_F(IteratorTest, ConcurrentInsertNotVisible) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + const int N = 500; + std::vector docs; + for (int i = 0; i < N; i++) docs.push_back(TestHelper::CreateDoc(i, *schema)); + ASSERT_TRUE(collection->Insert(docs).has_value()); + collection->Flush(); + + auto iter_result = collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + // Consume one doc to establish the snapshot, then insert concurrently. + auto first = iter->Next(); + ASSERT_TRUE(first.has_value()); + ASSERT_NE(first.value(), nullptr); + + std::atomic writer_failed{false}; + std::thread writer([&]() { + std::vector more; + for (int i = N; i < N + 200; i++) { + more.push_back(TestHelper::CreateDoc(i, *schema)); + } + if (!collection->Insert(more).has_value()) writer_failed = true; + collection->Flush(); + }); + + int count = 1; // already consumed one + while (true) { + auto r = iter->Next(); + ASSERT_TRUE(r.has_value()) << r.error().message(); + if (r.value() == nullptr) break; + count++; + } + writer.join(); + + EXPECT_FALSE(writer_failed); + // Snapshot was taken at creation → only the original N are visible. + EXPECT_EQ(count, N); + + // A fresh iterator sees all N + 200. + auto iter2 = collection->CreateIterator().value(); + int count2 = 0; + while (true) { + auto r = iter2->Next(); + ASSERT_TRUE(r.has_value()); + if (r.value() == nullptr) break; + count2++; + } + EXPECT_EQ(count2, N + 200); + + collection->Destroy(); +} + +// Test 10: Concurrency — Optimize during iteration must not crash. +TEST_F(IteratorTest, ConcurrentOptimizeNoCrash) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + const int N = 500; + // Insert in two batches with a flush between, so Optimize has >1 segment. + std::vector b1, b2; + for (int i = 0; i < N / 2; i++) + b1.push_back(TestHelper::CreateDoc(i, *schema)); + for (int i = N / 2; i < N; i++) + b2.push_back(TestHelper::CreateDoc(i, *schema)); + ASSERT_TRUE(collection->Insert(b1).has_value()); + collection->Flush(); + ASSERT_TRUE(collection->Insert(b2).has_value()); + collection->Flush(); + + auto iter_result = collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()); + auto iter = iter_result.value(); + + // Kick off Optimize concurrently (destroys old segments after compaction). + std::thread optimizer([&]() { collection->Optimize(); }); + + int count = 0; + while (true) { + auto r = iter->Next(); + ASSERT_TRUE(r.has_value()) << r.error().message(); + if (r.value() == nullptr) break; + EXPECT_FALSE(r.value()->pk().empty()); + count++; + // Slow the consumer slightly so Optimize overlaps with iteration. + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + optimizer.join(); + + // Snapshot isolation: all N docs remain visible despite concurrent Optimize. + EXPECT_EQ(count, N); + collection->Destroy(); +} + +// Test 11: Performance — iterate 100k docs; memory should stay bounded (one +// batch at a time). Reports elapsed time; asserts correctness of the count. +TEST_F(IteratorTest, Performance100k) { + auto schema = TestHelper::CreateNormalSchema(); + CollectionOptions options; + options.read_only_ = false; + options.enable_mmap_ = true; + + auto result = Collection::CreateAndOpen(iter_test_path, *schema, options); + ASSERT_TRUE(result.has_value()); + auto collection = std::move(result.value()); + + const int N = 100000; + const int kBatch = 1000; // max write batch size is 1024 (constants.h) + for (int start = 0; start < N; start += kBatch) { + std::vector docs; + docs.reserve(kBatch); + for (int i = start; i < start + kBatch; i++) { + docs.push_back(TestHelper::CreateDoc(i, *schema)); + } + auto ins = collection->Insert(docs); + ASSERT_TRUE(ins.has_value()) + << "insert failed at start=" << start << ": " << ins.error().message(); + } + collection->Flush(); + + // include_vector=false to isolate scan+scalar throughput. + IteratorOptions iter_opts; + iter_opts.include_vector_ = false; + auto iter = collection->CreateIterator(iter_opts).value(); + + auto t0 = std::chrono::steady_clock::now(); + int count = 0; + while (true) { + auto r = iter->Next(); + ASSERT_TRUE(r.has_value()); + if (r.value() == nullptr) break; + count++; + } + auto t1 = std::chrono::steady_clock::now(); + auto ms = + std::chrono::duration_cast(t1 - t0).count(); + + EXPECT_EQ(count, N); + std::cout << "[perf] iterated " << count << " docs (no vector) in " << ms + << " ms (" << (ms > 0 ? count / ms : count) << " docs/ms)" + << std::endl; + + collection->Destroy(); +}