feat(iterator): add DocIterator for full collection traversal - #597
feat(iterator): add DocIterator for full collection traversal#597YongqiYin wants to merge 19 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces a full-collection document iterator (“DocIterator”) across the C++ core, C API, and Python bindings to enable streaming traversal/export without relying on large topk queries (relates to #380).
Changes:
- Add C++
Collection::CreateIterator()andDocIteratorwith snapshot isolation, segment concatenation, and delete filtering. - Expose the iterator via the C API (
zvec_collection_create_iterator/next/close) and Python (Collection.iter_docs()generator). - Add C++/C/Python tests covering basic iteration, delete filtering, field selection, and concurrency/isolation.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/db/iterator_test.cc | New C++ unit/integration/concurrency/perf tests for full traversal. |
| tests/c/c_api_test.c | Adds C API iterator tests and option handling coverage. |
| src/include/zvec/db/options.h | Introduces IteratorOptions (output_fields/include_vector). |
| src/include/zvec/db/doc_iterator.h | Public C++ DocIterator interface. |
| src/include/zvec/db/collection.h | Adds Collection::CreateIterator() API. |
| src/include/zvec/c_api.h | Adds public C iterator and iterator-options API. |
| src/db/index/segment/filtering_reader.h | New Arrow reader wrapper to filter deleted docs. |
| src/db/index/segment/concatenating_reader.h | New Arrow reader to concatenate readers across segments. |
| src/db/doc_iterator.cc | Implements row-by-row Doc materialization + optional vector prefetch. |
| src/db/doc_iterator_internal.h | Internal DocIterator::Impl definition (lifetime ordering, caches). |
| src/db/collection.cc | Implements iterator creation, snapshot scan, and reader chain construction. |
| src/binding/python/model/python_collection.cc | Exposes _DocIterator + Collection.CreateIterator() to Python. |
| src/binding/python/include/python_collection.h | Declares bind_iterator() hook. |
| src/binding/c/c_api.cc | Implements iterator options + iterator handles for the C API. |
| python/zvec/model/collection.py | Adds Collection.iter_docs() streaming generator. |
| python/tests/test_iter_docs.py | New Python tests for iterator behavior and isolation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
b596564 to
4aca7cb
Compare
|
@copilot resolve the merge conflicts in this pull request |
4aca7cb to
feef1d7
Compare
96273fc to
4f542e0
Compare
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 alibaba#380
4f542e0 to
16c0bd2
Compare
- 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 alibaba#618) - drop stale comments and unused includes
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.
| template <typename ArrowArrayT, typename T> | ||
| Status SetListField(const std::shared_ptr<arrow::Array> &array, int64_t row, | ||
| const std::string &name, Doc *doc) { | ||
| auto list_array = std::dynamic_pointer_cast<arrow::ListArray>(array); |
| impl_->vector_cache_.clear(); | ||
| impl_->segments.clear(); | ||
| impl_->delete_store.reset(); | ||
| impl_->schema.reset(); |
There was a problem hiding this comment.
可以直接reset impl_吗?简单一点
如果不行的话,最后也reset一下impl_吧,当前虽然没问题,但万一后面逻辑有修改引入问题
There was a problem hiding this comment.
已在最后impl_.reset()
没有直接impl_.reset()是将现在的析构顺序固化,相对更安全
若进一步保证未来安全可以加一个单测,我觉得目前不需要做到这程度
There was a problem hiding this comment.
没有直接impl_.reset()是将现在的析构顺序固化,相对更安全
最好是让Impl本身的析构安全,然后这里直接调用.reset
There was a problem hiding this comment.
Impl 的成员声明顺序本身已保证析构安全,已简化为直接 impl.reset()
Extend the reviewer's shared_ptr-reduction feedback to the remaining hot paths: - ConvertArrowRowToDocField now takes a raw Array* (callers pass batch.columns()[i].get() / chunk.get()), removing per-row atomic ref-count traffic. - DocIterator PK/doc_id/row-id extraction uses type_id() + static_cast instead of per-row dynamic_pointer_cast. - Update the Scan comment after alibaba#614: snapshot consistency comes from the write_mtx_ atomic snapshot + shared_ptr keep-alive, not from blocking Optimize (which now takes the schema lock in shared mode).
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.
Suppressed comments (12)
tests/db/iterator_test.cc:183
- The iterator is still alive when
collection->Destroy()is called, which can leave Arrow/segment resources open during destruction and cause flaky cleanup. Close the iterator before destroying the collection.
EXPECT_EQ(count, N - static_cast<int>(pks_to_delete.size()));
collection->Destroy();
tests/db/iterator_test.cc:271
- The iterator is still alive when
collection->Destroy()is called, which can keep segment/Arrow resources open during destruction. Close the iterator before destroying the collection to avoid flaky cleanup.
EXPECT_EQ(count, N);
collection->Destroy();
tests/db/iterator_test.cc:317
- The iterator is still alive when
collection->Destroy()is called, which can keep segment/Arrow resources open during destruction. Close the iterator before destroying the collection to avoid flaky cleanup.
EXPECT_EQ(count, 5);
collection->Destroy();
tests/db/iterator_test.cc:391
- The iterator is not closed before
collection->Destroy(). Keeping the iterator alive can keep file handles open and makeDestroy()flaky on some platforms. Close the iterator before destroying the collection.
EXPECT_EQ((*a_s)[0], "value_" + std::to_string(kId));
collection->Destroy();
tests/db/iterator_test.cc:447
- The iterator is not closed before
collection->Destroy(). If the iterator still holds segment/Arrow resources, destroying the collection can be flaky (especially on Windows). Close the iterator before destroying the collection.
EXPECT_EQ(count, N);
EXPECT_EQ(seen_pks.size(), (size_t)N);
collection->Destroy();
tests/db/iterator_test.cc:509
- Both iterators (
iteranditer2) are still alive whencollection->Destroy()is called. Close them before destroying the collection to ensure underlying segment/Arrow resources are released and cleanup is reliable.
EXPECT_EQ(count2, N + 200);
collection->Destroy();
tests/db/iterator_test.cc:556
- The iterator is still alive when
collection->Destroy()is called, which can keep resources open during destruction and cause flaky cleanup. Close the iterator before destroying the collection.
EXPECT_EQ(count, N);
collection->Destroy();
}
tests/db/iterator_test.cc:607
- The iterator is still alive when
collection->Destroy()is called. Close it before destroying the collection to ensure segment/Arrow resources are released and test cleanup is reliable.
<< std::endl;
collection->Destroy();
}
tests/db/iterator_test.cc:16
std::coutis used later in this test file (Performance100k), but<iostream>is not included. This can break compilation depending on transitive includes.
#include <chrono>
tests/db/iterator_test.cc:126
- The iterator is still alive when
collection->Destroy()is called. On platforms with strict file-handle semantics (and per DocIterator's internal comment about releasing Arrow handles before segment cleanup), this can makeDestroy()fail or become flaky. Close the iterator (or let it go out of scope) before destroying the collection.
This issue also appears in the following locations of the same file:
- line 181
- line 270
- line 316
- line 389
- line 445
- ...and 3 more
EXPECT_EQ(r.value(), nullptr) << "Expected EOF on empty collection";
collection->Destroy();
src/binding/c/c_api.cc:7333
zvec_collection_create_iteratoralways maps C++ iterator creation failures toZVEC_ERROR_INTERNAL_ERROR, even when the underlyingStatusisINVALID_ARGUMENT(e.g., unknown/duplicate output fields). This loses actionable error information for C callers; use the existingstatus_to_error_code()helper for consistent mapping.
SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR,
"Failed to create iterator: " +
result.error().message());
return ZVEC_ERROR_INTERNAL_ERROR;
src/binding/c/c_api.cc:7361
zvec_doc_iterator_nextconverts all iterator failures intoZVEC_ERROR_INTERNAL_ERROR, which hides the underlyingStatusCodeand is inconsistent with other C API wrappers that usestatus_to_error_code(). Map the error code fromresult.error()so callers can distinguish invalid arguments vs internal failures.
if (!result.has_value()) {
SET_LAST_ERROR(ZVEC_ERROR_INTERNAL_ERROR,
"Iterator next failed: " + result.error().message());
return ZVEC_ERROR_INTERNAL_ERROR;
}
…ue patterns - Gate per-element IsNull checks behind an O(1) null_count() precheck (the common null-free case skips validity checks entirely, matching the optimization the reviewer asked to preserve). - Unify element extraction on Value() for all list element types (string/binary Value returns string_view, emplace_back converts it to std::string in place), removing the if-constexpr type branch.
Regression test for the deferred-snapshot issue: creating the iterator without consuming it must already freeze the snapshot, so documents inserted between iter_docs() and the first next() stay invisible. Verified red on the lazy-generator version, green on the fix.
- Fetch reads Scalars directly again and drops MakeArrayFromScalar boxing; its LIST branch now shares value extraction with the iterator/SQL-engine path via ExtractTypedArrayValues . - Open one segment reader lazily and release it when the segment is exhausted, so at most one segment's files stay open. - Prefetch vectors in bounded windows (kIteratorVectorPrefetchWindow) so a large Parquet row group cannot cache a million vectors. - Skip empty batches defensively; fail instead of silently emitting docs when uid/g_doc_id columns are missing; cache column indices per batch to avoid per-row GetFieldIndex lookups.
…pshot side effects - output_fields accepts only forward fields and rejects duplicates; validation runs before the writing segment is sealed so invalid options fail fast with no side effect. - Map CreateIterator/Next failures through status_to_error_code in the C API so callers see INVALID_ARGUMENT instead of INTERNAL_ERROR. - Document the seal side effect and the close-before-destroy contract in the C++/C/Python API docs.
- Add OutputFieldsSelection and InvalidOutputFieldsRejected. - Add ParquetVectorPrefetchWindows: a row group larger than the prefetch window verifies window-refill alignment. - Close iterators before Destroy() and add the missing <iostream> include.
| const auto &requested = *output_fields; | ||
| std::unordered_set<std::string> seen; | ||
| for (const auto &name : requested) { | ||
| if (schema_->get_forward_field(name) == nullptr) { |
There was a problem hiding this comment.
可以直接把get_forward_field的结果放到columns,后面不用再遍历检查一次了
There was a problem hiding this comment.
已采纳,校验通过即直接放入 columns
| impl_->vector_cache_.clear(); | ||
| impl_->segments.clear(); | ||
| impl_->delete_store.reset(); | ||
| impl_->schema.reset(); |
There was a problem hiding this comment.
没有直接impl_.reset()是将现在的析构顺序固化,相对更安全
最好是让Impl本身的析构安全,然后这里直接调用.reset
| 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) { |
There was a problem hiding this comment.
已采纳。
新增了列级 ConvertArrowColumnToDocFields,sqlengine中改为单次列级调用;行级接口保留给 DocIterator(逐行消费);底层 SetScalarValue 一份实现两级共享
| } | ||
|
|
||
| arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch> *batch) override { | ||
| while (true) { |
There was a problem hiding this comment.
已采纳,放到了filtering_reader.cc
… reviews - Add column-level ConvertArrowColumnToDocFields (one type dispatch and one null_count() check per column) and use it in fill_doc_field; the row-level converter stays for DocIterator. - Drop the redundant per-row null re-check in SetScalarField and the element-level null checks in ExtractTypedArrayValues (zvec array fields never store null elements); null_count() is only used as the row-level fast path. - Move FilteringReader::ReadNext into filtering_reader.cc and keep the header declaration-only. - Simplify Close() to a single impl_.reset() (Impl's member order guarantees teardown safety) and drop the redundant closed flag. - Push validated output_fields into columns in one pass and trim stale comments.
| } | ||
| { | ||
| const auto &col = batch.columns()[impl_->uid_col]; | ||
| if (col->type_id() != arrow::Type::STRING) { |
There was a problem hiding this comment.
不需要每个batch做一次检查,在创建record batch reader的时候做下检查就够了,然后可以保存下相关的index,后续每个batch直接使用
There was a problem hiding this comment.
已改,检查移到每个 segment reader 创建时执行一次,索引保存在了 DocIterator::Impl 的成员里
| } | ||
| } | ||
|
|
||
| impl_->current_row++; |
There was a problem hiding this comment.
这里还是逐行处理每个列的逻辑,效率比较低,建议改为按列处理每个batch
There was a problem hiding this comment.
整体复用ConvertArrowColumnToDocFields,而非只复用最里面的SetScalarValue
There was a problem hiding this comment.
已改为按列批量处理且标量/数组列整体复用 ConvertArrowColumnToDocFields
0d6ab14 to
9d9a66f
Compare
Replace the reopen-and-swap concurrency solution (reverted) with the simple model suggested in review: the iterator takes the schema lock (shared) at creation and holds it until Close(), so the schema and the segment set are frozen for its whole lifetime. - Operations needing the exclusive lock (create/drop index, add/alter/ drop column, Optimize, Flush, Close, Destroy) fail fast with PermissionDenied while any iterator is open, via an active-iterator count; this also prevents a single-threaded caller from deadlocking against its own iterator. - The iterator keeps the collection alive (Impl holds a shared_ptr<Collection>; CollectionImpl now enable_shared_from_this), so dropping the collection handle during iteration is safe and "next after collection close" can no longer happen. - Stats/Schema/Options switch from the exclusive lock to shared so they remain callable while iterators are open. - Concurrent writes and queries are unaffected; snapshot isolation is unchanged. Iterator vs Optimize concurrency is left as future work. Tests: OptimizeRejectedWhileIteratorOpen, DdlRejectedWhileIteratorOpen, CloseRejectedWhileIteratorOpen (drop-handle safety + reopen), MultipleIteratorsShareTheLock; python tests updated to the new flush contract. Full regression: iterator 18/18, collection 85/85, c_api 73/73, forward_recall 35/35, segment_helper 13/13.
9d9a66f to
ec45b71
Compare
Address review comments: resolve and validate column indices once per segment reader instead of per batch; materialize each batch into docs column by column via the shared ConvertArrowColumnToDocFields (type dispatch and null_count once per column), with vectors fetched per field in one pass — Next() hands out materialized docs one at a time and peak memory stays one batch. A materialization failure is sticky so callers can never iterate a partially filled batch. Drop the now-unused row-level converter API and the vector prefetch window machinery.
kMaxRecordBatchNumRows only caps MemForwardStore batches; a Parquet scan returns a whole row group per ReadNext (up to ~1M rows), so materializing an entire batch could peak at gigabytes with vectors included. Materialize docs window by window instead (at most 4096 rows at a time, scalar columns zero-copy sliced, vectors fetched per window), releasing the Arrow batch right after its last window — restoring the constant-memory contract while keeping column-level conversion and per-reader column resolution. Rename the Parquet test to reflect what it exercises: several materialization windows within one large row group.
|
|
||
| std::lock_guard maintenance_lock(maintenance_mtx_); | ||
|
|
||
| auto iter_check = check_no_active_iterators("Destroy"); |
There was a problem hiding this comment.
可以用try_lock实现,当前实现先检查再加锁。检查完成在加写锁之前其他地方先加了读锁,这里一样会hang住。
|
|
||
| class CollectionImpl : public Collection { | ||
| class CollectionImpl : public Collection, | ||
| public std::enable_shared_from_this<CollectionImpl> { |
There was a problem hiding this comment.
没必要enable_shared_from_this,读锁就能避免析构,Close会加写锁
| CollectionSchema::Ptr schema; | ||
| std::vector<std::string> scan_columns; | ||
| IndexFilter::Ptr filter; | ||
| auto s = Scan(options, segments, delete_store, schema, scan_columns, filter); |
There was a problem hiding this comment.
Scan含义模糊,实现和名称差别较大,建议改为类似PrepareIterate的命名,返回ResultDocIterator::Impl;

Add streaming full-collection traversal across C++/C/Python (relates to #380).
API
Collection::CreateIterator(IteratorOptions)→DocIterator(Next()returnsResult<Doc::Ptr>: error /nullptrEOF / doc;Close()).zvec_doc_iterator_t+zvec_iterator_options_thandles;zvec_collection_create_iterator/zvec_doc_iterator_next/zvec_doc_iterator_close; errors mapped tozvec_error_code_t.collection.iter_docs()generator (constant memory; releases native resources infinally, also on early break).Snapshot semantics
CreateIterator()seals the current writing segment on writable collections (read-only collections scan directly, including their writing segment — no flush, no data loss); the snapshot captures the segment set, a deep copy of the delete bitmap, and the schema. Writes after creation are invisible to the iterator; deletions after creation do not affect it.IteratorOptions.output_fields_selects forward fields (unknown/duplicate names rejected with an error);include_vector_controls vector materialization.Concurrency (schema read lock for the iterator lifetime)
Close(); the collection is kept alive by the iterator's own reference, so dropping the collection handle during iteration is safe.Implementation notes
FilteringReader(src/db/index/segment/filtering_reader.*)._zvec_row_id_, correct for compacted segments). A materialization failure is sticky (the error keeps being returned; no partial docs are ever handed out).src/db/index/common/doc_field_converter.*also serveSegmentImpl::Fetchand the SQL engine (ConvertVectorDataBufferToDocField/ConvertArrowColumnToDocFields/ExtractTypedArrayValues).Result/error code/exception — nothing is logged-and-continued (exception:SegmentImpl::Fetchkeeps its pre-existing lenient per-field contract).Tests
iterator_test: 18 tests — basic/empty/deleted/close-then-next, include/exclude vector, output_fields selection + rejection, scalar type mapping, 1000-doc integration, read-only collection, performance (100k docs, constant memory), Parquet multi-batch vector alignment, and concurrency: snapshot isolation under writes, Optimize/DDL/Flush/Close rejected while open, close-handle safety after the collection handle is dropped, multiple iterators.c_api_test: 5 iterator tests (basic/empty/exclude-vector/output-fields/null-args) inside the 73-test C API suite.test_iter_docs.py: 9 tests (basic fields/vectors, deletion filtering, output fields, isolation, snapshot-at-call-time, generator protocol).