diff --git a/nsparse/id_map_index.cpp b/nsparse/id_map_index.cpp index f95d606..509f17f 100644 --- a/nsparse/id_map_index.cpp +++ b/nsparse/id_map_index.cpp @@ -9,7 +9,15 @@ #include "nsparse/id_map_index.h" +#include +#include +#include +#include #include +#include +#include +#include +#include #include "nsparse/id_selector.h" #include "nsparse/io/index_io.h" @@ -78,6 +86,96 @@ void IDMapIndex::add_with_ids(idx_t n, const idx_t* indptr, external_to_internal_[ids[i]] = old_size + i; } } + +std::vector IDMapIndex::read_id_file(const char* id_path) { + check_if_file_valid(id_path, "id map file"); + + std::ifstream in(id_path, std::ios::binary); + if (!in.is_open()) { + throw std::runtime_error(std::string("cannot open id map file: ") + + id_path); + } + + int64_t count = 0; + in.read(reinterpret_cast(&count), sizeof(count)); + if (!in) { + throw std::runtime_error(std::string("truncated id map file: ") + + id_path); + } + if (count < 0) { + throw std::invalid_argument(std::string("negative id map count in: ") + + id_path); + } + const auto map_size = static_cast(count); + + // Guard the byte-size arithmetic against wraparound before relying on it. + if (map_size > (std::numeric_limits::max() - sizeof(int64_t)) / + sizeof(idx_t)) { + throw std::invalid_argument(std::string("id map count is too large: ") + + id_path); + } + + // Reject a truncated or oversized file up front. + const size_t expected_bytes = sizeof(int64_t) + map_size * sizeof(idx_t); + if (std::filesystem::file_size(id_path) != expected_bytes) { + throw std::invalid_argument( + std::string( + "id map file size does not match its count (expected ") + + std::to_string(expected_bytes) + " bytes): " + id_path); + } + + std::vector internal_to_external(map_size); + if (map_size > 0) { + in.read(reinterpret_cast(internal_to_external.data()), + static_cast(map_size * sizeof(idx_t))); + if (!in) { + throw std::runtime_error(std::string("truncated id map file: ") + + id_path); + } + } + return internal_to_external; +} + +void IDMapIndex::read_csr_and_ids(const char* csr_path, const char* id_path, + Residency residency) { + if (delegate_ == nullptr) { + throw std::logic_error("IDMapIndex has no delegate index"); + } + check_if_file_valid(csr_path, "csr file"); + + // Fully validate and load the id file BEFORE ingesting the CSR, so a + // missing/malformed/truncated id file leaves this index untouched. Only the + // count-vs-CSR-row check below can fail once the delegate has ingested; on + // ANY throw from this method the half-built index must be discarded. + std::vector internal_to_external = read_id_file(id_path); + + // The id file is known-good; now ingest the vectors (borrowed from the + // mapping when residency == kMmap). Afterward num_vectors() reflects the + // CSR rows. + delegate_->read_csr(csr_path, residency); + + // The map is row-aligned with the CSR, so its count must equal the vectors + // the delegate just ingested. + const size_t delegate_size = delegate_->num_vectors(); + if (internal_to_external.size() != delegate_size) { + throw std::invalid_argument( + "id map count (" + std::to_string(internal_to_external.size()) + + ") does not match the CSR vector count (" + + std::to_string(delegate_size) + "): " + id_path); + } + + set_id_map(std::move(internal_to_external)); +} + +void IDMapIndex::set_id_map(std::vector&& internal_to_external) { + internal_to_external_ = std::move(internal_to_external); + external_to_internal_.clear(); + external_to_internal_.reserve(internal_to_external_.size()); + for (size_t i = 0; i < internal_to_external_.size(); ++i) { + external_to_internal_[internal_to_external_[i]] = static_cast(i); + } +} + void IDMapIndex::write_index(IOWriter* io_writer) { // Write internal_to_external_ vector size_t map_size = internal_to_external_.size(); @@ -89,23 +187,18 @@ void IDMapIndex::write_index(IOWriter* io_writer) { nsparse::detail::write_index(delegate_.get(), io_writer, true); } -void IDMapIndex::read_index(IOReader* io_reader, const IndexHeader& header, +void IDMapIndex::read_index(IOReader* io_reader, const IndexHeader& /*header*/, int io_flags) { - // Read internal_to_external_ vector + // Read the id map into a local vector, then load the delegate. size_t map_size = 0; io_reader->read(&map_size, sizeof(size_t), 1); - internal_to_external_.resize(map_size); + std::vector internal_to_external(map_size); if (map_size > 0) { - io_reader->read(internal_to_external_.data(), sizeof(idx_t), map_size); + io_reader->read(internal_to_external.data(), sizeof(idx_t), map_size); } delegate_.reset(nsparse::detail::read_index(io_reader, true, io_flags)); - // Rebuild external_to_internal_ from internal_to_external_ - external_to_internal_.clear(); - external_to_internal_.reserve(map_size); - for (size_t i = 0; i < map_size; ++i) { - external_to_internal_[internal_to_external_[i]] = static_cast(i); - } + set_id_map(std::move(internal_to_external)); } } // namespace nsparse \ No newline at end of file diff --git a/nsparse/id_map_index.h b/nsparse/id_map_index.h index 6db9715..d448c42 100644 --- a/nsparse/id_map_index.h +++ b/nsparse/id_map_index.h @@ -101,6 +101,10 @@ class IDMapIndex : public Index, public IndexIO { void add_with_ids(idx_t n, const idx_t* indptr, const term_t* indices, const float* values, const idx_t* ids) override; + + void read_csr_and_ids(const char* csr_path, const char* id_path, + Residency residency = Residency::kInMemory); + [[nodiscard]] uint32_t format_version() const override { return kFormatVersion; } @@ -109,6 +113,10 @@ class IDMapIndex : public Index, public IndexIO { int io_flags = 0) override; private: + std::vector read_id_file(const char* id_path); + + void set_id_map(std::vector&& internal_to_external); + // Owns the wrapped delegate index. Using unique_ptr ensures the delegate is // freed when the IDMapIndex is destroyed (previously a raw pointer with a // defaulted destructor, which leaked the delegate and everything it owned). diff --git a/nsparse/utils/checks.h b/nsparse/utils/checks.h index 150a2b1..b851d50 100644 --- a/nsparse/utils/checks.h +++ b/nsparse/utils/checks.h @@ -11,8 +11,10 @@ #define COMMON_H #include +#include #include #include +#include namespace nsparse { @@ -24,6 +26,16 @@ T* throw_if_null(T* ptr, const char* msg = "unexpected nullptr") { return ptr; } +// Rejects a file path that is null or does not exist, so a caller gets a clear +// error before trying to open it. `what` names the file in the message. +inline void check_if_file_valid(const char* path, const char* what = "file") { + throw_if_null(path, "file path must not be null"); + if (!std::filesystem::exists(path)) { + throw std::invalid_argument(std::string(what) + + " does not exist: " + path); + } +} + template T throw_if_not_positive(T value, const char* msg = "value must be positive") { if (value <= 0) { diff --git a/python_tests/test_seismic_index.py b/python_tests/test_seismic_index.py index 11fa7fd..7fcc5df 100644 --- a/python_tests/test_seismic_index.py +++ b/python_tests/test_seismic_index.py @@ -67,6 +67,51 @@ def test_with_id_map(corpus, queries, oracle, doc_ids): assert recall_at_k(labels, want_external) >= RECALL_FLOOR +def _write_interchange_csr(path, corpus): + """Corpus as an interchange CSR: int64 header {n, dim, nnz}, int64 indptr, + int32 indices, float32 values -- the layout nsparse.convert consumes.""" + with open(path, "wb") as out: + np.array( + [corpus.n, corpus.dim, corpus.indices.size], dtype=np.int64 + ).tofile(out) + corpus.indptr.astype(np.int64).tofile(out) + corpus.indices.astype(np.int32).tofile(out) + corpus.values.astype(np.float32).tofile(out) + + +def _write_id_map(path, external_ids): + """The id-map file read_csr_and_ids reads: int64 count, then int32 ids, + row-aligned with the CSR.""" + with open(path, "wb") as out: + np.array([external_ids.size], dtype=np.int64).tofile(out) + external_ids.astype(np.int32).tofile(out) + + +def test_id_map_from_csr_and_id_files(corpus, queries, oracle, doc_ids, tmp_path): + """read_csr_and_ids builds an idmap from a native CSR (borrowed via mmap) + plus a separate id file -- the memory-saving build path -- and must return + the caller's external ids, matching the in-RAM add_with_ids path.""" + interchange = tmp_path / "corpus.csr" + native = tmp_path / "corpus.mcsr" + id_file = tmp_path / "ids.bin" + _write_interchange_csr(interchange, corpus) + nsparse.convert(str(interchange), str(native)) + _write_id_map(id_file, doc_ids) + + index = nsparse.index_factory(corpus.dim, f"idmap,{SPEC}") + index.read_csr_and_ids(str(native), str(id_file), nsparse.Residency_kMmap) + index.build() + + _, labels = search(index, queries) + returned = labels[labels >= 0] + assert returned.size > 0, "every query should return at least one hit" + assert np.isin(returned, doc_ids).all(), "returned ids must be caller ids" + + want_labels, _ = oracle + want_external = np.where(want_labels >= 0, doc_ids[want_labels], -1) + assert recall_at_k(labels, want_external) >= RECALL_FLOOR + + def test_exact_match(index, queries, oracle): """An enumerable selector of size <= k switches search to the exact path. diff --git a/tests/csr_interchange_test_util.h b/tests/csr_interchange_test_util.h index 034d64f..976181a 100644 --- a/tests/csr_interchange_test_util.h +++ b/tests/csr_interchange_test_util.h @@ -18,6 +18,8 @@ #include #include +#include "nsparse/types.h" + // Shared helpers for the mmap-CSR build path, used by both the regular and the // disk-resident index suites: write a corpus as an interchange CSR (the layout // csr_layout::convert consumes) and manage the interchange + native temp files. @@ -25,8 +27,8 @@ namespace nsparse::csr_test { // Writes a corpus as an interchange CSR: int64 header {rows, num_cols, nnz}, // int64 indptr[rows + 1], int32 indices[nnz], float values[nnz]. Templated on -// the corpus struct (any type exposing .n / .indptr / .indices / .values), so it -// serves any test corpus. The values are written verbatim, so a convert + +// the corpus struct (any type exposing .n / .indptr / .indices / .values), so +// it serves any test corpus. The values are written verbatim, so a convert + // read_csr(kMmap) build sees the exact same vectors as add(). template void write_interchange_csr(const std::string& path, const Corpus& c, @@ -47,12 +49,26 @@ void write_interchange_csr(const std::string& path, const Corpus& c, static_cast(c.values.size() * sizeof(float))); } +// Writes the id-map file that IDMapIndex::read_csr_and_ids reads: +// [int64 count][idx_t external_id x count]. Row-aligned with the CSR, so +// external_ids[i] is the external id of CSR row i. +inline void write_id_map_file(const std::string& path, + const std::vector& external_ids) { + std::ofstream out(path, std::ios::binary); + const int64_t count = static_cast(external_ids.size()); + out.write(reinterpret_cast(&count), sizeof(count)); + out.write( + reinterpret_cast(external_ids.data()), + static_cast(external_ids.size() * sizeof(idx_t))); +} + // An interchange CSR temp file and the native path convert writes it to, both // removed on destruction. class TempCsrFiles { public: explicit TempCsrFiles(const std::string& stem) - : interchange_(std::filesystem::temp_directory_path() / (stem + ".csr")), + : interchange_(std::filesystem::temp_directory_path() / + (stem + ".csr")), native_(std::filesystem::temp_directory_path() / (stem + ".mcsr")) { std::error_code ignored; std::filesystem::remove(interchange_, ignored); diff --git a/tests/id_map_index_test.cpp b/tests/id_map_index_test.cpp index 0adb36e..11c121f 100644 --- a/tests/id_map_index_test.cpp +++ b/tests/id_map_index_test.cpp @@ -11,6 +11,12 @@ #include +#include +#include +#include +#include +#include +#include #include #include "nsparse/id_selector.h" @@ -20,6 +26,8 @@ #include "nsparse/io/index_io.h" #include "nsparse/seismic_index.h" #include "nsparse/types.h" +#include "nsparse/utils/csr_layout.h" +#include "tests/csr_interchange_test_util.h" namespace { @@ -370,3 +378,235 @@ TEST_F(IDMapIndexTest, search_with_not_id_selector) { EXPECT_TRUE(label == 100 || label == 300 || label == -1); } } + +namespace { + +using nsparse::idx_t; +using nsparse::term_t; + +// A CSR corpus exposing the fields csr_test::write_interchange_csr needs +// (.n / .indptr / .indices / .values). +struct Corpus { + idx_t n = 0; + std::vector indptr; + std::vector indices; + std::vector values; +}; + +// Reproducible random corpus: distinct ascending terms per row (CSR +// convention) and values in (0, 1]. +Corpus make_corpus(idx_t rows, int dim, unsigned seed) { + std::mt19937 rng(seed); + std::uniform_int_distribution nnz_dist(3, 10); + std::uniform_int_distribution term_dist(0, dim - 1); + std::uniform_real_distribution val_dist(0.05F, 1.0F); + Corpus c; + c.n = rows; + c.indptr.push_back(0); + for (idx_t r = 0; r < rows; ++r) { + std::set terms; + const int nnz = nnz_dist(rng); + while (static_cast(terms.size()) < nnz) { + terms.insert(term_dist(rng)); + } + for (const int t : terms) { + c.indices.push_back(static_cast(t)); + c.values.push_back(val_dist(rng)); + } + c.indptr.push_back(static_cast(c.indices.size())); + } + return c; +} + +// External ids distinct from the internal row indices, to prove the search +// output is translated through the id map rather than echoing internal ids. +std::vector make_external_ids(idx_t rows) { + std::vector ids(static_cast(rows)); + for (idx_t i = 0; i < rows; ++i) { + ids[static_cast(i)] = 1000 + i * 7; + } + return ids; +} + +// Batch search over the corpus rows as queries; returns per-query {labels, +// scores} for a bit-exact comparison. +std::pair, std::vector> search_corpus( + nsparse::Index& index, const Corpus& queries, int k) { + std::vector labels(static_cast(queries.n) * k, + nsparse::detail::INVALID_IDX); + std::vector distances(static_cast(queries.n) * k, -1.0F); + index.search(queries.n, queries.indptr.data(), queries.indices.data(), + queries.values.data(), k, distances.data(), labels.data()); + return {labels, distances}; +} + +// A temp file removed on destruction (for the id-map file). +class TempIdFile { +public: + explicit TempIdFile(const std::string& stem) + : path_((std::filesystem::temp_directory_path() / stem).string()) { + std::error_code ignored; + std::filesystem::remove(path_, ignored); + } + ~TempIdFile() { + std::error_code ignored; + std::filesystem::remove(path_, ignored); + } + TempIdFile(const TempIdFile&) = delete; + TempIdFile& operator=(const TempIdFile&) = delete; + const std::string& path() const { return path_; } + +private: + std::string path_; +}; + +constexpr int kDim = 100; +const nsparse::SeismicClusterParameters kClusterParams{ + .lambda = 10, .beta = 2, .alpha = 0.5F, .seed = 42}; + +} // namespace + +// The single-function file build (delegate read_csr(kMmap) + id map from a +// file) must produce a search result bit-exact to the in-RAM add_with_ids +// build: same corpus, same external ids, same fixed cluster seed. This is the +// mmap-CSR memory-saving path, and its labels must be the external ids. +TEST(IDMapReadCsrAndId, MatchesAddWithIdsBuild) { + const Corpus corpus = make_corpus(300, kDim, /*seed=*/1); + const Corpus queries = make_corpus(20, kDim, /*seed=*/2); + const std::vector ids = make_external_ids(corpus.n); + constexpr int k = 10; + + // Reference: add_with_ids() then build(). + nsparse::IDMapIndex added(new nsparse::SeismicIndex(kDim, kClusterParams)); + added.add_with_ids(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data(), ids.data()); + added.build(); + const auto expected = search_corpus(added, queries, k); + + // Under test: build the delegate from a native CSR borrowed via mmap and + // read the id map from a file, through the single entry point. + nsparse::csr_test::TempCsrFiles csr("nsparse_idmap_src"); + nsparse::csr_test::write_interchange_csr(csr.interchange(), corpus, kDim); + nsparse::csr_layout::convert(csr.interchange(), csr.native()); + TempIdFile idfile("nsparse_idmap_src.ids"); + nsparse::csr_test::write_id_map_file(idfile.path(), ids); + + nsparse::IDMapIndex mapped(new nsparse::SeismicIndex(kDim, kClusterParams)); + mapped.read_csr_and_ids(csr.native().c_str(), idfile.path().c_str(), + nsparse::Residency::kMmap); + ASSERT_EQ(mapped.num_vectors(), static_cast(corpus.n)); + mapped.build(); + const auto got = search_corpus(mapped, queries, k); + + EXPECT_EQ(got.first, expected.first) << "external-id labels differ"; + ASSERT_EQ(got.second.size(), expected.second.size()); + for (size_t i = 0; i < got.second.size(); ++i) { + EXPECT_FLOAT_EQ(got.second[i], expected.second[i]) << "score at " << i; + } + // Sanity: the labels are external ids, not internal row indices. + bool saw_external = false; + for (const idx_t label : got.first) { + if (label >= 1000) { + saw_external = true; + break; + } + } + EXPECT_TRUE(saw_external) << "labels should be translated to external ids"; +} + +// The id map is row-aligned with the CSR, so a count that disagrees with the +// delegate's vector count is rejected. +TEST(IDMapReadCsrAndId, CountMismatchThrows) { + const Corpus corpus = make_corpus(50, kDim, /*seed=*/3); + nsparse::csr_test::TempCsrFiles csr("nsparse_idmap_mismatch"); + nsparse::csr_test::write_interchange_csr(csr.interchange(), corpus, kDim); + nsparse::csr_layout::convert(csr.interchange(), csr.native()); + + // One too few ids for the 50 CSR rows. + const std::vector short_ids = make_external_ids(corpus.n - 1); + TempIdFile idfile("nsparse_idmap_mismatch.ids"); + nsparse::csr_test::write_id_map_file(idfile.path(), short_ids); + + nsparse::IDMapIndex mapped(new nsparse::SeismicIndex(kDim, kClusterParams)); + EXPECT_THROW( + mapped.read_csr_and_ids(csr.native().c_str(), idfile.path().c_str(), + nsparse::Residency::kMmap), + std::invalid_argument); +} + +// A file whose byte size does not match its declared count is malformed. +TEST(IDMapReadCsrAndId, MalformedFileThrows) { + const Corpus corpus = make_corpus(5, kDim, /*seed=*/4); + nsparse::csr_test::TempCsrFiles csr("nsparse_idmap_malformed"); + nsparse::csr_test::write_interchange_csr(csr.interchange(), corpus, kDim); + nsparse::csr_layout::convert(csr.interchange(), csr.native()); + + // Header claims 5 ids but only 3 follow -> size mismatch. + TempIdFile idfile("nsparse_idmap_malformed.ids"); + { + std::ofstream out(idfile.path(), std::ios::binary); + const int64_t bogus_count = 5; + out.write(reinterpret_cast(&bogus_count), + sizeof(bogus_count)); + const std::vector only_three = {1, 2, 3}; + out.write( + reinterpret_cast(only_three.data()), + static_cast(only_three.size() * sizeof(idx_t))); + } + + nsparse::IDMapIndex mapped(new nsparse::SeismicIndex(kDim, kClusterParams)); + EXPECT_THROW( + mapped.read_csr_and_ids(csr.native().c_str(), idfile.path().c_str(), + nsparse::Residency::kMmap), + std::invalid_argument); +} + +// The reverse map (external -> internal) must be populated by the file build, +// so an id-selector filter over external ids still works. +TEST(IDMapReadCsrAndId, IdSelectorFilterWorksAfterFileBuild) { + const Corpus corpus = make_corpus(200, kDim, /*seed=*/5); + const std::vector ids = make_external_ids(corpus.n); + + nsparse::csr_test::TempCsrFiles csr("nsparse_idmap_filter"); + nsparse::csr_test::write_interchange_csr(csr.interchange(), corpus, kDim); + nsparse::csr_layout::convert(csr.interchange(), csr.native()); + TempIdFile idfile("nsparse_idmap_filter.ids"); + nsparse::csr_test::write_id_map_file(idfile.path(), ids); + + nsparse::IDMapIndex mapped(new nsparse::SeismicIndex(kDim, kClusterParams)); + mapped.read_csr_and_ids(csr.native().c_str(), idfile.path().c_str(), + nsparse::Residency::kMmap); + mapped.build(); + + // Allow only the first two external ids (internal rows 0 and 1). + const std::vector allowed = {ids[0], ids[1]}; + nsparse::SetIDSelector selector(allowed.size(), allowed.data()); + nsparse::SeismicSearchParameters params; + params.set_id_selector(&selector); + + // Query with the corpus rows themselves, so rows 0 and 1 (the allowed docs) + // match their own vectors and are guaranteed to score. For an allowed doc + // to survive the external-id filter and be returned, the reverse map + // (external -> internal) must be populated by the file build -- an empty + // reverse map would filter everything out and still pass a "no disallowed + // leak" check, so this asserts a positive hit too. + constexpr int k = 5; + std::vector labels(static_cast(corpus.n) * k, + nsparse::detail::INVALID_IDX); + std::vector distances(static_cast(corpus.n) * k, -1.0F); + mapped.search(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data(), k, distances.data(), labels.data(), + ¶ms); + + bool saw_allowed = false; + for (const idx_t label : labels) { + EXPECT_TRUE(label == ids[0] || label == ids[1] || + label == nsparse::detail::INVALID_IDX) + << "filter must restrict to allowed external ids, got " << label; + if (label == ids[0] || label == ids[1]) { + saw_allowed = true; + } + } + EXPECT_TRUE(saw_allowed) << "reverse map must be populated: an allowed doc " + "should actually be returned"; +}