feat(ivf): add balanced k-means clustering for IVF partitioning

docs/docs/en/src/indexes/ivf.md

@@ -33,6 +33,64 @@ splits the space into two orthogonal sets of centroids
 (`first_order_buckets_count` × `second_order_buckets_count`) for even finer
 partitioning on very large corpora.
 
+### Balanced k-means training (`ivf_train_type: "balanced_kmeans"`)
+
+Standard k-means assigns every vector to its nearest centroid without any
+capacity constraint, so skewed, duplicated, or unevenly distributed corpora can
+produce buckets that differ in size by orders of magnitude — and even empty
+ones. Oversized buckets raise search latency and weaken partition pruning.
+
+`balanced_kmeans` keeps the same squared-L2 objective and mean centroid update
+as `kmeans`, but every assignment step is capacity-constrained: with `n`
+training vectors and `buckets_count` clusters, each final cluster holds either
+`floor(n / k)` or `ceil(n / k)` vectors. The guarantee extends to the initial
+bucket assignment of `Build`, so the stored buckets inherit the balance, and no
+cluster is ever empty as long as `1 <= k <= n`.
+
+How it works, per iteration:
+
+1. Each vector keeps only its `num_candidates` (8 by default) nearest clusters;
+   distances are evaluated one vector at a time in fixed parallel blocks, so no
+   `n × k` distance matrix is ever materialized and peak memory stays
+   `O(n · candidates + k · dim)` instead of `O(n · k)`.
+2. The `n mod k` clusters with the highest unconstrained nearest-centroid demand
+   receive `ceil(n / k)` seats, the others `floor(n / k)`.
+3. Vectors are then placed most-constrained first, in decreasing order of their
+   *regret* (the cost gap between their best and second-best candidate). Each
+   vector takes the nearest cluster that still has a free seat; a vector whose
+   remembered candidates are all full recomputes distances against the clusters
+   that still have seats. Because the capacities sum to exactly `n`, a seat is
+   always available.
+4. Centroids are re-set to the mean of their assigned vectors and the cycle
+   repeats until the loss stops improving (`kmeans_iter_count` bounds it).
+
+Complexity and trade-offs. Per iteration the trainer pays the same
+`O(n · k · dim)` distance work as standard k-means, plus an `O(n log n)`
+ordering; memory is bounded as above. Balancing is a heuristic (regret-greedy
+insertion), not an exact transportation solve, so the loss can be somewhat
+above an unconstrained k-means: the gap is small on uniform or evenly clustered
+data and grows on strongly skewed mixtures, where equal-size buckets must split
+tight clusters. Run `cluster_balance_bench` to measure both sides on your own
+shapes. What the extra loss buys is a hard bound on bucket size: every bucket
+holds at most `ceil(n / k)` vectors, the per-query scanned work is bounded by
+`scan_buckets_count × ceil(n / k)`, and tail latency stops depending on how
+skewed the corpus happens to be.
+
+Notes:
+
+- `random` and `kmeans` are unchanged; the default remains `kmeans`.
+- Cosine indexes normalize the trained centroids exactly as in the `kmeans`
+  path, and the initial bucket assignment uses the metric's own routing order
+  (`1 - <x, c>` for IP/cosine, squared L2 for L2).
+- `balanced_kmeans` is available only with `partition_strategy_type: "ivf"`;
+  GNO-IMI always trains plain k-means.
+- The trainer runs on the host; `enable_gpu_build` does not accelerate balanced
+  training or the balanced initial assignment.
+- Results are deterministic for a given input and seed, independent of the
+  thread count.
+- Compare balance, loss, time, and memory yourself with the reproducible
+  `cluster_balance_bench` tool (`tools/cluster_balance_bench`).
+
 ## Quick start
 
 ```cpp
@@ -79,8 +137,8 @@ Build-time parameters live under `index_param`. See
 | `buckets_count` | int | `10` | Number of inverted lists (effective for `ivf`) |
 | `first_order_buckets_count` | int | `10` | First-level count (effective for `gno_imi`) |
 | `second_order_buckets_count` | int | `10` | Second-level count (effective for `gno_imi`) |
-| `ivf_train_type` | string | `"kmeans"` | Centroid training: `kmeans` or `random` |
-| `kmeans_iter_count` | int | `25` (`ivf`), `30` (`gno_imi`) | KMeans training iterations for centroid training (effective with `ivf_train_type: "kmeans"`). Higher → better bucket quality, longer build. |
+| `ivf_train_type` | string | `"kmeans"` | Centroid training: `kmeans`, `balanced_kmeans`, or `random` |
+| `kmeans_iter_count` | int | `25` (`ivf`), `30` (`gno_imi`) | KMeans training iterations for centroid training (effective with `ivf_train_type: "kmeans"` or `"balanced_kmeans"`). Higher → better bucket quality, longer build. |
 | `route_max_degree` | int | `64` | Routing HGraph maximum degree (effective for `ivf`) |
 | `route_ef_construction` | int | `300` | Routing HGraph construction search breadth (effective for `ivf`) |
 | `enable_gpu_build` | bool | `false` | Use the CUDA backend while building (effective for `ivf`); see [GPU-accelerated build](#gpu-accelerated-build) |

docs/docs/zh/src/indexes/ivf.md

@@ -28,6 +28,49 @@ IVF 通常是一个更合适的默认选择。
 正交中心划分（`first_order_buckets_count` × `second_order_buckets_count`），在超大规模
 语料上能得到更精细的分区。
 
+### 均衡 k-means 训练（`ivf_train_type: "balanced_kmeans"`）
+
+普通 k-means 把每条向量分配给最近的中心，不限制桶容量，因此偏斜、包含大量重复或
+分布不均的数据集会产出大小悬殊甚至为空的桶：过大的桶会拉高检索延迟、削弱分区裁剪。
+
+`balanced_kmeans` 保持与 `kmeans` 相同的平方 L2 目标与均值中心更新，但每一步分配都
+受容量约束：`n` 条训练向量、`buckets_count` 个桶时，每个桶最终恰好容纳
+`floor(n / k)` 或 `ceil(n / k)` 条向量。该保证还延伸到 `Build` 的首次桶分配，
+存储的桶继承均衡性；只要 `1 <= k <= n`，就不会出现空桶。
+
+每次迭代的流程：
+
+1. 每条向量只保留最近的 `num_candidates`（默认 8）个候选桶；距离按固定分块并行逐条
+   计算，永远不会物化 `n × k` 距离矩阵，峰值内存为 `O(n · candidates + k · dim)`
+   而非 `O(n · k)`。
+2. 无约束最近分配吸引票数最多的 `n mod k` 个桶获得 `ceil(n / k)` 个名额，其余桶
+   各得 `floor(n / k)` 个。
+3. 向量按“后悔值”（首选与次选候选的代价差）从大到小依次落位——最容易被挤出的向量
+   优先保住首选名额；每条向量选择仍有空位且距离最近的候选桶；若候选全部满员，则按需
+   重算到所有尚有空位桶的距离。由于总容量恰好等于 `n`，名额总是够用。
+4. 每个中心更新为其成员的均值，重复上述过程直到损失不再下降（由
+   `kmeans_iter_count` 限制迭代次数）。
+
+复杂度与取舍。每次迭代与标准 k-means 同样需要 `O(n · k · dim)` 次距离计算，另加
+`O(n log n)` 排序；内存如上所述有界。均衡分配采用后悔值贪心启发式而非精确的运输问题
+求解，因此损失可能略高于无约束 k-means：在均匀或簇大小均衡的数据上差距很小，在偏斜
+严重的数据上差距更大（等大的桶必须拆分紧密的簇）。可用 `cluster_balance_bench`
+在自家数据形状上实测两边。换来的收益是桶大小的硬上界：每个桶至多容纳
+`ceil(n / k)` 条向量，单次查询的扫描量被 `scan_buckets_count × ceil(n / k)` 封顶，
+长尾延迟不再取决于语料有多偏斜。
+
+补充说明：
+
+- `random` 与 `kmeans` 的行为不变，默认仍是 `kmeans`。
+- cosine 索引下训练后的中心归一化与 `kmeans` 路径完全一致；首次桶分配沿用各度量自身
+  的路由顺序（IP/cosine 用 `1 - <x, c>`，L2 用平方距离）。
+- `balanced_kmeans` 仅支持 `partition_strategy_type: "ivf"`；GNO-IMI 始终训练普通
+  k-means。
+- 均衡训练与首次分配在 CPU 上执行，`enable_gpu_build` 不会加速它们。
+- 给定输入与随机种子时结果确定，且与线程数无关。
+- 用可复现的 `cluster_balance_bench` 工具（`tools/cluster_balance_bench`）对比
+  均衡度、损失、时间与内存。
+
 ## 快速开始
 
 ```cpp
@@ -73,8 +116,8 @@ auto result = index->KnnSearch(
 | `buckets_count` | int | `10` | 倒排列表数量（`ivf` 策略下生效） |
 | `first_order_buckets_count` | int | `10` | 第一级桶数（`gno_imi` 策略下生效） |
 | `second_order_buckets_count` | int | `10` | 第二级桶数（`gno_imi` 策略下生效） |
-| `ivf_train_type` | string | `"kmeans"` | 中心训练方式：`kmeans` 或 `random` |
-| `kmeans_iter_count` | int | `25`（`ivf`）、`30`（`gno_imi`） | KMeans 质心训练的迭代次数（`ivf_train_type: "kmeans"` 时生效）。越大桶划分质量越好，构建越慢。 |
+| `ivf_train_type` | string | `"kmeans"` | 中心训练方式：`kmeans`、`balanced_kmeans` 或 `random` |
+| `kmeans_iter_count` | int | `25`（`ivf`）、`30`（`gno_imi`） | KMeans 质心训练的迭代次数（`ivf_train_type: "kmeans"` 或 `"balanced_kmeans"` 时生效）。越大桶划分质量越好，构建越慢。 |
 | `route_max_degree` | int | `64` | 路由 HGraph 的最大度数（`ivf` 策略下生效） |
 | `route_ef_construction` | int | `300` | 路由 HGraph 的构建搜索宽度（`ivf` 策略下生效） |
 | `enable_gpu_build` | bool | `false` | 构建时使用 CUDA 后端（`ivf` 策略下生效），参见 [GPU 加速构建](#gpu-加速构建) |

src/algorithm/ivf/ivf_build.cpp

@@ -193,8 +193,15 @@ IVF::add(const DatasetPtr& base, bool try_optimized_build) {
     const auto* attr_sets = base->GetAttributeSets();
     const auto* extra_info = base->GetExtraInfos();
     const auto extra_info_size = base->GetExtraInfoSize();
+    // The first Add of a fresh Build is the initial bucket assignment the
+    // partition strategies can refine (issue #2870); every other Add routes to
+    // the nearest buckets.
+    const bool initial_bucket_assignment = try_optimized_build and this->total_elements_ == 0;
     auto buckets =
-        partition_strategy_->ClassifyDatas(vectors, num_element, buckets_per_data_, nullptr);
+        initial_bucket_assignment
+            ? partition_strategy_->ClassifyDatasForBuild(
+                  vectors, num_element, buckets_per_data_, nullptr)
+            : partition_strategy_->ClassifyDatas(vectors, num_element, buckets_per_data_, nullptr);
 
     int64_t current_num;
     bool need_cal_memory_usage = false;

src/algorithm/ivf/ivf_nearest_partition.cpp

@@ -30,6 +30,7 @@
 #include "gpu/gpu_plan.h"
 #include "impl/allocator/safe_allocator.h"
 #include "impl/blas/blas_function.h"
+#include "impl/cluster/balanced_kmeans_cluster.h"
 #include "impl/cluster/kmeans_cluster.h"
 #include "inner_string_params.h"
 #include "query_context.h"
@@ -120,6 +121,19 @@ IVFNearestPartition::Train(const DatasetPtr dataset) {
                 dataset->GetNumElements(),
                 ivf_partition_strategy_param_->GetKMeansIterCount());
         memcpy(data.data(), cls.k_centroids_, dim * this->bucket_count_ * sizeof(float));
+    } else if (ivf_partition_strategy_param_->partition_train_type ==
+               IVFNearestPartitionTrainerType::BalancedKMeansTrainer) {
+        // Same sample, iteration budget and distance handling as the kmeans
+        // branch above; only the assignment step is capacity-constrained. The
+        // trainer always runs on the host: enable_gpu_build accelerates the
+        // plain kmeans passes, and a balanced training stays on the host even
+        // when the flag is on.
+        BalancedKMeansCluster cls(static_cast<int32_t>(dim), this->allocator_, this->thread_pool_);
+        cls.Run(this->bucket_count_,
+                dataset->GetFloat32Vectors(),
+                dataset->GetNumElements(),
+                ivf_partition_strategy_param_->GetKMeansIterCount());
+        memcpy(data.data(), cls.k_centroids_, dim * this->bucket_count_ * sizeof(float));
     } else if (ivf_partition_strategy_param_->partition_train_type ==
                IVFNearestPartitionTrainerType::RandomTrainer) {
         auto selected = select_k_numbers(dataset->GetNumElements(), this->bucket_count_);
@@ -177,6 +191,64 @@ IVFNearestPartition::ClassifyDatas(const void* datas,
     return classify_datas_on_host(datas, count, buckets_per_data, ctx);
 }
 
+Vector<BucketIdType>
+IVFNearestPartition::ClassifyDatasForBuild(const void* datas,
+                                           int64_t count,
+                                           BucketIdType buckets_per_data,
+                                           QueryContext* ctx) const {
+    // Balancing the first assignment is defined for one bucket per vector over a
+    // 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 ==
+                                     IVFNearestPartitionTrainerType::BalancedKMeansTrainer and
+                                 buckets_per_data == 1 and count > 0 and datas != nullptr and
+                                 this->is_trained_ and this->bucket_count_ > 0;
+    if (not balance_buckets) {
+        return this->ClassifyDatas(datas, count, buckets_per_data, ctx);
+    }
+
+    // The layout decides where the centroids are, the same question GetCentroid
+    // asks; the balanced assignment needs them all at once.
+    const uint64_t centroid_bytes = static_cast<uint64_t>(this->bucket_count_) *
+                                    static_cast<uint64_t>(this->dim_) * sizeof(float);
+    Vector<float> centroids(this->allocator_);
+    if (this->use_route_graph_) {
+        centroids.resize(static_cast<uint64_t>(this->bucket_count_) * this->dim_);
+        for (BucketIdType b = 0; b < this->bucket_count_; ++b) {
+            this->route_index_ptr_->GetCodeByInnerId(
+                b,
+                reinterpret_cast<uint8_t*>(centroids.data() +
+                                           static_cast<uint64_t>(b) * this->dim_));
+        }
+    } else {
+        // A trained scan-layout partition stores exactly this many centroids;
+        // anything else cannot be balanced and is routed instead.
+        if (this->centroids_.size() * sizeof(float) != centroid_bytes) {
+            return this->ClassifyDatas(datas, count, buckets_per_data, ctx);
+        }
+        centroids = this->centroids_;
+    }
+
+    // Assign under the index metric: squared L2 for L2SQR indexes, and the
+    // scan routing key 1 - <x, c> for IP and cosine (cosine centroids were
+    // normalized during training), the same orders the plain routing would
+    // follow.
+    BalancedKMeansCluster cls(
+        static_cast<int32_t>(this->dim_), this->allocator_, this->thread_pool_);
+    const auto labels = cls.Assign(this->bucket_count_,
+                                   centroids.data(),
+                                   static_cast<const float*>(datas),
+                                   static_cast<uint64_t>(count),
+                                   this->metric_type_);
+
+    Vector<BucketIdType> result(count, INVALID_BUCKET_ID, this->allocator_);
+    for (int64_t i = 0; i < count; ++i) {
+        result[i] = static_cast<BucketIdType>(labels[i]);
+    }
+    return result;
+}
+
 Vector<BucketIdType>
 IVFNearestPartition::ClassifyDatasForSearch(const void* datas,
                                             int64_t count,

src/algorithm/ivf/ivf_nearest_partition.h

@@ -45,6 +45,17 @@ class IVFNearestPartition : public IVFPartitionStrategy {
                   BucketIdType buckets_per_data,
                   QueryContext* ctx) const override;
 
+    /// With `ivf_train_type: "balanced_kmeans"` the buckets of the first Add of
+    /// a fresh Build are balanced over the trained centroids instead of routed
+    /// to the nearest one, so the capacity guarantee of the trainer extends
+    /// from the training sample to the stored index (issue #2870). Every other
+    /// case uses the plain routing.
+    Vector<BucketIdType>
+    ClassifyDatasForBuild(const void* datas,
+                          int64_t count,
+                          BucketIdType buckets_per_data,
+                          QueryContext* ctx) const override;
+
     Vector<BucketIdType>
     ClassifyDatasForSearch(const void* datas,
                            int64_t count,

src/algorithm/ivf/ivf_partition_strategy.h

@@ -61,6 +61,18 @@ class IVFPartitionStrategy {
                   BucketIdType buckets_per_data,
                   QueryContext* ctx) const = 0;
 
+    /// Bucket assignment for the first Add of a fresh Build, where the index is
+    /// still empty and a strategy may do more than route each vector to its
+    /// nearest bucket (issue #2870: the balanced k-means trainer balances the
+    /// initial buckets instead). The default is the plain routing.
+    virtual Vector<BucketIdType>
+    ClassifyDatasForBuild(const void* datas,
+                          int64_t count,
+                          BucketIdType buckets_per_data,
+                          QueryContext* ctx) const {
+        return this->ClassifyDatas(datas, count, buckets_per_data, ctx);
+    }
+
     virtual Vector<BucketIdType>
     ClassifyDatasForSearch(const void* datas,
                            int64_t count,