seqtree 1.2.0__tar.gz → 1.3.1__tar.gz

{seqtree-1.2.0 → seqtree-1.3.1}/CHANGELOG.md

@@ -3,6 +3,24 @@
 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.1] — 2026-10-19
+
+### Changed
+
+- Gap-prior lookup preparation evaluates pure callbacks only for input lengths,
+  avoiding cubic callback work for absent lengths while preserving exact scores,
+  count reductions and lane-specific priors. Equal-length/empty axes do not
+  evaluate unused priors. The dense native lookup layout is unchanged.
+
+## [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.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: seqtree
-Version: 1.2.0
+Version: 1.3.1
 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.1}/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.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "seqtree"
-version = "1.2.0"
+version = "1.3.1"
 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.1}/python/seqtree/gapblock.py

@@ -499,19 +499,22 @@ def gapblock_score(
     return best + gap_cost(d, gap_open, gap_extend), best_i
 
 
-def _prior_cube(prior: GapPrior, width: int) -> list[int]:
+def _prior_cube(prior: GapPrior, width: int, lengths=None) -> list[int]:
     """Flatten ``prior`` to ``[m][d][i]`` with stride ``width + 1``, for the C++ kernel.
 
     Only reachable cells are filled: ``d >= 1`` (the prior never applies to equal lengths) and
     ``i <= m - d`` (the block start indexes the shorter sequence). The rest stay zero and are
-    never read.
+    never read. Input lengths restrict callback evaluation to possible full-length pairs;
+    the dense native layout is unchanged. None retains the full-width internal oracle.
     """
     w1 = width + 1
     cube = [0] * (w1 * w1 * w1)
-    for m in range(w1):
-        for d in range(1, m + 1):
+    lengths = range(w1) if lengths is None else sorted(set(lengths))
+    for row, m in enumerate(lengths):
+        for n in reversed(lengths[:row]):
+            d = m - n
             base = (m * w1 + d) * w1
-            for i in range(m - d + 1):
+            for i in range(n + 1):
                 v = int(prior(i, d, m))
                 if v < 0:
                     raise ValueError(f"gap_prior({i}, {d}, {m}) = {v} is negative")
@@ -561,7 +564,8 @@ def score_matrix(
             leaving this at 1 with a real matrix makes gaps ~14x cheaper than substitutions.
         gap_extend: Cost per additional gap column.
         gap_prior: :data:`GapPrior`, materialized once into a lookup cube and then read from C++.
-            ``None`` lets the score alone choose the block position.
+            Evaluated only for unequal input lengths and valid block starts; it must be
+            a pure nonnegative function. ``None`` lets the score alone choose the block position.
         alphabet: ``"aa"``, ``"nt"``, or ``"iupac"``. Symbols outside it raise.
         threads: Worker threads; ``0`` means one per core. Rows are disjoint, so this scales.
 
@@ -586,8 +590,9 @@ def score_matrix(
         gap_open = _default_gap_open(matrix)
     if gap_open < 0 or gap_extend < 0:
         raise ValueError("gap_open and gap_extend must be >= 0")
-    width = max((len(s) for s in q + r), default=0)
-    cube = _prior_cube(gap_prior, width) if gap_prior is not None else []
+    lengths = {len(s) for s in q + r}
+    width = max(lengths, default=0)
+    cube = _prior_cube(gap_prior, width, lengths) if gap_prior is not None else []
     return _gapblock_matrix(q, r, alphabet, matrix, gap_open, gap_extend, cube, width, threads,
                             _position_weights(position_weights_by_length, q, r),
                             *_group_options(query_group_ids, reference_group_ids, group_distances, q, r),
@@ -611,8 +616,9 @@ def _batch_options(axes, matrix, gap_open, gap_extend, gap_prior, threads):
     go = _batch_integer(_default_gap_open(matrix) if gap_open is None else gap_open, "gap_open")
     ge = _batch_integer(gap_extend, "gap_extend")
     threads = _batch_integer(threads, "threads")
-    width = max((len(s) for axis in axes for s in axis), default=0)
-    cube = _prior_cube(gap_prior, width) if gap_prior is not None else []
+    lengths = {len(s) for axis in axes for s in axis}
+    width = max(lengths, default=0)
+    cube = _prior_cube(gap_prior, width, lengths) if gap_prior is not None else []
     return go, ge, cube, width, threads
 
 
@@ -915,11 +921,11 @@ def _additional_geometry(geometry, features):
         geometry.additional_features = terms
 
 
-def _lane_prior(geometry, prior, width):
+def _lane_prior(geometry, prior, width, queries, refs):
     if prior is not None:
         if not hasattr(geometry, "gap_prior"):
             raise RuntimeError("lane priors require native grouped gap_prior support")
-        geometry.gap_prior = _prior_cube(prior, width)
+        geometry.gap_prior = _prior_cube(prior, width, {len(s) for s in queries + refs})
 
 
 def _alignment_options(geometry, q, r, mode, scale, left, right, matrix, prior,
@@ -1099,8 +1105,8 @@ def grouped_paired_sum_count_batch(queries_a, queries_b, refs_a, refs_b, thresho
         reference_position_weights_beta, feature_offsets_beta, score_quantum_beta)
     _additional_geometry(alpha, additional_feature_offsets_alpha)
     _additional_geometry(beta, additional_feature_offsets_beta)
-    _lane_prior(alpha, gap_prior_alpha, width)
-    _lane_prior(beta, gap_prior_beta, width)
+    _lane_prior(alpha, gap_prior_alpha, width, qa, ra)
+    _lane_prior(beta, gap_prior_beta, width, qb, rb)
     _alignment_options(alpha, qa, ra, alignment_mode, alignment_scale, trim_left, trim_right,
         matrix, gap_prior if gap_prior_alpha is None else gap_prior_alpha, position_weights_by_length_alpha, query_position_weights_alpha,
         reference_position_weights_alpha)

{seqtree-1.2.0 → seqtree-1.3.1}/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.1}/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 {