diff --git a/include/rabitqlib/index/ivf/cluster.hpp b/include/rabitqlib/index/ivf/cluster.hpp index e20eae60..e343ec76 100644 --- a/include/rabitqlib/index/ivf/cluster.hpp +++ b/include/rabitqlib/index/ivf/cluster.hpp @@ -9,11 +9,21 @@ namespace rabitqlib::ivf { // Non-owning view of one IVF cluster's codes, reranking data and point IDs. class Cluster { private: + friend class IVF; + size_t num_; // Num of vectors in this cluster char* batch_data_ = nullptr; // RaBitQ code and factors char* ex_data_ = nullptr; // Extra-bit codes/factors or original float32 vectors PID* ids_ = nullptr; // PID of vectors + // The owning index refreshes only views whose storage or count changed. + void reset(size_t num, char* batch_data, char* ex_data, PID* ids) noexcept { + num_ = num; + batch_data_ = batch_data; + ex_data_ = ex_data; + ids_ = ids; + } + public: explicit Cluster(size_t, char*, char*, PID*); Cluster(const Cluster&) = default; diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index 54ee3323..fa1a3714 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include @@ -392,19 +393,32 @@ void IVF::add( ); } - std::vector> rows_of(num_clusters_); - for (size_t i = 0; i < n; ++i) { - rows_of[assigned[i]].push_back(static_cast(i)); - } - std::vector new_sizes(num_clusters_); - std::vector touched; - for (size_t i = 0; i < num_clusters_; ++i) { - new_sizes[i] = clusters_[i].num() + rows_of[i].size(); - if (!rows_of[i].empty()) { - touched.push_back(i); + struct ClusterAddition { + PID cluster_id; + std::vector rows; + size_t new_size; + std::optional staged{}; + ByteStorage tail{}; // Repacked tail for in-place growth. + }; + std::vector additions; + { + // Group only incoming rows, preserving their order within each cluster. + // A small append must not allocate or scan num_clusters_ entries. + std::unordered_map positions; + for (size_t i = 0; i < n; ++i) { + const auto [position, inserted] = + positions.try_emplace(assigned[i], additions.size()); + if (inserted) { + additions.push_back({assigned[i], {}, clusters_[assigned[i]].num()}); + } + auto& addition = additions[position->second]; + addition.rows.push_back(static_cast(i)); + ++addition.new_size; } } + const size_t workers = std::min(resolve_num_threads(num_threads), additions.size()); + static_cast(workers); // Consumed by the OpenMP clause below. quant::RabitqConfig config; if (faster) { config = quant::faster_config(padded_dim_, ex_bits_ + 1); @@ -415,42 +429,43 @@ void IVF::add( const PID first_pid = static_cast(max_elements_); /* Rebuild clusters without reserved capacity on the side; the rest grow in place */ - std::vector> staged(touched.size()); - std::vector tails(touched.size()); // repacked tails of in-place clusters - for (size_t t = 0; t < touched.size(); ++t) { - const size_t i = touched[t]; - if (has_capacity(cluster_storage_[i], new_sizes[i])) { - const size_t full_batches = clusters_[i].num() / fastscan::kBatchSize; - tails[t].resize(batch_data_bytes(new_sizes[i]) - (full_batches * batch_bytes)); + for (auto& addition : additions) { + const PID cluster_id = addition.cluster_id; + if (has_capacity(cluster_storage_[cluster_id], addition.new_size)) { + const size_t full_batches = clusters_[cluster_id].num() / fastscan::kBatchSize; + addition.tail.resize( + batch_data_bytes(addition.new_size) - (full_batches * batch_bytes) + ); } else { - staged[t].emplace(raw_reranking_); - resize_storage(*staged[t], new_sizes[i]); + addition.staged.emplace(raw_reranking_); + resize_storage(*addition.staged, addition.new_size); } } // Growing within capacity cannot reallocate or throw, so old cluster views stay valid - for (size_t t = 0; t < touched.size(); ++t) { - if (!staged[t]) { - resize_storage(cluster_storage_[touched[t]], new_sizes[touched[t]]); + for (const auto& addition : additions) { + if (!addition.staged) { + resize_storage(cluster_storage_[addition.cluster_id], addition.new_size); } } // An exception must not leave an OpenMP region, so keep the first one and rethrow it. index_detail::ParallelException failure; -#pragma omp parallel for schedule(dynamic) num_threads(resolve_num_threads(num_threads)) - for (std::ptrdiff_t index = 0; index < static_cast(touched.size()); +#pragma omp parallel for schedule(dynamic) num_threads(workers) + for (std::ptrdiff_t index = 0; index < static_cast(additions.size()); ++index) { - const auto t = static_cast(index); - const size_t i = touched[t]; + auto& addition = additions[static_cast(index)]; + const PID cluster_id = addition.cluster_id; try { - ClusterStorage& cs = staged[t] ? *staged[t] : cluster_storage_[i]; - const size_t full_batches = clusters_[i].num() / fastscan::kBatchSize; - char* tail = staged[t] ? cs.batch_data() + (full_batches * batch_bytes) - : reinterpret_cast(tails[t].data()); + ClusterStorage& cs = + addition.staged ? *addition.staged : cluster_storage_[cluster_id]; + const size_t full_batches = clusters_[cluster_id].num() / fastscan::kBatchSize; + char* tail = addition.staged ? cs.batch_data() + (full_batches * batch_bytes) + : reinterpret_cast(addition.tail.data()); extend_cluster( - clusters_[i], - rows_of[i], + clusters_[cluster_id], + addition.rows, data, - initializer_->centroid(static_cast(i)), + initializer_->centroid(cluster_id), first_pid, cs.batch_data(), rerank_bytes > 0 ? cs.ex_data() : nullptr, @@ -462,31 +477,37 @@ void IVF::add( } if (failure.failed()) { // In-place clusters were only written past their old size; shrinking drops that - for (size_t t = 0; t < touched.size(); ++t) { - if (!staged[t]) { - resize_storage(cluster_storage_[touched[t]], clusters_[touched[t]].num()); + for (const auto& addition : additions) { + if (!addition.staged) { + const PID cluster_id = addition.cluster_id; + resize_storage(cluster_storage_[cluster_id], clusters_[cluster_id].num()); } } failure.rethrow(); } /* Commit. Nothing below can throw. */ - for (size_t t = 0; t < touched.size(); ++t) { - const size_t i = touched[t]; - if (staged[t]) { - cluster_storage_[i] = std::move(*staged[t]); + for (auto& addition : additions) { + const PID cluster_id = addition.cluster_id; + auto& cs = cluster_storage_[cluster_id]; + if (addition.staged) { + cs = std::move(*addition.staged); } else { - const size_t full_batches = clusters_[i].num() / fastscan::kBatchSize; + const size_t full_batches = clusters_[cluster_id].num() / fastscan::kBatchSize; std::copy_n( - tails[t].data(), - tails[t].size(), - cluster_storage_[i].batch_storage.data() + (full_batches * batch_bytes) + addition.tail.data(), + addition.tail.size(), + cs.batch_storage.data() + (full_batches * batch_bytes) ); } + clusters_[cluster_id].reset( + addition.new_size, + cs.batch_data(), + rerank_bytes > 0 ? cs.ex_data() : nullptr, + cs.ids() + ); } max_elements_ = new_num_points; - clusters_.clear(); - init_clusters(new_sizes); } /** diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index da258836..a79862c9 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -725,6 +725,49 @@ TEST(IvfAddTest, AddingToOneClusterLeavesOtherClustersUntouched) { } } +TEST(IvfAddTest, SparseInterleavedClustersPreserveRowOrderAndMatchRebuild) { + constexpr size_t kOld = 3, kCount = 9, kClusters = 257; + const auto points = random_rows(kCount, 744); + const auto centroids = random_rows(kClusters, 745); + const std::array assigned{0, 128, 256, 256, 0, 256, 128, 0, 256}; + const std::string path = ::testing::TempDir() + "rabitq_ivf_sparse_add.index"; + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {1U, 4U, 32U}) { + IVF index(kOld, kDynDim, kClusters, bits, metric); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + index.reserve(256, 32); + const Cluster reserved = IvfTestAccess::cluster_view(index, 256); + index.add( + points.data() + kOld * kDynDim, + kCount - kOld, + assigned.data() + kOld, + false, + 4 + ); + EXPECT_EQ(IvfTestAccess::cluster_view(index, 256).ids(), reserved.ids()); + for (size_t i = 0; i < kClusters; ++i) { + std::vector expected; + for (size_t row = 0; row < kCount; ++row) { + if (assigned[row] == i) { + expected.push_back(static_cast(row)); + } + } + const Cluster cluster = IvfTestAccess::cluster_view(index, i); + ASSERT_EQ(cluster.num(), expected.size()); + for (size_t row = 0; row < expected.size(); ++row) { + EXPECT_EQ(cluster.ids()[row], expected[row]); + } + } + const auto expected = saved_bytes(index, path); + IVF rebuilt; + rebuilt.load(path.c_str()); + rebuilt.construct(points.data(), centroids.data(), assigned.data(), false, 1); + EXPECT_EQ(saved_bytes(rebuilt, path), expected); + } + } + std::remove(path.c_str()); +} + TEST(IvfAddTest, UnreservedClustersGrowToExactlyTheirNewSize) { constexpr size_t kOld = 40; const auto points = random_rows(kOld + 1, 851);