From 1f815b9d1b96821f4b903936ab3f98eeb2349b01 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Wed, 21 Oct 2026 13:53:02 +0800 Subject: [PATCH] Limit IVF append metadata work to affected clusters Group incoming IVF rows by touched cluster and keep each cluster's staging state together. Update only changed views and bound workers to the touched cluster count. Cover sparse interleaved appends and reserved storage with a rebuild-equivalence regression. --- include/rabitqlib/index/ivf/cluster.hpp | 10 +++ src/index/ivf.cpp | 113 ++++++++++++++---------- tests/unit/rabitqlib/index/ivf_test.cpp | 43 +++++++++ 3 files changed, 120 insertions(+), 46 deletions(-) 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);