diff --git a/docs/docs/en/src/indexes/ivf.md b/docs/docs/en/src/indexes/ivf.md index 640f6f14ab..6dade85ae2 100644 --- a/docs/docs/en/src/indexes/ivf.md +++ b/docs/docs/en/src/indexes/ivf.md @@ -33,6 +33,107 @@ 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. With `buckets_per_data: 1`, the guarantee +extends to the initial bucket assignment of `Build`, so the stored buckets inherit +the balance, and no cluster is empty when balancing applies with `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 - ` 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. +- Initial `Build` assignment is balanced only with `buckets_per_data: 1`. Other + values use ordinary nearest-centroid routing; balanced centroid training still + applies. +- Compare balance, loss, time, and memory yourself with the reproducible + `cluster_balance_bench` tool (`tools/cluster_balance_bench`). + +#### Benchmark tool + +`cluster_balance_bench` trains `KMeansCluster` and `BalancedKMeansCluster` on +seeded synthetic datasets: uniform, deliberately skewed, a skewed set whose size +is not divisible by `k`, and a duplicated-vector set. For each trainer and +dataset it reports: + +- minimum, maximum, and mean cluster size, size standard deviation, and + coefficient of variation; +- clustering loss (mean squared distance from each vector to its assigned + centroid, recomputed from the returned labels with the same yardstick for both + trainers); +- training time; and +- peak memory allocated by the trainer beyond the dataset, measured through a + recording allocator, with a dense `n × k` float distance matrix shown for + reference. + +On Linux, build the tool with `ENABLE_TOOLS=ON` and `ENABLE_CXX11_ABI=ON`: + +```bash +make release EXTRA_DEFINED="-DENABLE_TOOLS=ON -DENABLE_CXX11_ABI=ON" +# Or build just this target in an existing configured build directory: +cmake --build --target cluster_balance_bench +``` + +Usage: + +```text +cluster_balance_bench [--dataset all|uniform|skewed|indivisible|duplicates] + [--n 50000] [--dim 64] [--k 256] [--iter 25] + [--seed 42] [--threads N] +``` + +Every run is fully seeded, including deterministic plain-k-means reductions. +Repeated invocations with the same arguments reproduce the cluster assignments, +centroids, and clustering loss; elapsed time can vary. To preserve this guarantee, +the tool rejects `k >= 10000`: plain k-means would otherwise use approximate HGraph +assignment, whose parallel graph construction is not deterministic. Use a smaller +`k` for reproducible comparisons. + ## Quick start ```cpp @@ -79,8 +180,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) | diff --git a/docs/docs/zh/src/indexes/ivf.md b/docs/docs/zh/src/indexes/ivf.md index e24235d921..69b748cc31 100644 --- a/docs/docs/zh/src/indexes/ivf.md +++ b/docs/docs/zh/src/indexes/ivf.md @@ -28,6 +28,85 @@ 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)` 条向量。当 `buckets_per_data: 1` 时，该保证还延伸到 +`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 - `，L2 用平方距离）。 +- `balanced_kmeans` 仅支持 `partition_strategy_type: "ivf"`；GNO-IMI 始终训练普通 + k-means。 +- 均衡训练与首次分配在 CPU 上执行，`enable_gpu_build` 不会加速它们。 +- 给定输入与随机种子时结果确定，且与线程数无关。 +- 仅当 `buckets_per_data: 1` 时，`Build` 的首次分配才会均衡；其他取值仍按普通的 + 最近中心路由分配，但均衡中心训练仍然生效。 +- 用可复现的 `cluster_balance_bench` 工具（`tools/cluster_balance_bench`）对比 + 均衡度、损失、时间与内存。 + +#### 基准工具 + +`cluster_balance_bench` 在带随机种子的合成数据集上分别训练 `KMeansCluster` 和 +`BalancedKMeansCluster`：均匀数据、刻意偏斜的数据、大小不能整除 `k` 的偏斜数据，以及 +含重复向量的数据。工具针对每种训练器和数据集报告： + +- 桶大小的最小值、最大值、平均值、标准差和变异系数； +- 聚类损失（每条向量到其分配中心的平方距离均值；根据返回的标签重新计算，两种训练器 + 使用相同口径）； +- 训练耗时；以及 +- 训练器在数据集之外分配的峰值内存（通过记录分配器测量），并给出稠密 `n × k` 浮点 + 距离矩阵的大小作为参照。 + +在 Linux 上，使用 `ENABLE_TOOLS=ON` 和 `ENABLE_CXX11_ABI=ON` 构建工具： + +```bash +make release EXTRA_DEFINED="-DENABLE_TOOLS=ON -DENABLE_CXX11_ABI=ON" +# 或在已有的构建目录中只构建此目标： +cmake --build --target cluster_balance_bench +``` + +用法： + +```text +cluster_balance_bench [--dataset all|uniform|skewed|indivisible|duplicates] + [--n 50000] [--dim 64] [--k 256] [--iter 25] + [--seed 42] [--threads N] +``` + +每次运行都完全由种子控制，普通 k-means 也使用确定性归约。相同参数重复运行时，聚类 +分配、中心和聚类损失均可复现；耗时会因运行环境而异。为保证这一点，工具会拒绝 +`k >= 10000`：否则普通 k-means 会使用近似 HGraph 分配，其并行图构建过程不具有确定性。 +如需可复现的对比，请使用更小的 `k`。 + ## 快速开始 ```cpp @@ -73,8 +152,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-加速构建) | diff --git a/src/algorithm/ivf/ivf_build.cpp b/src/algorithm/ivf/ivf_build.cpp index 21ba07b711..d682ddc784 100644 --- a/src/algorithm/ivf/ivf_build.cpp +++ b/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; diff --git a/src/algorithm/ivf/ivf_nearest_partition.cpp b/src/algorithm/ivf/ivf_nearest_partition.cpp index 4c28f5ac26..17c6ce0926 100644 --- a/src/algorithm/ivf/ivf_nearest_partition.cpp +++ b/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(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,70 @@ IVFNearestPartition::ClassifyDatas(const void* datas, return classify_datas_on_host(datas, count, buckets_per_data, ctx); } +bool +IVFNearestPartition::HasDistinctBuildAssignment() const { + return this->ivf_partition_strategy_param_->partition_train_type == + IVFNearestPartitionTrainerType::BalancedKMeansTrainer; +} + +Vector +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(this->bucket_count_) * + static_cast(this->dim_) * sizeof(float); + Vector centroids(this->allocator_); + if (this->use_route_graph_) { + centroids.resize(static_cast(this->bucket_count_) * this->dim_); + for (BucketIdType b = 0; b < this->bucket_count_; ++b) { + this->route_index_ptr_->GetCodeByInnerId( + b, + reinterpret_cast(centroids.data() + + static_cast(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 - for IP and cosine (cosine centroids were + // normalized during training), the same orders the plain routing would + // follow. + BalancedKMeansCluster cls( + static_cast(this->dim_), this->allocator_, this->thread_pool_); + const auto labels = cls.Assign(this->bucket_count_, + centroids.data(), + static_cast(datas), + static_cast(count), + this->metric_type_); + + Vector result(count, INVALID_BUCKET_ID, this->allocator_); + for (int64_t i = 0; i < count; ++i) { + result[i] = static_cast(labels[i]); + } + return result; +} + Vector IVFNearestPartition::ClassifyDatasForSearch(const void* datas, int64_t count, diff --git a/src/algorithm/ivf/ivf_nearest_partition.h b/src/algorithm/ivf/ivf_nearest_partition.h index 90dc2de877..088ef6793e 100644 --- a/src/algorithm/ivf/ivf_nearest_partition.h +++ b/src/algorithm/ivf/ivf_nearest_partition.h @@ -45,6 +45,20 @@ 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 + ClassifyDatasForBuild(const void* datas, + int64_t count, + BucketIdType buckets_per_data, + QueryContext* ctx) const override; + + [[nodiscard]] bool + HasDistinctBuildAssignment() const override; + Vector ClassifyDatasForSearch(const void* datas, int64_t count, diff --git a/src/algorithm/ivf/ivf_partition_strategy.h b/src/algorithm/ivf/ivf_partition_strategy.h index 85c3e20489..da5dfc15e7 100644 --- a/src/algorithm/ivf/ivf_partition_strategy.h +++ b/src/algorithm/ivf/ivf_partition_strategy.h @@ -61,6 +61,25 @@ 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 + ClassifyDatasForBuild(const void* datas, + int64_t count, + BucketIdType buckets_per_data, + QueryContext* ctx) const { + return this->ClassifyDatas(datas, count, buckets_per_data, ctx); + } + + /// Whether Build routing can differ from nearest-centroid routing. Residual + /// training uses this to keep ordinary strategies on the single-route path. + [[nodiscard]] virtual bool + HasDistinctBuildAssignment() const { + return false; + } + virtual Vector ClassifyDatasForSearch(const void* datas, int64_t count, diff --git a/src/algorithm/ivf/ivf_partition_strategy_parameter.cpp b/src/algorithm/ivf/ivf_partition_strategy_parameter.cpp index 909665f3ff..1ac43b62b6 100644 --- a/src/algorithm/ivf/ivf_partition_strategy_parameter.cpp +++ b/src/algorithm/ivf/ivf_partition_strategy_parameter.cpp @@ -68,10 +68,25 @@ IVFPartitionStrategyParameters::IVFPartitionStrategyParameters() = default; void IVFPartitionStrategyParameters::FromJson(const JsonType& json) { - if (json[IVF_TRAIN_TYPE_KEY].GetString() == IVF_TRAIN_TYPE_KMEANS) { + const std::string train_type = json[IVF_TRAIN_TYPE_KEY].GetString(); + if (train_type == IVF_TRAIN_TYPE_KMEANS) { this->partition_train_type = IVFNearestPartitionTrainerType::KMeansTrainer; - } else if (json[IVF_TRAIN_TYPE_KEY].GetString() == IVF_TRAIN_TYPE_RANDOM) { + } else if (train_type == IVF_TRAIN_TYPE_RANDOM) { this->partition_train_type = IVFNearestPartitionTrainerType::RandomTrainer; + } else if (train_type == IVF_TRAIN_TYPE_BALANCED_KMEANS) { + this->partition_train_type = IVFNearestPartitionTrainerType::BalancedKMeansTrainer; + } else { + // Unknown trainer names used to be silently treated as the default. They + // are refused now (issue #2870): a typo must not train an index the user + // asked not to build. Serialized indexes only ever store the names above. + throw VsagException(ErrorType::INVALID_ARGUMENT, + fmt::format("unknown {}: \"{}\", expected one of: " + "\"{}\", \"{}\", \"{}\"", + IVF_TRAIN_TYPE_KEY, + train_type, + IVF_TRAIN_TYPE_KMEANS, + IVF_TRAIN_TYPE_RANDOM, + IVF_TRAIN_TYPE_BALANCED_KMEANS)); } if (json[IVF_PARTITION_STRATEGY_TYPE_KEY].GetString() == IVF_PARTITION_STRATEGY_TYPE_NEAREST) { @@ -80,6 +95,16 @@ IVFPartitionStrategyParameters::FromJson(const JsonType& json) { IVF_PARTITION_STRATEGY_TYPE_GNO_IMI) { this->partition_strategy_type = IVFPartitionStrategyType::GNO_IMI; } + // The two-level GNO-IMI strategy trains plain k-means at both levels; the + // balanced constraint is defined for the single-level IVF strategy. + CHECK_ARGUMENT( + not(this->partition_strategy_type == IVFPartitionStrategyType::GNO_IMI and + this->partition_train_type == IVFNearestPartitionTrainerType::BalancedKMeansTrainer), + fmt::format("{}: \"{}\" is only supported with {} \"{}\"", + IVF_TRAIN_TYPE_KEY, + IVF_TRAIN_TYPE_BALANCED_KMEANS, + IVF_PARTITION_STRATEGY_TYPE_KEY, + IVF_PARTITION_STRATEGY_TYPE_NEAREST)); if (json.Contains(IVF_ROUTE_MAX_DEGREE_KEY)) { this->route_max_degree = static_cast(json[IVF_ROUTE_MAX_DEGREE_KEY].GetInt()); CHECK_ARGUMENT(this->route_max_degree > 0, "route_max_degree must be positive"); @@ -138,6 +163,9 @@ IVFPartitionStrategyParameters::ToJson() const { json[IVF_TRAIN_TYPE_KEY].SetString(IVF_TRAIN_TYPE_KMEANS); } else if (this->partition_train_type == IVFNearestPartitionTrainerType::RandomTrainer) { json[IVF_TRAIN_TYPE_KEY].SetString(IVF_TRAIN_TYPE_RANDOM); + } else if (this->partition_train_type == + IVFNearestPartitionTrainerType::BalancedKMeansTrainer) { + json[IVF_TRAIN_TYPE_KEY].SetString(IVF_TRAIN_TYPE_BALANCED_KMEANS); } if (this->partition_strategy_type == IVFPartitionStrategyType::IVF) { diff --git a/src/algorithm/ivf/ivf_partition_strategy_parameter.h b/src/algorithm/ivf/ivf_partition_strategy_parameter.h index 57a6356269..e646c5b5c4 100644 --- a/src/algorithm/ivf/ivf_partition_strategy_parameter.h +++ b/src/algorithm/ivf/ivf_partition_strategy_parameter.h @@ -27,6 +27,7 @@ namespace vsag { enum class IVFNearestPartitionTrainerType { RandomTrainer = 0, KMeansTrainer = 1, + BalancedKMeansTrainer = 2, }; enum class IVFPartitionStrategyType { @@ -69,7 +70,8 @@ class IVFPartitionStrategyParameters : public Parameter { // are computed, not what the index looks like afterwards, which is why // CheckCompatibility ignores them. bool enable_gpu_build{false}; - // Iterations used by the kmeans centroid training (ivf_train_type == "kmeans"). + // Iterations used by the centroid training (ivf_train_type == "kmeans" or + // "balanced_kmeans"). // Unset (-1) means "use the per-strategy default", which keeps the historical values: // 25 for the single-level IVF strategy and 30 for GNO-IMI. int32_t kmeans_iter_count{-1}; diff --git a/src/algorithm/ivf/ivf_partition_strategy_parameter_test.cpp b/src/algorithm/ivf/ivf_partition_strategy_parameter_test.cpp index 2217b2d33b..db5b8db848 100644 --- a/src/algorithm/ivf/ivf_partition_strategy_parameter_test.cpp +++ b/src/algorithm/ivf/ivf_partition_strategy_parameter_test.cpp @@ -201,3 +201,72 @@ TEST_CASE("IVF Partition Strategy GPU Settings Do Not Affect Compatibility", REQUIRE(host->CheckCompatibility(device)); REQUIRE(device->CheckCompatibility(host)); } + +TEST_CASE("IVF Partition Strategy Balanced KMeans Train Type", + "[ut][IVFPartitionStrategyParameters]") { + // Issue #2870: "balanced_kmeans" is a first-class ivf_train_type value that + // parses, serializes back to its own name, and keeps the defaults of the + // existing trainers untouched. + auto param = std::make_shared(); + param->FromString(R"({"partition_strategy_type": "ivf", "ivf_train_type": "balanced_kmeans"})"); + REQUIRE(param->partition_train_type == + vsag::IVFNearestPartitionTrainerType::BalancedKMeansTrainer); + + const auto dumped = param->ToJson(); + REQUIRE(dumped[vsag::IVF_TRAIN_TYPE_KEY].GetString() == "balanced_kmeans"); + + // The dump round-trips: a reader sees the same trainer. + auto reloaded = std::make_shared(); + reloaded->FromJson(dumped); + REQUIRE(reloaded->partition_train_type == + vsag::IVFNearestPartitionTrainerType::BalancedKMeansTrainer); + vsag::ParameterTest::TestToJson(param); + + // The trainer is a training choice, not a stored-layout difference, so it + // does not split compatibility - as for kmeans vs random. + auto plain = std::make_shared(); + plain->FromString(R"({"partition_strategy_type": "ivf", "ivf_train_type": "kmeans"})"); + auto balanced = std::make_shared(); + balanced->FromString( + R"({"partition_strategy_type": "ivf", "ivf_train_type": "balanced_kmeans"})"); + REQUIRE(plain->CheckCompatibility(balanced)); + REQUIRE(balanced->CheckCompatibility(plain)); + + // The resolved iteration count applies to the balanced trainer as well: + // it shares the kmeans_iter_count knob with the plain kmeans trainer. + auto iter_param = std::make_shared(); + iter_param->FromString( + R"({"partition_strategy_type": "ivf", "ivf_train_type": "balanced_kmeans", + "kmeans_iter_count": 9})"); + REQUIRE(iter_param->partition_train_type == + vsag::IVFNearestPartitionTrainerType::BalancedKMeansTrainer); + REQUIRE(iter_param->GetKMeansIterCount() == 9); + REQUIRE(iter_param->ToJson()[vsag::IVF_KMEANS_ITER_COUNT_KEY].GetInt() == 9); +} + +TEST_CASE("IVF Partition Strategy Rejects Unknown and Incompatible Train Types", + "[ut][IVFPartitionStrategyParameters]") { + // An unknown trainer name must be refused with a clear invalid-argument + // error instead of silently training with the default (issue #2870). + auto param = std::make_shared(); + REQUIRE_THROWS( + param->FromString(R"({"partition_strategy_type": "ivf", "ivf_train_type": "kmeanz"})")); + + // GNO-IMI always trains plain k-means at both levels; the balanced + // constraint belongs to the single-level IVF strategy. + REQUIRE_THROWS(param->FromString(R"( + { + "partition_strategy_type": "gno_imi", + "ivf_train_type": "balanced_kmeans", + "gno_imi": { + "first_order_buckets_count": 20, + "second_order_buckets_count": 10 + } + })")); + + // The existing names keep parsing with both strategies exactly as before. + REQUIRE_NOTHROW(param->FromString(R"({"partition_strategy_type": "ivf", + "ivf_train_type": "kmeans"})")); + REQUIRE_NOTHROW(param->FromString(R"({"partition_strategy_type": "ivf", + "ivf_train_type": "random"})")); +} diff --git a/src/datacell/bucket_datacell.h b/src/datacell/bucket_datacell.h index 08f50bdd26..1bdde8bda6 100644 --- a/src/datacell/bucket_datacell.h +++ b/src/datacell/bucket_datacell.h @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -625,24 +626,178 @@ void BucketDataCell::Train(const void* data, uint64_t count) { Vector train_data_buffer(allocator_); if (use_residual_) { - auto data_ptr = static_cast(data); - auto dim = this->quantizer_->GetDim(); - train_data_buffer.resize(count * dim); + const auto* original_data_ptr = static_cast(data); + auto data_ptr = original_data_ptr; + const uint64_t dim = static_cast(this->quantizer_->GetDim()); + const bool train_nearest = strategy_->HasDistinctBuildAssignment(); + bool prefix_limited = false; + bool scalar_limited = false; + if (train_nearest) { + 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 + quantizer_name == QUANTIZATION_TYPE_VALUE_SQ4 or + quantizer_name == QUANTIZATION_TYPE_VALUE_SQ8_UNIFORM or + quantizer_name == QUANTIZATION_TYPE_VALUE_SQ4_UNIFORM; + } + const uint64_t buffer_factor = train_nearest ? 2 : 1; + if (dim == 0 or count > std::numeric_limits::max() / dim / buffer_factor) { + throw VsagException(ErrorType::INVALID_ARGUMENT, "residual training size overflow"); + } + if (metric_ == MetricType::METRIC_TYPE_COSINE or not train_nearest) { + // Cosine classification needs a normalized corpus. Other single-route training can + // write its residuals immediately; dual-route L2/IP allocates after classification. + train_data_buffer.resize(count * dim); + } if (metric_ == MetricType::METRIC_TYPE_COSINE) { - for (int i = 0; i < count; ++i) { + for (uint64_t i = 0; i < count; ++i) { Normalize(data_ptr + i * dim, train_data_buffer.data() + i * dim, dim); } data_ptr = train_data_buffer.data(); } Vector centroid(this->quantizer_->GetDim(), allocator_); - auto buckets = strategy_->ClassifyDatas(data_ptr, count, 1, nullptr); - for (int i = 0; i < count; ++i) { - strategy_->GetCentroid(buckets[i], centroid); - for (int j = 0; j < dim; ++j) { - train_data_buffer[i * dim + j] = data_ptr[i * dim + j] - centroid[j]; + auto build_buckets = strategy_->ClassifyDatasForBuild(data_ptr, count, 1, nullptr); + uint64_t train_count = count; + if (train_nearest) { + // Balanced Build uses capacity-constrained assignments, while later Add + // calls use nearest-centroid routing. Only this path needs both residual + // populations and the corresponding second classification pass. + auto nearest_buckets = strategy_->ClassifyDatas(data_ptr, count, 1, nullptr); + Vector nearest_rows(allocator_); + if (prefix_limited or scalar_limited) { + for (uint64_t i = 0; i < count; ++i) { + if (build_buckets[i] != nearest_buckets[i]) { + nearest_rows.emplace_back(i); + } + } + const uint64_t max_train_count = prefix_limited ? 65536 : 100000; + const bool sample_populations = count + nearest_rows.size() > max_train_count; + uint64_t build_samples = count; + uint64_t nearest_samples = static_cast(nearest_rows.size()); + if (sample_populations) { + // PQ/PQFS retain only their first 65,536 rows, while scalar quantizers + // sample 100,000 rows. Budget both populations before either limit can + // discard appended nearest residuals. + build_samples = std::min(count, max_train_count / 2); + nearest_samples = std::min(static_cast(nearest_rows.size()), + max_train_count - build_samples); + build_samples = std::min(count, max_train_count - nearest_samples); + } + train_count = build_samples + nearest_samples; + // The sampled rows are re-derived from the original input. Release + // any normalized classification copy before allocating only the rows + // that the sample-limited quantizer will actually consume. + Vector(allocator_).swap(train_data_buffer); + train_data_buffer.resize(train_count * dim); + Vector normalized_row(dim, allocator_); + auto write_residual = + [&](uint64_t destination, uint64_t source, BucketIdType bucket) { + const float* src = original_data_ptr + source * dim; + if (metric_ == MetricType::METRIC_TYPE_COSINE) { + Normalize(src, normalized_row.data(), dim); + src = normalized_row.data(); + } + strategy_->GetCentroid(bucket, centroid); + for (uint64_t j = 0; j < dim; ++j) { + train_data_buffer[destination * dim + j] = src[j] - centroid[j]; + } + }; + std::mt19937_64 random(0x524553494455414cULL); + auto for_each_sample = [&](bool should_sample, + uint64_t sample_count, + uint64_t population_count, + auto write_sample) { + for (uint64_t i = 0; i < sample_count; ++i) { + write_sample(i, i); + } + if (not should_sample) { + return; + } + // Use the residual buffer itself as a reservoir. Every subset has the + // same probability, without fixed-stride phase aliases, duplicate rows, + // or any sampling allocation beyond the existing training budget. + for (uint64_t i = sample_count; i < population_count; ++i) { + std::uniform_int_distribution select(0, i); + const uint64_t selected = select(random); + if (selected < sample_count) { + write_sample(selected, i); + } + } + }; + // Without nearest-only rows, retain the original PQ prefix. Scalar training + // already sampled a lone population, so reservoir-sample it here instead of + // silently reducing it to a prefix. + const bool sample_build = + sample_populations && (nearest_samples != 0 or scalar_limited); + for_each_sample( + sample_build, build_samples, count, [&](uint64_t destination, uint64_t source) { + write_residual(destination, source, build_buckets[source]); + }); + for_each_sample(sample_populations, + nearest_samples, + nearest_rows.size(), + [&](uint64_t destination, uint64_t nearest) { + const uint64_t source = nearest_rows[nearest]; + write_residual(build_samples + destination, + source, + nearest_buckets[source]); + }); + } else { + // Keep complete populations when they fit. Allocate only rows whose routes + // disagree in addition to the Build corpus; cosine retains its normalized prefix. + for (uint64_t i = 0; i < count; ++i) { + if (build_buckets[i] != nearest_buckets[i]) { + ++train_count; + } + } + const bool regrow_normalized = + metric_ == MetricType::METRIC_TYPE_COSINE and train_count > count; + if (regrow_normalized) { + // Do not let vector growth hold the old and new corpus allocations at once. + Vector(allocator_).swap(train_data_buffer); + } + train_data_buffer.resize(train_count * dim); + if (metric_ == MetricType::METRIC_TYPE_COSINE) { + if (regrow_normalized) { + for (uint64_t i = 0; i < count; ++i) { + Normalize(original_data_ptr + i * dim, + train_data_buffer.data() + i * dim, + dim); + } + } + data_ptr = train_data_buffer.data(); + } + // Each nearest extra is written before its normalized source row is overwritten + // by the Build residual. + uint64_t nearest_offset = count; + for (uint64_t i = 0; i < count; ++i) { + if (build_buckets[i] != nearest_buckets[i]) { + strategy_->GetCentroid(nearest_buckets[i], centroid); + for (uint64_t j = 0; j < dim; ++j) { + train_data_buffer[nearest_offset * dim + j] = + data_ptr[i * dim + j] - centroid[j]; + } + ++nearest_offset; + } + strategy_->GetCentroid(build_buckets[i], centroid); + for (uint64_t j = 0; j < dim; ++j) { + train_data_buffer[i * dim + j] = data_ptr[i * dim + j] - centroid[j]; + } + } + } + } else { + // Ordinary kmeans/random strategies have identical Build and Add routes: + // retain the previous single-classification, single-buffer behavior. + for (uint64_t i = 0; i < count; ++i) { + strategy_->GetCentroid(build_buckets[i], centroid); + for (uint64_t j = 0; j < dim; ++j) { + train_data_buffer[i * dim + j] = data_ptr[i * dim + j] - centroid[j]; + } } } data = train_data_buffer.data(); + count = train_count; } this->quantizer_->Train(reinterpret_cast(data), count); } diff --git a/src/datacell/bucket_datacell_test.cpp b/src/datacell/bucket_datacell_test.cpp index 242f6366c8..8f6a3a814b 100644 --- a/src/datacell/bucket_datacell_test.cpp +++ b/src/datacell/bucket_datacell_test.cpp @@ -46,7 +46,26 @@ class BoundedBucketAllocator : public DefaultAllocator { void* Allocate(uint64_t size) override { check_size(size); - return DefaultAllocator::Allocate(size); + auto* pointer = DefaultAllocator::Allocate(size); + if (size == capture_size) { + // Classification arrays can have the same size as the training buffer, which is + // allocated last. Ignore earlier matches instead of depending on destruction order. + captured_pointers.clear(); + captured_pointers.emplace_back(pointer); + } + return pointer; + } + + void + Deallocate(void* pointer) override { + auto captured = std::find(captured_pointers.begin(), captured_pointers.end(), pointer); + if (captured != captured_pointers.end()) { + REQUIRE(captured_pointers.size() == 1); + const auto* values = static_cast(pointer); + captured_values.assign(values, values + capture_size / sizeof(float)); + captured_pointers.erase(captured); + } + DefaultAllocator::Deallocate(pointer); } void* @@ -56,12 +75,17 @@ class BoundedBucketAllocator : public DefaultAllocator { } uint64_t limit{std::numeric_limits::max()}; + uint64_t rejected_size{0}; uint64_t rejected{0}; + uint64_t capture_size{0}; + std::vector captured_values; private: + std::vector captured_pointers; + void check_size(uint64_t size) { - if (size > limit) { + if (size > limit or (rejected_size != 0 and size == rejected_size)) { ++rejected; throw std::bad_alloc(); } @@ -83,6 +107,7 @@ class FixedCentroidPartitionStrategy : public IVFPartitionStrategy { Vector ClassifyDatas(const void*, int64_t count, BucketIdType, QueryContext*) const override { + classify_requests_.fetch_add(1, std::memory_order_relaxed); return Vector(count, 0, allocator_); } @@ -99,11 +124,173 @@ class FixedCentroidPartitionStrategy : public IVFPartitionStrategy { return centroid_requests_.load(std::memory_order_relaxed); } + [[nodiscard]] uint64_t + GetClassifyRequestCount() const { + return classify_requests_.load(std::memory_order_relaxed); + } + private: std::atomic centroid_requests_{0}; + mutable std::atomic classify_requests_{0}; float centroid_scale_{1.0F}; }; +class DualRouteFixedCentroidPartitionStrategy : public FixedCentroidPartitionStrategy { +public: + using FixedCentroidPartitionStrategy::FixedCentroidPartitionStrategy; + + [[nodiscard]] bool + HasDistinctBuildAssignment() const override { + return true; + } +}; + +class BuildAssignedPartitionStrategy : public FixedCentroidPartitionStrategy { +public: + using FixedCentroidPartitionStrategy::FixedCentroidPartitionStrategy; + + Vector + ClassifyDatasForBuild(const void*, int64_t count, BucketIdType, QueryContext*) const override { + return Vector(count, 1, allocator_); + } + + [[nodiscard]] bool + HasDistinctBuildAssignment() const override { + return true; + } +}; + +class OrderedBalancedResidualPartitionStrategy : public IVFPartitionStrategy { +public: + OrderedBalancedResidualPartitionStrategy(const IndexCommonParam& common_param, + BucketIdType bucket_count) + : IVFPartitionStrategy(common_param, bucket_count) { + is_trained_ = true; + } + + void + Train(const DatasetPtr) override { + } + + Vector + ClassifyDatas(const void* datas, int64_t count, BucketIdType, QueryContext*) const override { + const auto* values = static_cast(datas); + Vector result(count, 0, allocator_); + for (int64_t i = 0; i < count; ++i) { + result[i] = values[i] > 1.0F ? 1 : 0; + } + return result; + } + + Vector + ClassifyDatasForBuild(const void*, int64_t count, BucketIdType, QueryContext*) const override { + Vector result(count, 0, allocator_); + for (int64_t i = 0; i < count; ++i) { + result[i] = i < count / 2 ? 0 : 1; + } + return result; + } + + [[nodiscard]] bool + HasDistinctBuildAssignment() const override { + return true; + } + + void + GetCentroid(BucketIdType bucket_id, Vector& centroid) override { + centroid[0] = bucket_id == 0 ? 0.0F : 185530.0F / 65536.0F; + } +}; + +class MixedResidualPartitionStrategy : public IVFPartitionStrategy { +public: + MixedResidualPartitionStrategy(const IndexCommonParam& common_param, BucketIdType bucket_count) + : IVFPartitionStrategy(common_param, bucket_count) { + is_trained_ = true; + } + + void + Train(const DatasetPtr) override { + } + + Vector + ClassifyDatas(const void* datas, int64_t count, BucketIdType, QueryContext*) const override { + const auto* values = static_cast(datas); + Vector result(count, 0, allocator_); + for (int64_t i = 0; i < count; ++i) { + result[i] = values[i] == 100.0F ? 1 : 0; + } + return result; + } + + Vector + ClassifyDatasForBuild(const void* datas, + int64_t count, + BucketIdType, + QueryContext*) const override { + const auto* values = static_cast(datas); + Vector result(count, 0, allocator_); + for (int64_t i = 0; i < count; ++i) { + result[i] = values[i] == 0.0F ? 0 : 1; + } + return result; + } + + [[nodiscard]] bool + HasDistinctBuildAssignment() const override { + return true; + } + + void + GetCentroid(BucketIdType bucket_id, Vector& centroid) override { + centroid[0] = bucket_id == 0 ? 0.0F : 28.0F; + } +}; + +class PeriodicResidualPartitionStrategy : public IVFPartitionStrategy { +public: + PeriodicResidualPartitionStrategy(const IndexCommonParam& common_param, + BucketIdType bucket_count) + : IVFPartitionStrategy(common_param, bucket_count) { + is_trained_ = true; + } + + void + Train(const DatasetPtr) override { + } + + Vector + ClassifyDatas(const void*, int64_t count, BucketIdType, QueryContext*) const override { + Vector result(count, 0, allocator_); + for (int64_t i = 0; i < count; ++i) { + // Half the rows mismatch Build, exceeding the nearest-only sampling budget. + result[i] = (i % 4 == 0 or i % 4 == 3) ? 1 : 0; + } + return result; + } + + Vector + ClassifyDatasForBuild(const void*, int64_t count, BucketIdType, QueryContext*) const override { + Vector result(count, 0, allocator_); + for (int64_t i = 0; i < count; ++i) { + // Each [0, 100, 100, 100] period forces its first 100 into centroid 50. + result[i] = i % 4 < 2 ? 0 : 1; + } + return result; + } + + [[nodiscard]] bool + HasDistinctBuildAssignment() const override { + return true; + } + + void + GetCentroid(BucketIdType bucket_id, Vector& centroid) override { + std::fill(centroid.begin(), centroid.end(), 0.0F); + centroid[0] = bucket_id == 0 ? 50.0F : 100.0F; + } +}; + class TrackingReader : public Reader { public: explicit TrackingReader(std::shared_ptr data) : data_(std::move(data)) { @@ -795,6 +982,315 @@ TEST_CASE("BucketDataCell batch query", "[ut][BucketDataCell]") { } } +TEST_CASE("BucketDataCell ordinary residual training keeps one route and buffer", + "[ut][BucketDataCell]") { + auto allocator = std::make_shared(); + constexpr int64_t dim = 64; + constexpr uint64_t count = 1024; + constexpr uint64_t bucket_count = 1; + const std::vector vectors(count * dim, 1.0F); + auto param = std::make_shared(); + param->FromJson(JsonType::Parse(R"({ + "io_params": { "type": "memory_io" }, + "quantization_params": { "type": "fp32" }, + "buckets_count": 1, + "use_residual": true + })")); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + auto strategy = std::make_shared(common_param, bucket_count); + bucket->SetStrategy(strategy); + // One FP32 corpus is 256 KiB and the old dual residual buffer is 512 KiB. + allocator->limit = 300 * 1024; + REQUIRE_NOTHROW(bucket->Train(vectors.data(), count)); + REQUIRE(allocator->rejected == 0); + REQUIRE(strategy->GetClassifyRequestCount() == 1); +} + +TEST_CASE("BucketDataCell residual training uses build and nearest assignments", + "[ut][BucketDataCell]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + constexpr int64_t dim = 1; + constexpr uint64_t bucket_count = 2; + const std::vector vectors = {50.0F, 100.0F}; + auto param_json = JsonType::Parse(R"({ + "io_params": { "type": "memory_io" }, + "quantization_params": { "type": "sq8" }, + "buckets_count": 2, + "use_residual": true + })"); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + // Centroids are 50 and 100. Build assigns both rows to bucket 1, producing + // residuals [-50, 0], while nearest routing assigns both to bucket 0, + // producing [0, 50]. Each one-sided training set clips the other path. + auto strategy = std::make_shared( + common_param, bucket_count, /*centroid_scale=*/5000.0F); + bucket->SetStrategy(strategy); + bucket->Train(vectors.data(), vectors.size()); + bucket->InsertVector(vectors.data(), 1, 0); + bucket->InsertVector(vectors.data() + dim, 0, 1); + + auto build_computer = bucket->FactoryComputer(vectors.data()); + REQUIRE(std::abs(bucket->QueryOneById(build_computer, 1, 0)) < 1e-4F); + auto nearest_computer = bucket->FactoryComputer(vectors.data() + dim); + REQUIRE(std::abs(bucket->QueryOneById(nearest_computer, 0, 0)) < 1e-4F); +} + +TEST_CASE("BucketDataCell samples scalar build and nearest residual populations", + "[ut][BucketDataCell]") { + auto allocator = std::make_shared(); + constexpr int64_t dim = 1; + constexpr uint64_t count = 131072; + constexpr uint64_t zero_count = 65536; + constexpr uint64_t outlier_count = 6; + constexpr uint64_t nearest_population = count - zero_count - outlier_count; + constexpr uint64_t train_count = count + nearest_population; + constexpr uint64_t max_scalar_samples = 100000; + constexpr uint64_t bucket_count = 2; + static_assert(nearest_population == 65530); + static_assert(2147483647UL % train_count == 1); + const std::string quantizer_type = GENERATE("sq8", "sq4", "sq8_uniform", "sq4_uniform"); + CAPTURE(quantizer_type); + + std::vector vectors(count, 1.0F); + std::fill_n(vectors.begin(), zero_count, 0.0F); + std::fill(vectors.end() - outlier_count, vectors.end(), 20000.0F); + + JsonType param_json; + JsonType io_json; + io_json["type"].SetString("memory_io"); + param_json["io_params"].SetJson(io_json); + JsonType quantizer_json; + quantizer_json["type"].SetString(quantizer_type); + param_json["quantization_params"].SetJson(quantizer_json); + param_json["buckets_count"].SetInt(bucket_count); + param_json["use_residual"].SetBool(true); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + bucket->SetStrategy( + std::make_shared(common_param, bucket_count)); + allocator->capture_size = max_scalar_samples * dim * sizeof(float); + allocator->rejected_size = train_count * dim * sizeof(float); + bucket->Train(vectors.data(), count); + REQUIRE(allocator->rejected == 0); + REQUIRE(allocator->captured_values.size() == max_scalar_samples); + + // Half the scalar budget represents the 65,530 nearest-only residuals. The old + // stride of one sampled only the Build prefix, making the upper bound zero. + for (uint64_t i = max_scalar_samples / 2; i < max_scalar_samples; ++i) { + REQUIRE(allocator->captured_values[i] == 1.0F); + } + + const float added = 1.0F; + bucket->InsertVector(&added, 0, 0); + auto computer = bucket->FactoryComputer(&added); + // Training only the Build prefix clips this +1 residual against a zero upper bound. + REQUIRE(std::abs(bucket->QueryOneById(computer, 0, 0)) < 1e-3F); +} + +TEST_CASE("BucketDataCell residual training preserves the PQ sample prefix", + "[ut][BucketDataCell]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + constexpr int64_t dim = 1; + constexpr uint64_t count = 131072; + constexpr uint64_t bucket_count = 1; + const std::string quantizer_type = GENERATE("pq", "pqfs"); + CAPTURE(quantizer_type); + std::vector vectors(count); + for (uint64_t i = 0; i < count; ++i) { + vectors[i] = i % 2 == 0 ? 0.0F : 100.0F; + } + + JsonType param_json; + JsonType io_json; + io_json["type"].SetString("memory_io"); + param_json["io_params"].SetJson(io_json); + param_json["quantization_params"].SetJson( + JsonType::Parse(quantizer_type == "pq" ? R"({"type": "pq", "pq_dim": 1, "pq_bits": 8})" + : R"({"type": "pqfs", "pq_dim": 1})")); + param_json["buckets_count"].SetInt(bucket_count); + param_json["use_residual"].SetBool(true); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + // The sole centroid is 50, so alternating rows have residuals -50 and +50. + // Both Build and nearest routing agree, leaving nearest_samples at zero. + // Fixed-stride sampling used to select only even rows, filling PQ's + // 65,536-row prefix with just the -50 residual. + auto strategy = std::make_shared( + common_param, bucket_count, /*centroid_scale=*/5000.0F); + bucket->SetStrategy(strategy); + bucket->Train(vectors.data(), count); + REQUIRE(strategy->GetClassifyRequestCount() == 2); + bucket->InsertVector(vectors.data() + count - 1, 0, count - 1); + bucket->Package(); + + auto computer = bucket->FactoryComputer(vectors.data() + count - 1); + REQUIRE(std::abs(bucket->QueryOneById(computer, 0, 0)) < 1e-3F); +} + +TEST_CASE("BucketDataCell residual sampling avoids periodic population aliases", + "[ut][BucketDataCell]") { + auto allocator = std::make_shared(); + constexpr int64_t dim = 2; + constexpr uint64_t count = 131072; + constexpr uint64_t build_samples = 32768; + constexpr uint64_t nearest_samples = 32768; + constexpr uint64_t nearest_population = count / 2; + static_assert(nearest_samples < nearest_population); + constexpr uint64_t bucket_count = 2; + const std::string quantizer_type = GENERATE("pq", "pqfs"); + CAPTURE(quantizer_type); + std::vector vectors(count * dim, 0.0F); + for (uint64_t i = 0; i < count; ++i) { + vectors[i * dim] = i % 4 == 0 ? 0.0F : 100.0F; + // The second dimension only identifies source rows for uniqueness checks. + vectors[i * dim + 1] = static_cast(i); + } + + JsonType param_json; + JsonType io_json; + io_json["type"].SetString("memory_io"); + param_json["io_params"].SetJson(io_json); + param_json["quantization_params"].SetJson( + JsonType::Parse(quantizer_type == "pq" ? R"({"type": "pq", "pq_dim": 1, "pq_bits": 8})" + : R"({"type": "pqfs", "pq_dim": 1})")); + param_json["buckets_count"].SetInt(bucket_count); + param_json["use_residual"].SetBool(true); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + bucket->SetStrategy( + std::make_shared(common_param, bucket_count)); + allocator->capture_size = (build_samples + nearest_samples) * dim * sizeof(float); + allocator->rejected_size = 2 * count * dim * sizeof(float); + bucket->Train(vectors.data(), count); + REQUIRE(allocator->rejected == 0); + REQUIRE(allocator->captured_values.size() == (build_samples + nearest_samples) * dim); + + std::vector build_seen(count, false); + std::vector nearest_seen(count, false); + uint64_t negative_build_residuals = 0; + uint64_t positive_build_residuals = 0; + uint64_t negative_nearest_residuals = 0; + uint64_t positive_nearest_residuals = 0; + for (uint64_t i = 0; i < build_samples; ++i) { + const float residual = allocator->captured_values[i * dim]; + const float source_value = allocator->captured_values[i * dim + 1]; + REQUIRE(source_value >= 0.0F); + REQUIRE(source_value < static_cast(count)); + REQUIRE(source_value == std::floor(source_value)); + const auto source = static_cast(source_value); + REQUIRE_FALSE(build_seen[source]); + build_seen[source] = true; + + const float value = source % 4 == 0 ? 0.0F : 100.0F; + const float centroid = source % 4 < 2 ? 50.0F : 100.0F; + REQUIRE(residual == value - centroid); + negative_build_residuals += residual == -50.0F ? 1 : 0; + positive_build_residuals += residual == 50.0F ? 1 : 0; + } + // Fixed stride 4*i selected only the -50 Build residuals from this corpus. + REQUIRE(negative_build_residuals > 0); + REQUIRE(positive_build_residuals > 0); + + 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(count)); + REQUIRE(source_value == std::floor(source_value)); + const auto source = static_cast(source_value); + REQUIRE_FALSE(nearest_seen[source]); + nearest_seen[source] = true; + REQUIRE(source % 2 == 0); + + const float value = source % 4 == 0 ? 0.0F : 100.0F; + const float centroid = source % 4 < 2 ? 100.0F : 50.0F; + REQUIRE(residual == value - centroid); + negative_nearest_residuals += residual < 0.0F ? 1 : 0; + positive_nearest_residuals += residual > 0.0F ? 1 : 0; + } + // Fixed stride over the 65,536 mismatches selected only -100 nearest residuals. + REQUIRE(negative_nearest_residuals > 0); + REQUIRE(positive_nearest_residuals > 0); +} + +TEST_CASE("BucketDataCell samples nearest residuals within the PQ prefix", "[ut][BucketDataCell]") { + auto allocator = std::make_shared(); + constexpr int64_t dim = 1; + constexpr uint64_t count = 65536; + constexpr uint64_t bucket_count = 2; + const std::string quantizer_type = GENERATE("pq", "pqfs"); + CAPTURE(quantizer_type); + std::vector vectors(count, 0.0F); + for (uint64_t i = 0; i < count; ++i) { + const uint64_t kind = i % 10; + vectors[i] = kind < 5 ? 0.0F : (kind < 9 ? 10.0F : 100.0F); + } + + JsonType param_json; + JsonType io_json; + io_json["type"].SetString("memory_io"); + param_json["io_params"].SetJson(io_json); + param_json["quantization_params"].SetJson( + JsonType::Parse(quantizer_type == "pq" ? R"({"type": "pq", "pq_dim": 1, "pq_bits": 8})" + : R"({"type": "pqfs", "pq_dim": 1})")); + param_json["buckets_count"].SetInt(bucket_count); + param_json["use_residual"].SetBool(true); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + bucket->SetStrategy( + std::make_shared(common_param, bucket_count)); + // The old temporary held 2 * count rows although PQ/PQFS consume at most + // count rows here. Reject that exact allocation without bounding trainer workspaces. + allocator->rejected_size = 2 * count * dim * sizeof(float); + bucket->Train(vectors.data(), count); + REQUIRE(allocator->rejected == 0); + + // Build assigns 10 to centroid 28 (residual -18), but Add routes it to + // centroid 0 (residual +10). The latter used to start after the PQ prefix. + const float added = 10.0F; + bucket->InsertVector(&added, 0, 0); + bucket->Package(); + auto computer = bucket->FactoryComputer(&added); + REQUIRE(std::abs(bucket->QueryOneById(computer, 0, 0)) < 1e-3F); +} + TEST_CASE("BucketDataCell batch query preserves residual correction", "[ut][BucketDataCell]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); constexpr int64_t dim = 8; @@ -1382,6 +1878,59 @@ TEST_CASE("RaBitQ split bucket supports optimized build", REQUIRE_FALSE(optimized->IsOptimizedBuildActive()); } +TEST_CASE("RaBitQ split residual training uses build and nearest assignments", + "[ut][BucketDataCell][rabitq_split][residual]") { + auto allocator = std::make_shared(); + constexpr uint64_t dim = 16; + constexpr uint64_t training_count = 64; + constexpr uint64_t bucket_count = 2; + auto vectors = fixtures::generate_vectors(training_count, dim); + auto param = std::make_shared(); + param->FromJson(JsonType::Parse(R"({ + "io_params": { "type": "memory_io" }, + "quantization_params": { + "type": "rabitq", + "rabitq_version": "split", + "rabitq_bits_per_dim_query": 32, + "rabitq_bits_per_dim_base": 8, + "rabitq_bits_per_dim_filter": 1 + }, + "buckets_count": 2, + "use_residual": true + })")); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + auto strategy = std::make_shared( + common_param, bucket_count, /*centroid_scale=*/5000.0F); + bucket->SetStrategy(strategy); + allocator->capture_size = 2 * training_count * dim * sizeof(float); + bucket->Train(vectors.data(), training_count); + + // Every row has a different Build and nearest route, so training must fetch + // both centroids per row. Rebuilding residual query transforms may add more. + REQUIRE(strategy->GetClassifyRequestCount() == 1); + REQUIRE(strategy->GetCentroidRequestCount() >= 2 * training_count); + + // Check the actual quantizer input, not just the classification counters: Build residuals + // occupy the prefix and each nearest residual is appended at count + i against centroid 0. + Vector nearest_centroid(dim, allocator.get()); + Vector build_centroid(dim, allocator.get()); + strategy->GetCentroid(0, nearest_centroid); + strategy->GetCentroid(1, build_centroid); + std::vector expected(2 * training_count * dim); + for (uint64_t i = 0; i < training_count; ++i) { + for (uint64_t d = 0; d < dim; ++d) { + expected[i * dim + d] = vectors[i * dim + d] - build_centroid[d]; + expected[(training_count + i) * dim + d] = vectors[i * dim + d] - nearest_centroid[d]; + } + } + REQUIRE(allocator->captured_values == expected); +} + TEST_CASE("RaBitQ split bucket scan spills computed masks above 256 blocks", "[ut][BucketDataCell][rabitq_split][residual]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); diff --git a/src/datacell/rabitq_split_bucket_datacell.cpp b/src/datacell/rabitq_split_bucket_datacell.cpp index 0b5337ff0b..4d9e5a531c 100644 --- a/src/datacell/rabitq_split_bucket_datacell.cpp +++ b/src/datacell/rabitq_split_bucket_datacell.cpp @@ -1348,31 +1348,74 @@ RaBitQSplitBucketDataCell::Train(const void* data, uint64_t count) { return; } + CHECK_ARGUMENT(count <= static_cast(std::numeric_limits::max()), + "RaBitQ split training count is too large"); const auto* input = static_cast(data); - Vector train_data(count * this->dim_, this->allocator_); - Vector normalized(this->dim_, this->allocator_); + const bool train_nearest = this->strategy_->HasDistinctBuildAssignment(); + const uint64_t buffer_factor = train_nearest ? 2 : 1; + CHECK_ARGUMENT(this->dim_ != 0 and + count <= std::numeric_limits::max() / this->dim_ / buffer_factor, + "RaBitQ split residual training size overflow"); Vector centroid(this->dim_, this->allocator_); const float* classify_data = input; - Vector normalized_data(this->allocator_); + Vector train_data(this->allocator_); if (this->metric_ == MetricType::METRIC_TYPE_COSINE) { - normalized_data.resize(count * this->dim_); + // Normalize into the residual buffer so the corpus is not duplicated for classification. + train_data.resize(count * this->dim_); + for (uint64_t i = 0; i < count; ++i) { + Normalize(input + i * this->dim_, train_data.data() + i * this->dim_, this->dim_); + } + classify_data = train_data.data(); + } + // Balanced Build uses capacity-constrained assignments, while later Add calls + // use nearest-centroid routing. Train on both residual populations so either + // route remains inside the population used to fit the RaBitQ codes. + auto build_buckets = this->strategy_->ClassifyDatasForBuild( + classify_data, static_cast(count), 1, nullptr); + Vector nearest_buckets(this->allocator_); + uint64_t train_count = count; + if (train_nearest) { + nearest_buckets = + this->strategy_->ClassifyDatas(classify_data, static_cast(count), 1, nullptr); for (uint64_t i = 0; i < count; ++i) { - Normalize(input + i * this->dim_, normalized_data.data() + i * this->dim_, this->dim_); + if (build_buckets[i] != nearest_buckets[i]) { + ++train_count; + } } - classify_data = normalized_data.data(); } - CHECK_ARGUMENT(count <= static_cast(std::numeric_limits::max()), - "RaBitQ split training count is too large"); - auto buckets = - this->strategy_->ClassifyDatas(classify_data, static_cast(count), 1, nullptr); + const bool regrow_normalized = + this->metric_ == MetricType::METRIC_TYPE_COSINE and train_count > count; + if (regrow_normalized) { + // Release the classification corpus before growing so vector reallocation cannot briefly + // retain both the old buffer and the complete dual-population buffer. + Vector(this->allocator_).swap(train_data); + } + train_data.resize(train_count * this->dim_); + if (this->metric_ == MetricType::METRIC_TYPE_COSINE) { + if (regrow_normalized) { + for (uint64_t i = 0; i < count; ++i) { + Normalize(input + i * this->dim_, train_data.data() + i * this->dim_, this->dim_); + } + } + classify_data = train_data.data(); + } + uint64_t nearest_offset = count; for (uint64_t i = 0; i < count; ++i) { - this->strategy_->GetCentroid(buckets[i], centroid); + if (train_nearest and build_buckets[i] != nearest_buckets[i]) { + this->strategy_->GetCentroid(nearest_buckets[i], centroid); + FP32Sub(classify_data + i * this->dim_, + centroid.data(), + train_data.data() + nearest_offset * this->dim_, + this->dim_); + ++nearest_offset; + } + this->strategy_->GetCentroid(build_buckets[i], centroid); FP32Sub(classify_data + i * this->dim_, centroid.data(), train_data.data() + i * this->dim_, this->dim_); } - this->codes_->Train(train_data.data(), count); + this->codes_->Train(train_data.data(), train_count); this->rebuild_residual_centroid_transforms(); } diff --git a/src/impl/cluster/CMakeLists.txt b/src/impl/cluster/CMakeLists.txt index e5da301671..2eb57a2767 100644 --- a/src/impl/cluster/CMakeLists.txt +++ b/src/impl/cluster/CMakeLists.txt @@ -14,5 +14,9 @@ # limitations under the License. -add_library (cluster OBJECT kmeans_cluster.cpp nearest_centroid_assign.cpp kmedoids_cluster.cpp) +add_library (cluster OBJECT + kmeans_cluster.cpp + balanced_kmeans_cluster.cpp + nearest_centroid_assign.cpp + kmedoids_cluster.cpp) target_link_libraries (cluster PRIVATE coverage_config vsag_src_common) diff --git a/src/impl/cluster/balanced_kmeans_cluster.cpp b/src/impl/cluster/balanced_kmeans_cluster.cpp new file mode 100644 index 0000000000..a365d53d0e --- /dev/null +++ b/src/impl/cluster/balanced_kmeans_cluster.cpp @@ -0,0 +1,597 @@ + +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "balanced_kmeans_cluster.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "impl/logger/logger.h" +#include "simd/fp32_simd.h" +#include "simd/normalize.h" +#include "utils/util_functions.h" +#include "vsag_exception.h" + +namespace vsag { + +namespace { + +/// The largest finite float. Distances are clamped into [-FLT_MAX, FLT_MAX] +/// before they are stored or subtracted: a vector pair whose squared L2 or +/// dot product overflows to infinity would otherwise make two infinite +/// distances subtract to NaN, and a NaN in the ordering comparator of the +/// assignment would leave the sort order undefined. Inputs are checked to be +/// finite, so reaching a clamp means the data magnitudes themselves overflow, +/// which is a pathological but not invalid input. +constexpr float K_MAX_DISTANCE = std::numeric_limits::max(); + +float +clamp_distance(float distance) { + if (distance != distance) { + return K_MAX_DISTANCE; + } + return std::min(std::max(distance, -K_MAX_DISTANCE), K_MAX_DISTANCE); +} + +/// The distance a balanced assignment orders by. Squared L2 for L2SQR, and +/// the IVF scan routing key `1 - ` for IP and cosine, matching the +/// constant term the scan layout uses (norms_[b] = 1.0). Cosine centroids are +/// normalized by the IVF training path before assignment, so this key orders +/// the same way as the true cosine distance. +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)); + } + return clamp_distance(FP32ComputeL2Sqr(x, c, dim)); +} + +[[noreturn]] void +throw_invalid_argument(const std::string& message) { + throw VsagException(ErrorType::INVALID_ARGUMENT, message); +} + +} // namespace + +BalancedKMeansCluster::BalancedKMeansCluster(int32_t dim, + Allocator* allocator, + SafeThreadPoolPtr thread_pool) + : dim_(dim), allocator_(allocator) { + if (dim <= 0) { + throw_invalid_argument("dim must be positive"); + } + if (allocator_ == nullptr) { + throw_invalid_argument("allocator cannot be null"); + } + if (thread_pool == nullptr) { + this->thread_pool_ = SafeThreadPool::FactoryDefaultThreadPool(); + } else { + this->thread_pool_ = std::move(thread_pool); + } +} + +BalancedKMeansCluster::~BalancedKMeansCluster() { + if (k_centroids_ != nullptr) { + allocator_->Deallocate(k_centroids_); + k_centroids_ = nullptr; + } +} + +void +BalancedKMeansCluster::check_arguments(const float* datas, + uint64_t count, + uint32_t k, + bool allow_k_gt_count) const { + if (k == 0) { + throw_invalid_argument("k must be positive"); + } + if (count == 0) { + throw_invalid_argument("count must be positive"); + } + if (datas == nullptr) { + throw_invalid_argument("datas cannot be null"); + } + if (not allow_k_gt_count and k > count) { + throw_invalid_argument("k cannot be larger than count"); + } + // The assignment math is only well-defined for finite input. Checking + // every value costs one pass over the data, negligible next to the + // distance computations of even a single iteration. + const uint64_t total = count * static_cast(dim_); + for (uint64_t i = 0; i < total; ++i) { + if (datas[i] != datas[i] or datas[i] > std::numeric_limits::max() or + datas[i] < -std::numeric_limits::max()) { + throw_invalid_argument("datas must contain only finite values"); + } + } +} + +Vector +BalancedKMeansCluster::Run(uint32_t k, + const float* datas, + uint64_t count, + int iter, + double* loss, + float threshold, + KMeansInitMethod init_method, + std::optional random_seed, + uint32_t num_candidates) { + this->check_arguments(datas, count, k, /*allow_k_gt_count=*/false); + if (iter <= 0) { + throw_invalid_argument("iter must be positive"); + } + if (num_candidates == 0) { + throw_invalid_argument("num_candidates must be positive"); + } + + if (k_centroids_ != nullptr) { + allocator_->Deallocate(k_centroids_); + k_centroids_ = nullptr; + } + const uint64_t centroid_values = static_cast(k) * static_cast(dim_); + k_centroids_ = static_cast(allocator_->Allocate(centroid_values * sizeof(float))); + std::fill(k_centroids_, k_centroids_ + centroid_values, 0.0F); + + std::random_device rd; + std::mt19937 gen(random_seed.has_value() ? *random_seed : rd()); + + if (init_method == KMeansInitMethod::KMEANS_PLUS_PLUS) { + select_initial_centroids_kmeans_plus_plus(datas, count, k, gen); + } else { + select_initial_centroids_random(datas, count, k, gen); + } + + logger::trace("BalancedKMeansCluster::Run k: {}, count: {}, iter: {}", k, count, iter); + + Vector labels(count, -1, allocator_); + double current_loss = 0.0; + double last_loss = std::numeric_limits::max(); + int32_t iterations_run = 0; + + for (int it = 0; it < iter; ++it) { + ++iterations_run; + // The capacity-constrained assignment step of one Lloyd iteration. + // Training always orders by squared L2, the same objective the plain + // k-means trainer optimizes regardless of the index metric. + current_loss = this->balanced_assign(k_centroids_, + datas, + count, + k, + MetricType::METRIC_TYPE_L2SQR, + num_candidates, + labels.data()); + // Centroids become the mean of the vectors assigned to them. With + // k <= count every cluster holds at least floor(count / k) >= 1 + // vectors by construction, so no centroid loses its support and no + // empty-cluster re-seeding is needed. + this->update_centroids(datas, labels.data(), count, k); + + logger::trace("[{}] BalancedKMeansCluster::Run iter: {}/{} finished, cur loss is {}", + get_current_time(), + it, + iter, + current_loss); + if (it > 0 and std::fabs(last_loss - current_loss) / static_cast(count) < + static_cast(threshold)) { + break; + } + last_loss = current_loss; + } + + this->record_stats(k, count, labels.data(), current_loss, iterations_run); + if (loss != nullptr) { + *loss = current_loss; + } + return labels; +} + +Vector +BalancedKMeansCluster::Assign(uint32_t k, + const float* centroids, + const float* datas, + uint64_t count, + MetricType metric, + uint32_t num_candidates) { + this->check_arguments(datas, count, k, /*allow_k_gt_count=*/true); + if (centroids == nullptr) { + throw_invalid_argument("centroids cannot be null"); + } + if (num_candidates == 0) { + throw_invalid_argument("num_candidates must be positive"); + } + // The same finiteness bound as the data: a non-finite centroid would make + // every distance to it meaningless. + const uint64_t centroid_values = static_cast(k) * static_cast(dim_); + for (uint64_t i = 0; i < centroid_values; ++i) { + if (centroids[i] != centroids[i] or centroids[i] > std::numeric_limits::max() or + centroids[i] < -std::numeric_limits::max()) { + throw_invalid_argument("centroids must contain only finite values"); + } + } + + Vector labels(count, -1, allocator_); + const double loss = + this->balanced_assign(centroids, datas, count, k, metric, num_candidates, labels.data()); + this->record_stats(k, count, labels.data(), loss, /*iters=*/0); + return labels; +} + +double +BalancedKMeansCluster::balanced_assign(const float* centroids, + const float* datas, + uint64_t count, + uint32_t k, + MetricType metric, + uint32_t num_candidates, + int32_t* labels) { + const uint32_t c = std::min(num_candidates, k); + + // Step 1: per-vector candidate lists. Distances are evaluated one vector + // against all k centroids in fixed blocks, so peak memory never grows + // beyond the per-vector lists: no n x k distance matrix is materialized + // for any n and k. Per-vector work is independent, so the result does not + // depend on the thread schedule. + const uint64_t cand_total = count * static_cast(c); + Vector cand_dist(cand_total, K_MAX_DISTANCE, allocator_); + Vector cand_id(cand_total, 0U, allocator_); + Vector demand(k, 0ULL, allocator_); + std::mutex demand_mutex; + + constexpr uint64_t block_size = 1024; + auto scan_block = [&](uint64_t begin, uint64_t end) { + Vector local_demand(k, 0ULL, allocator_); + // The small top-c buffers are reused across the block's vectors; c is + // runtime, so they are heap buffers rather than variable-length arrays. + // Infinity is only an uninitialized sentinel. A real distance may + // legitimately clamp to K_MAX_DISTANCE and must still be selectable. + Vector best_dist(c, std::numeric_limits::infinity(), allocator_); + Vector best_id(c, 0U, allocator_); + Vector normalized_row(allocator_); + if (metric == MetricType::METRIC_TYPE_COSINE) { + normalized_row.resize(dim_); + } + for (uint64_t i = begin; i < end; ++i) { + const float* x = datas + i * static_cast(dim_); + if (metric == MetricType::METRIC_TYPE_COSINE) { + Normalize(x, normalized_row.data(), dim_); + x = normalized_row.data(); + } + std::fill(best_dist.begin(), best_dist.end(), std::numeric_limits::infinity()); + std::fill(best_id.begin(), best_id.end(), 0U); + for (uint32_t j = 0; j < k; ++j) { + const float dist = + metric_distance(metric, x, centroids + j * static_cast(dim_), dim_); + 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(c) + j] = best_dist[j]; + cand_id[i * static_cast(c) + j] = best_id[j]; + } + ++local_demand[best_id[0]]; + } + std::lock_guard 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> futures; + std::future pending_future; + auto wait_for_submissions = [&]() { + // If emplace_back throws, the just-submitted block is still owned by + // pending_future rather than by futures. Drain it as well. + if (pending_future.valid()) { + pending_future.wait(); + } + for (auto& future : futures) { + future.wait(); + } + }; + try { + 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 { + pending_future = thread_pool_->GeneralEnqueue(scan_block, begin, end); + futures.emplace_back(std::move(pending_future)); + } + } + } catch (...) { + // Submission may fail after earlier blocks have started. Those blocks + // reference stack state, so none may outlive this function. + wait_for_submissions(); + throw; + } + std::exception_ptr scan_exception; + for (auto& future : futures) { + try { + future.get(); + } catch (...) { + if (scan_exception == nullptr) { + scan_exception = std::current_exception(); + } + } + } + // Every submitted block has finished before the captured stack state can + // unwind, and no placeholder candidate data is consumed after a failure. + if (scan_exception != nullptr) { + std::rethrow_exception(scan_exception); + } + + // 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(k); + const uint64_t extra = count % static_cast(k); + Vector remaining(k, base, allocator_); + { + Vector 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 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(c) + 1] - + cand_dist[i * static_cast(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(c)]; + const float db = cand_dist[b * static_cast(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. The same sequential scratch row is reused by rescans and by + // the loss pass below; neither needs a normalized copy of the whole corpus. + Vector normalized_row(allocator_); + if (metric == MetricType::METRIC_TYPE_COSINE) { + normalized_row.resize(dim_); + } + for (uint64_t i : point_order) { + const uint64_t row = i * static_cast(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(dim_); + if (metric == MetricType::METRIC_TYPE_COSINE) { + Normalize(x, normalized_row.data(), dim_); + x = normalized_row.data(); + } + float best = std::numeric_limits::infinity(); + for (uint32_t j = 0; j < k; ++j) { + if (remaining[j] == 0) { + continue; + } + const float dist = + metric_distance(metric, x, centroids + j * static_cast(dim_), dim_); + if (dist < best) { + best = dist; + chosen = j; + } + } + // Unreachable: some cluster always holds a free seat while vectors + // remain, and the loop above visits every such cluster. + if (chosen == k) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "balanced assignment ran out of capacity"); + } + } + labels[i] = static_cast(chosen); + --remaining[chosen]; + } + + // The loss is recomputed from the final (vector, centroid) pairs with the + // same distance as the assignment, in a fixed order. + double total_loss = 0.0; + for (uint64_t i = 0; i < count; ++i) { + const float* x = datas + i * static_cast(dim_); + if (metric == MetricType::METRIC_TYPE_COSINE) { + Normalize(x, normalized_row.data(), dim_); + x = normalized_row.data(); + } + const float* centroid = + centroids + static_cast(labels[i]) * static_cast(dim_); + total_loss += static_cast(metric_distance(metric, x, centroid, dim_)); + } + return total_loss; +} + +void +BalancedKMeansCluster::update_centroids(const float* datas, + const int32_t* labels, + uint64_t count, + uint32_t k) { + const uint64_t centroid_values = static_cast(k) * static_cast(dim_); + Vector sums(centroid_values, 0.0, allocator_); + Vector sizes(k, 0ULL, allocator_); + // Double accumulation is order-dependent too, so keep one sequential pass + // in fixed vector order. Besides making the result independent of the + // assignment workers' schedule, this prevents large clusters from reaching + // the binary32 point where adding another finite vector has no effect. + for (uint64_t i = 0; i < count; ++i) { + const uint32_t label = static_cast(labels[i]); + const float* src = datas + i * static_cast(dim_); + double* dest = sums.data() + static_cast(label) * static_cast(dim_); + ++sizes[label]; + for (int32_t d = 0; d < dim_; ++d) { + dest[d] += static_cast(src[d]); + } + } + for (uint32_t j = 0; j < k; ++j) { + if (sizes[j] == 0) { + // Cannot happen for Run with k <= count (every cluster keeps at + // least floor(count / k) >= 1 vectors by construction) nor for the + // clusters a count < k Assign fills; a cluster that lost all its + // vectors keeps its previous centroid. + continue; + } + const double scale = 1.0 / static_cast(sizes[j]); + float* centroid = k_centroids_ + static_cast(j) * static_cast(dim_); + const double* sum = sums.data() + static_cast(j) * static_cast(dim_); + for (int32_t d = 0; d < dim_; ++d) { + centroid[d] = static_cast(sum[d] * scale); + } + } +} + +void +BalancedKMeansCluster::select_initial_centroids_random(const float* datas, + uint64_t count, + uint32_t k, + std::mt19937& gen) { + std::uniform_int_distribution dis(0, count - 1); + for (uint32_t i = 0; i < k; ++i) { + const uint64_t index = dis(gen); + for (int32_t d = 0; d < dim_; ++d) { + k_centroids_[static_cast(i) * static_cast(dim_) + + static_cast(d)] = + datas[static_cast(index) * static_cast(dim_) + + static_cast(d)]; + } + } +} + +void +BalancedKMeansCluster::select_initial_centroids_kmeans_plus_plus(const float* datas, + uint64_t count, + uint32_t k, + std::mt19937& gen) { + // Standard D^2 seeding, identical in spirit to the plain trainer's, so the + // two trainers can be compared on equal initialization footing. + std::uniform_int_distribution first_dis(0, count - 1); + const uint64_t first_idx = first_dis(gen); + for (int32_t d = 0; d < dim_; ++d) { + k_centroids_[static_cast(d)] = + datas[static_cast(first_idx) * static_cast(dim_) + + static_cast(d)]; + } + + Vector min_distances(count, K_MAX_DISTANCE, allocator_); + for (uint32_t ci = 1; ci < k; ++ci) { + const float* centroid = + k_centroids_ + static_cast(ci - 1) * static_cast(dim_); + for (uint64_t i = 0; i < count; ++i) { + const float dist = + FP32ComputeL2Sqr(datas + i * static_cast(dim_), centroid, dim_); + min_distances[i] = std::min(min_distances[i], clamp_distance(dist)); + } + double total_weight = 0.0; + for (uint64_t i = 0; i < count; ++i) { + total_weight += static_cast(min_distances[i]); + } + uint64_t selected_idx = 0; + if (total_weight <= 0.0) { + // Every vector coincides with an already chosen centroid; any + // vector is then a valid pick, so keep the generator moving and + // take a uniform one. + std::uniform_int_distribution dis(0, count - 1); + selected_idx = dis(gen); + } else { + std::uniform_real_distribution prob_dis(0.0, total_weight); + const double target = prob_dis(gen); + double cumulative = 0.0; + selected_idx = count - 1; + for (uint64_t i = 0; i < count; ++i) { + cumulative += static_cast(min_distances[i]); + if (cumulative >= target) { + selected_idx = i; + break; + } + } + } + for (int32_t d = 0; d < dim_; ++d) { + k_centroids_[static_cast(ci) * static_cast(dim_) + + static_cast(d)] = + datas[static_cast(selected_idx) * static_cast(dim_) + + static_cast(d)]; + } + } +} + +void +BalancedKMeansCluster::record_stats( + uint32_t k, uint64_t count, const int32_t* labels, double loss, int32_t iters) { + Vector sizes(k, 0ULL, allocator_); + for (uint64_t i = 0; i < count; ++i) { + ++sizes[static_cast(labels[i])]; + } + const auto extremes = std::minmax_element(sizes.begin(), sizes.end()); + stats_.cluster_count = k; + stats_.vector_count = count; + stats_.min_cluster_size = *extremes.first; + stats_.max_cluster_size = *extremes.second; + stats_.loss = loss; + stats_.iterations_run = iters; +} + +} // namespace vsag diff --git a/src/impl/cluster/balanced_kmeans_cluster.h b/src/impl/cluster/balanced_kmeans_cluster.h new file mode 100644 index 0000000000..0529a4a36f --- /dev/null +++ b/src/impl/cluster/balanced_kmeans_cluster.h @@ -0,0 +1,214 @@ + +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +#include "impl/thread_pool/safe_thread_pool.h" +#include "kmeans_cluster.h" +#include "metric_type.h" +#include "typing.h" + +namespace vsag { +class Allocator; + +/// Summary of the most recent clustering or assignment done by +/// `BalancedKMeansCluster`, so that tests, tools and internal callers can +/// evaluate balance and loss without recomputing them from the returned +/// labels (issue #2870 observability requirement). +struct BalancedKMeansStats { + /// The number of clusters the last call was asked for. + uint32_t cluster_count{0}; + /// The number of vectors the last call assigned. + uint64_t vector_count{0}; + /// Smallest and largest final cluster size, derived from the returned + /// labels rather than from the capacity targets. + uint64_t min_cluster_size{0}; + uint64_t max_cluster_size{0}; + /// Total assignment cost of the last assignment step: the sum over all + /// vectors of the distance to the centroid they were assigned to, using + /// the same distance as the assignment itself. + double loss{0.0}; + /// Number of assignment/centroid-update iterations `Run` executed; 0 for + /// a single `Assign` call, which does not iterate. + int32_t iterations_run{0}; +}; + +/// Balanced k-means clustering for IVF partitioning (issue #2870). +/// +/// Standard k-means (`KMeansCluster`) assigns every vector to its nearest +/// centroid, which can produce heavily skewed or even empty clusters on +/// skewed or duplicated data. This trainer keeps the same squared-L2 +/// objective and mean centroid update, but every assignment step is +/// capacity-constrained: with `n` vectors and `k` clusters each cluster +/// finally holds either `floor(n / k)` or `ceil(n / k)` vectors, so the +/// difference between the largest and the smallest cluster is at most one +/// and no cluster is empty when `1 <= k <= n`. +/// +/// The balancing strategy is a regret-greedy capacitated assignment, which +/// is a different point in the design space from a price-based auction or +/// an exact transportation simplex: +/// 1. Each vector keeps only its `num_candidates` nearest clusters +/// (bounded per-vector state; no `n x k` distance matrix is ever +/// materialized; distances are computed one vector at a time in fixed +/// parallel blocks and the reduced statistics are order-independent). +/// 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 placed one at a time, 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 on its candidate +/// list that still has a free seat; if every candidate is full, the +/// distances to all clusters with a free seat are recomputed on demand. +/// Because the capacities sum to exactly `n`, a seat is always available +/// and every vector is placed. +/// 4. Centroids are re-set to the mean of their assigned vectors and the +/// cycle repeats until the loss stops improving or the iteration budget +/// is exhausted. +/// +/// Expected complexity per iteration: `O(n * k * dim)` distance evaluations +/// plus `O(n log n)` ordering - the same asymptotic work as one standard +/// k-means iteration plus a sort - and `O(n * num_candidates + k * dim)` +/// working memory. The assignment is a heuristic for the capacity-constrained +/// step (it does not price seats or search augmenting swaps), so it can pay +/// a somewhat larger loss than a solver-exact assignment on strongly skewed +/// data; the balance guarantee itself holds for every input. +/// +/// Results are deterministic for a given input, iteration count, candidate +/// count and seed, independent of the thread pool size: no shared accumulation +/// happens in parallel code and the assignment order is a total order. +/// +/// Like `KMeansCluster`, training always minimizes squared L2 regardless of +/// the index metric; cosine index handling (centroid normalization) stays in +/// the IVF layer above. `Assign` additionally supports the IP and cosine +/// routing keys so that a balanced bucket assignment can mirror the ordering +/// IVF search uses. This trainer always runs on the host; the CUDA build +/// backend only accelerates `KMeansCluster`. +class BalancedKMeansCluster { +public: + /// Per-vector candidates kept during assignment. Larger values reduce the + /// number of on-demand full rescans (better loss) at a linear memory cost + /// of `8 * num_candidates` bytes per vector. Values above the cluster + /// count are clamped, and `k <= num_candidates` makes every assignment + /// exact per-iteration since no vector then needs a rescan. + static constexpr uint32_t DEFAULT_NUM_CANDIDATES = 8; + + explicit BalancedKMeansCluster(int32_t dim, + Allocator* allocator, + SafeThreadPoolPtr thread_pool = nullptr); + + ~BalancedKMeansCluster(); + + /// Runs balanced k-means and returns the final vector-to-cluster + /// assignment (`labels[i]` is the cluster of vector `i`), which also + /// serves as the observability handle on the balanced partitioning. The + /// trained centroids are left in `k_centroids_` and summary statistics in + /// `Stats()`. + /// + /// Throws `VsagException(INVALID_ARGUMENT)` when `k == 0`, `count == 0`, + /// `datas == nullptr`, `k > count`, `iter <= 0`, `num_candidates == 0`, + /// or `datas` contains a non-finite value. + Vector + Run(uint32_t k, + const float* datas, + uint64_t count, + int iter = 25, + double* loss = nullptr, + float threshold = 1e-6F, + KMeansInitMethod init_method = KMeansInitMethod::KMEANS_PLUS_PLUS, + std::optional random_seed = std::nullopt, + uint32_t num_candidates = DEFAULT_NUM_CANDIDATES); + + /// Assigns `count` vectors to fixed centroids under the balanced capacity + /// constraint, without updating the centroids. `metric` selects the + /// ordering key: squared L2 for `METRIC_TYPE_L2SQR`, `1 - ` for IP + /// and cosine, matching the IVF scan routing key. Cosine input rows are + /// normalized internally before their assignment costs are compared. + /// Unlike `Run`, this + /// accepts `k > count`: then `count` clusters receive one vector each and + /// the remaining clusters stay empty, which is the best achievable + /// balance when there are fewer vectors than clusters. + /// + /// Throws `VsagException(INVALID_ARGUMENT)` on the same invalid inputs as + /// `Run`, except that `k > count` is allowed. + Vector + Assign(uint32_t k, + const float* centroids, + const float* datas, + uint64_t count, + MetricType metric = MetricType::METRIC_TYPE_L2SQR, + uint32_t num_candidates = DEFAULT_NUM_CANDIDATES); + + /// Statistics of the most recent `Run` or `Assign` call. + const BalancedKMeansStats& + Stats() const { + return stats_; + } + +public: + /// Trained centroids from the most recent `Run`, `k * dim` floats + /// (cluster-major). Unset after `Assign`, which does not train. + float* k_centroids_{nullptr}; + +private: + /// Validates the shared arguments of `Run` and `Assign`, including the + /// finiteness of the input data. + void + check_arguments(const float* datas, uint64_t count, uint32_t k, bool allow_k_gt_count) const; + + /// Places `count` vectors over `centroids` with exact capacity balance and + /// returns the total assignment cost; fills `labels`. + double + balanced_assign(const float* centroids, + const float* datas, + uint64_t count, + uint32_t k, + MetricType metric, + uint32_t num_candidates, + int32_t* labels); + + /// Recomputes every centroid as the mean of the vectors assigned to it. + /// Clusters with no member keep their previous centroid. + void + update_centroids(const float* datas, const int32_t* labels, uint64_t count, uint32_t k); + + void + select_initial_centroids_random(const float* datas, + uint64_t count, + uint32_t k, + std::mt19937& gen); + + void + select_initial_centroids_kmeans_plus_plus(const float* datas, + uint64_t count, + uint32_t k, + std::mt19937& gen); + + /// Fills `stats_` from the final labels of a completed assignment. + void + record_stats(uint32_t k, uint64_t count, const int32_t* labels, double loss, int32_t iters); + +private: + Allocator* const allocator_{nullptr}; + + SafeThreadPoolPtr thread_pool_{nullptr}; + + const int32_t dim_{0}; + + BalancedKMeansStats stats_{}; +}; + +} // namespace vsag diff --git a/src/impl/cluster/balanced_kmeans_cluster_test.cpp b/src/impl/cluster/balanced_kmeans_cluster_test.cpp new file mode 100644 index 0000000000..cdbbf05f68 --- /dev/null +++ b/src/impl/cluster/balanced_kmeans_cluster_test.cpp @@ -0,0 +1,742 @@ + +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "balanced_kmeans_cluster.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "allocator/memory_record_allocator.h" +#include "impl/allocator/default_allocator.h" +#include "impl/allocator/safe_allocator.h" +#include "impl/thread_pool/default_thread_pool.h" +#include "unittest.h" + +namespace { + +/// Well separated spherical clusters of equal size, with the vector order +/// shuffled. `truth[i]` is the generating cluster of vector i. The centroid +/// axes are spaced far apart relative to the noise, so the nearest-centroid +/// answer is never in doubt. +std::vector +make_separated(uint32_t k, + uint64_t per_cluster, + int32_t dim, + std::mt19937& gen, + std::vector& truth) { + REQUIRE(k <= static_cast(dim)); + const uint64_t n = static_cast(k) * per_cluster; + std::vector centroids(static_cast(k) * dim, 0.0F); + for (uint32_t c = 0; c < k; ++c) { + centroids[static_cast(c) * dim + c] = 1000.0F * (1.0F + static_cast(c)); + } + std::normal_distribution noise(0.0F, 0.01F); + std::vector data(n * static_cast(dim)); + truth.resize(n); + for (uint64_t i = 0; i < n; ++i) { + truth[i] = static_cast(i / per_cluster); + } + std::shuffle(truth.begin(), truth.end(), gen); + for (uint64_t i = 0; i < n; ++i) { + const uint32_t c = truth[i]; + for (int32_t d = 0; d < dim; ++d) { + data[i * dim + d] = + centroids[static_cast(c) * dim + d] + static_cast(noise(gen)); + } + } + return data; +} + +/// A skewed mixture: cluster sizes drop geometrically, so a free nearest-centroid +/// assignment would produce strongly uneven clusters (issue #2870 scenario). +std::vector +make_skewed(uint32_t k, uint64_t largest, int32_t dim, std::mt19937& gen, uint64_t& count_out) { + std::vector sizes(k); + uint64_t total = 0; + uint64_t cur = largest; + for (uint32_t c = 0; c < k; ++c) { + sizes[c] = std::max(1, cur); + total += sizes[c]; + cur = cur / 3 + 1; + } + count_out = total; + std::vector centroids(static_cast(k) * dim); + std::uniform_real_distribution coord(-100.0F, 100.0F); + std::normal_distribution noise(0.0F, 0.05F); + for (uint32_t c = 0; c < k; ++c) { + for (int32_t d = 0; d < dim; ++d) { + centroids[static_cast(c) * dim + d] = coord(gen); + } + } + std::vector data(total * static_cast(dim)); + uint64_t i = 0; + for (uint32_t c = 0; c < k; ++c) { + for (uint64_t j = 0; j < sizes[c]; ++j, ++i) { + for (int32_t d = 0; d < dim; ++d) { + data[i * dim + d] = + centroids[static_cast(c) * dim + d] + static_cast(noise(gen)); + } + } + } + return data; +} + +std::vector +cluster_sizes(const vsag::Vector& labels, uint32_t k) { + std::vector sizes(k, 0); + for (const int label : labels) { + REQUIRE(label >= 0); + REQUIRE(label < static_cast(k)); + ++sizes[static_cast(label)]; + } + return sizes; +} + +/// The balance guarantee of the issue: every size is floor(n/k) or ceil(n/k), +/// the spread is at most one, and no cluster is empty when 1 <= k <= n. +void +require_exact_balance(const std::vector& sizes, uint64_t n, uint32_t k) { + const uint64_t base = n / k; + const uint64_t extra = n % k; + REQUIRE(sizes.size() == k); + uint64_t big = 0; + for (const uint64_t size : sizes) { + REQUIRE((size == base or size == base + 1)); + big += (size == base + 1) ? 1 : 0; + if (base >= 1) { + REQUIRE(size >= 1); + } + } + REQUIRE(big == extra); +} + +class WorkerFailAllocator : public vsag::DefaultAllocator { +public: + WorkerFailAllocator() : owner_(std::this_thread::get_id()) { + } + + void* + Allocate(uint64_t size) override { + if (std::this_thread::get_id() != owner_) { + throw std::bad_alloc(); + } + return vsag::DefaultAllocator::Allocate(size); + } + + void* + Reallocate(void* pointer, uint64_t size) override { + if (std::this_thread::get_id() != owner_) { + throw std::bad_alloc(); + } + return vsag::DefaultAllocator::Reallocate(pointer, size); + } + +private: + std::thread::id owner_; +}; + +class SubmitFailThreadPool : public vsag::ThreadPool { +public: + SubmitFailThreadPool() : worker_(1), release_future_(release_.get_future()) { + } + + std::future + Enqueue(std::function task) override { + if (submission_count_++ != 0) { + rejected_.set_value(); + throw std::runtime_error("submission rejected by test"); + } + return worker_.Enqueue([this, task = std::move(task)]() mutable { + release_future_.wait(); + task(); + }); + } + + void + WaitUntilEmpty() override { + worker_.WaitUntilEmpty(); + } + + void + SetQueueSizeLimit(std::uint64_t limit) override { + worker_.SetQueueSizeLimit(limit); + } + + void + SetPoolSize(std::uint64_t limit) override { + worker_.SetPoolSize(limit); + } + + void + WaitForRejectedSubmission() { + rejected_.get_future().wait(); + } + + void + ReleaseWorker() { + release_.set_value(); + } + +private: + vsag::DefaultThreadPool worker_; + uint64_t submission_count_{0}; + std::promise rejected_; + std::promise release_; + std::shared_future release_future_; +}; + +} // namespace + +TEST_CASE("BalancedKMeans exact sizes when n divides by k", "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 8; + constexpr int32_t dim = 16; + constexpr uint64_t per_cluster = 150; + std::mt19937 gen(12345); + std::vector truth; + auto data = make_separated(k, per_cluster, dim, gen, truth); + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + const auto labels = cluster.Run(k, + data.data(), + k * per_cluster, + 10, + nullptr, + 1e-6F, + vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, + 42U); + + REQUIRE(labels.size() == k * per_cluster); + const auto sizes = cluster_sizes(labels, k); + require_exact_balance(sizes, k * per_cluster, k); + // On perfectly separated equal-size clusters the balance costs nothing: + // every generating cluster occupies exactly one bucket. + for (uint32_t c = 0; c < k; ++c) { + REQUIRE(sizes[c] == per_cluster); + } + std::vector label_of_cluster(k, -1); + for (uint64_t i = 0; i < labels.size(); ++i) { + const uint32_t truth_c = truth[i]; + if (label_of_cluster[truth_c] == -1) { + label_of_cluster[truth_c] = labels[i]; + } else { + REQUIRE(label_of_cluster[truth_c] == labels[i]); + } + } + REQUIRE(std::unique(label_of_cluster.begin(), label_of_cluster.end()) == + label_of_cluster.end()); +} + +TEST_CASE("BalancedKMeans exact sizes when n does not divide by k", "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 7; + constexpr int32_t dim = 8; + constexpr uint64_t per_cluster = 100; + std::mt19937 gen(999); + std::vector truth; + auto data = make_separated(k, per_cluster, dim, gen, truth); + // Drop 9 vectors so n % k == 4 with n = 691. + const uint64_t n = k * per_cluster - 9; + data.resize(n * static_cast(dim)); + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + double loss = 0.0; + const auto labels = cluster.Run( + k, data.data(), n, 12, &loss, 1e-6F, vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, 7U); + + const auto sizes = cluster_sizes(labels, k); + require_exact_balance(sizes, n, k); + REQUIRE(loss >= 0.0); + REQUIRE(std::isfinite(loss)); + const auto& stats = cluster.Stats(); + REQUIRE(stats.cluster_count == k); + REQUIRE(stats.vector_count == n); + REQUIRE(stats.min_cluster_size == n / k); + REQUIRE(stats.max_cluster_size == n / k + 1); + REQUIRE(stats.iterations_run >= 1); + REQUIRE(stats.iterations_run <= 12); +} + +TEST_CASE("BalancedKMeans constrains a heavily skewed dataset", "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 16; + constexpr int32_t dim = 12; + std::mt19937 gen(2718); + uint64_t n = 0; + auto data = make_skewed(k, 20000, dim, gen, n); + REQUIRE(n % k != 0); // The dataset is deliberately not divisible. + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + double loss = -1.0; + const auto labels = + cluster.Run(k, data.data(), n, 8, &loss, 1e-6F, vsag::KMeansInitMethod::RANDOM, 11U); + + const auto sizes = cluster_sizes(labels, k); + require_exact_balance(sizes, n, k); + // The whole point of the feature: the largest bucket is ceil(n/k), where an + // unconstrained k-means on this mixture would have left a multiple of it. + REQUIRE(*std::max_element(sizes.begin(), sizes.end()) == n / k + 1); + REQUIRE(std::isfinite(loss)); +} + +TEST_CASE("BalancedKMeans splits duplicated vectors evenly", "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 10; + constexpr int32_t dim = 4; + constexpr uint64_t n = 1000; + // Every vector identical: nearest-centroid assignment would pile them into + // one bucket. Duplicates are the extreme case of the issue's scenario. + std::vector data(n * static_cast(dim), 1.5F); + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + const auto labels = + cluster.Run(k, data.data(), n, 5, nullptr, 1e-6F, vsag::KMeansInitMethod::RANDOM, 3U); + + const auto sizes = cluster_sizes(labels, k); + require_exact_balance(sizes, n, k); + for (const uint64_t size : sizes) { + REQUIRE(size == n / k); + } + // Every centroid is the mean of a cluster of vectors that are all exactly + // 1.5, so the trained centroids coincide with the data. + for (uint64_t j = 0; j < static_cast(k) * dim; ++j) { + REQUIRE(std::abs(cluster.k_centroids_[j] - 1.5F) < 1e-5F); + } +} + +TEST_CASE("BalancedKMeans uses stable accumulation for cluster means", + "[ut][BalancedKMeansCluster]") { + // This compact input has the same sequential-binary32 failure mode as + // 2^25 rows of 1.0F: after the first value, both unit contributions are + // rounded away. A double accumulator retains the exact sum 16,777,218. + const std::vector data = {16777216.0F, 1.0F, 1.0F}; + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(1, allocator.get()); + cluster.Run(1, data.data(), data.size(), 1, nullptr, 1e-6F, vsag::KMeansInitMethod::RANDOM, 0U); + + REQUIRE(cluster.k_centroids_[0] == 5592406.0F); +} + +TEST_CASE("BalancedKMeans k equals 1 and k equals n", "[ut][BalancedKMeansCluster]") { + constexpr int32_t dim = 6; + constexpr uint64_t n = 100; + std::mt19937 gen(555); + std::normal_distribution noise(0.0F, 1.0F); + std::vector data(n * static_cast(dim)); + for (auto& v : data) { + v = noise(gen); + } + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + + SECTION("k == 1 puts everything in one cluster") { + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + const auto labels = cluster.Run(1, data.data(), n, 3); + for (const int label : labels) { + REQUIRE(label == 0); + } + REQUIRE(cluster.Stats().min_cluster_size == n); + REQUIRE(cluster.Stats().max_cluster_size == n); + } + SECTION("k == n gives every cluster exactly one vector") { + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + const auto labels = cluster.Run(static_cast(n), data.data(), n, 2); + const auto sizes = cluster_sizes(labels, static_cast(n)); + require_exact_balance(sizes, n, n); + for (const uint64_t size : sizes) { + REQUIRE(size == 1); + } + } +} + +TEST_CASE("BalancedKMeans rejects invalid arguments", "[ut][BalancedKMeansCluster]") { + constexpr int32_t dim = 4; + constexpr uint64_t n = 64; + std::vector data(n * static_cast(dim), 0.5F); + std::vector centroids(4 * static_cast(dim), 0.5F); + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + + SECTION("Run") { + // k == 0 + REQUIRE_THROWS(cluster.Run(0, data.data(), n, 1)); + // n == 0 + REQUIRE_THROWS(cluster.Run(4, data.data(), 0, 1)); + // null data + REQUIRE_THROWS(cluster.Run(4, nullptr, n, 1)); + // k > n + REQUIRE_THROWS(cluster.Run(static_cast(n + 1), data.data(), n, 1)); + // non-positive iterations + REQUIRE_THROWS(cluster.Run(4, data.data(), n, 0)); + REQUIRE_THROWS(cluster.Run(4, data.data(), n, -1)); + // zero candidates + REQUIRE_THROWS(cluster.Run( + 4, data.data(), n, 1, nullptr, 1e-6F, vsag::KMeansInitMethod::RANDOM, 1U, 0U)); + // non-finite values + auto bad = data; + bad[13] = std::numeric_limits::quiet_NaN(); + REQUIRE_THROWS(cluster.Run(4, bad.data(), n, 1)); + bad[13] = std::numeric_limits::infinity(); + REQUIRE_THROWS(cluster.Run(4, bad.data(), n, 1)); + } + SECTION("Assign") { + REQUIRE_THROWS(cluster.Assign(0, centroids.data(), data.data(), n)); + REQUIRE_THROWS(cluster.Assign(4, nullptr, data.data(), n)); + REQUIRE_THROWS(cluster.Assign(4, centroids.data(), nullptr, n)); + REQUIRE_THROWS(cluster.Assign(4, centroids.data(), data.data(), 0)); + REQUIRE_THROWS(cluster.Assign( + 4, centroids.data(), data.data(), n, vsag::MetricType::METRIC_TYPE_L2SQR, 0U)); + auto bad_centroids = centroids; + bad_centroids[5] = std::numeric_limits::quiet_NaN(); + REQUIRE_THROWS(cluster.Assign(4, bad_centroids.data(), data.data(), n)); + auto bad_data = data; + bad_data[9] = std::numeric_limits::quiet_NaN(); + REQUIRE_THROWS(cluster.Assign(4, centroids.data(), bad_data.data(), n)); + } + SECTION("Constructor") { + REQUIRE_THROWS(vsag::BalancedKMeansCluster(0, allocator.get())); + REQUIRE_THROWS(vsag::BalancedKMeansCluster(-4, allocator.get())); + } +} + +TEST_CASE("BalancedKMeans is reproducible across thread counts", "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 5; + constexpr int32_t dim = 20; + std::mt19937 gen(8121); + uint64_t n = 0; + auto data = make_skewed(k, 3000, dim, gen, n); + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + auto single_thread_pool = vsag::SafeThreadPool::FactoryDefaultThreadPool(); + single_thread_pool->SetPoolSize(1); + auto multi_thread_pool = vsag::SafeThreadPool::FactoryDefaultThreadPool(); + multi_thread_pool->SetPoolSize(4); + + vsag::BalancedKMeansCluster single_thread(dim, allocator.get(), single_thread_pool); + vsag::BalancedKMeansCluster multi_thread(dim, allocator.get(), multi_thread_pool); + + const uint32_t iterations = 6; + const auto single_labels = single_thread.Run(k, + data.data(), + n, + iterations, + nullptr, + 1e-6F, + vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, + 0x2870U); + const auto multi_labels = multi_thread.Run(k, + data.data(), + n, + iterations, + nullptr, + 1e-6F, + vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, + 0x2870U); + + // No shared float accumulation happens in parallel code, so a run is + // bit-identical for any thread count under the same seed. + REQUIRE(single_labels == multi_labels); + const uint64_t centroid_values = static_cast(k) * dim; + REQUIRE(std::equal(single_thread.k_centroids_, + single_thread.k_centroids_ + centroid_values, + multi_thread.k_centroids_)); + REQUIRE(single_thread.Stats().loss == multi_thread.Stats().loss); +} + +TEST_CASE("BalancedKMeans Assign balances fixed centroids", "[ut][BalancedKMeansCluster]") { + // Two centroids, two vectors, capacity one each: the vector with the + // stronger preference keeps its favourite and the other takes the rest, + // under the IP routing key (cost 1 - ) IVF uses. + constexpr uint32_t k = 2; + constexpr int32_t dim = 2; + const std::vector centroids = {0.0F, 0.0F, 10.0F, 0.0F}; + const std::vector data = {10.0F, 0.0F, 0.5F, 0.0F}; + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + + const auto labels = + cluster.Assign(k, centroids.data(), data.data(), 2, vsag::MetricType::METRIC_TYPE_IP, 2U); + REQUIRE(labels[0] == 1); + REQUIRE(labels[1] == 0); + REQUIRE(cluster.Stats().min_cluster_size == 1); + REQUIRE(cluster.Stats().max_cluster_size == 1); +} + +TEST_CASE("BalancedKMeans cosine assignment normalizes input rows", "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 2; + constexpr int32_t dim = 2; + const std::vector centroids = {1.0F, 0.0F, 0.0F, 1.0F}; + const std::vector data = {100.0F, 99.0F, 1.0F, 0.0F}; + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + + const auto labels = cluster.Assign( + k, centroids.data(), data.data(), 2, vsag::MetricType::METRIC_TYPE_COSINE, 2U); + // Raw magnitudes make the regrets tie and give the first row centroid 0. + // Cosine regrets instead give the exactly aligned second row priority. + REQUIRE(labels[0] == 1); + REQUIRE(labels[1] == 0); +} + +TEST_CASE("BalancedKMeans cosine assignment avoids a full normalized corpus", + "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 2; + constexpr int32_t dim = 64; + constexpr uint64_t n = 4096; + std::vector centroids(static_cast(k) * dim, 0.0F); + centroids[0] = 1.0F; + centroids[dim] = 1.0F; + std::vector data(n * static_cast(dim), 0.0F); + for (uint64_t i = 0; i < n; ++i) { + data[i * static_cast(dim)] = static_cast(i + 1); + } + auto allocator = std::make_shared(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + + // One candidate forces a fallback pass for half the rows. The loss pass also + // visits every row; all three passes must use row-local normalization. + const auto labels = cluster.Assign( + k, centroids.data(), data.data(), n, vsag::MetricType::METRIC_TYPE_COSINE, 1U); + require_exact_balance(cluster_sizes(labels, k), n, k); + REQUIRE(std::isfinite(cluster.Stats().loss)); + const uint64_t normalized_corpus_bytes = n * static_cast(dim) * sizeof(float); + REQUIRE(allocator->GetMemoryPeak() < normalized_corpus_bytes / 2); +} + +TEST_CASE("BalancedKMeans saturated distances still use available seats", + "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 2; + constexpr int32_t dim = 1; + const std::vector centroids = {-1e20F, 1e20F}; + const std::vector data = {0.0F, 0.0F}; + const uint32_t candidates = GENERATE(1U, 2U); + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + + const auto labels = cluster.Assign( + k, centroids.data(), data.data(), 2, vsag::MetricType::METRIC_TYPE_L2SQR, candidates); + REQUIRE(labels[0] != labels[1]); + require_exact_balance(cluster_sizes(labels, k), 2, k); +} + +TEST_CASE("BalancedKMeans propagates candidate worker exceptions", "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 2; + constexpr int32_t dim = 1; + constexpr uint64_t count = 2048; // More than one candidate-scan block. + const std::vector centroids = {0.0F, 1.0F}; + const std::vector data(count, 0.0F); + auto pool = vsag::SafeThreadPool::FactoryDefaultThreadPool(); + pool->SetPoolSize(2); + WorkerFailAllocator allocator; + vsag::BalancedKMeansCluster cluster(dim, &allocator, pool); + + REQUIRE_THROWS_AS(cluster.Assign(k, centroids.data(), data.data(), count), std::bad_alloc); +} + +TEST_CASE("BalancedKMeans drains workers when submission fails", "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 2; + constexpr int32_t dim = 1; + constexpr uint64_t count = 2048; // The second block is rejected. + const std::vector centroids = {0.0F, 1.0F}; + const std::vector data(count, 0.0F); + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + auto rejecting_pool = std::make_shared(); + auto pool = std::make_shared(rejecting_pool); + vsag::BalancedKMeansCluster cluster(dim, allocator.get(), pool); + + auto assign_result = std::async(std::launch::async, [&]() { + try { + cluster.Assign(k, centroids.data(), data.data(), count); + } catch (const std::runtime_error&) { + return true; + } + return false; + }); + rejecting_pool->WaitForRejectedSubmission(); + // The first task remains gated. Assign must be draining it, rather than + // unwinding the stack state referenced by that task. + const bool waited_for_worker = + assign_result.wait_for(std::chrono::milliseconds(50)) == std::future_status::timeout; + rejecting_pool->ReleaseWorker(); + const bool threw_submission_error = assign_result.get(); + REQUIRE(waited_for_worker); + REQUIRE(threw_submission_error); +} + +TEST_CASE("BalancedKMeans Assign with fewer vectors than clusters spreads them", + "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 10; + constexpr int32_t dim = 3; + constexpr uint64_t n = 6; + std::mt19937 gen(77); + std::uniform_real_distribution coord(-1.0F, 1.0F); + std::vector centroids(static_cast(k) * dim); + for (auto& v : centroids) { + v = coord(gen); + } + std::vector data(n * static_cast(dim)); + for (auto& v : data) { + v = coord(gen); + } + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + const auto labels = cluster.Assign(k, centroids.data(), data.data(), n); + const auto sizes = cluster_sizes(labels, k); + // floor(6/10) == 0 and ceil == 1: no bucket holds more than one vector. + for (const uint64_t size : sizes) { + REQUIRE(size <= 1); + } + REQUIRE(std::accumulate(sizes.begin(), sizes.end(), 0ULL) == n); +} + +TEST_CASE("BalancedKMeans stays exact with one candidate and with full candidates", + "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 4; + constexpr int32_t dim = 8; + constexpr uint64_t per_cluster = 60; + std::mt19937 gen(313); + std::vector truth; + auto data = make_separated(k, per_cluster, dim, gen, truth); + // Skew by deleting so n % k == 2. + const uint64_t n = k * per_cluster - 2; + data.resize(n * static_cast(dim)); + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + + SECTION("num_candidates == 1 forces on-demand rescans") { + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + const auto labels = cluster.Run(k, + data.data(), + n, + 5, + nullptr, + 1e-6F, + vsag::KMeansInitMethod::RANDOM, + 5U, + /*num_candidates=*/1U); + const auto sizes = cluster_sizes(labels, k); + require_exact_balance(sizes, n, k); + } + SECTION("num_candidates >= k keeps every cluster as a candidate") { + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + const auto labels = cluster.Run(k, + data.data(), + n, + 5, + nullptr, + 1e-6F, + vsag::KMeansInitMethod::RANDOM, + 5U, + /*num_candidates=*/1024U); + const auto sizes = cluster_sizes(labels, k); + require_exact_balance(sizes, n, k); + } +} + +TEST_CASE("BalancedKMeans assignment matches centroids and reported loss", + "[ut][BalancedKMeansCluster]") { + constexpr uint32_t k = 6; + constexpr int32_t dim = 10; + std::mt19937 gen(4242); + std::normal_distribution noise(0.0F, 1.0F); + constexpr uint64_t n = 500; + std::vector data(n * static_cast(dim)); + for (auto& v : data) { + v = noise(gen); + } + + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::BalancedKMeansCluster cluster(dim, allocator.get()); + double reported = 0.0; + const auto labels = cluster.Run(k, data.data(), n, 4, &reported); + REQUIRE(reported > 0.0); + + // Recompute the final assignment loss against the centroids the labels were + // produced with: the mean of each cluster's members is the stored centroid, + // so recompute the cost of the stored (labels, centroids) pair. + std::vector sums(static_cast(k) * dim, 0.0); + std::vector sizes(k, 0); + for (uint64_t i = 0; i < n; ++i) { + const uint32_t label = static_cast(labels[i]); + ++sizes[label]; + for (int32_t d = 0; d < dim; ++d) { + sums[static_cast(label) * dim + d] += data[i * dim + d]; + } + } + double recomputed = 0.0; + for (uint32_t j = 0; j < k; ++j) { + REQUIRE(sizes[j] >= 1); + for (int32_t d = 0; d < dim; ++d) { + const float mean = + static_cast(sums[static_cast(j) * dim + d] / sizes[j]); + REQUIRE(std::abs(mean - cluster.k_centroids_[static_cast(j) * dim + d]) < + 1e-3F); + } + } + // The reported loss is the cost of the last assignment step; it must be no + // worse than the cost of placing every vector at its own centroid mean. + for (uint64_t i = 0; i < n; ++i) { + const uint32_t label = static_cast(labels[i]); + double dist = 0.0; + for (int32_t d = 0; d < dim; ++d) { + const double diff = + static_cast(data[i * dim + d]) - + static_cast(cluster.k_centroids_[static_cast(label) * dim + d]); + dist += diff * diff; + } + recomputed += dist; + } + REQUIRE(recomputed <= reported + 1.0); +} + +TEST_CASE("BalancedKMeans avoids a dense n-by-k distance matrix", "[ut][BalancedKMeansCluster]") { + // The assignment is required to bound its working set (issue #2870): the + // trainer keeps per-vector candidate lists, never an n x k matrix. A + // recording allocator measures the peak against that matrix's size. + constexpr uint32_t k = 200; + constexpr int32_t dim = 32; + constexpr uint64_t n = 4000; + std::mt19937 gen(12); + std::uniform_real_distribution coord(-1.0F, 1.0F); + std::vector data(n * static_cast(dim)); + for (auto& v : data) { + v = coord(gen); + } + const uint64_t dense_matrix_bytes = n * static_cast(k) * sizeof(float); + + auto allocator = std::make_shared(); + // The recording allocator only measures what the trainer allocates through + // the allocator it is handed; the dataset itself lives outside of it. + vsag::BalancedKMeansCluster cluster(dim, allocator.get(), nullptr); + const auto labels = cluster.Run(k, data.data(), n, 3); + + const uint64_t peak = allocator->GetMemoryPeak(); + REQUIRE(peak < dense_matrix_bytes / 4); + // Sanity: the run still balances. + const auto sizes = cluster_sizes(labels, k); + require_exact_balance(sizes, n, k); +} diff --git a/src/impl/cluster/kmeans_cluster.h b/src/impl/cluster/kmeans_cluster.h index 5f6bb8b9f6..18293dc968 100644 --- a/src/impl/cluster/kmeans_cluster.h +++ b/src/impl/cluster/kmeans_cluster.h @@ -50,6 +50,8 @@ struct KMeansGpuConfig { class KMeansCluster { public: + static constexpr uint64_t THRESHOLD_FOR_HGRAPH = 10000ULL; + explicit KMeansCluster(int32_t dim, Allocator* allocator, SafeThreadPoolPtr thread_pool = nullptr, @@ -99,8 +101,6 @@ class KMeansCluster { const int32_t dim_{0}; const KMeansGpuConfig gpu_config_{}; - - static constexpr uint64_t THRESHOLD_FOR_HGRAPH = 10000ULL; }; } // namespace vsag diff --git a/src/inner_string_params.h b/src/inner_string_params.h index e5e279e99b..95cd7531fa 100644 --- a/src/inner_string_params.h +++ b/src/inner_string_params.h @@ -176,6 +176,7 @@ const char* const BUCKET_USE_RESIDUAL_KEY = "use_residual"; const char* const IVF_TRAIN_TYPE_KEY = "ivf_train_type"; const char* const IVF_TRAIN_TYPE_RANDOM = "random"; const char* const IVF_TRAIN_TYPE_KMEANS = "kmeans"; +const char* const IVF_TRAIN_TYPE_BALANCED_KMEANS = "balanced_kmeans"; const char* const TRAIN_SAMPLE_COUNT_KEY = "train_sample_count"; // used after v0.18 for both Hgraph and IVF diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 9ea7dd6d54..49f7234997 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -23,6 +23,8 @@ #include #include #include +#include +#include #include #include #include @@ -4354,3 +4356,117 @@ TEST_CASE("IVF builds with the documented GPU configuration", "[ft][ivf][pr]") { REQUIRE(ids_of(reloaded_result.value(), topk) == ids); } } + +TEST_CASE("IVF balanced_kmeans builds balanced buckets", "[ft][ivf][balanced_kmeans][pr]") { + // Issue #2870: with ivf_train_type "balanced_kmeans" both the centroid + // training and the initial bucket assignment of a Build are + // capacity-constrained. On a dataset whose vectors pile onto one point, + // plain k-means would leave a bucket holding nearly everything, while the + // balanced trainer must spread them so every bucket size is floor(n/k) or + // ceil(n/k). + const int64_t dim = 16; + const int64_t count = 1200; + const int buckets_count = 10; + + const std::string build_param = fmt::format(R"({{ + "dtype": "float32", + "metric_type": "l2", + "dim": {}, + "index_param": {{ + "buckets_count": {}, + "base_quantization_type": "fp32", + "partition_strategy_type": "ivf", + "ivf_train_type": "balanced_kmeans", + "kmeans_iter_count": 10 + }} + }})", + dim, + buckets_count); + auto index = vsag::Factory::CreateIndex("ivf", build_param); + REQUIRE(index.has_value()); + + // 1000 exact duplicates of one vector plus 200 spread-out vectors: the + // pathological skew the feature exists for. + std::vector vectors(static_cast(count) * dim, 0.0F); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + std::mt19937 gen(2870); + std::uniform_real_distribution coord(-50.0F, 50.0F); + for (int64_t i = 1000; i < count; ++i) { + for (int64_t d = 0; d < dim; ++d) { + vectors[static_cast(i) * dim + d] = coord(gen); + } + } + auto base = vsag::Dataset::Make(); + base->NumElements(count) + ->Dim(dim) + ->Ids(ids.data()) + ->Float32Vectors(vectors.data()) + ->Owner(false); + REQUIRE(index.value()->Build(base).has_value()); + + const auto stats = vsag::JsonType::Parse(index.value()->GetStats()); + const float min_bucket = stats["bucket_num"]["min"].GetFloat(); + const float max_bucket = stats["bucket_num"]["max"].GetFloat(); + CAPTURE(min_bucket, max_bucket); + // n divides by k exactly here, so every bucket must hold n / k = 120. + REQUIRE(max_bucket - min_bucket <= 1.0F); + REQUIRE(min_bucket > 0.0F); + + // Search works and a serialize round trip preserves the answers. + auto query = vsag::Dataset::Make(); + std::vector query_vector(static_cast(dim), 0.0F); + query->NumElements(1)->Dim(dim)->Float32Vectors(query_vector.data())->Owner(false); + const std::string search_param = R"({"ivf":{"scan_buckets_count":10}})"; + auto result = index.value()->KnnSearch(query, 10, search_param); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 10); + std::vector found(result.value()->GetIds(), + result.value()->GetIds() + result.value()->GetDim()); + + std::stringstream stream; + REQUIRE(index.value()->Serialize(stream).has_value()); + auto reloaded = vsag::Factory::CreateIndex("ivf", build_param); + REQUIRE(reloaded.has_value()); + REQUIRE(reloaded.value()->Deserialize(stream).has_value()); + auto reloaded_result = reloaded.value()->KnnSearch(query, 10, search_param); + REQUIRE(reloaded_result.has_value()); + std::vector found_after( + reloaded_result.value()->GetIds(), + reloaded_result.value()->GetIds() + reloaded_result.value()->GetDim()); + REQUIRE(found_after == found); +} + +TEST_CASE("IVF rejects unknown trainers and unbalanced combinations", + "[ft][ivf][balanced_kmeans][pr]") { + const std::string unknown_train = R"({ + "dtype": "float32", + "metric_type": "l2", + "dim": 8, + "index_param": { + "buckets_count": 4, + "base_quantization_type": "fp32", + "ivf_train_type": "kmeanss" + } + })"; + auto index = vsag::Factory::CreateIndex("ivf", unknown_train); + REQUIRE_FALSE(index.has_value()); + REQUIRE(index.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + // GNO-IMI trains plain k-means at both levels; balanced_kmeans belongs to + // the single-level strategy. + const std::string gno_imi_balanced = R"({ + "dtype": "float32", + "metric_type": "l2", + "dim": 8, + "index_param": { + "buckets_count": 4, + "base_quantization_type": "fp32", + "partition_strategy_type": "gno_imi", + "ivf_train_type": "balanced_kmeans" + } + })"; + auto gno_index = vsag::Factory::CreateIndex("ivf", gno_imi_balanced); + REQUIRE_FALSE(gno_index.has_value()); + REQUIRE(gno_index.error().type == vsag::ErrorType::INVALID_ARGUMENT); +} diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 30408efaee..7cc4adc359 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -19,3 +19,4 @@ add_subdirectory (autotune) add_subdirectory (check_compatibility) add_subdirectory (analyze_index) add_subdirectory (visualize_index) +add_subdirectory (cluster_balance_bench) diff --git a/tools/README.md b/tools/README.md index d98103d7e8..0536fb9cce 100644 --- a/tools/README.md +++ b/tools/README.md @@ -9,6 +9,7 @@ nearby entry points for local development and AI coding agents. | `check_compatibility` | [check_compatibility/README.md](check_compatibility/README.md) | [docs/docs/en/src/resources/check_compatibility.md](../docs/docs/en/src/resources/check_compatibility.md) | | | `eval_performance` | [eval/README.md](eval/README.md) | [docs/docs/en/src/resources/eval.md](../docs/docs/en/src/resources/eval.md) | | | `autotune` | [autotune/README.md](autotune/README.md) | [docs/docs/en/src/resources/autotune.md](../docs/docs/en/src/resources/autotune.md) | | +| `cluster_balance_bench` | [cluster_balance_bench/README.md](cluster_balance_bench/README.md) | — | — | Chinese documentation is available under `docs/docs/zh/src/resources/` and on the Chinese website paths under . diff --git a/tools/cluster_balance_bench/CMakeLists.txt b/tools/cluster_balance_bench/CMakeLists.txt new file mode 100644 index 0000000000..debf0bbb89 --- /dev/null +++ b/tools/cluster_balance_bench/CMakeLists.txt @@ -0,0 +1,16 @@ +# Copyright 2024-present the vsag project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +add_executable (cluster_balance_bench cluster_balance_bench.cpp) +target_link_libraries (cluster_balance_bench PRIVATE vsag simd) diff --git a/tools/cluster_balance_bench/README.md b/tools/cluster_balance_bench/README.md new file mode 100644 index 0000000000..c28e66935c --- /dev/null +++ b/tools/cluster_balance_bench/README.md @@ -0,0 +1,7 @@ +# cluster_balance_bench + +The canonical build and usage guide is in the +[IVF documentation](../../docs/docs/en/src/indexes/ivf.md#benchmark-tool), under +“Balanced k-means training → Benchmark tool”. + +[中文说明](../../docs/docs/zh/src/indexes/ivf.md#基准工具) diff --git a/tools/cluster_balance_bench/cluster_balance_bench.cpp b/tools/cluster_balance_bench/cluster_balance_bench.cpp new file mode 100644 index 0000000000..e9024ae562 --- /dev/null +++ b/tools/cluster_balance_bench/cluster_balance_bench.cpp @@ -0,0 +1,463 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// cluster_balance_bench: reproducible comparison of the plain k-means and the +// balanced k-means IVF trainers (issue #2870). +// +// It trains both on synthetic datasets - uniform, deliberately skewed, a +// skewed one whose size is not divisible by k, and a duplicated-vector set - +// and reports, per trainer and dataset: +// * minimum, maximum and average cluster size, and the size standard +// deviation and coefficient of variation; +// * the clustering loss (mean squared distance from each vector to its +// assigned centroid, recomputed from the returned labels and the final +// centroids so the two trainers are measured with the same yardstick); +// * the training time; +// * the peak memory the trainer allocated beyond the dataset itself, +// measured through a recording allocator, alongside the size of a dense +// n x k float distance matrix for reference. +// +// Usage: +// cluster_balance_bench [--dataset all|uniform|skewed|indivisible|duplicates] +// [--n 50000] [--dim 64] [--k 256] [--iter 25] +// [--seed 42] [--threads N] +// +// Every run is fully seeded; repeated invocations with the same arguments +// reproduce the reported numbers up to timing noise. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "impl/cluster/balanced_kmeans_cluster.h" +#include "impl/cluster/kmeans_cluster.h" +#include "impl/thread_pool/safe_thread_pool.h" +#include "typing.h" +#include "vsag/allocator.h" + +namespace { + +// Keep the numeric threshold in the en/zh IVF documentation in sync as well. +static_assert(vsag::KMeansCluster::THRESHOLD_FOR_HGRAPH == 10000ULL, + "Update the en/zh IVF documentation for the new HGraph threshold"); + +/// Allocator that tracks currently-live bytes and records the peak, so the +/// benchmark can report trainer memory without process-level instrumentation. +class RecordingAllocator : public vsag::Allocator { +public: + std::string + Name() override { + return "RecordingAllocator"; + } + + void* + Allocate(uint64_t size) override { + void* ptr = std::malloc(size); + std::lock_guard lock(mutex_); + live_[ptr] = size; + current_ += size; + peak_ = std::max(peak_, current_); + return ptr; + } + + void + Deallocate(void* p) override { + { + std::lock_guard lock(mutex_); + auto it = live_.find(p); + if (it != live_.end()) { + current_ -= it->second; + live_.erase(it); + } + } + std::free(p); + } + + void* + Reallocate(void* p, uint64_t size) override { + Deallocate(p); + return Allocate(size); + } + + uint64_t + Peak() const { + return peak_; + } + + void + Reset() { + std::lock_guard lock(mutex_); + peak_ = current_; + } + +private: + std::mutex mutex_; + std::unordered_map live_; + uint64_t current_{0}; + uint64_t peak_{0}; +}; + +/// Geometric-mixture dataset: cluster sizes drop by ~3x and well-separated +/// centers make plain k-means pile vectors into few buckets. The duplicated +/// mode collapses every point of the largest component onto one vector. +std::vector +make_dataset(const std::string& kind, + uint32_t k, + uint64_t n, + int32_t dim, + uint32_t seed, + uint64_t& actual_n) { + std::mt19937 gen(seed); + std::vector data; + if (kind == "uniform") { + actual_n = n; + data.resize(n * static_cast(dim)); + std::uniform_real_distribution coord(0.0F, 1.0F); + for (auto& v : data) { + v = coord(gen); + } + return data; + } + if (kind == "duplicates") { + actual_n = n; + data.assign(n * static_cast(dim), 0.0F); + // The whole set coincides except for a few outliers: an unconstrained + // trainer has nowhere sensible to put 90% of the mass. + for (uint64_t i = 0; i < n / 20; ++i) { + for (int32_t d = 0; d < dim; ++d) { + data[(n - 1 - i) * static_cast(dim) + d] = 100.0F + static_cast(i); + } + } + return data; + } + // skewed (and indivisible = skewed trimmed to n % k != 0) + std::vector sizes(k); + uint64_t total = 0; + // Geometric tail: with cur as the head component, the sizes sum to roughly + // 1.5 * cur + k, so pick cur to land near the requested n. + uint64_t cur = (n > 2 * static_cast(k)) ? (n - k) * 2 / 3 : n; + for (uint32_t c = 0; c < k; ++c) { + sizes[c] = std::max(1, cur); + total += sizes[c]; + cur = cur / 3 + 1; + } + actual_n = total; + // With one cluster every size is divisible; in particular, do not turn the + // valid n=1, k=1 case into an empty dataset. + if (kind == "indivisible" and k > 1 and total % k == 0) { + actual_n = total - 1; + } + std::vector centroids(static_cast(k) * dim); + std::uniform_real_distribution coord(-100.0F, 100.0F); + for (auto& v : centroids) { + v = coord(gen); + } + std::normal_distribution noise(0.0F, 0.05F); + data.assign(actual_n * static_cast(dim), 0.0F); + uint64_t i = 0; + for (uint32_t c = 0; c < k and i < actual_n; ++c) { + for (uint64_t j = 0; j < sizes[c] and i < actual_n; ++j, ++i) { + for (int32_t d = 0; d < dim; ++d) { + data[i * static_cast(dim) + d] = + centroids[static_cast(c) * static_cast(dim) + d] + + static_cast(noise(gen)); + } + } + } + return data; +} + +struct Result { + std::string trainer; + std::string dataset; + uint64_t n{0}; + uint32_t k{0}; + uint64_t min_size{0}; + uint64_t max_size{0}; + double mean_size{0.0}; + double std_size{0.0}; + double cv_size{0.0}; + double loss{0.0}; + double seconds{0.0}; + uint64_t peak_bytes{0}; +}; + +Result +measure(const std::string& trainer, + const std::string& dataset, + const vsag::Vector& labels, + const std::vector& data, + const float* centroids, + uint32_t k, + int32_t dim, + uint64_t n, + double seconds, + uint64_t peak_bytes) { + Result r; + r.trainer = trainer; + r.dataset = dataset; + r.n = n; + r.k = k; + std::vector sizes(k, 0); + for (uint64_t i = 0; i < n; ++i) { + ++sizes[static_cast(labels[i])]; + } + r.min_size = *std::min_element(sizes.begin(), sizes.end()); + r.max_size = *std::max_element(sizes.begin(), sizes.end()); + r.mean_size = static_cast(n) / static_cast(k); + double var = 0.0; + for (uint64_t s : sizes) { + var += (static_cast(s) - r.mean_size) * (static_cast(s) - r.mean_size); + } + var /= static_cast(k); + r.std_size = std::sqrt(var); + r.cv_size = r.std_size / r.mean_size; + // Loss measured identically for both trainers: mean squared distance from + // each vector to the centroid it was assigned to. + double loss = 0.0; + for (uint64_t i = 0; i < n; ++i) { + const float* x = data.data() + i * static_cast(dim); + const float* c = centroids + static_cast(labels[i]) * static_cast(dim); + double sum = 0.0; + for (int32_t d = 0; d < dim; ++d) { + const double diff = static_cast(x[d]) - static_cast(c[d]); + sum += diff * diff; + } + loss += sum; + } + r.loss = loss / static_cast(n); + r.seconds = seconds; + r.peak_bytes = peak_bytes; + return r; +} + +void +run_bench(const std::string& dataset, + uint64_t n, + int32_t dim, + uint32_t k, + int iter, + uint32_t seed, + int threads, + std::vector& results) { + uint64_t actual_n = 0; + const auto data = make_dataset(dataset, k, n, dim, seed, actual_n); + + auto pool = vsag::SafeThreadPool::FactoryDefaultThreadPool(); + if (threads > 0) { + pool->SetPoolSize(static_cast(threads)); + } + + // Plain k-means, seeded identically to the balanced trainer. + { + RecordingAllocator allocator; + vsag::KMeansCluster cluster(dim, &allocator, pool); + const auto start = std::chrono::steady_clock::now(); + const auto labels = cluster.Run(k, + data.data(), + actual_n, + iter, + nullptr, + false, + 1e-6F, + vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, + seed, + /*deterministic_reduction=*/true); + const auto end = std::chrono::steady_clock::now(); + results.push_back(measure("kmeans", + dataset, + labels, + data, + cluster.k_centroids_, + k, + dim, + actual_n, + std::chrono::duration(end - start).count(), + allocator.Peak())); + } + + // Balanced k-means. + { + RecordingAllocator allocator; + vsag::BalancedKMeansCluster cluster(dim, &allocator, pool); + const auto start = std::chrono::steady_clock::now(); + const auto labels = cluster.Run(k, + data.data(), + actual_n, + iter, + nullptr, + 1e-6F, + vsag::KMeansInitMethod::KMEANS_PLUS_PLUS, + seed); + const auto end = std::chrono::steady_clock::now(); + results.push_back(measure("balanced_kmeans", + dataset, + labels, + data, + cluster.k_centroids_, + k, + dim, + actual_n, + std::chrono::duration(end - start).count(), + allocator.Peak())); + } +} + +void +print_usage(const char* argv0) { + std::printf( + "usage: %s [--dataset all|uniform|skewed|indivisible|duplicates] [--n N] [--dim D]\n" + " [--k K] [--iter I] [--seed S] [--threads T]\n", + argv0); +} + +bool +parse(const int argc, + char** argv, + std::string& dataset, + uint64_t& n, + int32_t& dim, + uint32_t& k, + int& iter, + uint32_t& seed, + int& threads) { + for (int i = 1; i < argc; ++i) { + const std::string key = argv[i]; + auto next_value = [&](const char* name) -> std::string { + if (i + 1 >= argc) { + std::printf("missing value for %s\n", name); + std::exit(2); + } + return argv[++i]; + }; + if (key == "--dataset") { + dataset = next_value("--dataset"); + } else if (key == "--n") { + n = std::strtoull(next_value("--n").c_str(), nullptr, 10); + } else if (key == "--dim") { + dim = static_cast(std::strtol(next_value("--dim").c_str(), nullptr, 10)); + } else if (key == "--k") { + k = static_cast(std::strtoul(next_value("--k").c_str(), nullptr, 10)); + } else if (key == "--iter") { + iter = static_cast(std::strtol(next_value("--iter").c_str(), nullptr, 10)); + } else if (key == "--seed") { + seed = static_cast(std::strtoul(next_value("--seed").c_str(), nullptr, 10)); + } else if (key == "--threads") { + threads = static_cast(std::strtol(next_value("--threads").c_str(), nullptr, 10)); + } else if (key == "--help" or key == "-h") { + print_usage(argv[0]); + std::exit(0); + } else { + std::printf("unknown argument: %s\n", key.c_str()); + print_usage(argv[0]); + return false; + } + } + if (k >= vsag::KMeansCluster::THRESHOLD_FOR_HGRAPH) { + // KMeansCluster switches to non-deterministic parallel HGraph assignment here. + std::printf("--k must be less than %llu for reproducible HGraph-free assignment\n", + static_cast(vsag::KMeansCluster::THRESHOLD_FOR_HGRAPH)); + return false; + } + return k > 0 and n > 0 and dim > 0 and iter > 0 and (k <= n); +} + +} // namespace + +int +main(int argc, char** argv) { + std::string dataset = "all"; + uint64_t n = 50000; + int32_t dim = 64; + uint32_t k = 256; + int iter = 25; + uint32_t seed = 42; + int threads = 0; + if (not parse(argc, argv, dataset, n, dim, k, iter, seed, threads)) { + print_usage(argv[0]); + return 2; + } + + const std::vector all = {"uniform", "skewed", "indivisible", "duplicates"}; + std::vector datasets; + if (dataset == "all") { + datasets = all; + } else { + datasets.push_back(dataset); + } + + std::vector results; + for (const auto& name : datasets) { + run_bench(name, n, dim, k, iter, seed, threads, results); + } + + std::printf("cluster_balance_bench: n=%llu dim=%d k=%u iter=%d seed=%u threads=%d\n", + static_cast(n), + dim, + k, + iter, + seed, + threads); + std::printf("dense n*k*4 reference: %.1f MiB (balanced_kmeans must stay far below)\n", + static_cast(n) * static_cast(k) * 4.0 / 1048576.0); + std::printf("\n%-16s %-12s %-12s %-13s %-9s %-9s %-8s %-9s %-8s %-11s\n", + "trainer", + "dataset", + "n", + "sizes", + "mean", + "std", + "CV", + "loss(MSE)", + "time", + "peak-mem"); + for (const auto& r : results) { + char sizes[64]; + std::snprintf(sizes, + sizeof(sizes), + "%llu-%llu", + static_cast(r.min_size), + static_cast(r.max_size)); + char mem[32]; + std::snprintf(mem, sizeof(mem), "%.1f MiB", static_cast(r.peak_bytes) / 1048576.0); + std::printf("%-16s %-12s %-12llu %-13s %-9.1f %-9.1f %-8.3f %-9.4f %-7.2fs %s\n", + r.trainer.c_str(), + r.dataset.c_str(), + static_cast(r.n), + sizes, + r.mean_size, + r.std_size, + r.cv_size, + r.loss, + r.seconds, + mem); + if (r.trainer == "balanced_kmeans") { + const bool exact = + r.max_size - r.min_size <= 1 and r.min_size >= 1; // valid input has k <= n + std::printf(" balance check: %s\n", exact ? "PASS" : "FAIL"); + if (not exact) { + return 1; + } + } + } + return 0; +}