seqtree 1.2.0__tar.gz → 1.3.0__tar.gz

{seqtree-1.2.0 → seqtree-1.3.0}/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

{seqtree-1.2.0 → seqtree-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: seqtree
-Version: 1.2.0
+Version: 1.3.0
 Summary: Fast fuzzy search over biological sequences (C++ core, Python bindings)
 Keywords: sequence-search,fuzzy-matching,CDR3,immunology,bioinformatics,trie
 Author-Email: ISALGO laboratory <mikhail.shugay@gmail.com>

{seqtree-1.2.0 → seqtree-1.3.0}/include/seqtree/seqtree.hpp

@@ -2,6 +2,7 @@
 #include "seqtree/types.hpp"
 
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <optional>
 #include <string>
@@ -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 @@ public:
     std::vector<std::vector<Hit>> search_batch(const std::vector<std::string>& 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<std::string>& queries, const SearchParams&,
+        const std::function<void(size_t, size_t, CountBuffer&&)>& 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<std::string>& queries, const std::vector<std::string>& other_queries,
+        const SearchParams&,
+        const std::function<void(size_t, size_t, CountBuffer&&, CountBuffer&&)>& 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<std::vector<uint64_t>> edit_histogram_batch(
         const std::vector<std::string>& queries, const SearchParams&,

{seqtree-1.2.0 → seqtree-1.3.0}/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"

{seqtree-1.2.0 → seqtree-1.3.0}/src/_bindings.cpp

@@ -12,6 +12,7 @@
 #include <nanobind/stl/vector.h>
 
 #include <algorithm>
+#include <exception>
 #include <optional>
 #include <string>
 #include <vector>
@@ -186,6 +187,53 @@ nb::list py_search_batch(const Index& idx, const std::vector<std::string>& queri
     return out;
 }
 
+void py_search_batch_consume(const Index& idx, const std::vector<std::string>& 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<std::string>& queries, const std::vector<std::string>& 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<std::vector<uint64_t>> py_edit_histogram_batch(
         const Index& idx, const std::vector<std::string>& 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,

{seqtree-1.2.0 → seqtree-1.3.0}/src/index.cpp

@@ -5,9 +5,12 @@
 #include "count_unit.hpp"
 
 #include <algorithm>
+#include <atomic>
 #include <cstring>
 #include <fstream>
+#include <iterator>
 #include <limits>
+#include <mutex>
 #include <stdexcept>
 #include <vector>
 
@@ -188,6 +191,131 @@ std::vector<std::vector<Hit>> Index::search_batch(const std::vector<std::string>
     return results;
 }
 
+void Index::search_batch_consume(const std::vector<std::string>& queries,
+        const SearchParams& p,
+        const std::function<void(size_t, size_t, CountBuffer&&)>& 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<uint32_t>::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<bool> cancelled{false};
+    std::mutex callback_mutex;
+    struct Local { Searcher searcher; std::vector<std::vector<Hit>> hits; };
+    parallel_for(blocks, threads,
+        [&] { return Local{Searcher(*this), std::vector<std::vector<Hit>>(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<std::mutex> 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<std::string>& queries, const std::vector<std::string>& other_queries,
+        const SearchParams& p,
+        const std::function<void(size_t, size_t, CountBuffer&&, CountBuffer&&)>& 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<uint32_t>::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<bool> cancelled{false};
+    std::mutex callback_mutex;
+    struct Local { Searcher searcher, other_searcher; std::vector<std::vector<Hit>> hits, other_hits; };
+    parallel_for(blocks, threads,
+        [&] { return Local{Searcher(*this), Searcher(other), std::vector<std::vector<Hit>>(chunk_size), std::vector<std::vector<Hit>>(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<std::mutex> 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<std::vector<uint64_t>> Index::edit_histogram_batch(
         const std::vector<std::string>& queries, const SearchParams& p,
         int threads, bool exclude_exact, int max_indels) const {