diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e93a5e..e6b740b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,15 @@ All notable changes to `seqtree`. Dates are release dates. From 1.0.0 the project follows semantic versioning: breaking changes need a **major** bump. +## [1.3.0] — 2026-10-19 + +### Added + +- Bounded complete-query hit consumers for single and coordinated two-index + searches, with owned numeric buffers, explicit byte budgets and worker cleanup. +- Optional empty query axes in coordinated search retain original query offsets + and supported marginal evidence without creating a query for the missing axis. + ## [1.2.0] — 2026-10-17 ### Fixed diff --git a/docs/api.rst b/docs/api.rst index f4a940c..18bcd95 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -47,6 +47,30 @@ and exact constrained dynamic programming; other profiles use the generic search worker. Existing eligible count paths are retained when the bound is redundant. Explicit ``seqtrie`` remains unsupported. +For full hits with bounded output, use +``Index.search_batch_consume(queries, params, consume, threads=0, chunk_size=32, max_hit_bytes=64*1024**2)``. +It searches the complete population with one native pool. Serialized callbacks, +which may run on native worker threads, +receive ``(start, stop, owned_buffer)`` with read-only uint32 columns +``query_index, ref_id, score, n_subs, n_ins, n_dels``. Every query is complete; +block order may vary across threads, while hit order within each query is unchanged. +Each output buffer obeys the explicit byte ceiling. An oversized single query +raises rather than truncating its hits. Workers also retain the current block's +native hit scratch. Consume and release buffers to bound total memory; buffers +explicitly retained by a callback remain owned. Parameters and matrices are +snapshotted, callbacks are serialized, and workers join before errors propagate. + +For coordinated searches over two same-alphabet indexes, use +``Index.search_paired_batch_consume(other, queries, other_queries, params, consume, threads=0, chunk_size=32, max_hit_bytes=64*1024**2, skip_empty=False)``. +Both equal-length query populations run in one pool with two reusable Searchers +per worker. Serialized callbacks receive ``(start, stop, owned_buffer, +other_owned_buffer)`` for the same complete query block. The byte ceiling applies +to the sum of both output buffers; ownership, snapshots and error cleanup follow +the single-index consumer contract. This API provides no pairing or score fusion; +the caller owns those interpretations. ``skip_empty=True`` emits no hits for an +empty query axis while retaining searches on its partner; the default preserves +ordinary empty-string search semantics. + Floating alignment target reductions ------------------------------------ diff --git a/include/seqtree/seqtree.hpp b/include/seqtree/seqtree.hpp index 663a228..56598dc 100644 --- a/include/seqtree/seqtree.hpp +++ b/include/seqtree/seqtree.hpp @@ -2,6 +2,7 @@ #include "seqtree/types.hpp" #include +#include #include #include #include @@ -13,6 +14,7 @@ namespace seqtree { struct Trie; // internal, src/trie.hpp struct Scratch; // internal, src/engines.hpp +struct CountBuffer; // Maps characters to small integer codes (<=32) and back. Case-insensitive on // input. Invalid characters encode to kInvalid. @@ -134,6 +136,19 @@ class Index { std::vector> search_batch(const std::vector& queries, const SearchParams&, int threads = 0) const; + // One pool, serialized owned primitive blocks; each query is emitted completely. + // Columns: query index, ref_id, score, n_subs, n_ins, n_dels (all nonnegative). + void search_batch_consume(const std::vector& queries, const SearchParams&, + const std::function& consume, + int threads = 0, size_t chunk_size = 32, size_t max_hit_bytes = 64*1024*1024) const; + + // Coordinate two indexes in one pool; both owned buffers cover the same queries. + void search_paired_batch_consume(const Index& other, + const std::vector& queries, const std::vector& other_queries, + const SearchParams&, + const std::function& consume, + int threads = 0, size_t chunk_size = 32, size_t max_hit_bytes = 64*1024*1024, bool skip_empty = false) const; + // Exact edit-count bins 0..max_total_edits; one reusable hit buffer per worker. std::vector> edit_histogram_batch( const std::vector& queries, const SearchParams&, diff --git a/pyproject.toml b/pyproject.toml index b13195d..9a00cc6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "seqtree" -version = "1.2.0" +version = "1.3.0" description = "Fast fuzzy search over biological sequences (C++ core, Python bindings)" readme = "README.md" requires-python = ">=3.10" diff --git a/src/_bindings.cpp b/src/_bindings.cpp index a30c163..8d23730 100644 --- a/src/_bindings.cpp +++ b/src/_bindings.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -186,6 +187,53 @@ nb::list py_search_batch(const Index& idx, const std::vector& queri return out; } +void py_search_batch_consume(const Index& idx, const std::vector& queries, + const PyParams& pp, nb::object consume, int threads, size_t chunk_size, size_t max_hit_bytes) { + if (!PyCallable_Check(consume.ptr())) throw nb::value_error("consumer must be callable"); + // A callback may mutate the original Python parameters; native workers read + // only these owned matrix/positional snapshots until the whole pool joins. + const PyParams owned = pp; + auto mat = make_matrix(owned, idx.alphabet()); + SearchParams cp = to_cpp(owned, mat ? &*mat : nullptr); + std::exception_ptr error; + { + nb::gil_scoped_release release; + try { + idx.search_batch_consume(queries, cp, + [&](size_t start, size_t stop, CountBuffer&& buffer) { + nb::gil_scoped_acquire acquire; + consume(start, stop, nb::cast(std::move(buffer))); + }, threads, chunk_size, max_hit_bytes); + } catch (...) { error = std::current_exception(); } + } + // Restore the caller's GIL before Python exceptions are restored/destroyed. + if (error) std::rethrow_exception(error); +} + +void py_search_paired_batch_consume(const Index& idx, const Index& other, + const std::vector& queries, const std::vector& other_queries, + const PyParams& pp, nb::object consume, int threads, size_t chunk_size, size_t max_hit_bytes, bool skip_empty) { + if (!PyCallable_Check(consume.ptr())) throw nb::value_error("consumer must be callable"); + // A callback may mutate the original Python parameters; native workers read + // only these owned matrix/positional snapshots until the whole pool joins. + const PyParams owned = pp; + auto mat = make_matrix(owned, idx.alphabet()); + SearchParams cp = to_cpp(owned, mat ? &*mat : nullptr); + std::exception_ptr error; + { + nb::gil_scoped_release release; + try { + idx.search_paired_batch_consume(other, queries, other_queries, cp, + [&](size_t start, size_t stop, CountBuffer&& buffer, CountBuffer&& other_buffer) { + nb::gil_scoped_acquire acquire; + consume(start, stop, nb::cast(std::move(buffer)), nb::cast(std::move(other_buffer))); + }, threads, chunk_size, max_hit_bytes, skip_empty); + } catch (...) { error = std::current_exception(); } + } + // Restore the caller's GIL before Python exceptions are restored/destroyed. + if (error) std::rethrow_exception(error); +} + std::vector> py_edit_histogram_batch( const Index& idx, const std::vector& queries, const PyParams& pp, int threads, bool exclude_exact, int max_indels) { @@ -718,6 +766,24 @@ NB_MODULE(_core, m) { nb::arg("threads") = 0, "Search many queries in parallel (releases the GIL). ``threads=0`` uses all " "cores. Returns one hit list per query, in input order.") + .def("search_batch_consume", &py_search_batch_consume, + nb::arg("queries"), nb::arg("params"), nb::arg("consume"), nb::arg("threads") = 0, + nb::arg("chunk_size") = 32, nb::arg("max_hit_bytes") = 64*1024*1024, + "One population/pool, serialized consume(start, stop, owned_buffer) callbacks. " + "Read-only uint32 rows: query_index, ref_id, score, n_subs, n_ins, n_dels. " + "Every query is complete; block order may vary across threads. Output bytes per " + "block never exceed max_hit_bytes; an oversized query raises, never truncates. " + "Native worker scratch also holds the current block of hits. Consumer-retained " + "buffers remain owned; callers control their retention. Parameters are snapshotted.") + .def("search_paired_batch_consume", &py_search_paired_batch_consume, + nb::arg("other"), nb::arg("queries"), nb::arg("other_queries"), + nb::arg("params"), nb::arg("consume"), nb::arg("threads") = 0, + nb::arg("chunk_size") = 32, nb::arg("max_hit_bytes") = 64*1024*1024, nb::arg("skip_empty") = false, + "One pool with two Searchers; consume(start, stop, owned_buffer, other_owned_buffer). " + "Both six-column buffers cover the same complete query block. Combined output " + "bytes are bounded by max_hit_bytes; block order may vary. Same-alphabet indexes " + "and equal-length query populations required. Parameters are snapshotted. " + "skip_empty=True emits no hits for empty query axes; other axes remain searched.") .def("edit_histogram_batch", &py_edit_histogram_batch, nb::arg("queries"), nb::arg("params"), nb::arg("threads") = 0, nb::arg("exclude_exact") = false, nb::arg("max_indels") = -1, diff --git a/src/index.cpp b/src/index.cpp index 559e033..7d2f2bc 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -5,9 +5,12 @@ #include "count_unit.hpp" #include +#include #include #include +#include #include +#include #include #include @@ -188,6 +191,131 @@ std::vector> Index::search_batch(const std::vector return results; } +void Index::search_batch_consume(const std::vector& queries, + const SearchParams& p, + const std::function& consume, + int threads, size_t chunk_size, size_t max_hit_bytes) const { + constexpr size_t columns = 6; + const size_t max_rows = max_hit_bytes / (columns * sizeof(uint32_t)); + if (!chunk_size || !max_rows || !consume) + throw std::invalid_argument("positive chunk_size/max_hit_bytes and a consumer required"); + if (queries.size() > uint64_t(std::numeric_limits::max()) + 1) + throw std::overflow_error("query indices exceed the uint32 output buffer"); + // Bound both the primitive output and the per-worker Hit vectors. A single + // query may still exceed the explicit output cap; detect that without truncation. + chunk_size = std::min(chunk_size, std::max(size_t(1), max_rows / std::max(size_t(1), size_t(size())))); + const size_t blocks = queries.size() / chunk_size + (queries.size() % chunk_size != 0); + std::atomic cancelled{false}; + std::mutex callback_mutex; + struct Local { Searcher searcher; std::vector> hits; }; + parallel_for(blocks, threads, + [&] { return Local{Searcher(*this), std::vector>(chunk_size)}; }, + [&](size_t block, Local& worker) { + if (cancelled.load()) return; + try { + const size_t start = block * chunk_size; + const size_t stop = start + std::min(chunk_size, queries.size() - start); + size_t rows = 0; + for (size_t i = start; i < stop; ++i) { + if (cancelled.load()) return; + auto& hits = worker.hits[i-start]; + worker.searcher.search_into(queries[i], p, hits); + if (hits.size() > max_rows - rows) + throw std::overflow_error("complete query hits exceed max_hit_bytes; increase the explicit buffer limit"); + rows += hits.size(); + } + CountBuffer buffer; + buffer.shape = {rows, columns}; + buffer.data.resize(rows * columns); + size_t offset = 0; + for (size_t i = start; i < stop; ++i) + for (const auto& hit : worker.hits[i-start]) { + const uint32_t values[] = {uint32_t(i), hit.ref_id, uint32_t(hit.score), + hit.n_subs, hit.n_ins, hit.n_dels}; + std::copy(std::begin(values), std::end(values), buffer.data.begin() + offset); + offset += columns; + } + std::lock_guard lock(callback_mutex); + if (!cancelled.load()) { + try { consume(start, stop, std::move(buffer)); } + catch (...) { cancelled.store(true); throw; } + } + } catch (...) { + cancelled.store(true); + throw; + } + }); +} + +void Index::search_paired_batch_consume(const Index& other, + const std::vector& queries, const std::vector& other_queries, + const SearchParams& p, + const std::function& consume, + int threads, size_t chunk_size, size_t max_hit_bytes, bool skip_empty) const { + if (queries.size() != other_queries.size()) + throw std::invalid_argument("paired query populations must have equal lengths"); + if (alphabet() != other.alphabet()) + throw std::invalid_argument("paired indexes must use the same alphabet"); + constexpr size_t columns = 6; + const size_t max_rows = max_hit_bytes / (columns * sizeof(uint32_t)); + if (!chunk_size || !max_rows || !consume) + throw std::invalid_argument("positive chunk_size/max_hit_bytes and a consumer required"); + if (queries.size() > uint64_t(std::numeric_limits::max()) + 1) + throw std::overflow_error("query indices exceed the uint32 output buffer"); + // Bound both the primitive output and the per-worker Hit vectors. A single + // query may still exceed the explicit output cap; detect that without truncation. + chunk_size = std::min(chunk_size, std::max(size_t(1), max_rows / std::max(size_t(1), size_t(size()) + size_t(other.size())))); + const size_t blocks = queries.size() / chunk_size + (queries.size() % chunk_size != 0); + std::atomic cancelled{false}; + std::mutex callback_mutex; + struct Local { Searcher searcher, other_searcher; std::vector> hits, other_hits; }; + parallel_for(blocks, threads, + [&] { return Local{Searcher(*this), Searcher(other), std::vector>(chunk_size), std::vector>(chunk_size)}; }, + [&](size_t block, Local& worker) { + if (cancelled.load()) return; + try { + const size_t start = block * chunk_size; + const size_t stop = start + std::min(chunk_size, queries.size() - start); + size_t rows = 0; + for (size_t i = start; i < stop; ++i) { + if (cancelled.load()) return; + auto& hits = worker.hits[i-start]; + if (skip_empty && queries[i].empty()) hits.clear(); + else worker.searcher.search_into(queries[i], p, hits); + auto& other_hits = worker.other_hits[i-start]; + if (skip_empty && other_queries[i].empty()) other_hits.clear(); + else worker.other_searcher.search_into(other_queries[i], p, other_hits); + if (hits.size() > max_rows - rows || other_hits.size() > max_rows - rows - hits.size()) + throw std::overflow_error("complete query hits exceed max_hit_bytes; increase the explicit buffer limit"); + rows += hits.size() + other_hits.size(); + } + auto pack = [&](const auto& groups) { + CountBuffer buffer; + size_t count = 0; + for (size_t i = start; i < stop; ++i) count += groups[i-start].size(); + buffer.shape = {count, columns}; + buffer.data.reserve(count * columns); + for (size_t i = start; i < stop; ++i) + for (const auto& hit : groups[i-start]) { + const uint32_t values[] = {uint32_t(i), hit.ref_id, uint32_t(hit.score), + hit.n_subs, hit.n_ins, hit.n_dels}; + buffer.data.insert(buffer.data.end(), std::begin(values), std::end(values)); + } + return buffer; + }; + auto buffer = pack(worker.hits), other_buffer = pack(worker.other_hits); + std::lock_guard lock(callback_mutex); + if (!cancelled.load()) { + try { consume(start, stop, std::move(buffer), std::move(other_buffer)); } + catch (...) { cancelled.store(true); throw; } + } + } catch (...) { + cancelled.store(true); + throw; + } + }); +} + std::vector> Index::edit_histogram_batch( const std::vector& queries, const SearchParams& p, int threads, bool exclude_exact, int max_indels) const { diff --git a/tests/python/test_paired_consume.py b/tests/python/test_paired_consume.py new file mode 100644 index 0000000..3e504c9 --- /dev/null +++ b/tests/python/test_paired_consume.py @@ -0,0 +1,98 @@ +"""Coordinated full-population search preserves each chain's ordinary hits.""" +import pytest +from seqtree import Index, SearchParams, PositionalMatrix, SubstitutionMatrix + +@pytest.mark.parametrize('threads',[1,2]) +@pytest.mark.parametrize('weighted',[False,True]) +def test_paired_buffers_equal_legacy(threads,weighted): + a=Index.build(['AAAA','AAAC','AAAA','CCCC'],alphabet='nt') + b=Index.build(['CCCC','CCCA','AAAA'],alphabet='nt') + qa=['AAAA','GGGG','AAAC']*25;qb=['CCCC','GGGG','AAAA']*25 + p=SearchParams(max_subs=2,max_ins=1,max_dels=1,max_total_edits=2) + if weighted: + p.matrix=SubstitutionMatrix.unit(4);p.gap_open=p.gap_extend=2 + expected=[[[tuple(h) for h in row] for row in index.search_batch(q,p,threads)] for index,q in [(a,qa),(b,qb)]] + blocks=[] + a.search_paired_batch_consume(b,qa,qb,p,lambda *args:blocks.append(args),threads=threads,chunk_size=2,max_hit_bytes=24*7*2) + actual=[[[] for _ in qa] for _ in range(2)];covered=[] + for start,stop,*buffers in sorted(blocks): + covered.extend(range(start,stop)) + assert 0