Lift the uint16 assembly-count limit with global record offsets

cpp/include/seqwin/graph.hpp

@@ -15,10 +15,8 @@ namespace seqwin {
 struct Kmer {
     /** 0-based position of the minimizer within its FASTA record. */
     std::uint32_t pos;
-    /** 0-based index of the FASTA record within the assembly. */
-    std::uint16_t record_idx;
-    /** Assembly index assigned to the source assembly. */
-    std::uint16_t assembly_idx;
+    /** 0-based global index of the FASTA record. */
+    std::uint32_t record_idx;
 };
 
 /**
@@ -70,29 +68,29 @@ struct Graph {
     NoInitArray<Node> nodes;
     /** Sorted by hash. */
     NoInitArray<Edge> edges;
+    /** Cumulative global FASTA record offsets by assembly. */
+    std::vector<std::uint64_t> record_offsets;
     /** FASTA record IDs of each assembly. */
     std::vector<std::vector<std::string>> ids_by_assembly;
 };
 
 /**
  * @brief Build a minimizer graph from assembly FASTA files.
  *
- * `assembly_paths`, `assembly_indices`, and `is_targets` are parallel lists.
+ * `assembly_paths` and `is_targets` are parallel lists.
  *
  * @param assembly_paths Paths to input assemblies in FASTA format (plain or gzipped).
  * @param kmerlen K-mer length for minimizer sketch.
  * @param windowsize Window size for minimizer sketch.
- * @param assembly_indices Assembly indices, expected to be `0..N-1`.
  * @param is_targets Whether each assembly is a target assembly.
  * @param n_cpu Number of worker threads to use.
  * @return Minimizer graph.
- * @throws `std::runtime_error` If input sizes are inconsistent or indices exceed supported ranges.
+ * @throws `std::runtime_error` If input sizes are inconsistent or counts exceed supported ranges.
  */
 Graph build(
     const std::vector<std::string>& assembly_paths,
     std::size_t kmerlen,
     std::size_t windowsize,
-    const std::vector<std::size_t>& assembly_indices,
     const std::vector<bool>& is_targets,
     std::size_t n_cpu
 );

cpp/include/seqwin/helpers.hpp

@@ -30,6 +30,8 @@ struct ThreadGraph {
     NoInitArray<Node> nodes;
     /** Unsorted. */
     NoInitArray<Edge> edges;
+    /** Thread-local cumulative FASTA record offsets by assembly. */
+    std::vector<std::uint64_t> record_offsets;
     /** FASTA record IDs of each assembly. */
     std::vector<std::vector<std::string>> ids_by_assembly;
     /** Total number of minimizers generated by this worker. */

cpp/src/bindings/python_bindings.cpp

@@ -1,6 +1,8 @@
 #include <cstdint>
 #include <memory>
 #include <string>
+#include <type_traits>
+#include <utility>
 #include <vector>
 
 #include <pybind11/numpy.h>
@@ -12,8 +14,32 @@
 
 namespace py = pybind11;
 
+namespace {
+
+template <typename Array>
+auto array_to_numpy(Array&& values) {
+    using Owner = typename std::decay<Array>::type;
+    using T = typename std::remove_cv<
+        typename std::remove_pointer<decltype(std::declval<Owner&>().data())>::type
+    >::type;
+
+    auto* owner = new Owner(std::forward<Array>(values));
+    auto capsule = py::capsule(owner, [](void* ptr) {
+        delete static_cast<Owner*>(ptr);
+    });
+
+    return py::array_t<T>(
+        {static_cast<py::ssize_t>(owner->size())},
+        {static_cast<py::ssize_t>(sizeof(T))},
+        owner->data(),
+        capsule
+    );
+}
+
+}  // namespace
+
 PYBIND11_MODULE(_core, m) {
-    PYBIND11_NUMPY_DTYPE(seqwin::Kmer, pos, record_idx, assembly_idx);
+    PYBIND11_NUMPY_DTYPE(seqwin::Kmer, pos, record_idx);
     PYBIND11_NUMPY_DTYPE(seqwin::Node, hash, start, stop, n_tar, n_neg, penalty);
     PYBIND11_NUMPY_DTYPE(seqwin::Edge, first, second, weight);
 
@@ -23,7 +49,6 @@ PYBIND11_MODULE(_core, m) {
         [](const std::vector<std::string>& assembly_paths,
            std::size_t kmerlen,
            std::size_t windowsize,
-           const std::vector<std::size_t>& assembly_indices,
            const std::vector<bool>& is_targets,
            std::size_t n_cpu
         ) {
@@ -34,58 +59,31 @@ PYBIND11_MODULE(_core, m) {
                     assembly_paths,
                     kmerlen,
                     windowsize,
-                    assembly_indices,
                     is_targets,
                     n_cpu
                 );
             }
-            auto owner = std::make_shared<seqwin::Graph>(std::move(graph));
-            auto capsule = py::capsule(
-                new std::shared_ptr<seqwin::Graph>(owner),
-                [](void* ptr) {
-                    delete static_cast<std::shared_ptr<seqwin::Graph>*>(ptr);
-                }
-            );
-
-            auto kmers = py::array_t<seqwin::Kmer>(
-                {static_cast<py::ssize_t>(owner->kmers.size())},
-                {static_cast<py::ssize_t>(sizeof(seqwin::Kmer))},
-                owner->kmers.data(),
-                capsule
-            );
-            auto idx = py::array_t<std::uint64_t>(
-                {static_cast<py::ssize_t>(owner->idx.size())},
-                {static_cast<py::ssize_t>(sizeof(std::uint64_t))},
-                owner->idx.data(),
-                capsule
-            );
-            auto nodes = py::array_t<seqwin::Node>(
-                {static_cast<py::ssize_t>(owner->nodes.size())},
-                {static_cast<py::ssize_t>(sizeof(seqwin::Node))},
-                owner->nodes.data(),
-                capsule
-            );
-            auto edges = py::array_t<seqwin::Edge>(
-                {static_cast<py::ssize_t>(owner->edges.size())},
-                {static_cast<py::ssize_t>(sizeof(seqwin::Edge))},
-                owner->edges.data(),
-                capsule
-            );
+
+            auto kmers = array_to_numpy(std::move(graph.kmers));
+            auto idx = array_to_numpy(std::move(graph.idx));
+            auto nodes = array_to_numpy(std::move(graph.nodes));
+            auto edges = array_to_numpy(std::move(graph.edges));
+            auto record_offsets = array_to_numpy(std::move(graph.record_offsets));
+
             py::list ids_by_assembly;
-            for (const auto& ids : owner->ids_by_assembly) {
+            for (const auto& ids : graph.ids_by_assembly) {
                 py::tuple ids_tuple(ids.size());
                 for (std::size_t i = 0; i < ids.size(); ++i) {
                     ids_tuple[i] = ids[i];
                 }
                 ids_by_assembly.append(ids_tuple);
             }
 
-            return py::make_tuple(kmers, idx, nodes, edges, ids_by_assembly);
+            return py::make_tuple(kmers, idx, nodes, edges, record_offsets, ids_by_assembly);
         },
         py::arg("assembly_paths"),
         py::arg("kmerlen"),
         py::arg("windowsize"),
-        py::arg("assembly_indices"),
         py::arg("is_targets"),
         py::arg("n_cpu") = 1
     );
@@ -116,32 +114,10 @@ PYBIND11_MODULE(_core, m) {
                     used_hashes
                 );
             }
-            auto owner = std::make_shared<seqwin::Graph>(std::move(graph));
-            auto capsule = py::capsule(
-                new std::shared_ptr<seqwin::Graph>(owner),
-                [](void* ptr) {
-                    delete static_cast<std::shared_ptr<seqwin::Graph>*>(ptr);
-                }
-            );
-
-            auto kmers_new = py::array_t<seqwin::Kmer>(
-                {static_cast<py::ssize_t>(owner->kmers.size())},
-                {static_cast<py::ssize_t>(sizeof(seqwin::Kmer))},
-                owner->kmers.data(),
-                capsule
-            );
-            auto idx_new = py::array_t<std::uint64_t>(
-                {static_cast<py::ssize_t>(owner->idx.size())},
-                {static_cast<py::ssize_t>(sizeof(std::uint64_t))},
-                owner->idx.data(),
-                capsule
-            );
-            auto nodes_new = py::array_t<seqwin::Node>(
-                {static_cast<py::ssize_t>(owner->nodes.size())},
-                {static_cast<py::ssize_t>(sizeof(seqwin::Node))},
-                owner->nodes.data(),
-                capsule
-            );
+
+            auto kmers_new = array_to_numpy(std::move(graph.kmers));
+            auto idx_new = array_to_numpy(std::move(graph.idx));
+            auto nodes_new = array_to_numpy(std::move(graph.nodes));
 
             return py::make_tuple(kmers_new, idx_new, nodes_new);
         },

cpp/src/seqwin/graph.cpp

@@ -47,7 +47,6 @@ ThreadGraph build_worker(
     const std::vector<std::string>& assembly_paths,
     std::size_t kmerlen,
     std::size_t windowsize,
-    const std::vector<std::size_t>& assembly_indices,
     const std::vector<bool>& is_targets,
     std::size_t start_assembly,
     std::size_t end_assembly,
@@ -64,7 +63,9 @@ ThreadGraph build_worker(
 
     ThreadGraph graph;
     graph.kmers.reserve(n_kmers_est);
-    graph.ids_by_assembly.resize(end_assembly - start_assembly);
+    graph.record_offsets.reserve(end_assembly - start_assembly + 1);
+    graph.record_offsets.push_back(0);
+    graph.ids_by_assembly.reserve(end_assembly - start_assembly);
     graph.start_assembly = start_assembly;
 
     std::vector<std::uint64_t> hashes; // Used to build ThreadGraph.idx
@@ -78,31 +79,29 @@ ThreadGraph build_worker(
     edge_map.reserve(n_map_entries_est);
 
     for (std::size_t assembly_i = start_assembly; assembly_i < end_assembly; ++assembly_i) {
-        const auto assembly_idx = assembly_indices[assembly_i];
         const bool is_target = is_targets[assembly_i];
 
         auto records = seqwin::read_fasta(assembly_paths[assembly_i]);
-        auto& record_ids = graph.ids_by_assembly[assembly_i - start_assembly];
-        record_ids.resize(records.size());
-
-        for (std::size_t record_idx = 0; record_idx < records.size(); ++record_idx) {
-            if (record_idx > std::numeric_limits<std::uint16_t>::max()) {
-                throw std::runtime_error("Number of FASTA records exceeds uint16 range");
+        std::vector<std::string> record_ids;
+        record_ids.reserve(records.size());
+
+        for (std::size_t record_i = 0; record_i < records.size(); ++record_i) {
+            const auto record_idx = graph.record_offsets.back() + record_i;
+            auto& record = records[record_i];
+            if (record.sequence.size() > std::numeric_limits<std::uint32_t>::max()) {
+                throw std::runtime_error(
+                    "Sequence length exceeds uint32 range for record " +
+                    record.id + " in assembly " + assembly_paths[assembly_i]);
             }
-            auto& record = records[record_idx];
-            record_ids[record_idx] = std::move(record.id);
+            record_ids.push_back(std::move(record.id));
 
             // Generate minimizers for the current record
             const auto mins = btllib::minimize_sequence(record.sequence, kmerlen, windowsize);
 
             for (const auto& m : mins) {
-                if (m.pos > std::numeric_limits<std::uint32_t>::max()) {
-                    throw std::runtime_error("Minimizer position exceeds uint32 range");
-                }
                 graph.kmers.push_back(Kmer{
                     static_cast<std::uint32_t>(m.pos),
-                    static_cast<std::uint16_t>(record_idx),
-                    static_cast<std::uint16_t>(assembly_idx)
+                    static_cast<std::uint32_t>(record_idx)
                 });
                 hashes.push_back(m.out_hash);
 
@@ -140,6 +139,8 @@ ThreadGraph build_worker(
                 }
             }
         }
+        graph.record_offsets.push_back(graph.record_offsets.back() + records.size());
+        graph.ids_by_assembly.push_back(std::move(record_ids));
     }
 
     // Materialize edges first to reduce peak memory
@@ -164,6 +165,7 @@ ThreadGraph build_worker(
         auto node_it = node_map.find(hash);
         graph.idx[node_it->second.cursor++] = static_cast<std::uint64_t>(i);
     }
+    std::vector<std::uint64_t>().swap(hashes);
 
     graph.nodes = NoInitArray<Node>(node_map.size());
     std::size_t node_i = 0;
@@ -188,19 +190,14 @@ Graph build(
     const std::vector<std::string>& assembly_paths,
     std::size_t kmerlen,
     std::size_t windowsize,
-    const std::vector<std::size_t>& assembly_indices,
     const std::vector<bool>& is_targets,
     std::size_t n_cpu
 ) {
-    if (assembly_paths.size() != assembly_indices.size() ||
-        assembly_paths.size() != is_targets.size()) {
-        throw std::runtime_error(
-            "assembly_paths, assembly_indices, and is_targets must have the same length");
+    if (assembly_paths.size() != is_targets.size()) {
+        throw std::runtime_error("assembly_paths and is_targets must have the same length");
     }
-    for (std::size_t i = 0; i < assembly_indices.size(); ++i) {
-        if (assembly_indices[i] > std::numeric_limits<std::uint16_t>::max()) {
-            throw std::runtime_error("Assembly index exceeds uint16 range");
-        }
+    if (assembly_paths.size() > std::numeric_limits<std::uint32_t>::max()) {
+        throw std::runtime_error("Number of input assemblies exceeds uint32 range");
     }
 
     const auto n_assemblies = assembly_paths.size();
@@ -223,7 +220,6 @@ Graph build(
                 assembly_paths,
                 kmerlen,
                 windowsize,
-                assembly_indices,
                 is_targets,
                 chunk_start,
                 chunk_end,