Widen CSR nnz-offset to 64-bit (offset_t) so one segment can hold >2.1B nnz

benchmarks/batched_build_mem_bench.cpp

@@ -216,10 +216,10 @@ void streaming_add(nsparse::Index* index, const std::string& path) {
         const int64_t bnnz = indptr64[row_end] - indptr64[row_start];
         const int64_t boff = indptr64[row_start];
 
-        std::vector<nsparse::idx_t> bindptr(brows + 1);
+        std::vector<nsparse::offset_t> bindptr(brows + 1);
         for (int64_t i = 0; i <= brows; ++i) {
             bindptr[i] =
-                static_cast<nsparse::idx_t>(indptr64[row_start + i] - boff);
+                static_cast<nsparse::offset_t>(indptr64[row_start + i] - boff);
         }
         std::vector<nsparse::term_t> bindices(bnnz);
         {

benchmarks/index_build_benchmark.cpp

@@ -48,7 +48,7 @@ struct CSRMatrix {
     int64_t nrow;
     int64_t ncol;
     int64_t nnz;
-    std::vector<nsparse::idx_t> indptr;
+    std::vector<nsparse::offset_t> indptr;
     std::vector<nsparse::term_t> indices;
     std::vector<float> data;
 };
@@ -71,7 +71,7 @@ CSRMatrix read_csr(const std::string& path) {
            static_cast<std::streamsize>((m.nrow + 1) * sizeof(int64_t)));
     m.indptr.resize(m.nrow + 1);
     for (int64_t i = 0; i <= m.nrow; ++i) {
-        m.indptr[i] = static_cast<nsparse::idx_t>(indptr64[i]);
+        m.indptr[i] = static_cast<nsparse::offset_t>(indptr64[i]);
     }
 
     std::vector<int32_t> indices32(m.nnz);
@@ -112,7 +112,7 @@ const CSRMatrix& shared_data() {
 }
 
 // Seismic cluster parameters comparable to the search benchmark's index.
-constexpr nsparse::SeismicClusterParameters kParams = {
+const nsparse::SeismicClusterParameters kParams = {
     .lambda = 6000, .beta = 400, .alpha = 0.4F};
 
 // Builds a fresh SeismicIndex from the shared corpus and times only build().

benchmarks/index_search_benchmark.cpp

@@ -29,7 +29,7 @@ struct CSRMatrix {
     int64_t nrow;
     int64_t ncol;
     int64_t nnz;
-    std::vector<nsparse::idx_t> indptr;
+    std::vector<nsparse::offset_t> indptr;
     std::vector<nsparse::term_t> indices;
     std::vector<float> data;
 };
@@ -52,7 +52,7 @@ CSRMatrix read_csr(const std::string& path) {
            static_cast<std::streamsize>((m.nrow + 1) * sizeof(int64_t)));
     m.indptr.resize(m.nrow + 1);
     for (int64_t i = 0; i <= m.nrow; ++i) {
-        m.indptr[i] = static_cast<nsparse::idx_t>(indptr64[i]);
+        m.indptr[i] = static_cast<nsparse::offset_t>(indptr64[i]);
     }
 
     std::vector<int32_t> indices32(m.nnz);

benchmarks/sq_residency_bench.cpp

@@ -48,7 +48,7 @@ struct CSRMatrix {
     int64_t nrow = 0;
     int64_t ncol = 0;
     int64_t nnz = 0;
-    std::vector<nsparse::idx_t> indptr;
+    std::vector<nsparse::offset_t> indptr;
     std::vector<nsparse::term_t> indices;
     std::vector<float> data;
 };
@@ -70,7 +70,7 @@ CSRMatrix read_csr(const std::string& path) {
            static_cast<std::streamsize>((m.nrow + 1) * sizeof(int64_t)));
     m.indptr.resize(m.nrow + 1);
     for (int64_t i = 0; i <= m.nrow; ++i) {
-        m.indptr[i] = static_cast<nsparse::idx_t>(indptr64[i]);
+        m.indptr[i] = static_cast<nsparse::offset_t>(indptr64[i]);
     }
 
     std::vector<int32_t> indices32(m.nnz);
@@ -266,9 +266,9 @@ int do_search(int argc, char** argv) {
     std::vector<double> per_query_ms;
     per_query_ms.reserve(static_cast<size_t>(n_queries));
     for (int qi = 0; qi < n_queries; ++qi) {
-        const nsparse::idx_t start = query.indptr[qi];
-        const nsparse::idx_t end = query.indptr[qi + 1];
-        std::vector<nsparse::idx_t> q_indptr = {0, end - start};
+        const nsparse::offset_t start = query.indptr[qi];
+        const nsparse::offset_t end = query.indptr[qi + 1];
+        std::vector<nsparse::offset_t> q_indptr = {0, end - start};
         const auto t0 = std::chrono::steady_clock::now();
         index->search(1, q_indptr.data(), query.indices.data() + start,
                       query.data.data() + start, k, distances.data(),

nsparse/brutal_index.cpp

@@ -24,7 +24,7 @@ namespace nsparse {
 
 BrutalIndex::BrutalIndex(int dim) : Index(dim) {}
 
-void BrutalIndex::add(idx_t n, const idx_t* indptr, const term_t* indices,
+void BrutalIndex::add(idx_t n, const offset_t* indptr, const term_t* indices,
                       const float* values) {
     throw_if_not_positive(n);
     throw_if_any_null(indptr, indices, values);
@@ -42,7 +42,7 @@ void BrutalIndex::add(idx_t n, const idx_t* indptr, const term_t* indices,
                           nnz * element_size);
 }
 
-auto BrutalIndex::search(idx_t n, const idx_t* indptr, const term_t* indices,
+auto BrutalIndex::search(idx_t n, const offset_t* indptr, const term_t* indices,
                          const float* values, int k,
                          SearchParameters* search_parameters)
     -> pair_of_score_id_vectors_t {
@@ -84,7 +84,7 @@ auto BrutalIndex::single_query(const std::vector<float>& dense, int k)
     const auto& [indptr, indices, values] = vectors_->get_all_data();
 
     for (size_t i = 0; i < num_docs; ++i) {
-        const idx_t start = indptr[i];
+        const offset_t start = indptr[i];
         const size_t len = indptr[i + 1] - start;
         float score = detail::dot_product_float_dense(
             indices + start, values + start, len, dense.data());

nsparse/brutal_index.h

@@ -25,14 +25,14 @@ class BrutalIndex : public Index {
 
     BrutalIndex(const BrutalIndex&) = delete;
     BrutalIndex& operator=(const BrutalIndex&) = delete;
-    void add(idx_t n, const idx_t* indptr, const term_t* indices,
+    void add(idx_t n, const offset_t* indptr, const term_t* indices,
              const float* values) override;
 
     std::array<char, 4> id() const override { return name; }
     static constexpr std::array<char, 4> name = {'B', 'R', 'U', 'T'};
 
 protected:
-    auto search(idx_t n, const idx_t* indptr, const term_t* indices,
+    auto search(idx_t n, const offset_t* indptr, const term_t* indices,
                 const float* values, int k,
                 SearchParameters* search_parameters = nullptr)
         -> pair_of_score_id_vectors_t override;

nsparse/cluster/inverted_list_clusters.cpp

@@ -93,9 +93,9 @@ SparseVectors summarize_with_cpu_(const SparseVectors* vectors,
         auto doc_ids = std::span<const idx_t>(
             group_of_doc_ids.data() + offsets[i], offsets[i + 1] - offsets[i]);
         for (const auto& doc_id : doc_ids) {
-            int start = indptr_data[doc_id];
-            int end = indptr_data[doc_id + 1];
-            for (size_t j = start; j < end; ++j) {
+            offset_t start = indptr_data[doc_id];
+            offset_t end = indptr_data[doc_id + 1];
+            for (offset_t j = start; j < end; ++j) {
                 const term_t term = indices_data[j];
                 // j is element index, need byte offset for T access
                 const T v =
@@ -280,9 +280,9 @@ void InvertedListClusters::build_transpose(const SparseVectors& summaries) {
     std::vector<uint8_t> csc_value(nnz * esz);
     std::vector<idx_t> cursor(term_ptr.begin(), term_ptr.end() - 1);
     for (size_t cluster = 0; cluster < n_clusters_; ++cluster) {
-        const idx_t start = indptr[cluster];
-        const idx_t end = indptr[cluster + 1];
-        for (idx_t j = start; j < end; ++j) {
+        const offset_t start = indptr[cluster];
+        const offset_t end = indptr[cluster + 1];
+        for (offset_t j = start; j < end; ++j) {
             const size_t col = term_column(indices[j]);
             const idx_t pos = cursor[col]++;
             csc_cluster[pos] = static_cast<cluster_id_t>(cluster);

nsparse/cluster/kmeans_utils.cpp

@@ -52,7 +52,7 @@ template <class T>
 CentroidIndex<T> build_centroid_index(
     const SparseVectors* vectors,
     const std::vector<std::vector<idx_t>>& clusters) {
-    const idx_t* indptr = vectors->indptr_data();
+    const offset_t* indptr = vectors->indptr_data();
     const term_t* indices = vectors->indices_data();
     const T* values = vectors->typed_values_data<T>();
     const size_t n_clusters = clusters.size();
@@ -63,7 +63,7 @@ CentroidIndex<T> build_centroid_index(
     size_t nnz = 0;
     for (const auto& cluster : clusters) {
         const idx_t centroid = cluster.at(0);
-        for (idx_t j = indptr[centroid]; j < indptr[centroid + 1]; ++j) {
+        for (offset_t j = indptr[centroid]; j < indptr[centroid + 1]; ++j) {
             max_term = std::max<size_t>(max_term, indices[j]);
             ++nnz;
         }
@@ -74,7 +74,7 @@ CentroidIndex<T> build_centroid_index(
     index.term_ptr.assign(index.n_cols + 1, 0);
     for (const auto& cluster : clusters) {
         const idx_t centroid = cluster.at(0);
-        for (idx_t j = indptr[centroid]; j < indptr[centroid + 1]; ++j) {
+        for (offset_t j = indptr[centroid]; j < indptr[centroid + 1]; ++j) {
             index.term_ptr[indices[j] + 1]++;
         }
     }
@@ -87,7 +87,7 @@ CentroidIndex<T> build_centroid_index(
     std::vector<idx_t> cursor(index.term_ptr.begin(), index.term_ptr.end() - 1);
     for (size_t c = 0; c < n_clusters; ++c) {
         const idx_t centroid = clusters[c].at(0);
-        for (idx_t j = indptr[centroid]; j < indptr[centroid + 1]; ++j) {
+        for (offset_t j = indptr[centroid]; j < indptr[centroid + 1]; ++j) {
             const idx_t pos = cursor[indices[j]]++;
             index.cluster[pos] = static_cast<local_cluster_id_t>(c);
             index.weight[pos] = values[j];
@@ -108,7 +108,7 @@ template <class T>
 void map_docs_to_clusters_typed(const SparseVectors* vectors,
                                 const std::vector<idx_t>& docs,
                                 std::vector<std::vector<idx_t>>& clusters) {
-    const idx_t* indptr = vectors->indptr_data();
+    const offset_t* indptr = vectors->indptr_data();
     const term_t* indices = vectors->indices_data();
     const T* values = vectors->typed_values_data<T>();
     const size_t n_clusters = clusters.size();
@@ -131,7 +131,7 @@ void map_docs_to_clusters_typed(const SparseVectors* vectors,
             continue;
         }
         std::ranges::fill(similarities, acc_t(0));
-        for (idx_t j = indptr[doc_id]; j < indptr[doc_id + 1]; ++j) {
+        for (offset_t j = indptr[doc_id]; j < indptr[doc_id + 1]; ++j) {
             const size_t term = indices[j];
             if (term >= index.n_cols) {
                 continue;  // no centroid carries this term

nsparse/disk_seismic_index_base.cpp

@@ -42,7 +42,7 @@ DiskSeismicIndexBase::DiskSeismicIndexBase(int dim,
                                            SeismicClusterParameters parameter)
     : MmapIndex(dim), cluster_parameter_(parameter) {}
 
-void DiskSeismicIndexBase::add(idx_t n, const idx_t* indptr,
+void DiskSeismicIndexBase::add(idx_t n, const offset_t* indptr,
                                const term_t* indices, const float* values) {
     throw_if_not_positive(n);
     throw_if_any_null(indptr, indices, values);
@@ -70,7 +70,7 @@ void DiskSeismicIndexBase::build() {
         &batch_spill_);
 }
 
-auto DiskSeismicIndexBase::search(idx_t n, const idx_t* indptr,
+auto DiskSeismicIndexBase::search(idx_t n, const offset_t* indptr,
                                   const term_t* indices, const float* values,
                                   int k, SearchParameters* search_parameters)
     -> pair_of_score_id_vectors_t {
@@ -174,7 +174,7 @@ auto DiskSeismicIndexBase::search(idx_t n, const idx_t* indptr,
 
 #pragma omp for schedule(dynamic, 64)
         for (idx_t query_idx = 0; query_idx < n; ++query_idx) {
-            const idx_t start = indptr[query_idx];
+            const offset_t start = indptr[query_idx];
             const size_t len = indptr[query_idx + 1] - start;
             const term_t* query_indices = indices + start;
             const uint8_t* query_codes =
@@ -251,17 +251,17 @@ void DiskSeismicIndexBase::write_doc_directory(
     SparseVectors remainder(
         {.element_size = element_size,
          .dimension = static_cast<size_t>(get_dimension())});
-    const idx_t* indptr = vectors.indptr_data();
+    const offset_t* indptr = vectors.indptr_data();
     const term_t* indices = vectors.indices_data();
     const uint8_t* values = vectors.values_data();
     uint32_t remainder_row = 0;
     for (size_t doc_id = 0; doc_id < num_docs; ++doc_id) {
         if (covered[doc_id]) {
             continue;
         }
-        const idx_t start = indptr[doc_id];
+        const offset_t start = indptr[doc_id];
         const size_t nnz = static_cast<size_t>(indptr[doc_id + 1] - start);
-        const idx_t row_indptr[2] = {0, static_cast<idx_t>(nnz)};
+        const offset_t row_indptr[2] = {0, static_cast<offset_t>(nnz)};
         remainder.add_vectors(
             row_indptr, 2, indices + start, nnz,
             values + static_cast<size_t>(start) * element_size,
@@ -334,8 +334,8 @@ auto DiskSeismicIndexBase::get_doc(idx_t doc_id, size_t element_size) const
             throw std::runtime_error(
                 "DiskSeismic exact match: remainder locator out of range");
         }
-        const idx_t* r_indptr = remainder_.indptr_data();
-        const idx_t r_start = r_indptr[loc.block];
+        const offset_t* r_indptr = remainder_.indptr_data();
+        const offset_t r_start = r_indptr[loc.block];
         return {remainder_.indices_data() + r_start,
                 remainder_.values_data() +
                     static_cast<size_t>(r_start) * element_size,
@@ -352,7 +352,7 @@ auto DiskSeismicIndexBase::get_doc(idx_t doc_id, size_t element_size) const
 }
 
 auto DiskSeismicIndexBase::exact_match_directory(
-    idx_t n, const idx_t* indptr, const term_t* indices, const float* values,
+    idx_t n, const offset_t* indptr, const term_t* indices, const float* values,
     int k, const IDSelectorEnumerable& selector,
     const SearchParameters* search_parameters) const
     -> pair_of_score_id_vectors_t {
@@ -374,7 +374,7 @@ auto DiskSeismicIndexBase::exact_match_directory(
         std::vector<uint8_t> dense(dense_bytes, 0);
 #pragma omp for schedule(dynamic, 64)
         for (idx_t query_idx = 0; query_idx < n; ++query_idx) {
-            const idx_t start = indptr[query_idx];
+            const offset_t start = indptr[query_idx];
             const size_t len =
                 static_cast<size_t>(indptr[query_idx + 1] - start);
             const term_t* q_indices = indices + start;
@@ -396,8 +396,8 @@ auto DiskSeismicIndexBase::exact_match_directory(
                 const DocSlice doc = get_doc(doc_id, element_size);
                 // Dot the doc's slice against the dense query via a 2-entry
                 // indptr.
-                const idx_t slice_indptr[2] = {0,
-                                               static_cast<idx_t>(doc.nnz)};
+                const offset_t slice_indptr[2] = {
+                    0, static_cast<offset_t>(doc.nnz)};
                 const float score = detail::compute_similarity(
                     0, slice_indptr, doc.comps, doc.vals, dense.data(),
                     element_size);

nsparse/disk_seismic_index_base.h

@@ -56,7 +56,7 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO {
     // Persisted, since a mapped index has no in-RAM vectors_ to derive it from.
     size_t num_vectors() const override { return num_vectors_; }
 
-    void add(idx_t n, const idx_t* indptr, const term_t* indices,
+    void add(idx_t n, const offset_t* indptr, const term_t* indices,
              const float* values) override;
     void build() override;
 
@@ -123,7 +123,7 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO {
     detail::InlineForwardIndex fwd_;
 
 private:
-    auto search(idx_t n, const idx_t* indptr, const term_t* indices,
+    auto search(idx_t n, const offset_t* indptr, const term_t* indices,
                 const float* values, int k,
                 SearchParameters* search_parameters = nullptr)
         -> pair_of_score_id_vectors_t override;
@@ -152,7 +152,7 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO {
     // Scores every selected doc directly through the doc-locator directory, for
     // a mapped index. Requires doc_locators_ populated.
     [[nodiscard]] auto exact_match_directory(
-        idx_t n, const idx_t* indptr, const term_t* indices,
+        idx_t n, const offset_t* indptr, const term_t* indices,
         const float* values, int k, const IDSelectorEnumerable& selector,
         const SearchParameters* search_parameters) const
         -> pair_of_score_id_vectors_t;

nsparse/disk_seismic_search.cpp

@@ -77,11 +77,11 @@ void score_block(const InlineForwardIndex* fwd, const SparseVectors* vectors,
                       element_size, id_selector, heap, visited);
         }
     } else if (vectors != nullptr) {
-        const idx_t* const indptr = vectors->indptr_data();
+        const offset_t* const indptr = vectors->indptr_data();
         const term_t* const indices = vectors->indices_data();
         const uint8_t* const values = vectors->values_data();
         for (const idx_t doc_id : clusters[pl].get_docs(cid)) {
-            const idx_t start = indptr[doc_id];
+            const offset_t start = indptr[doc_id];
             const size_t len = indptr[doc_id + 1] - start;
             score_doc(doc_id, indices + start,
                       values + static_cast<size_t>(start) * element_size, len,

nsparse/gpu/gpu_cluster_assigner.cu

@@ -183,7 +183,7 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors,
         return;
     }
 
-    const idx_t* indptr = vectors->indptr_data();
+    const offset_t* indptr = vectors->indptr_data();
     const size_t dim = vectors->get_dimension();
 
     // Centroids are clusters[j].front(); collect them and record which input

nsparse/gpu/gpu_common.cuh

@@ -89,17 +89,29 @@ public:
     // are a cheap identity check.
     const DeviceCorpus& ensure_resident(const SparseVectors* vectors) {
         const size_t n_vectors = vectors->num_vectors();
-        const idx_t* indptr = vectors->indptr_data();
+        const offset_t* indptr = vectors->indptr_data();
         const term_t* indices = vectors->indices_data();
         const float* values = vectors->values_data_float();
         const int64_t nnz = indptr[n_vectors];
 
+        // The device stores CSR offsets as int32 (and cuSPARSE's CSR API is
+        // 32-bit), so this path cannot represent nnz > INT32_MAX; rebuild such a
+        // corpus with the CPU path.
+        if (nnz > INT32_MAX) {
+            throw std::runtime_error(
+                "GPU build path requires nnz <= INT32_MAX; use the CPU path");
+        }
+
         std::lock_guard<std::mutex> lock(mutex_);
         if (corpus_.matches(vectors, n_vectors, nnz)) {
             return corpus_;
         }
         free_locked();
 
+        std::vector<int32_t> indptr32(n_vectors + 1);
+        for (size_t i = 0; i <= n_vectors; ++i) {
+            indptr32[i] = static_cast<int32_t>(indptr[i]);
+        }
         std::vector<int32_t> indices32(static_cast<size_t>(nnz));
         for (int64_t i = 0; i < nnz; ++i) {
             indices32[i] = static_cast<int32_t>(indices[i]);
@@ -110,7 +122,7 @@ public:
                    "cudaMalloc(corpus.indices)");
         check_cuda(cudaMalloc(&corpus_.values, nnz * sizeof(float)),
                    "cudaMalloc(corpus.values)");
-        check_cuda(cudaMemcpy(corpus_.indptr, indptr,
+        check_cuda(cudaMemcpy(corpus_.indptr, indptr32.data(),
                               (n_vectors + 1) * sizeof(int32_t),
                               cudaMemcpyHostToDevice),
                    "cudaMemcpy(corpus.indptr)");

nsparse/gpu/gpu_summarizer.cu

@@ -165,7 +165,7 @@ bool summarize_list_impl(const SparseVectors* vectors, const idx_t* docs,
                          const idx_t* offsets, size_t n_clusters,
                          std::vector<GpuSummarizer::ClusterSummary>& out) {
     const size_t dim = vectors->get_dimension();
-    const idx_t* indptr = vectors->indptr_data();
+    const offset_t* indptr = vectors->indptr_data();
 
     const size_t n_docs = static_cast<size_t>(offsets[n_clusters] - offsets[0]);
     if (n_docs == 0) {