Support IVFPQ training (#2078)

### What problem does this PR solve? Support IVFPQ training Issue link:#2077 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
infiniflow · Oct 21, 2024 · 86f0c8b · 86f0c8b
1 parent fe3c181
commit 86f0c8b
Show file tree

Hide file tree

Showing 5 changed files with 830 additions and 282 deletions.
diff --git a/src/storage/definition/index_ivf.cpp b/src/storage/definition/index_ivf.cpp
@@ -276,11 +276,12 @@ void CheckIndexIVFStorageOption(IndexIVFStorageOption &storage_option, const Dat
             } else if (storage_option.type_ == IndexIVFStorageOption::Type::kProductQuantization) {
                 // check product_quantization_subspace_num_, product_quantization_subspace_bits_
                 if (const auto product_quantization_subspace_num = storage_option.product_quantization_subspace_num_;
-                    product_quantization_subspace_num > embedding_dimension) {
-                    RecoverableError(
-                        Status::InvalidIndexDefinition(std::format("product_quantization_subspace_num ({}) is larger than embedding_dimension({}).",
-                                                                   product_quantization_subspace_num,
-                                                                   embedding_dimension)));
+                    product_quantization_subspace_num >= embedding_dimension || product_quantization_subspace_num <= 0 ||
+                    embedding_dimension % product_quantization_subspace_num != 0) {
+                    RecoverableError(Status::InvalidIndexDefinition(
+                        std::format("product_quantization_subspace_num ({}) is not divisor of embedding_dimension({})!",
+                                    product_quantization_subspace_num,
+                                    embedding_dimension)));
                 }
                 if (const auto product_quantization_subspace_bits = storage_option.product_quantization_subspace_bits_;
                     product_quantization_subspace_bits < 4 || product_quantization_subspace_bits > 16) {

diff --git a/src/storage/definition/index_ivf.cppm b/src/storage/definition/index_ivf.cppm
@@ -46,7 +46,7 @@ export struct IndexIVFStorageOption {
     // kScalarQuantization
     u32 scalar_quantization_bits_ = 0; // 4 or 8
     // kProductQuantization
-    u32 product_quantization_subspace_num_ = 0;  // in range [1, dim]
+    u32 product_quantization_subspace_num_ = 0;  // divisor of embedding dimension
     u32 product_quantization_subspace_bits_ = 0; // in range [4, 16]
     bool operator==(const IndexIVFStorageOption &other) const = default;
     String ToString() const;