Repository navigation
feat(ivf): add balanced k-means clustering for IVF partitioning - #3260
DoubleCore wants to merge 14 commits into
Conversation
Adds BalancedKMeansCluster, a reusable capacity-constrained trainer for IVF partitioning (issue antgroup#2870), plus a new optional ivf_train_type value "balanced_kmeans". Unlike a price-based auction solver, the balancing step uses a regret-greedy capacitated assignment: every vector keeps only its top-C nearest clusters (no dense n x k distance matrix), the n mod k clusters with the highest unconstrained nearest-centroid demand receive ceil(n/k) seats, and vectors are placed most-constrained first - in decreasing order of their regret - each into its nearest cluster that still has a free seat, with an on-demand full rescan for the rare vector whose candidates are all full. Each cluster therefore holds exactly floor(n/k) or ceil(n/k) vectors, no cluster is empty while k <= n, and results are deterministic for a given seed regardless of thread count. - src/impl/cluster/balanced_kmeans_cluster.{h,cpp}: the trainer, its metric-aware fixed-centroid Assign (used by the IVF build-time bucket assignment) and a BalancedKMeansStats observability handle. - IVF integration: ivf_train_type "balanced_kmeans" parses and serializes, unknown trainer names now fail with INVALID_ARGUMENT, GNO-IMI rejects the combination, defaults and existing kmeans/random behavior are unchanged, and the first bucket assignment of IVF::Build balances through a new IVFPartitionStrategy::ClassifyDatasForBuild hook. - Tests: src/impl/cluster/balanced_kmeans_cluster_test.cpp (exact balance for divisible, indivisible, skewed and duplicated data, k == 1, k == n, invalid arguments, cross-thread determinism, Assign), parameter tests for parsing/serialization/rejection, and functional tests in tests/test_ivf.cpp checking bucket balance through GetStats, search, serialization round-trip and error paths. - tools/cluster_balance_bench: seeded reproducible comparison of the two trainers on uniform, skewed, indivisible and duplicated datasets, reporting cluster-size statistics, loss, training time and peak memory against a dense n x k distance matrix. - docs: English and Chinese IVF pages describe the trainer, its complexity and its balance/quality trade-off. Signed-off-by: DoubleCore <851678303@qq.com> Assisted-by: OpenCode:qwen3.8-flash
Merger | Merge ReadinessHead: Updated: 2026-11-06 07:29:49 UTC Status: BLOCKED Merger decision: HOLD (shadow assessment; no merge executed).
Next owner: Author. Required next step: Resolve the GitHub merge conditions, then reassess. Merge reliability: 0.00 / 1.00 (WITHHELD)
Evidence subtotal: 0.67; applicable ceiling: 0.00. Final score = min(subtotal, ceiling). Policy:
AI review: COMPLETE. Static review: patch is correct; CI is reported separately. Review findings: 0; P0/P1 findings in that review: 0; coverage: 24/24 changed files. Findings are not a count of currently unresolved threads. CI summary: 1 neutral, 5 success. Check source: Next action: Merge commit freshness: Waiting for GitHub to refresh the merge commit for the current head and base. Only current head checks are shown.; PR CI Gate: Waiting for required result from the expected source; Human review: GitHub: REVIEW_REQUIRED; approvals >= 1; Code Owner approval required; stale approvals dismissed; Conversations: 44 unresolved; resolution required; GitHub merge state: BLOCKED CI check details
This is an uncalibrated evidence score, not a probability, correctness guarantee or approval. Zero means the assessment is withheld or lacks positive evidence, not a proven-bad patch. Automatic merging is disabled. New commits or base changes invalidate earlier evidence; never authorize a merge from this score alone. |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Reviewer | Review completeAutomated pull request review completed. Model: No actionable inline findings were found. |
vsag-bot
added
kind/feature
Merge Protections🟢 All 3 merge protections satisfied — ready to merge. Show 3 satisfied protections🟢 Require kind label
🟢 Require version label
🟢 Require linked issue for feature/bug PRs
|
|
Jev-driven classification Applied: Option probabilities
Model: |
mergify
Bot
added
area/docs
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Reviewer | Code review
Automated pull request review completed.
Model: gpt-6.1-sol; reasoning effort: high.
Review scope: high (2247 changed lines across 20 files).
Submitted 5 inline comments.
Review coverage checks completed; this does not prove absence of defects.
Static review only; compilation and runtime tests are outside this review's scope.
File coverage: 20/20; observed passing test runs: 0.
Reviewed commit a08dcb7.
| auto buckets = | ||
| partition_strategy_->ClassifyDatas(vectors, num_element, buckets_per_data_, nullptr); | ||
| initial_bucket_assignment | ||
| ? partition_strategy_->ClassifyDatasForBuild( |
There was a problem hiding this comment.
[P2] [P2] Train residual quantizers with the balanced assignment
This changes storage assignment, but BucketDataCell::Train still calls ClassifyDatas, so residual quantizers learn a different distribution. For L2 data [0,0,0,100] with k=2, balanced centroids are 0 and 50. SQ8 trains on residuals [0,0,0,50], while Build must encode one zero as -50. That residual clips to 0, reconstructing the point as 50 and producing incorrect distances even when all buckets are scanned. Residual training must account for the balanced assignments used during Build.
| float | ||
| metric_distance(MetricType metric, const float* x, const float* c, int32_t dim) { | ||
| if (metric == MetricType::METRIC_TYPE_IP or metric == MetricType::METRIC_TYPE_COSINE) { | ||
| return clamp_distance(1.0F - FP32ComputeIP(x, c, dim)); |
There was a problem hiding this comment.
[P2] [P2] Normalize cosine rows before comparing assignment regrets
ClassifyDatasForBuild passes raw vectors here. Normalized centroids preserve each row's nearest-centroid order, but raw row norms change the regrets compared across vectors. With centroids (1,0)/(0,1) and vectors (100,99)/(1,0), both raw regrets equal 1, so the first vector takes centroid 0 and forces (1,0) into centroid 1. Normalized cosine costs give (1,0) priority instead. Normalize cosine inputs before computing costs so vector magnitude does not distort placement and recall.
| } | ||
| const float dist = | ||
| metric_distance(metric, x, centroids + j * static_cast<uint64_t>(dim_), dim_); | ||
| if (dist < best) { |
There was a problem hiding this comment.
[P2] [P2] Allow saturated distances to select an available seat
metric_distance clamps overflow to K_MAX_DISTANCE, which is also the initial best value. For Assign with two zero vectors and finite centroids -1e20/+1e20 in one dimension, every distance clamps to this value. Candidate lists retain only their default bucket 0; after its seat fills, the fallback rejects bucket 1 because its distance is not strictly smaller than best and throws 'ran out of capacity'. Use a distinct uninitialized sentinel or explicitly accept the first available cluster, including in candidate construction.
| false, | ||
| 1e-6F, | ||
| vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, | ||
| seed); |
There was a problem hiding this comment.
[P3] [P3] Make the seeded baseline reduction reproducible
This leaves KMeansCluster::Run's final deterministic_reduction argument at false. With multiple workers, block centroid sums merge in completion order, so floating-point centroids, subsequent labels, and reported loss can vary despite identical arguments and seed. This contradicts the tool's reproducibility contract. Enable deterministic reduction for the reproducible comparison or explicitly qualify the reported results.
| } | ||
| actual_n = total; | ||
| if (kind == "indivisible" and total % k == 0) { | ||
| actual_n = total - 1; |
There was a problem hiding this comment.
[P3] [P3] Handle the single-point indivisible dataset
The parser accepts --dataset indivisible --n 1 --k 1, but this subtraction changes actual_n from 1 to 0. run_bench then passes an empty dataset to KMeansCluster::Run, which throws an uncaught invalid-argument exception. Handle k=1 explicitly or reject this unsupported dataset combination before generating it.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a08dcb7ec5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| const auto labels = cls.Assign(this->bucket_count_, | ||
| centroids.data(), | ||
| static_cast<const float*>(datas), | ||
| static_cast<uint64_t>(count), | ||
| this->metric_type_); |
There was a problem hiding this comment.
With balanced_kmeans, use_residual: true, and SQ8 base codes, Build now uses capacity-constrained labels, but BucketDataCell::Train still computes training residuals using nearest-centroid ClassifyDatas (src/datacell/bucket_datacell.h:638). For 90 zero vectors and ten vectors equal to 100 with two buckets, the centroids become 0 and 20: SQ8 trains on residuals in [0, 80], while Build generates residual -20 for 40 zero vectors. SQ8 clamps those residuals, reconstructing the affected zeros as 20. Without exact reorder, the resulting distance errors persist even when every bucket is scanned. Train the residual codec using assignments consistent with balanced Build, or reject this combination.
Useful? React with 👍 / 👎.
| The binary is part of the tools build (`-DENABLE_TOOLS=ON`, requires | ||
| `-DENABLE_CXX11_ABI=ON`, Linux): |
There was a problem hiding this comment.
The new benchmark's build and CLI guide lives in this standalone README, while a repository-wide rg search for cluster_balance_bench finds only brief mentions in the English and Chinese website sources. Root AGENTS.md requires user-facing documentation under docs/docs/{en,zh}/src/ and says, “Tool-level READMEs should stay as local pointers to those canonical pages.” Publish the benchmark guide in those website sources and make this README point there rather than establishing a separate documentation source.
AGENTS.md reference: AGENTS.md:L101-L103
Useful? React with 👍 / 👎.
| } | ||
| } | ||
| for (auto& future : futures) { | ||
| future.wait(); |
There was a problem hiding this comment.
If an allocation inside scan_block fails, GeneralEnqueue stores the exception in its future, but wait() never observes it. Assignment then proceeds with the failed rows' placeholder distances and cluster IDs and can return success with incorrect labels. For example, failing the worker allocations when assigning [10,10,0,0] to centroids [0,10] returns labels [0,0,1,1] and loss 400 instead of reporting the allocation failure. Retrieve each future's result, drain all outstanding tasks, and propagate any recorded exception before consuming the candidate buffers.
Useful? React with 👍 / 👎.
| } | ||
| const float dist = | ||
| metric_distance(metric, x, centroids + j * static_cast<uint64_t>(dim_), dim_); | ||
| if (dist < best) { |
There was a problem hiding this comment.
When finite inputs overflow the FP32 distance calculation, metric_distance clamps the result to K_MAX_DISTANCE, which is also the initial value of best. This strict comparison therefore never selects a cluster if every remaining distance reaches that bound, despite seats being available. For example, Assign with centroids [-1e20,1e20] and vectors [0,0] throws “balanced assignment ran out of capacity” for the second vector. Select the first feasible cluster even when its distance equals the sentinel so the advertised finite-input handling preserves the capacity guarantee.
Useful? React with 👍 / 👎.
| vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, | ||
| seed); |
There was a problem hiding this comment.
With multiple workers and more than one 1,024-vector reduction block, this call leaves KMeansCluster::Run's final deterministic_reduction argument at its default false. The CPU implementation consequently merges floating-point centroid sums in worker-completion order (src/impl/cluster/kmeans_cluster.cpp:345-354), so a fixed seed does not determine the resulting centroids and subsequent assignments. This undermines the tool's advertised reproducible clustering comparison. Pass deterministic_reduction=true, as the existing seeded reproducibility test does.
Useful? React with 👍 / 👎.
vsag-lgtm
left a comment
vsag-lgtm
left a comment
There was a problem hiding this comment.
Automated review @ a08dcb7e — 4 inline comment(s).
| if (dist < best_dist[c - 1]) { | ||
| uint32_t pos = c - 1; | ||
| while (pos > 0 and best_dist[pos - 1] > dist) { | ||
| best_dist[pos] = best_dist[pos - 1]; | ||
| best_id[pos] = best_id[pos - 1]; | ||
| --pos; | ||
| } | ||
| best_dist[pos] = dist; | ||
| best_id[pos] = j; | ||
| } | ||
| } | ||
| for (uint32_t j = 0; j < c; ++j) { | ||
| cand_dist[i * static_cast<uint64_t>(c) + j] = best_dist[j]; | ||
| cand_id[i * static_cast<uint64_t>(c) + j] = best_id[j]; | ||
| } | ||
| ++local_demand[best_id[0]]; | ||
| } | ||
| std::lock_guard<std::mutex> lock(demand_mutex); | ||
| for (uint32_t j = 0; j < k; ++j) { | ||
| // Exact integer sums are order-independent, so any merge | ||
| // schedule produces the same demand vector. | ||
| demand[j] += local_demand[j]; | ||
| } | ||
| }; | ||
|
|
||
| std::vector<std::future<void>> futures; | ||
| for (uint64_t i = 0; i < count; i += block_size) { | ||
| const uint64_t begin = i; | ||
| const uint64_t end = std::min(i + block_size, count); | ||
| if (thread_pool_->IsCurrentThread()) { | ||
| scan_block(begin, end); | ||
| } else { | ||
| futures.emplace_back(thread_pool_->GeneralEnqueue(scan_block, begin, end)); | ||
| } | ||
| } | ||
| for (auto& future : futures) { | ||
| future.wait(); | ||
| } | ||
|
|
||
| // Step 2: capacity targets. Exactly n mod k clusters hold ceil(n / k) | ||
| // vectors. The seats go to the clusters that the unconstrained nearest | ||
| // pass drew most, which keeps the constrained assignment closest to the | ||
| // unconstrained one. Ties break by cluster id, so the targets are | ||
| // deterministic. | ||
| const uint64_t base = count / static_cast<uint64_t>(k); | ||
| const uint64_t extra = count % static_cast<uint64_t>(k); | ||
| Vector<uint64_t> remaining(k, base, allocator_); | ||
| { | ||
| Vector<uint32_t> cluster_order(k, 0U, allocator_); | ||
| std::iota(cluster_order.begin(), cluster_order.end(), 0U); | ||
| std::sort(cluster_order.begin(), cluster_order.end(), [&](uint32_t a, uint32_t b) { | ||
| if (demand[a] != demand[b]) { | ||
| return demand[a] > demand[b]; | ||
| } | ||
| return a < b; | ||
| }); | ||
| for (uint64_t i = 0; i < extra; ++i) { | ||
| ++remaining[cluster_order[i]]; | ||
| } | ||
| } | ||
|
|
||
| // Step 3: regret-greedy placement. A vector's regret is the extra cost of | ||
| // losing its best candidate; the vectors that lose the most when pushed | ||
| // are placed first, while their best seats are still free. The order is | ||
| // a total order (regret, then nearest distance, then vector id), so the | ||
| // result does not depend on the thread count. | ||
| Vector<uint64_t> point_order(count, 0ULL, allocator_); | ||
| std::iota(point_order.begin(), point_order.end(), 0ULL); | ||
| auto regret = [&](uint64_t i) { | ||
| if (c >= 2) { | ||
| return cand_dist[i * static_cast<uint64_t>(c) + 1] - | ||
| cand_dist[i * static_cast<uint64_t>(c)]; | ||
| } | ||
| return 0.0F; | ||
| }; | ||
| std::sort(point_order.begin(), point_order.end(), [&](uint64_t a, uint64_t b) { | ||
| const float ra = regret(a); | ||
| const float rb = regret(b); | ||
| if (ra != rb) { | ||
| return ra > rb; | ||
| } | ||
| const float da = cand_dist[a * static_cast<uint64_t>(c)]; | ||
| const float db = cand_dist[b * static_cast<uint64_t>(c)]; | ||
| if (da != db) { | ||
| return da < db; | ||
| } | ||
| return a < b; | ||
| }); | ||
|
|
||
| // Total remaining capacity equals the number of unplaced vectors at every | ||
| // point of this loop, so the on-demand rescan always finds a cluster with | ||
| // a free seat. | ||
| for (uint64_t i : point_order) { | ||
| const uint64_t row = i * static_cast<uint64_t>(c); | ||
| uint32_t chosen = k; | ||
| for (uint32_t j = 0; j < c; ++j) { | ||
| const uint32_t id = cand_id[row + j]; | ||
| if (remaining[id] > 0) { | ||
| chosen = id; | ||
| break; | ||
| } | ||
| } | ||
| if (chosen == k) { | ||
| // Every remembered candidate is full: recompute distances against | ||
| // the clusters that still have seats instead of keeping the full | ||
| // row. Rare for c > 1, and bounded work even when it repeats. | ||
| const float* x = datas + i * static_cast<uint64_t>(dim_); | ||
| float best = K_MAX_DISTANCE; | ||
| for (uint32_t j = 0; j < k; ++j) { | ||
| if (remaining[j] == 0) { | ||
| continue; | ||
| } | ||
| const float dist = | ||
| metric_distance(metric, x, centroids + j * static_cast<uint64_t>(dim_), dim_); | ||
| if (dist < best) { |
There was a problem hiding this comment.
[critical] Strict < comparisons in the candidate insertion (line 271) and fallback rescan (line 385) treat a clamped distance equal to K_MAX_DISTANCE as unselected. If any vector-to-centroid distance overflows and is clamped to K_MAX_DISTANCE (or if the input maximum is reached), the candidate list keeps the default cluster 0 and the fallback never updates best, so chosen stays k and the function throws balanced assignment ran out of capacity. Use <= in both places (or a value-typed uninitialized sentinel instead of the valid max float) so clamped distances still participate in the selection and the capacity guarantee does not depend on avoiding overflow.
| std::uniform_int_distribution<uint64_t> dis(0, count - 1); | ||
| selected_idx = dis(gen); | ||
| } else { | ||
| std::uniform_real_distribution<double> prob_dis(0.0, total_weight); |
There was a problem hiding this comment.
[warning] std::uniform_real_distribution<double>(0.0, total_weight) has undefined behavior when total_weight is +inf or NaN. total_weight sums distances clamped to FLT_MAX; for very large n / extreme magnitudes that sum can overflow double, making the distribution bounds invalid. Guard with std::isfinite(total_weight) and, on overflow, fall back to the uniform random pick used for the total_weight <= 0.0 case.
| } | ||
| } | ||
| for (auto& future : futures) { | ||
| future.wait(); |
There was a problem hiding this comment.
[warning] future.wait() only blocks until the task finishes; it does not propagate exceptions stored in the future. If a worker throws (e.g., an allocation failure inside scan_block or metric_distance), the exception is swallowed and balanced_assign continues with uninitialized/placeholder candidate data, returning wrong labels and silently corrupted loss. Iterate the futures with future.get() (rethrowing the first failure) or explicitly collect and rethrow after all tasks complete.
| // trained partition; anything else (including a build that brings fewer | ||
| // vectors than buckets, where the trainer has already placed one vector per | ||
| // bucket best it can) takes the plain routing. | ||
| const bool balance_buckets = this->ivf_partition_strategy_param_->partition_train_type == |
There was a problem hiding this comment.
[suggestion] balance_buckets requires buckets_per_data == 1. That means if a user enables balanced_kmeans but sets n_probe > 1 (so buckets_per_data_ > 1), the index silently falls back to plain nearest-centroid routing for the first Build assignment. The behavior is internally consistent but surprising from a user's perspective; consider either documenting this limitation in the IVF docs or emitting a logger::warn trace when partition_train_type is BalancedKMeansTrainer and buckets_per_data != 1, so the balance guarantee does not appear to be silently dropped.
- train residual quantizers with the capacity-constrained build assignment - normalize cosine rows before comparing assignment regrets - use infinity as the candidate sentinel so saturated distances stay selectable - propagate worker exceptions instead of waiting on failed futures - make the seeded baseline reduction reproducible; keep n=1,k=1 datasets valid Signed-off-by: DoubleCore <DoubleCore@users.noreply.github.com>
mergify
Bot
added
the
module/datacell
|
本轮评审已处理，新提交 已修
未处理（如实说明）
验证： |
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Reviewer | Code review
Automated pull request review completed.
Review decision: Comment only. No approval was submitted.
Model: gpt-6.1-sol; reasoning effort: high.
Review scope: high (2412 changed lines across 23 files).
Submitted 2 inline comments.
Review coverage checks completed; this does not prove absence of defects.
Static review only; compilation and runtime tests are outside this review's scope.
File coverage: 23/23; observed passing test runs: 0.
Reviewed commit e81e314.
| if (thread_pool_->IsCurrentThread()) { | ||
| scan_block(begin, end); | ||
| } else { | ||
| futures.emplace_back(thread_pool_->GeneralEnqueue(scan_block, begin, end)); |
There was a problem hiding this comment.
[P2] [P2] Drain submitted scans when task submission fails
If GeneralEnqueue or futures.emplace_back throws after a scan has been queued, execution skips the draining loop below. Futures backed by packaged_task do not join on destruction, so workers can continue using cand_dist, cand_id, demand, and demand_mutex after those objects are destroyed. This turns an allocation failure or rejecting custom thread pool into use-after-free during training or Build. Guard submission as well as worker completion, and ensure every submitted task is joined before unwinding.
| float* dest = sums.data() + static_cast<uint64_t>(label) * static_cast<uint64_t>(dim_); | ||
| ++sizes[label]; | ||
| for (int32_t d = 0; d < dim_; ++d) { | ||
| dest[d] += src[d]; |
There was a problem hiding this comment.
[P2] [P2] Use a stable accumulator for large cluster means
Accumulating an entire cluster into one float loses contributions once the running sum becomes large. For dim=1, k=1, and 2^25 copies of 1.0F, the sum stops increasing at 2^24 and the returned centroid becomes 0.5 instead of 1.0. This valid input is also reachable with an explicitly enlarged train_sample_count; the existing trainer's 1024-vector block reduction preserves this example. Use double accumulation or a stable deterministic reduction so large clusters retain correct means.
vsag-lgtm
left a comment
vsag-lgtm
left a comment
There was a problem hiding this comment.
Automated review @ e81e314f — 4 inline comment(s).
| // Residuals must cover the same capacity-constrained assignments used by | ||
| // the initial Build; nearest-centroid training can otherwise clip the | ||
| // residuals of vectors that balancing moves to another bucket. | ||
| auto buckets = strategy_->ClassifyDatasForBuild(data_ptr, count, 1, nullptr); |
There was a problem hiding this comment.
[warning] The comment above promises "the same capacity-constrained assignments used by the initial Build", but the training batch is not the build batch, so that guarantee does not hold.
IVF::Train routes the sampled dataset here: sample_train_data(...) limits the set to train_sample_count_ (default MAX_TRAIN_COUNT) and bucket_->Train(data_ptr, sample_count) is called with that subset (src/algorithm/ivf/ivf_build.cpp:159-166). The build assignment, in contrast, runs on the entire first Add batch (src/algorithm/ivf/ivf_build.cpp:199-204). BalancedKMeansCluster::Assign is capacity-constrained, so its labels depend on the batch: with n_train != n_build the per-cluster targets ceil(n/k) / floor(n/k) differ and, via the regret ordering, a vector near a capacity boundary can end up in a different bucket than at build time. For such a vector this quantizer records the residual range against a centroid that encode never uses for it - i.e. exactly the clipping this change is meant to remove. Note also that the mismatch runs in both directions: balanced residuals are typically larger than nearest-centroid residuals, so many ranges become looser than needed and encode cannot produce them.
The Train(); Add(); API path is a second mismatch: Add calls add(base, false), so initial_bucket_assignment is always false and every vector is routed by nearest centroid, while the residual ranges here were trained from balanced labels.
No test covers either case - the new UT uses BuildAssignedPartitionStrategy, which returns bucket 1 for every row - so CI cannot detect the divergence. src/datacell/rabitq_split_bucket_datacell.cpp:1368 has the same call and the same exposure.
Suggestion: use the balanced assignment only when the datacell is trained on the exact batch that will later be assigned, and keep plain ClassifyDatas otherwise (or state the sampling caveat in the comment).
|
|
||
| Vector<float> normalized_datas(allocator_); | ||
| if (metric == MetricType::METRIC_TYPE_COSINE) { | ||
| normalized_datas.resize(count * static_cast<uint64_t>(dim_)); |
There was a problem hiding this comment.
[warning] For METRIC_TYPE_COSINE this materializes a second full copy of the input (count * dim_ * 4 bytes) and normalizes every row into it.
Both datacell call sites that now reach Assign through ClassifyDatasForBuild already hold normalized rows: BucketDataCell::Train normalizes into train_data_buffer (src/datacell/bucket_datacell.h:631-635) and RaBitQSplitBucketDataCell::Train into normalized_data (src/datacell/rabitq_split_bucket_datacell.cpp:1357-1363). There the copy is pure duplication and Normalize runs a second time over the same rows. On the IVF::add path the caller passes raw user vectors, so a normalization is genuinely required - but the extra dataset-sized allocation remains, and it scales with the whole initial batch (10M x 768 float rows is ~28 GiB), which is at odds with this function's documented bounded-memory behaviour.
Suggestion: normalize block-wise inside balanced_assign (the candidate scan already runs in 1024-row blocks and the loss pass can do the same), or let callers declare their rows already normalized so the copy is skipped. The former also removes the double normalization in both Train paths.
| if (metric == MetricType::METRIC_TYPE_COSINE) { | ||
| normalized_datas.resize(count * static_cast<uint64_t>(dim_)); | ||
| for (uint64_t i = 0; i < count; ++i) { | ||
| Normalize(datas + i * static_cast<uint64_t>(dim_), |
There was a problem hiding this comment.
[note] Normalize divides by the row norm without a zero-norm guard, so a zero vector becomes NaN in normalized_datas. clamp_distance maps NaN to K_MAX_DISTANCE, so such rows are ranked after every non-degenerate row and each one contributes FLT_MAX to the accumulated loss, making Stats().loss unusable for any batch containing zero vectors. Before this change the same row compared as 1 - 0 = 1.0 and stayed finite. Labels remain valid because balance is still enforced, so this is informational unless Stats().loss is consumed.
| 1e-6F, | ||
| vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, | ||
| seed, | ||
| /*deterministic_reduction=*/true); |
There was a problem hiding this comment.
[note] This makes the baseline arm deterministic while the production path is not: the IVF k-means trainer calls KMeansCluster::Run(k, data, n, iter) and keeps the default deterministic_reduction = false (src/algorithm/ivf/ivf_nearest_partition.cpp:117-122), and the balanced arm has no equivalent switch. Since the table prints both time columns side by side, any runtime difference introduced by the deterministic reduction makes the printed comparison non-like-for-like. Please confirm the difference is negligible, or state the asymmetry in the benchmark output.
Document the buckets_per_data limitation and format the new regression tests. Signed-off-by: DoubleCore <doublecore@users.noreply.github.com> Assisted-by: OpenCode:ark-code-latest
|
Validation for 50c0c97: C++ was not compiled or tested locally (CMake/build was deliberately not run); only static symbol checks were performed. clang-format 20.1.8 and git diff --check passed. This PR does not touch tools/dataset_profiler, so its Python pytest command was not applicable. C++ test or CI success is not claimed. |
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Reviewer | Code review
Automated pull request review completed.
Review decision: Comment only. No approval was submitted.
Model: gpt-6.1-sol; reasoning effort: high.
Review scope: high (3080 changed lines across 23 files).
Submitted 1 inline comment.
Review coverage checks completed; this does not prove absence of defects.
Static review only; compilation and runtime tests are outside this review's scope.
File coverage: 23/23; observed passing test runs: 0.
Reviewed commit 50c0c97.
| [](uint64_t sample, uint64_t sample_count, uint64_t population_count) { | ||
| const uint64_t step = population_count / sample_count; | ||
| const uint64_t remainder = population_count % sample_count; | ||
| return sample * step + sample * remainder / sample_count; |
There was a problem hiding this comment.
[P2] [P2] Avoid periodic aliasing when sampling PQ residuals
This fixed-stride sampler can discard an entire residual population. For an L2 balanced index with one bucket, residual PQ/PQFS, train_sample_count=131072, and 131072 scalar rows alternating 0 and 100, both routes agree and build_samples becomes 65536. This expression selects only rows 2*i, so the centroid is 50 but every training residual is -50. All codewords then represent -50, collapsing both input values to the same reconstruction and corrupting distances and recall. The previous 65536-row prefix contained both residual values. Use representative sampling that avoids periodic aliasing, or preserve the existing prefix when no extra population is needed.
There was a problem hiding this comment.
Fixed. When nearest_samples == 0, Build residuals now retain the original 65,536-row prefix; fixed-stride sampling is used only when a nearest-only residual population must share that prefix. I also updated the PQ/PQFS regression to use 131,072 alternating 0/100 rows, so it covers this exact aliasing case and checks reconstruction of the +50 residual. Validation: per task constraints, I did not compile or run C++ tests locally. Static review, clang-format 20.1.8 --dry-run --Werror, and git diff --check passed. No Python code was changed, so the dataset_profiler Python tests were not applicable.
Signed-off-by: DoubleCore <851678303@qq.com> Assisted-by: OpenCode:ark-code-latest
| // The sole centroid is 50, so alternating rows have residuals -50 and +50. | ||
| // Both Build and nearest routing agree. Fixed-stride sampling used to select | ||
| // only even rows, filling PQ's 65,536-row prefix with just the -50 residual. |
There was a problem hiding this comment.
[warning] This regression test cannot reach the sampling code it is meant to lock in, so the fix in bac458bb is currently unverified.
FixedCentroidPartitionStrategy (src/datacell/bucket_datacell_test.cpp:90) does not override HasDistinctBuildAssignment(), whose IVFPartitionStrategy base implementation returns false (src/algorithm/ivf/ivf_partition_strategy.h:79). BucketDataCell::Train therefore sets train_nearest = false and takes the single-route else branch (src/datacell/bucket_datacell.h:721); the whole prefix_limited block - including the new guard const bool sample_build = sample_prefix && nearest_samples != 0; (src/datacell/bucket_datacell.h:708) and the sample_index stride sampling - never executes for this strategy.
Consequences:
- The comment above is inaccurate: no fixed-stride sampling happens here. PQ simply truncates the first 65,536 rows of the single-route residual buffer.
- The aliasing scenario fixed by
bac458bb(dual routing, pq/pqfs,count > 65536, zero rows where Build and nearest assignments disagree) has no regression coverage. Reverting the guard to plainsample_prefixwould still leave this test green, with the build prefix collapsing onto every-even-row residue sampling again.
To cover the path, the strategy must return true from HasDistinctBuildAssignment() while ClassifyDatasForBuild and ClassifyDatas agree (so nearest_rows stays empty) on a pq/pqfs bucket with count > 65536. The sibling test "samples nearest residuals within the PQ prefix" uses MixedResidualPartitionStrategy, whose data always contains disagreeing rows, so it only exercises sample_build == true.
| return false; | ||
| } | ||
| } | ||
| if (k >= 10000) { |
There was a problem hiding this comment.
[note] Implicit coupling worth tracking: the magic 10000 duplicates the private KMeansCluster::THRESHOLD_FOR_HGRAPH = 10000ULL (src/impl/cluster/kmeans_cluster.h:103), which selects the deterministic BLAS path when k < THRESHOLD_FOR_HGRAPH (src/impl/cluster/kmeans_cluster.cpp). I verified the two agree today, so this guard and the en/zh ivf.md passages are accurate as of this push. Nothing ties them together at compile or test time, though: if the library threshold ever changes, the tool keeps rejecting k >= 10000 and silently diverges from the real HGraph switch point, voiding the reproducibility guarantee the docs advertise. Consider exposing the constant (public constexpr or accessor) or cross-checking it at tool startup.
Signed-off-by: DoubleCore <851678303@qq.com> Assisted-by: OpenCode:ark-code-latest
Validation: per the local task constraints, I did not run CMake, compile C++, or execute C++ tests. Static symbol and control-flow review, |
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Reviewer | Code review
Automated pull request review completed.
Review decision: Comment only. No approval was submitted.
Model: gpt-6.1-sol; reasoning effort: high.
Review scope: high (3106 changed lines across 24 files).
Submitted 1 inline comment.
Review coverage checks completed; this does not prove absence of defects.
Static review only; compilation and runtime tests are outside this review's scope.
File coverage: 24/24; observed passing test runs: 0.
Reviewed commit ccacd7b.
| const bool sample_build = sample_prefix && nearest_samples != 0; | ||
| for (uint64_t i = 0; i < build_samples; ++i) { | ||
| const uint64_t source = | ||
| sample_build ? sample_index(i, build_samples, count) : i; |
There was a problem hiding this comment.
[P2] [P2] Avoid periodic aliasing when sampling Build residuals
The fixed stride can remove an entire Build-residual population when nearest-route extras exist. For one-dimensional L2 residual PQ/PQFS, use k=2, n=131072, train_sample_count=n, and rows repeating [0,100,100,100]. Balanced training yields centroids 50 and 100, with 32768 rows of value 100 forced into centroid 50. Here build_samples and nearest_samples are both 32768, so sample_index selects only Build rows 4*i: their residual is always -50. Every nearest extra has residual 0, so the codebook never sees the +50 residual needed by a quarter of the corpus and encodes those vectors as reconstructed value 50 instead of 100. Use sampling that avoids periodic aliasing for both populations; the nearest_samples == 0 special case does not cover this case.
There was a problem hiding this comment.
Addressed in a33a9ee.
Replaced fixed-stride indexing for both Build rows and nearest-route extras with fixed-seed reservoir sampling. The existing residual training buffer serves as the reservoir, so each population keeps its exact build_samples/nearest_samples budget, has no duplicate source rows, and needs no additional sampling allocation. The nearest_samples == 0 lone-prefix path remains unchanged.
Added PQ/PQFS regression coverage using repeating [0,100,100,100] rows. It inspects the quantizer input directly and asserts both -50 and +50 Build residuals, exact population sizes, and per-population source uniqueness; existing lone-prefix coverage remains.
Validation: clang-format 20.1.8 --dry-run --Werror and git diff --check passed; a static reviewer-scenario model confirmed that the old stride selected only -50 Build residuals. C++ was not compiled and unit tests were not run locally.
There was a problem hiding this comment.
Addressed in a33a9ee. Both Build and nearest-extra populations now use fixed-seed reservoir sampling directly in the existing residual training buffer, replacing the fixed stride without duplicate rows, changing either sample budget, or adding sampling allocations. The nearest_samples == 0 lone-prefix path is unchanged. Added PQ/PQFS regression coverage using the [0,100,100,100] routes and inspecting the actual quantizer input: Build residuals must contain both -50 and +50, each population has exactly 32768 rows, and source-row IDs verify uniqueness within each population. Validation performed: clang-format 20.1.8 dry-run, git diff --check, static text checks, and a Python algorithm model (not a C++ test). No C++ compilation or C++ test execution was performed locally.
vsag-lgtm
left a comment
vsag-lgtm
left a comment
There was a problem hiding this comment.
Incremental review of bac458bb..ccacd7b3 (event SHA 50c0c979 is an ancestor of head): no new findings in this push; previously raised review thread(s) remain unresolved, so not approving.
Signed-off-by: DoubleCore <851678303@qq.com> Assisted-by: OpenCode:ark-code-latest
|
Review follow-up: a33a9ee replaces periodic residual sampling for both Build and nearest-extra populations with deterministic in-buffer reservoir sampling. Budgets and memory bounds stay unchanged, and the lone-prefix path is preserved. Added direct PQ/PQFS regression assertions for both -50/+50 Build residuals, population counts, and source uniqueness. Validation completed: clang-format 20.1.8 dry-run, git diff --check, and static reviewer-scenario checks. No local C++ compilation or unit-test execution was performed. |
|
Review follow-up in a33a9ee: Build and nearest-extra residuals now use fixed-seed reservoir sampling in the existing training buffer, removing periodic aliases without extra sampling allocations. The new PQ/PQFS regression directly checks quantizer input for both Build residual signs, exact population sizes, and source-row uniqueness; lone-prefix coverage remains. clang-format 20.1.8, diff/static checks, and a Python algorithm model passed. C++ compilation and C++ tests were not run locally. |
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Reviewer | Code review
Automated pull request review completed.
Review decision: Approved. No actionable findings were identified in this complete static review. CI and merge requirements remain separate.
Model: gpt-6.1-sol; reasoning effort: high.
Review scope: high (3252 changed lines across 24 files).
No actionable inline findings were found.
Review coverage checks completed; this does not prove absence of defects.
Static review only; compilation and runtime tests are outside this review's scope.
File coverage: 24/24; observed passing test runs: 0.
Reviewed commit a33a9ee.
vsag-lgtm
left a comment
vsag-lgtm
left a comment
There was a problem hiding this comment.
Automated review @ a33a9eea — 2 inline comment(s).
| for (uint64_t i = 0; i < nearest_samples; ++i) { | ||
| const uint64_t offset = (build_samples + i) * dim; | ||
| const float residual = allocator->captured_values[offset]; | ||
| const float source_value = allocator->captured_values[offset + 1]; | ||
| REQUIRE(source_value >= 0.0F); | ||
| REQUIRE(source_value < static_cast<float>(count)); | ||
| REQUIRE(source_value == std::floor(source_value)); | ||
| const auto source = static_cast<uint64_t>(source_value); | ||
| REQUIRE_FALSE(nearest_seen[source]); | ||
| nearest_seen[source] = true; | ||
| // Only the forced 100 rows differ between Build and nearest routing. | ||
| REQUIRE(source % 4 == 1); | ||
| REQUIRE(residual == 0.0F); |
There was a problem hiding this comment.
[warning] The nearest-extras reservoir path is never exercised by this regression test
With count = 131072 and exactly 32,768 mismatched rows, nearest_samples becomes min(32768, 65536 - 32768) == 32768, which equals nearest_rows.size(). The second call for_each_sample(sample_prefix, nearest_samples, nearest_rows.size(), ...) therefore executes the reservoir loop for (uint64_t i = sample_count; i < population_count; ++i) (src/datacell/bucket_datacell.h:715) zero times, so every nearest assertion in this block only validates the plain initial fill.
Only the Build population actually goes through the new reservoir code, even though the fix is described as replacing fixed-stride sampling for both the Build and the nearest-extra populations. Re-introducing stride-based selection for the nearest population alone (the aliasing problem this commit targets) would keep this test green.
Consider adding a case where the mismatched population exceeds the 32,768 budget, for example count = 131072 with every row mismatching so that nearest_samples = 32768 < nearest_rows.size() = 131072, and assert on the sampled nearest half (uniqueness plus both residual signs, as done for Build) so the reservoir behaviour is locked in for that population too.
| return DefaultAllocator::Allocate(size); | ||
| auto* pointer = DefaultAllocator::Allocate(size); | ||
| if (size == capture_size) { | ||
| captured_pointers.emplace_back(pointer); |
There was a problem hiding this comment.
[suggestion] Captured values depend on deallocation order and on an accidental byte-size match
captured_values is overwritten by every deallocation of a buffer whose byte size equals capture_size, so the assertions at the end of the test read whichever matching buffer happened to be freed last. In the new case capture_size is 65536 * 2 * sizeof(float) == 524288 bytes, and build_buckets / nearest_buckets are 131072 * sizeof(BucketIdType) == 524288 bytes as well, so the two classification arrays are tracked and captured too.
The residual assertions currently observe the training buffer only because train_data_buffer is the first local declared in Train and therefore destroyed last; reordering local declarations or changing the sampling budget would silently make these assertions compare against bucket IDs reinterpreted as floats. Consider making the assumption explicit, for example by keeping only the most recently allocated matching pointer (the training buffer is also the last such allocation) and asserting captured_pointers.size() == 1 at capture time.
Signed-off-by: DoubleCore <doublecore@users.noreply.github.com> Assisted-by: OpenCode:ark-code-latest
|
Addressed both new review comments in commit
Validation: |
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Reviewer | Code review
Automated pull request review completed.
Review decision: Comment only. No approval was submitted.
Model: gpt-6.1-sol; reasoning effort: high.
Review scope: high (3267 changed lines across 24 files).
Submitted 1 inline comment.
Review coverage checks completed; this does not prove absence of defects.
Static review only; compilation and runtime tests are outside this review's scope.
File coverage: 24/24; observed passing test runs: 0.
Reviewed commit 403ce4f.
| } | ||
| // Each nearest extra is written before its normalized source row is overwritten | ||
| // by the Build residual. | ||
| uint64_t nearest_offset = count; |
There was a problem hiding this comment.
[P2] [P2] Account for scalar quantizers' sampling of appended residuals
SQ8/SQ4 also subsample: ScalarQuantizationTrainer takes 100000 rows with stride 2147483647 % train_count. For a 131072-row corpus with 65530 differing routes, train_count=196602 makes that stride 1, so every appended nearest residual is skipped. A concrete L2 case is k=2, train_sample_count=131072, and ordered data containing 65536 zeros, 65530 ones, then six values of 20000. Balanced centroids are approximately 0 and 2.831; the sampled Build prefix contains only zero/negative residuals, giving an upper bound of zero. Later Add(1) routes to centroid zero and its residual +1 is clipped to zero, corrupting distances. The previous nearest-only training includes the +1 population. Apply population-aware sampling to scalar quantizers as well.
Signed-off-by: DoubleCore <doublecore@users.noreply.github.com> Assisted-by: OpenCode:ark-code-latest
|
Addressed discussion_r4478684635 in
Validation: clang-format 20.1.8, |
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Reviewer | Code review
Automated pull request review completed.
Review decision: Comment only. No approval was submitted.
Model: gpt-6.1-sol; reasoning effort: high.
Review scope: high (3372 changed lines across 24 files).
Submitted 1 inline comment.
Review coverage checks completed; this does not prove absence of defects.
Static review only; compilation and runtime tests are outside this review's scope.
File coverage: 24/24; observed passing test runs: 0.
Reviewed commit 7ac69e0.
| const auto quantizer_name = this->GetQuantizerName(); | ||
| prefix_limited = quantizer_name == QUANTIZATION_TYPE_VALUE_PQ or | ||
| quantizer_name == QUANTIZATION_TYPE_VALUE_PQFS; | ||
| scalar_limited = quantizer_name == QUANTIZATION_TYPE_VALUE_SQ8 or |
There was a problem hiding this comment.
[P2] [P2] Include uniform scalar quantizers in residual sampling
sq8_uniform and sq4_uniform also use ScalarQuantizationTrainer's 100000-row sampler, but this condition sends them through the unsampled dual-buffer branch. For the ordered 131072-row population in the new scalar test, the resulting 196602 residual rows give sampler stride 1: it reads only the Build prefix and drops every appended +1 nearest-route residual. Both uniform trainers then learn an upper bound of zero, so later Add vectors with that residual are clipped and become indistinguishable from zero residuals. Include both uniform quantizers in the population-budgeting path.
Signed-off-by: DoubleCore <doublecore@users.noreply.github.com> Assisted-by: OpenCode:ark-code-latest
|
Addressed vsag-bot review discussion_r4478974261 in commit
Validation performed:
No local CMake/C++ build or C++ test execution was performed in this environment. This comment does not claim that CI is green. |
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Reviewer | Code review
Automated pull request review completed.
Review decision: Approved. No actionable findings were identified in this complete static review. CI and merge requirements remain separate.
Model: gpt-6.1-sol; reasoning effort: high.
Review scope: high (3374 changed lines across 24 files).
No actionable inline findings were found.
Review coverage checks completed; this does not prove absence of defects.
Static review only; compilation and runtime tests are outside this review's scope.
File coverage: 24/24; observed passing test runs: 0.
Reviewed commit ebadeaa.
vsag-lgtm
left a comment
vsag-lgtm
left a comment
There was a problem hiding this comment.
Incremental review @ ebadeaab (increment a33a9eea..ebadeaab, includes 403ce4f/7ac69e0e; event SHA 403ce4f is an ancestor of head): scalar residual population budgeting, reservoir sampling and the new regression tests verified; the uniform-variant coverage added in ebadeaa matches the SQ8UniformQuantizer/SQ4UniformQuantizer names and the trainer's 100,000-row cap. No new findings.
Implements #2870.
Fixes: #2870.