From 16c0bd2aa3ebc739edaaea4547d6a4341c1d93dd 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 1/3] 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 | 137 ++++ src/db/doc_iterator.cc | 400 +++++++++++ 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 | 60 ++ src/include/zvec/db/options.h | 8 + tests/c/c_api_test.c | 211 ++++++ tests/db/iterator_test.cc | 657 ++++++++++++++++++ 16 files changed, 2225 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..2cd40ed86 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,121 @@ 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() + 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_); + + if (options_.read_only_) { + // Read-only: cannot flush (that writes to disk, violating read-only). + // Read the writing segment directly instead — SegmentImpl::scan() also + // reads the in-memory writing block, and a read-only collection has no + // concurrent writes, so the snapshot is stable and complete. + segments = get_all_segments(); // persist segments + writing segment + } else { + // Writable: flush the writing segment to seal it against concurrent + // writes, then snapshot the persist segments (now including it). + if (writing_segment_->doc_count() != 0) { + auto s = switch_to_new_segment_for_writing(); + CHECK_RETURN_STATUS_EXPECTED(s); + } + segments = get_all_persist_segments(); + } + + // Deep clone DeleteStore bitmap + 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); + } else { + LOG_ERROR("Segment::scan returned no reader; skipping a segment"); + } + } + // Vectors are not read here; DocIterator::Next() batch-prefetches them from + // each segment's vector indexer. + + // 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 = std::move(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 std::make_shared(DocIterator::Passkey{}, 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..10dae09cd --- /dev/null +++ b/src/db/doc_iterator.cc @@ -0,0 +1,400 @@ +// 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, 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(Passkey, 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); + bool found_segment = false; + for (const auto &seg : impl_->segments) { + if (first_gdoc >= seg->meta()->min_doc_id() && + first_gdoc <= seg->meta()->max_doc_id()) { + found_segment = true; + 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; + } + } + if (!found_segment) { + LOG_ERROR("vector prefetch: no segment owns g_doc_id %llu", + static_cast(first_gdoc)); + } + } + } + } + } + + 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 typed_array = + std::dynamic_pointer_cast(array); + if (typed_array) doc->set(field->name(), typed_array->Value(row)); + break; + } + case DataType::INT64: { + auto typed_array = + std::dynamic_pointer_cast(array); + if (typed_array) doc->set(field->name(), typed_array->Value(row)); + break; + } + case DataType::UINT32: { + auto typed_array = + std::dynamic_pointer_cast(array); + if (typed_array) doc->set(field->name(), typed_array->Value(row)); + break; + } + case DataType::UINT64: { + auto typed_array = + std::dynamic_pointer_cast(array); + if (typed_array) doc->set(field->name(), typed_array->Value(row)); + break; + } + case DataType::FLOAT: { + auto typed_array = + std::dynamic_pointer_cast(array); + if (typed_array) doc->set(field->name(), typed_array->Value(row)); + break; + } + case DataType::DOUBLE: { + auto typed_array = + std::dynamic_pointer_cast(array); + if (typed_array) doc->set(field->name(), typed_array->Value(row)); + break; + } + case DataType::BOOL: { + auto typed_array = + std::dynamic_pointer_cast(array); + if (typed_array) doc->set(field->name(), typed_array->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..d0aa06fd5 --- /dev/null +++ b/src/include/zvec/db/doc_iterator.h @@ -0,0 +1,60 @@ +// 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 { + // Pimpl: implementation holds Arrow types (RecordBatchReader, RecordBatch). + // Forward-declared here (private) so the public constructor can name it. + struct Impl; + + public: + using Ptr = std::shared_ptr; + + // Passkey idiom: Passkey's constructor is private and CollectionImpl is its + // friend, so only CollectionImpl can construct a DocIterator. This lets + // CollectionImpl use std::make_shared while keeping construction controlled. + struct Passkey { + private: + Passkey() {} + friend class CollectionImpl; + }; + + // Called by CollectionImpl::CreateIterator + DocIterator(Passkey, std::unique_ptr impl); + + // !has_value() → error + // has_value() && value() == nullptr → EOF + // has_value() && value() != nullptr → success + Result Next(); + + void Close(); + + ~DocIterator(); + + private: + // CollectionImpl builds the iterator (constructs Impl and fills its fields). + friend class CollectionImpl; + + 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..9fb6eff61 --- /dev/null +++ b/tests/db/iterator_test.cc @@ -0,0 +1,657 @@ +// 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(); +} + +// Test 12: Read-only collection iteration — reopen an existing collection in +// read-only mode and verify the full traversal works WITHOUT flushing (the +// read-only path reads the writing segment directly instead of flushing). +TEST_F(IteratorTest, ReadOnlyCollectionIteration) { + auto schema = TestHelper::CreateNormalSchema(); + const int N = 50; + + // 1. Create + insert + flush, then close in writable mode. + { + 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()); + 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(); + } // writable collection closed (dtor flushes + releases lock) + + // 2. Reopen the same collection in read-only mode. + CollectionOptions ro_options; + ro_options.read_only_ = true; + auto ro_result = Collection::Open(iter_test_path, ro_options); + ASSERT_TRUE(ro_result.has_value()) << ro_result.error().message(); + auto ro_collection = std::move(ro_result.value()); + + // 3. Iterate: must succeed (no disk write) and return every doc. + auto iter_result = ro_collection->CreateIterator(); + ASSERT_TRUE(iter_result.has_value()) << iter_result.error().message(); + auto iter = iter_result.value(); + + int count = 0; + std::set pks; + while (true) { + auto r = iter->Next(); + ASSERT_TRUE(r.has_value()) << r.error().message(); + if (r.value() == nullptr) break; + pks.insert(r.value()->pk()); + count++; + } + EXPECT_EQ(count, N); + EXPECT_EQ(pks.size(), N); + // No Destroy(): a read-only collection cannot be destroyed; TearDown cleans + // up. +} From ddfd640998a4292e51841d2d1d25003f24cdc159 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=90=AA=E9=9C=84?= Date: Wed, 29 Jul 2026 15:13:16 +0800 Subject: [PATCH 2/3] fix(iterator): address review feedback - iterate segment-by-segment: one FilteringReader per segment (drop ConcatenatingReader), so each batch's owning segment is known directly - hide CollectionImpl from the public doc_iterator.h (Pimpl fwd-decl + public ctor, no friend/Passkey in the public header) - propagate errors to callers instead of logging: Segment::scan failure, missing vector indexer, vector fetch failure, buffer conversion failure - fetch vectors via the segment-local row id column (LOCAL_ROW_ID) instead of g_doc_id - min_doc_id arithmetic (safe for compacted segments with non-contiguous doc ids) - use has_record() for the writable flush condition (aligns with #618) - drop stale comments and unused includes --- src/db/collection.cc | 106 ++++++++--------- src/db/doc_iterator.cc | 125 ++++++++++---------- src/db/doc_iterator_internal.h | 20 ++-- src/db/index/segment/concatenating_reader.h | 70 ----------- src/include/zvec/db/doc_iterator.h | 28 ++--- 5 files changed, 131 insertions(+), 218 deletions(-) delete mode 100644 src/db/index/segment/concatenating_reader.h diff --git a/src/db/collection.cc b/src/db/collection.cc index 982cd0766..10f985144 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -45,7 +45,6 @@ #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" @@ -175,17 +174,17 @@ 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 + // Columns for Segment::scan: system columns (+ LOCAL_ROW_ID when vectors + // are needed) plus the requested forward fields (all when nullopt). std::vector build_scan_columns( - const std::optional> &output_fields) const; + const IteratorOptions &options) 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); + // Isolated scan: Flush (writable) + snapshot + clone, then build one + // filtered reader per segment. segments/delete_store are filled for the + // caller to keep alive; the returned readers are parallel to `segments`. + Result> Scan( + const IteratorOptions &options, std::vector &segments, + DeleteStore::Ptr &delete_store, CollectionSchema::Ptr &schema); std::vector get_all_segments() const; @@ -2110,12 +2109,18 @@ std::vector CollectionImpl::get_all_persist_segments() const { } std::vector CollectionImpl::build_scan_columns( - const std::optional> &output_fields) const { - // Same logic as Segment::Fetch() + const IteratorOptions &options) const { std::vector columns; - columns.push_back(GLOBAL_DOC_ID); // _zvec_g_doc_id_ - columns.push_back(USER_ID); // _zvec_uid_ + columns.push_back(GLOBAL_DOC_ID); + columns.push_back(USER_ID); + if (options.include_vector_) { + // Segment-local row id, used by DocIterator to fetch vectors from the + // vector indexer (robust against non-contiguous g_doc_ids in compacted + // segments, unlike g_doc_id - min_doc_id arithmetic). + columns.push_back(LOCAL_ROW_ID); + } + const auto &output_fields = options.output_fields_; if (!output_fields.has_value()) { // nullopt = all forward fields for (const auto &field : schema_->forward_fields()) { @@ -2136,7 +2141,7 @@ std::vector CollectionImpl::build_scan_columns( return columns; } -Result CollectionImpl::Scan( +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) @@ -2150,80 +2155,69 @@ Result CollectionImpl::Scan( schema = schema_; { - // write_mtx_ blocks concurrent writes (Insert/Upsert/Update/Delete) - // auto-released when scope exits + // write_mtx_ blocks concurrent writes while taking the snapshot std::lock_guard write_lock(write_mtx_); if (options_.read_only_) { - // Read-only: cannot flush (that writes to disk, violating read-only). - // Read the writing segment directly instead — SegmentImpl::scan() also - // reads the in-memory writing block, and a read-only collection has no - // concurrent writes, so the snapshot is stable and complete. - segments = get_all_segments(); // persist segments + writing segment + // Read-only: flushing is not allowed, so read the writing segment + // directly (SegmentImpl::scan() also reads the in-memory writing block; + // no concurrent writes exist, so the snapshot is stable and complete). + segments = get_all_segments(); } else { - // Writable: flush the writing segment to seal it against concurrent - // writes, then snapshot the persist segments (now including it). - if (writing_segment_->doc_count() != 0) { + // Writable: seal the writing segment so concurrent writes cannot + // mutate the snapshot. has_record() also covers delete-only segments. + if (writing_segment_->has_record()) { auto s = switch_to_new_segment_for_writing(); CHECK_RETURN_STATUS_EXPECTED(s); } segments = get_all_persist_segments(); } - // Deep clone DeleteStore bitmap 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_); + auto scan_columns = build_scan_columns(options); + auto filter = delete_store->make_filter(); - std::vector> readers; + // Build one filtered reader per segment (NOT concatenated). DocIterator + // iterates segment-by-segment, so each batch's owning segment is known + // directly and vectors are fetched from it. readers[i] pairs segments[i]. + std::vector readers; + readers.reserve(segments.size()); for (const auto &seg : segments) { auto scalar_reader = seg->scan(scan_columns); - if (scalar_reader) { - readers.push_back(scalar_reader); - } else { - LOG_ERROR("Segment::scan returned no reader; skipping a segment"); + if (!scalar_reader) { + return tl::make_unexpected( + Status::InternalError("Segment::scan failed during collection scan")); } + readers.push_back(FilteringReader::Make(scalar_reader, filter)); } - // Vectors are not read here; DocIterator::Next() batch-prefetches them from - // each segment's vector indexer. - // 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; + return readers; } 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. + // Scan() takes shared_lock(schema_handle_mtx_) internally and captures the + // schema + one filtered reader per segment under that lock. 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()); + auto readers_result = Scan(options, segments, delete_store, schema); + if (!readers_result) { + return tl::make_unexpected(readers_result.error()); } - // Create DocIterator, keeping segments + delete_store alive + // Create DocIterator, keeping segments + delete_store alive. + // readers are parallel to segments (readers[i] scans segments[i]). auto impl = std::make_unique(); - impl->reader = std::move(reader_result.value()); + impl->readers = std::move(readers_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 std::make_shared(DocIterator::Passkey{}, std::move(impl)); + return std::make_shared(std::move(impl)); } } // namespace zvec diff --git a/src/db/doc_iterator.cc b/src/db/doc_iterator.cc index 10dae09cd..2277db3c4 100644 --- a/src/db/doc_iterator.cc +++ b/src/db/doc_iterator.cc @@ -13,7 +13,6 @@ // limitations under the License. #include -#include #include #include "db/common/constants.h" #include "db/doc_iterator_internal.h" @@ -119,8 +118,7 @@ static void SetNumericListField( // ── DocIterator implementation ── -DocIterator::DocIterator(Passkey, std::unique_ptr impl) - : impl_(std::move(impl)) {} +DocIterator::DocIterator(std::unique_ptr impl) : impl_(std::move(impl)) {} DocIterator::~DocIterator() { Close(); @@ -129,10 +127,10 @@ DocIterator::~DocIterator() { void DocIterator::Close() { if (impl_) { impl_->closed = true; - // Release the Arrow reader/batch first (they reference segment files), + // Release the Arrow readers/batch first (they reference segment files), // then release the kept-alive snapshot resources so a closed iterator // retains nothing. - impl_->reader.reset(); + impl_->readers.clear(); impl_->current_batch.reset(); impl_->vector_cache_.clear(); impl_->segments.clear(); @@ -146,68 +144,71 @@ Result DocIterator::Next() { return tl::make_unexpected(Status::InternalError("Iterator is closed")); } - // If current batch exhausted or not loaded, read next batch + // Load a new batch if the current one is exhausted, advancing across + // segments. Each segment has its own reader, so the batch's owning segment + // is always known directly via current_segment_index (no doc-id search). 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())); + bool loaded = false; + while (impl_->current_segment_index < impl_->readers.size()) { + auto &reader = impl_->readers[impl_->current_segment_index]; + auto status = reader->ReadNext(&impl_->current_batch); + if (!status.ok()) { + return tl::make_unexpected( + Status::InternalError("ReadNext failed: ", status.ToString())); + } + if (impl_->current_batch) { + loaded = true; + break; + } + // Current segment exhausted; move to the next one. + impl_->current_segment_index++; } - impl_->current_row = 0; - if (!impl_->current_batch) { - return Doc::Ptr(nullptr); // EOF + if (!loaded) { + return Doc::Ptr(nullptr); // EOF: all segments consumed } - // Pre-fetch vectors for the new batch (batch-level, not per-doc) + impl_->current_row = 0; + + // Pre-fetch vectors for the new batch from its owning segment. 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); - bool found_segment = false; - for (const auto &seg : impl_->segments) { - if (first_gdoc >= seg->meta()->min_doc_id() && - first_gdoc <= seg->meta()->max_doc_id()) { - found_segment = true; - 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; - } - } - if (!found_segment) { - LOG_ERROR("vector prefetch: no segment owns g_doc_id %llu", - static_cast(first_gdoc)); + // Segment-local row ids emitted by Segment::scan (requested in + // build_scan_columns). Unlike g_doc_id arithmetic, they stay valid for + // compacted segments whose g_doc_ids are non-contiguous. + int row_id_col = batch.schema()->GetFieldIndex(LOCAL_ROW_ID); + auto row_ids = row_id_col >= 0 + ? std::dynamic_pointer_cast( + batch.column(row_id_col)) + : nullptr; + if (!row_ids) { + return tl::make_unexpected(Status::InternalError( + "Iterator batch is missing the segment row-id column")); + } + const auto &seg = impl_->segments[impl_->current_segment_index]; + for (const auto &field : impl_->schema->vector_fields()) { + auto indexer = seg->get_combined_vector_indexer(field->name()); + if (!indexer) { + // A declared vector field must have an indexer; report the internal + // inconsistency instead of silently omitting the vector. + return tl::make_unexpected(Status::InternalError( + "vector indexer missing for field: ", field->name())); + } + std::vector bufs; + bufs.reserve(batch.num_rows()); + for (int64_t i = 0; i < batch.num_rows(); i++) { + uint32_t seg_doc_id = static_cast(row_ids->Value(i)); + auto fr = indexer->Fetch(seg_doc_id); + if (!fr.has_value()) { + // Propagate the failure instead of returning a doc with a + // silently missing vector. + return tl::make_unexpected(Status::InternalError( + "vector fetch failed, field: ", field->name(), ": ", + fr.error().message())); } + bufs.push_back(std::move(fr.value())); } + impl_->vector_cache_[field->name()] = std::move(bufs); } } } @@ -377,18 +378,16 @@ Result DocIterator::Next() { } } - // 4. Extract vector fields from pre-fetched cache + // 4. Extract vector fields from the 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()); + 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()); + return tl::make_unexpected(s); } } } diff --git a/src/db/doc_iterator_internal.h b/src/db/doc_iterator_internal.h index b9b387fda..926e6c689 100644 --- a/src/db/doc_iterator_internal.h +++ b/src/db/doc_iterator_internal.h @@ -16,7 +16,6 @@ // Shared by collection.cc and doc_iterator.cc #pragma once -#include #include #include #include @@ -32,25 +31,26 @@ 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. + // readers 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; + // Index of the segment currently being read (segments/readers are parallel: + // readers[i] was built from segments[i]). + size_t current_segment_index{0}; 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>> + // Pre-fetched vector data for the current batch. + // Key: field_name, Value: one VectorDataBuffer per row in current_batch. + std::unordered_map> vector_cache_; - RecordBatchReaderPtr reader; // destroyed before segments + // One reader per segment (parallel to `segments`); destroyed before segments + std::vector readers; }; } // namespace zvec diff --git a/src/db/index/segment/concatenating_reader.h b/src/db/index/segment/concatenating_reader.h deleted file mode 100644 index 708c3e211..000000000 --- a/src/db/index/segment/concatenating_reader.h +++ /dev/null @@ -1,70 +0,0 @@ -// 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/include/zvec/db/doc_iterator.h b/src/include/zvec/db/doc_iterator.h index d0aa06fd5..14345584a 100644 --- a/src/include/zvec/db/doc_iterator.h +++ b/src/include/zvec/db/doc_iterator.h @@ -19,27 +19,20 @@ namespace zvec { -class CollectionImpl; - class DocIterator { - // Pimpl: implementation holds Arrow types (RecordBatchReader, RecordBatch). - // Forward-declared here (private) so the public constructor can name it. + public: + // Pimpl: the implementation holds Arrow types (RecordBatchReader, + // RecordBatch). Only a forward declaration is exposed here; the definition + // lives in an internal header (src/db/doc_iterator_internal.h), so external + // code cannot construct or inspect it. struct Impl; - public: using Ptr = std::shared_ptr; - // Passkey idiom: Passkey's constructor is private and CollectionImpl is its - // friend, so only CollectionImpl can construct a DocIterator. This lets - // CollectionImpl use std::make_shared while keeping construction controlled. - struct Passkey { - private: - Passkey() {} - friend class CollectionImpl; - }; - - // Called by CollectionImpl::CreateIterator - DocIterator(Passkey, std::unique_ptr impl); + // Constructed only by the collection implementation (which can build an + // Impl via the internal header). Public so std::make_shared can be used; + // external code cannot call it because Impl is incomplete here. + explicit DocIterator(std::unique_ptr impl); // !has_value() → error // has_value() && value() == nullptr → EOF @@ -51,9 +44,6 @@ class DocIterator { ~DocIterator(); private: - // CollectionImpl builds the iterator (constructs Impl and fills its fields). - friend class CollectionImpl; - std::unique_ptr impl_; }; From 2785b751359bc943c2296e1aa7dc25d333e23893 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=90=AA=E9=9C=84?= Date: Thu, 30 Jul 2026 19:52:41 +0800 Subject: [PATCH 3/3] refactor(db): unify Arrow/vector-buffer to Doc field conversion Address review comments on shared conversion helpers: - Add db/index/common/doc_field_converter with two shared helpers: ConvertVectorDataBufferToDocField (dense + sparse vector buffers) and ConvertArrowRowToDocField (all 9 scalar + 9 array data types). - DocIterator, SegmentImpl::Fetch and sqlengine fill_doc_field now share the same row-level conversion, removing ~650 lines of duplicated type-switch code. - SegmentImpl::Fetch boxes its single-value scalars via arrow::MakeArrayFromScalar to reuse the row converter, keeping its lenient log-and-continue contract unchanged. - fill_doc_vector/fill_doc_sparse_vector stay in sqlengine: they decode a query-result-specific Arrow encoding with a single consumer. --- src/db/doc_iterator.cc | 239 +------------- src/db/index/common/doc_field_converter.cc | 230 ++++++++++++++ src/db/index/common/doc_field_converter.h | 43 +++ src/db/index/segment/segment.cc | 348 ++------------------- src/db/sqlengine/sqlengine_impl.cc | 117 ++----- 5 files changed, 325 insertions(+), 652 deletions(-) create mode 100644 src/db/index/common/doc_field_converter.cc create mode 100644 src/db/index/common/doc_field_converter.h diff --git a/src/db/doc_iterator.cc b/src/db/doc_iterator.cc index 2277db3c4..8da67b73d 100644 --- a/src/db/doc_iterator.cc +++ b/src/db/doc_iterator.cc @@ -16,106 +16,10 @@ #include #include "db/common/constants.h" #include "db/doc_iterator_internal.h" +#include "db/index/common/doc_field_converter.h" namespace zvec { -// ── VectorDataBuffer → Doc field conversion ── -// Same logic as SegmentImpl::ConvertVectorDataBufferToDocField, 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)) {} @@ -238,142 +142,18 @@ Result DocIterator::Next() { } } - // 3. Extract scalar fields from Arrow batch + // 3. Extract scalar/array fields from the Arrow batch via the shared + // row-level converter (same type coverage and null semantics as the + // SQL engine). 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 typed_array = - std::dynamic_pointer_cast(array); - if (typed_array) doc->set(field->name(), typed_array->Value(row)); - break; - } - case DataType::INT64: { - auto typed_array = - std::dynamic_pointer_cast(array); - if (typed_array) doc->set(field->name(), typed_array->Value(row)); - break; - } - case DataType::UINT32: { - auto typed_array = - std::dynamic_pointer_cast(array); - if (typed_array) doc->set(field->name(), typed_array->Value(row)); - break; - } - case DataType::UINT64: { - auto typed_array = - std::dynamic_pointer_cast(array); - if (typed_array) doc->set(field->name(), typed_array->Value(row)); - break; - } - case DataType::FLOAT: { - auto typed_array = - std::dynamic_pointer_cast(array); - if (typed_array) doc->set(field->name(), typed_array->Value(row)); - break; - } - case DataType::DOUBLE: { - auto typed_array = - std::dynamic_pointer_cast(array); - if (typed_array) doc->set(field->name(), typed_array->Value(row)); - break; - } - case DataType::BOOL: { - auto typed_array = - std::dynamic_pointer_cast(array); - if (typed_array) doc->set(field->name(), typed_array->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; + auto s = + ConvertArrowRowToDocField(batch.column(col), row, *field, doc.get()); + if (!s.ok()) { + return tl::make_unexpected(s); } } } @@ -385,7 +165,8 @@ Result DocIterator::Next() { if (it == impl_->vector_cache_.end()) continue; if (row >= static_cast(it->second.size())) continue; - auto s = SetVectorFieldFromBuffer(field, it->second[row], doc.get()); + auto s = + ConvertVectorDataBufferToDocField(field, it->second[row], doc.get()); if (!s.ok()) { return tl::make_unexpected(s); } diff --git a/src/db/index/common/doc_field_converter.cc b/src/db/index/common/doc_field_converter.cc new file mode 100644 index 000000000..1e0d26d53 --- /dev/null +++ b/src/db/index/common/doc_field_converter.cc @@ -0,0 +1,230 @@ +// 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 "doc_field_converter.h" +#include +#include + +namespace zvec { + +namespace { + +template +Status DenseVectorDataConverter( + const FieldSchema::Ptr &field, + const vector_column_params::DenseVectorBuffer &buffer, Doc *doc) { + const T *data_ptr = reinterpret_cast(buffer.data.data()); + size_t data_size = buffer.data.size() / sizeof(T); + std::vector vector_data(data_ptr, data_ptr + data_size); + doc->set(field->name(), vector_data); + return Status::OK(); +} + +template +Status SparseVectorDataConverter( + const FieldSchema::Ptr &field, + const vector_column_params::SparseVectorBuffer &buffer, Doc *doc) { + const IndexType *indices_ptr = + reinterpret_cast(buffer.indices.data()); + size_t indices_size = buffer.indices.size() / sizeof(IndexType); + std::vector indices_vector(indices_ptr, + indices_ptr + indices_size); + + const ValueType *values_ptr = + reinterpret_cast(buffer.values.data()); + size_t values_size = buffer.values.size() / sizeof(ValueType); + std::vector values_vector(values_ptr, values_ptr + values_size); + + std::pair, std::vector> sparse_vector_pair( + std::move(indices_vector), std::move(values_vector)); + doc->set(field->name(), sparse_vector_pair); + return Status::OK(); +} + +//! Set a scalar Doc field from row `row` of a typed Arrow array. +template +Status SetScalarField(const std::shared_ptr &array, int64_t row, + const std::string &name, Doc *doc) { + auto typed_array = std::dynamic_pointer_cast(array); + if (!typed_array) { + return Status::InternalError("Arrow array type mismatch for field: ", name); + } + if constexpr (std::is_same_v || + std::is_same_v) { + doc->set(name, typed_array->GetString(row)); + } else { + doc->set(name, typed_array->Value(row)); + } + return Status::OK(); +} + +//! Set an array Doc field from row `row` of an Arrow ListArray whose values +//! are of type ArrowArrayT. Null elements inside the list are skipped. +template +Status SetListField(const std::shared_ptr &array, int64_t row, + const std::string &name, Doc *doc) { + auto list_array = std::dynamic_pointer_cast(array); + if (!list_array) { + return Status::InternalError("Arrow list type mismatch for field: ", name); + } + auto values = + std::dynamic_pointer_cast(list_array->value_slice(row)); + if (!values) { + return Status::InternalError("Arrow list value type mismatch for field: ", + name); + } + std::vector vec; + vec.reserve(values->length()); + for (int64_t i = 0; i < values->length(); ++i) { + if (!values->IsValid(i)) { + continue; + } + if constexpr (std::is_same_v || + std::is_same_v) { + vec.push_back(values->GetString(i)); + } else { + vec.push_back(values->Value(i)); + } + } + doc->set(name, vec); + return Status::OK(); +} + +} // namespace + +Status ConvertVectorDataBufferToDocField( + const FieldSchema::Ptr &field, + const vector_column_params::VectorDataBuffer &buf, Doc *doc) { + Status status; + if (std::holds_alternative( + buf.vector_buffer)) { + const auto &dense_buffer = + std::get(buf.vector_buffer); + switch (field->data_type()) { + case DataType::VECTOR_BINARY32: { + status = DenseVectorDataConverter(field, dense_buffer, doc); + break; + } + case DataType::VECTOR_BINARY64: { + status = DenseVectorDataConverter(field, dense_buffer, doc); + break; + } + case DataType::VECTOR_FP16: { + status = DenseVectorDataConverter(field, dense_buffer, doc); + break; + } + case DataType::VECTOR_FP32: { + status = DenseVectorDataConverter(field, dense_buffer, doc); + break; + } + case DataType::VECTOR_FP64: { + status = DenseVectorDataConverter(field, dense_buffer, doc); + break; + } + case DataType::VECTOR_INT8: { + status = DenseVectorDataConverter(field, dense_buffer, doc); + break; + } + case DataType::VECTOR_INT16: { + status = DenseVectorDataConverter(field, dense_buffer, doc); + break; + } + default: + return Status::InvalidArgument( + "Unsupported dense vector element type: ", field->data_type()); + } + } else if (std::holds_alternative( + buf.vector_buffer)) { + const auto &sparse_buffer = + std::get(buf.vector_buffer); + switch (field->data_type()) { + case DataType::SPARSE_VECTOR_FP16: { + status = SparseVectorDataConverter( + field, sparse_buffer, doc); + break; + } + case DataType::SPARSE_VECTOR_FP32: { + status = SparseVectorDataConverter(field, + sparse_buffer, doc); + break; + } + default: + return Status::InvalidArgument( + "Unsupported sparse vector element type: ", field->data_type()); + } + } else { + return Status::InvalidArgument("Unsupported vector buffer type"); + } + + return status; +} + +Status ConvertArrowRowToDocField(const std::shared_ptr &array, + int64_t row, const FieldSchema &field, + Doc *doc) { + if (!array || row < 0 || row >= array->length()) { + return Status::InvalidArgument("Arrow row out of range for field: ", + field.name()); + } + if (array->IsNull(row)) { + return Status::OK(); // null value: leave the field unset + } + + const auto &name = field.name(); + switch (field.data_type()) { + case DataType::BINARY: + return SetScalarField(array, row, name, doc); + case DataType::STRING: + return SetScalarField(array, row, name, doc); + case DataType::BOOL: + return SetScalarField(array, row, name, doc); + case DataType::INT32: + return SetScalarField(array, row, name, doc); + case DataType::INT64: + return SetScalarField(array, row, name, doc); + case DataType::UINT32: + return SetScalarField(array, row, name, doc); + case DataType::UINT64: + return SetScalarField(array, row, name, doc); + case DataType::FLOAT: + return SetScalarField(array, row, name, doc); + case DataType::DOUBLE: + return SetScalarField(array, row, name, doc); + case DataType::ARRAY_BINARY: + return SetListField(array, row, name, + doc); + case DataType::ARRAY_STRING: + return SetListField(array, row, name, + doc); + case DataType::ARRAY_BOOL: + return SetListField(array, row, name, doc); + case DataType::ARRAY_INT32: + return SetListField(array, row, name, doc); + case DataType::ARRAY_INT64: + return SetListField(array, row, name, doc); + case DataType::ARRAY_UINT32: + return SetListField(array, row, name, doc); + case DataType::ARRAY_UINT64: + return SetListField(array, row, name, doc); + case DataType::ARRAY_FLOAT: + return SetListField(array, row, name, doc); + case DataType::ARRAY_DOUBLE: + return SetListField(array, row, name, doc); + default: + return Status::InvalidArgument("Unsupported data type for field: ", name, + ": ", field.data_type()); + } +} + +} // namespace zvec diff --git a/src/db/index/common/doc_field_converter.h b/src/db/index/common/doc_field_converter.h new file mode 100644 index 000000000..7e3f68100 --- /dev/null +++ b/src/db/index/common/doc_field_converter.h @@ -0,0 +1,43 @@ +// 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. + +// Shared row-level converters that fill Doc fields, used by SegmentImpl, +// DocIterator and the SQL engine so that field-type coverage and null/list +// semantics stay in one place. +#pragma once + +#include +#include +#include +#include +#include +#include "db/index/column/vector_column/vector_column_params.h" + +namespace zvec { + +//! Convert a VectorDataBuffer (dense or sparse) fetched from a vector indexer +//! into the corresponding Doc vector field. +Status ConvertVectorDataBufferToDocField( + const FieldSchema::Ptr &field, + const vector_column_params::VectorDataBuffer &buf, Doc *doc); + +//! Convert row `row` of an Arrow array into the corresponding Doc scalar or +//! array field. Null rows are skipped (the field is left unset). Covers every +//! scalar DataType (BINARY/STRING/BOOL/INT32/INT64/UINT32/UINT64/FLOAT/DOUBLE) +//! and every ARRAY_* DataType. +Status ConvertArrowRowToDocField(const std::shared_ptr &array, + int64_t row, const FieldSchema &field, + Doc *doc); + +} // namespace zvec diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index 984d73621..e1eb80230 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -50,6 +50,7 @@ #include "db/index/column/inverted_column/inverted_indexer.h" #include "db/index/column/vector_column/vector_column_indexer.h" #include "db/index/column/vector_column/vector_column_params.h" +#include "db/index/common/doc_field_converter.h" #include "db/index/common/index_filter.h" #include "db/index/common/meta.h" #include "db/index/segment/segment_helper.h" @@ -294,9 +295,6 @@ class SegmentImpl : public Segment, template Status InsertVector(VectorColumnIndexer::Ptr &indexer, const Doc &doc, const FieldSchema::Ptr &field); - Status ConvertVectorDataBufferToDocField( - const FieldSchema::Ptr &field, - const vector_column_params::VectorDataBuffer &buf, Doc *doc); Status insert_scalar_indexer(Doc &doc); Status insert_fts_indexer(Doc &doc); @@ -994,110 +992,6 @@ Status SegmentImpl::Delete(uint64_t g_doc_id) { return internal_delete(mutable_doc); } -template -Status DenseVectorDataConverter( - const FieldSchema::Ptr &field, - const vector_column_params::DenseVectorBuffer &buffer, Doc *doc) { - const T *data_ptr = reinterpret_cast(buffer.data.data()); - size_t data_size = buffer.data.size() / sizeof(T); - std::vector vector_data(data_ptr, data_ptr + data_size); - doc->set(field->name(), vector_data); - return Status::OK(); -} - -template -Status SparseVectorDataConverter( - const FieldSchema::Ptr &field, - const vector_column_params::SparseVectorBuffer &buffer, Doc *doc) { - const IndexType *indices_ptr = - reinterpret_cast(buffer.indices.data()); - size_t indices_size = buffer.indices.size() / sizeof(IndexType); - std::vector indices_vector(indices_ptr, - indices_ptr + indices_size); - - const ValueType *values_ptr = - reinterpret_cast(buffer.values.data()); - size_t values_size = buffer.values.size() / sizeof(ValueType); - std::vector values_vector(values_ptr, values_ptr + values_size); - - std::pair, std::vector> sparse_vector_pair( - std::move(indices_vector), std::move(values_vector)); - doc->set(field->name(), sparse_vector_pair); - return Status::OK(); -} - - -Status SegmentImpl::ConvertVectorDataBufferToDocField( - const FieldSchema::Ptr &field, - const vector_column_params::VectorDataBuffer &buf, Doc *doc) { - Status status; - if (std::holds_alternative( - buf.vector_buffer)) { - const auto &dense_buffer = - std::get(buf.vector_buffer); - switch (field->data_type()) { - case DataType::VECTOR_BINARY32: { - status = DenseVectorDataConverter(field, dense_buffer, doc); - break; - } - case DataType::VECTOR_BINARY64: { - status = DenseVectorDataConverter(field, dense_buffer, doc); - break; - } - case DataType::VECTOR_FP16: { - status = DenseVectorDataConverter(field, dense_buffer, doc); - break; - } - case DataType::VECTOR_FP32: { - status = DenseVectorDataConverter(field, dense_buffer, doc); - break; - } - case DataType::VECTOR_FP64: { - status = DenseVectorDataConverter(field, dense_buffer, doc); - break; - } - // case DataType::VECTOR_INT4: { - // status = DenseVectorDataConverter(field, dense_buffer, doc); - // break; - // } - case DataType::VECTOR_INT8: { - status = DenseVectorDataConverter(field, dense_buffer, doc); - break; - } - case DataType::VECTOR_INT16: { - status = DenseVectorDataConverter(field, dense_buffer, doc); - break; - } - default: - return Status::InvalidArgument( - "Unsupported dense vector element type: ", field->data_type()); - } - } else if (std::holds_alternative( - buf.vector_buffer)) { - const auto &sparse_buffer = - std::get(buf.vector_buffer); - switch (field->data_type()) { - case DataType::SPARSE_VECTOR_FP16: { - status = SparseVectorDataConverter( - field, sparse_buffer, doc); - break; - } - case DataType::SPARSE_VECTOR_FP32: { - status = SparseVectorDataConverter(field, - sparse_buffer, doc); - break; - } - default: - return Status::InvalidArgument( - "Unsupported sparse vector element type: ", field->data_type()); - } - } else { - return Status::InvalidArgument("Unsupported vector buffer type"); - } - - return status; -} - Doc::Ptr SegmentImpl::Fetch( uint64_t g_doc_id, @@ -1144,7 +1038,7 @@ Doc::Ptr SegmentImpl::Fetch( } } - // Build result schema + // Validate that every forward column has a convertible Arrow field std::vector> fields; for (size_t i = 0; i < forward_columns.size(); ++i) { const auto &col = forward_columns[i]; @@ -1164,7 +1058,6 @@ Doc::Ptr SegmentImpl::Fetch( fields.push_back(std::move(arrow_field)); } } - auto result_schema = std::make_shared(fields); // fetch forward columns auto exec_batch = fetch(forward_columns, segment_doc_id); @@ -1207,227 +1100,30 @@ Doc::Ptr SegmentImpl::Fetch( // other forward columns for (int col_idx = 2; col_idx < exec_batch->num_values(); ++col_idx) { auto column_name = forward_columns[col_idx]; - auto column = result_schema->GetFieldByName(column_name); auto &column_scalar = (*exec_batch)[col_idx].scalar(); if (column_scalar == nullptr || column_scalar->is_valid == false) { continue; } - switch (column->type()->id()) { - case arrow::Type::STRING: { - auto str_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, std::string(str_scalar->view())); - break; - } - case arrow::Type::INT32: { - auto int32_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, int32_scalar->value); - break; - } - case arrow::Type::INT64: { - auto int64_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, int64_scalar->value); - break; - } - case arrow::Type::UINT32: { - auto uint32_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, uint32_scalar->value); - break; - } - case arrow::Type::UINT64: { - auto uint64_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, uint64_scalar->value); - break; - } - case arrow::Type::DOUBLE: { - auto double_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, double_scalar->value); - break; - } - case arrow::Type::FLOAT: { - auto float_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, float_scalar->value); - break; - } - case arrow::Type::BOOL: { - auto bool_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, bool_scalar->value); - break; - } - case arrow::Type::BINARY: { - auto binary_scalar = - std::dynamic_pointer_cast(column_scalar); - doc->set(column_name, std::string(binary_scalar->view())); - break; - } - case arrow::Type::LIST: { - auto list_scalar = - std::dynamic_pointer_cast(column_scalar); - if (list_scalar && list_scalar->value) { - auto list_type = - std::dynamic_pointer_cast(column->type()); - if (list_type) { - auto value_type = list_type->value_type(); - switch (value_type->id()) { - case arrow::Type::BOOL: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->Value(i)); - } else { - LOG_ERROR( - "Invalid arrow::boolean array value at index %zu", - (size_t)i); - continue; - } - } - doc->set(column_name, values); - } - break; - } - case arrow::Type::INT32: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->Value(i)); - } - } - doc->set(column_name, values); - } - break; - } - case arrow::Type::INT64: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->Value(i)); - } - } - doc->set(column_name, values); - } - break; - } - case arrow::Type::UINT32: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->Value(i)); - } - } - doc->set(column_name, values); - } - break; - } - case arrow::Type::UINT64: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->Value(i)); - } - } - doc->set(column_name, values); - } - break; - } - case arrow::Type::FLOAT: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->Value(i)); - } - } - doc->set(column_name, values); - } - break; - } - case arrow::Type::DOUBLE: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->Value(i)); - } - } - doc->set(column_name, values); - } - break; - } - case arrow::Type::STRING: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->GetString(i)); - } - } - doc->set(column_name, values); - } - break; - } - case arrow::Type::BINARY: { - std::vector values; - auto array = std::dynamic_pointer_cast( - list_scalar->value); - if (array) { - values.reserve(array->length()); - for (int64_t i = 0; i < array->length(); ++i) { - if (array->IsValid(i)) { - values.push_back(array->GetString(i)); - } - } - doc->set(column_name, values); - } - break; - } - default: - LOG_WARN("Unsupported list element type: %s", - value_type->ToString().c_str()); - break; - } - } - } - break; - } - default: - LOG_ERROR("Unsupported type: %s", column_name.c_str()); - break; + auto *field = collection_schema_->get_field(column_name); + if (field == nullptr) { + LOG_ERROR("Field not found in schema: %s", column_name.c_str()); + continue; + } + // Box the single-value scalar as a one-row array and reuse the shared + // row-level converter, so the type coverage lives in one place + // (doc_field_converter) together with DocIterator and the SQL engine. + auto array_result = arrow::MakeArrayFromScalar(*column_scalar, 1); + if (!array_result.ok()) { + LOG_ERROR("Box scalar failed for column %s: %s", column_name.c_str(), + array_result.status().ToString().c_str()); + continue; + } + auto s = ConvertArrowRowToDocField(array_result.ValueOrDie(), 0, *field, + doc.get()); + if (!s.ok()) { + // Keep Fetch's lenient contract: log and continue with other fields. + LOG_ERROR("Convert field %s failed: %s", column_name.c_str(), + s.message().c_str()); } } diff --git a/src/db/sqlengine/sqlengine_impl.cc b/src/db/sqlengine/sqlengine_impl.cc index 8a84bb51d..f32069d05 100644 --- a/src/db/sqlengine/sqlengine_impl.cc +++ b/src/db/sqlengine/sqlengine_impl.cc @@ -23,6 +23,7 @@ #include "db/index/column/fts_column/fts_ast_rewriter.h" #include "db/index/column/fts_column/fts_pipeline.h" #include "db/index/column/fts_column/fts_query_ast.h" +#include "db/index/common/doc_field_converter.h" #include "db/sqlengine/analyzer/query_analyzer.h" #include "db/sqlengine/parser/select_info.h" #include "db/sqlengine/parser/sql_info_helper.h" @@ -426,117 +427,39 @@ Status fill_doc_vector(const arrow::BinaryArray *typed_arr, return Status::OK(); } -template -Status fill_doc_field(const arrow::Array *arr, const std::string &field_name, - DocPtrList::iterator doc_it) { - auto *typed_arr = static_cast(arr); - bool no_null = typed_arr->null_count() == 0; - for (int64_t i = 0; i < typed_arr->length(); ++i, ++doc_it) { - if (no_null || !typed_arr->IsNull(i)) { - if constexpr (std::is_same_v || - std::is_same_v || - std::is_same_v || - std::is_same_v) { - (*doc_it)->set(field_name, typed_arr->GetString(i)); - } else { - (*doc_it)->set(field_name, typed_arr->Value(i)); - } - } - } - return Status::OK(); -} - -template -Status fill_doc_array_field(const arrow::Array *arr, - const std::string &field_name, - DocPtrList::iterator doc_it) { - const auto *list_arr = static_cast(arr); - auto *typed_arr = - dynamic_cast(list_arr->values().get()); - bool has_null = list_arr->null_count() > 0; - for (int64_t i = 0; i < list_arr->length(); ++i, ++doc_it) { - if (has_null && list_arr->IsNull(i)) { - continue; - } - int64_t offset = list_arr->value_offset(i); - int64_t length = list_arr->value_length(i); - std::vector vec(length); - for (int64_t j = 0; j < length; ++j) { - vec[j] = typed_arr->Value(offset + j); - } - (*doc_it)->set(field_name, std::move(vec)); - } - return Status::OK(); -} - Status fill_doc_field(const std::shared_ptr &chunk, const FieldSchema &field_schema, DocPtrList::iterator doc_it) { switch (field_schema.data_type()) { + case DataType::BINARY: + case DataType::STRING: + case DataType::BOOL: case DataType::INT32: - return fill_doc_field(chunk.get(), field_schema.name(), - doc_it); - case DataType::UINT32: - return fill_doc_field(chunk.get(), - field_schema.name(), doc_it); case DataType::INT64: - return fill_doc_field(chunk.get(), field_schema.name(), - doc_it); + case DataType::UINT32: case DataType::UINT64: - return fill_doc_field(chunk.get(), - field_schema.name(), doc_it); case DataType::FLOAT: - return fill_doc_field(chunk.get(), field_schema.name(), - doc_it); case DataType::DOUBLE: - return fill_doc_field(chunk.get(), - field_schema.name(), doc_it); - case DataType::BOOL: - return fill_doc_field(chunk.get(), - field_schema.name(), doc_it); - case DataType::BINARY: - return fill_doc_field(chunk.get(), - field_schema.name(), doc_it); - - case DataType::STRING: - return fill_doc_field(chunk.get(), - field_schema.name(), doc_it); - + case DataType::ARRAY_BINARY: + case DataType::ARRAY_STRING: + case DataType::ARRAY_BOOL: case DataType::ARRAY_INT32: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); - case DataType::ARRAY_INT64: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); - case DataType::ARRAY_UINT32: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); - case DataType::ARRAY_UINT64: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); - case DataType::ARRAY_FLOAT: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); - - case DataType::ARRAY_DOUBLE: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); - - case DataType::ARRAY_STRING: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); - - case DataType::ARRAY_BINARY: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); - - case DataType::ARRAY_BOOL: - return fill_doc_array_field( - chunk.get(), field_schema.name(), doc_it); + case DataType::ARRAY_DOUBLE: { + // Scalar/array fields: shared row-level conversion (same type coverage + // and null semantics as DocIterator; see doc_field_converter.h). + for (int64_t i = 0; i < chunk->length(); ++i, ++doc_it) { + auto status = + ConvertArrowRowToDocField(chunk, i, field_schema, doc_it->get()); + if (!status.ok()) { + return status; + } + } + return Status::OK(); + } case DataType::VECTOR_FP32: return fill_doc_vector((arrow::BinaryArray *)chunk.get(),