Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 24 additions & 0 deletions docs/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
------------------------------------

Expand Down
15 changes: 15 additions & 0 deletions include/seqtree/seqtree.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#include "seqtree/types.hpp"

#include <cstdint>
#include <functional>
#include <memory>
#include <optional>
#include <string>
Expand All @@ -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.
Expand Down Expand Up @@ -134,6 +136,19 @@ class Index {
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&,
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
66 changes: 66 additions & 0 deletions src/_bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
#include <nanobind/stl/vector.h>

#include <algorithm>
#include <exception>
#include <optional>
#include <string>
#include <vector>
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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,
Expand Down
128 changes: 128 additions & 0 deletions src/index.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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>

Expand Down Expand Up @@ -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 {
Expand Down
Loading
Loading