diff --git a/cactus-engine/src/index.cpp b/cactus-engine/src/index.cpp index 940a476e4..a97127783 100644 --- a/cactus-engine/src/index.cpp +++ b/cactus-engine/src/index.cpp @@ -146,8 +146,14 @@ namespace index { if (mapped_index_ == MAP_FAILED || mapped_data_ == MAP_FAILED) throw std::runtime_error("Index is in a failed memory-mapped state"); validate_documents(documents); + std::vector> normalized_embeddings; + normalized_embeddings.reserve(documents.size()); size_t added_data_size = 0; for (const auto& doc : documents) { + std::vector<__fp16> normalized_embedding(embedding_dim_); + cactus_fp32_to_fp16(doc.embedding.data(), normalized_embedding.data(), embedding_dim_); + normalize(normalized_embedding.data(), embedding_dim_); + normalized_embeddings.emplace_back(std::move(normalized_embedding)); added_data_size += sizeof(DataEntry) + doc.content.size() + doc.metadata.size(); } @@ -211,12 +217,8 @@ namespace index { memcpy(index_write_pos, &index_entry, sizeof(IndexEntry)); index_write_pos += sizeof(IndexEntry); - std::vector<__fp16> normalized_embedding(embedding_dim_); - cactus_fp32_to_fp16(doc.embedding.data(), normalized_embedding.data(), embedding_dim_); - normalize(normalized_embedding.data(), embedding_dim_); - size_t embedding_bytes = embedding_dim_ * sizeof(__fp16); - memcpy(index_write_pos, normalized_embedding.data(), embedding_bytes); + memcpy(index_write_pos, normalized_embeddings[i].data(), embedding_bytes); index_write_pos += embedding_bytes; doc_id_map_[doc.id] = num_documents_ + static_cast(i); diff --git a/cactus-engine/tests/test_index.cpp b/cactus-engine/tests/test_index.cpp index 27eebcbe5..45b5cf5fe 100644 --- a/cactus-engine/tests/test_index.cpp +++ b/cactus-engine/tests/test_index.cpp @@ -81,6 +81,19 @@ class IndexFixture { return {r, content}; } + bool matches(int id, const char* content, const std::vector& embedding) { + char doc[64], meta[16]; + char* doc_ptr = doc; + char* meta_ptr = meta; + std::vector actual_embedding(dim_); + float* emb_ptr = actual_embedding.data(); + size_t doc_size = sizeof(doc), meta_size = sizeof(meta), emb_size = dim_; + int r = cactus_index_get(idx_, &id, 1, &doc_ptr, &doc_size, &meta_ptr, &meta_size, + &emb_ptr, &emb_size); + return r == 0 && std::strcmp(doc, content) == 0 && std::strcmp(meta, "meta") == 0 && + emb_size == embedding.size() && actual_embedding == embedding; + } + int del(int id) { return cactus_index_delete(idx_, &id, 1); } int del(const std::vector& ids) { return cactus_index_delete(idx_, ids.data(), ids.size()); } int compact() { return cactus_index_compact(idx_); } @@ -150,6 +163,53 @@ bool test_batch() { return true; } +bool test_rejected_batch(bool zero_first) { + IndexFixture f(zero_first ? "test_batch_zero_first" : "test_batch_valid_first", 4); + if (!f.init()) return false; + + const std::vector existing = {1, 0, 0, 0}; + const std::vector valid = {0, 1, 0, 0}; + const std::vector zero(4, 0.0f); + const std::vector later = {0, 0, 1, 0}; + if (f.add(1, "existing", existing) != 0) return false; + + auto read_file = [](const std::string& path) { + std::ifstream file(path, std::ios::binary); + std::ostringstream contents; + contents << file.rdbuf(); + return contents.str(); + }; + const std::string index_before = read_file(f.path() + "/index.bin"); + const std::string data_before = read_file(f.path() + "/data.bin"); + if (index_before.empty() || data_before.empty()) return false; + + int ids[] = {2, 3}; + const char* docs[] = {"rejected-first", "rejected-second"}; + const char* metas[] = {"meta", "meta"}; + const float* embeddings[] = {zero_first ? zero.data() : valid.data(), + zero_first ? valid.data() : zero.data()}; + if (cactus_index_add(f.get_idx(), ids, docs, metas, embeddings, 2, 4) != -1) return false; + if (std::string(cactus_get_last_error()) != "Cannot normalize zero vector") return false; + + if (read_file(f.path() + "/index.bin") != index_before || + read_file(f.path() + "/data.bin") != data_before) return false; + if (!f.matches(1, "existing", existing)) return false; + if (f.get(2).first == 0 || f.get(3).first == 0) return false; + if (f.query(existing, 10) != std::vector{1}) return false; + + if (f.add(4, "later-valid", later) != 0) return false; + auto check_documents = [&]() { + return f.matches(1, "existing", existing) && f.matches(4, "later-valid", later) && + f.get(2).first != 0 && f.get(3).first != 0 && + f.query(existing, 1) == std::vector{1} && + f.query(later, 1) == std::vector{4} && + f.query(existing, 10).size() == 2; + }; + if (!check_documents()) return false; + if (!f.reopen()) return false; + return check_documents(); +} + bool test_query() { IndexFixture f("test_query"); if (!f.init()) return false; @@ -270,6 +330,8 @@ int main() { runner.run_test("crud", test_crud()); runner.run_test("batch", test_batch()); + runner.run_test("batch_reject_valid_first", test_rejected_batch(false)); + runner.run_test("batch_reject_zero_first", test_rejected_batch(true)); runner.run_test("query", test_query()); runner.run_test("compact", test_compact()); runner.run_test("persistence", test_persistence());