faiss 0.2.4 → 0.2.6

Sign up to get free protection for your applications and to get access to all the features.
Files changed (178) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +10 -0
  3. data/LICENSE.txt +1 -1
  4. data/README.md +23 -21
  5. data/ext/faiss/extconf.rb +11 -0
  6. data/ext/faiss/index.cpp +17 -4
  7. data/ext/faiss/index_binary.cpp +6 -6
  8. data/ext/faiss/product_quantizer.cpp +4 -4
  9. data/lib/faiss/version.rb +1 -1
  10. data/vendor/faiss/faiss/AutoTune.cpp +13 -0
  11. data/vendor/faiss/faiss/IVFlib.cpp +101 -2
  12. data/vendor/faiss/faiss/IVFlib.h +26 -2
  13. data/vendor/faiss/faiss/Index.cpp +36 -3
  14. data/vendor/faiss/faiss/Index.h +43 -6
  15. data/vendor/faiss/faiss/Index2Layer.cpp +6 -2
  16. data/vendor/faiss/faiss/Index2Layer.h +6 -1
  17. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +219 -16
  18. data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +63 -5
  19. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +299 -0
  20. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h +199 -0
  21. data/vendor/faiss/faiss/IndexBinary.cpp +20 -4
  22. data/vendor/faiss/faiss/IndexBinary.h +18 -3
  23. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +9 -2
  24. data/vendor/faiss/faiss/IndexBinaryFlat.h +4 -2
  25. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +4 -1
  26. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +2 -1
  27. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +5 -1
  28. data/vendor/faiss/faiss/IndexBinaryHNSW.h +2 -1
  29. data/vendor/faiss/faiss/IndexBinaryHash.cpp +17 -4
  30. data/vendor/faiss/faiss/IndexBinaryHash.h +8 -4
  31. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +28 -13
  32. data/vendor/faiss/faiss/IndexBinaryIVF.h +10 -7
  33. data/vendor/faiss/faiss/IndexFastScan.cpp +626 -0
  34. data/vendor/faiss/faiss/IndexFastScan.h +145 -0
  35. data/vendor/faiss/faiss/IndexFlat.cpp +34 -21
  36. data/vendor/faiss/faiss/IndexFlat.h +7 -4
  37. data/vendor/faiss/faiss/IndexFlatCodes.cpp +35 -1
  38. data/vendor/faiss/faiss/IndexFlatCodes.h +12 -0
  39. data/vendor/faiss/faiss/IndexHNSW.cpp +66 -138
  40. data/vendor/faiss/faiss/IndexHNSW.h +4 -2
  41. data/vendor/faiss/faiss/IndexIDMap.cpp +247 -0
  42. data/vendor/faiss/faiss/IndexIDMap.h +107 -0
  43. data/vendor/faiss/faiss/IndexIVF.cpp +121 -33
  44. data/vendor/faiss/faiss/IndexIVF.h +35 -16
  45. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +84 -7
  46. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +63 -1
  47. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +590 -0
  48. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +171 -0
  49. data/vendor/faiss/faiss/IndexIVFFastScan.cpp +1290 -0
  50. data/vendor/faiss/faiss/IndexIVFFastScan.h +213 -0
  51. data/vendor/faiss/faiss/IndexIVFFlat.cpp +37 -17
  52. data/vendor/faiss/faiss/IndexIVFFlat.h +4 -2
  53. data/vendor/faiss/faiss/IndexIVFPQ.cpp +234 -50
  54. data/vendor/faiss/faiss/IndexIVFPQ.h +5 -1
  55. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +23 -852
  56. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +7 -112
  57. data/vendor/faiss/faiss/IndexIVFPQR.cpp +3 -3
  58. data/vendor/faiss/faiss/IndexIVFPQR.h +1 -1
  59. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +3 -1
  60. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +2 -1
  61. data/vendor/faiss/faiss/IndexLSH.cpp +4 -2
  62. data/vendor/faiss/faiss/IndexLSH.h +2 -1
  63. data/vendor/faiss/faiss/IndexLattice.cpp +7 -1
  64. data/vendor/faiss/faiss/IndexLattice.h +3 -1
  65. data/vendor/faiss/faiss/IndexNNDescent.cpp +4 -3
  66. data/vendor/faiss/faiss/IndexNNDescent.h +2 -1
  67. data/vendor/faiss/faiss/IndexNSG.cpp +37 -3
  68. data/vendor/faiss/faiss/IndexNSG.h +25 -1
  69. data/vendor/faiss/faiss/IndexPQ.cpp +106 -69
  70. data/vendor/faiss/faiss/IndexPQ.h +19 -5
  71. data/vendor/faiss/faiss/IndexPQFastScan.cpp +15 -450
  72. data/vendor/faiss/faiss/IndexPQFastScan.h +15 -78
  73. data/vendor/faiss/faiss/IndexPreTransform.cpp +47 -8
  74. data/vendor/faiss/faiss/IndexPreTransform.h +15 -3
  75. data/vendor/faiss/faiss/IndexRefine.cpp +8 -4
  76. data/vendor/faiss/faiss/IndexRefine.h +4 -2
  77. data/vendor/faiss/faiss/IndexReplicas.cpp +4 -2
  78. data/vendor/faiss/faiss/IndexReplicas.h +2 -1
  79. data/vendor/faiss/faiss/IndexRowwiseMinMax.cpp +438 -0
  80. data/vendor/faiss/faiss/IndexRowwiseMinMax.h +92 -0
  81. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +26 -15
  82. data/vendor/faiss/faiss/IndexScalarQuantizer.h +6 -7
  83. data/vendor/faiss/faiss/IndexShards.cpp +4 -1
  84. data/vendor/faiss/faiss/IndexShards.h +2 -1
  85. data/vendor/faiss/faiss/MetaIndexes.cpp +5 -178
  86. data/vendor/faiss/faiss/MetaIndexes.h +3 -81
  87. data/vendor/faiss/faiss/VectorTransform.cpp +43 -0
  88. data/vendor/faiss/faiss/VectorTransform.h +22 -4
  89. data/vendor/faiss/faiss/clone_index.cpp +23 -1
  90. data/vendor/faiss/faiss/clone_index.h +3 -0
  91. data/vendor/faiss/faiss/cppcontrib/SaDecodeKernels.h +300 -0
  92. data/vendor/faiss/faiss/cppcontrib/detail/CoarseBitType.h +24 -0
  93. data/vendor/faiss/faiss/cppcontrib/detail/UintReader.h +195 -0
  94. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-avx2-inl.h +2058 -0
  95. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +408 -0
  96. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-neon-inl.h +2147 -0
  97. data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMax-inl.h +460 -0
  98. data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMaxFP16-inl.h +465 -0
  99. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-avx2-inl.h +1618 -0
  100. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-inl.h +251 -0
  101. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-neon-inl.h +1452 -0
  102. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +1 -0
  103. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +0 -4
  104. data/vendor/faiss/faiss/gpu/GpuIndex.h +28 -4
  105. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +2 -1
  106. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +10 -8
  107. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +75 -14
  108. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +19 -32
  109. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +22 -31
  110. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +22 -28
  111. data/vendor/faiss/faiss/gpu/GpuResources.cpp +14 -0
  112. data/vendor/faiss/faiss/gpu/GpuResources.h +16 -3
  113. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +3 -3
  114. data/vendor/faiss/faiss/gpu/impl/IndexUtils.h +32 -0
  115. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +1 -0
  116. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +311 -75
  117. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +10 -0
  118. data/vendor/faiss/faiss/gpu/test/TestUtils.h +3 -0
  119. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +2 -2
  120. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +5 -4
  121. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +116 -47
  122. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +44 -13
  123. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +0 -54
  124. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +0 -76
  125. data/vendor/faiss/faiss/impl/DistanceComputer.h +64 -0
  126. data/vendor/faiss/faiss/impl/HNSW.cpp +123 -27
  127. data/vendor/faiss/faiss/impl/HNSW.h +19 -16
  128. data/vendor/faiss/faiss/impl/IDSelector.cpp +125 -0
  129. data/vendor/faiss/faiss/impl/IDSelector.h +135 -0
  130. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +6 -28
  131. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +6 -1
  132. data/vendor/faiss/faiss/impl/LookupTableScaler.h +77 -0
  133. data/vendor/faiss/faiss/impl/NNDescent.cpp +1 -0
  134. data/vendor/faiss/faiss/impl/NSG.cpp +1 -1
  135. data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp +383 -0
  136. data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.h +154 -0
  137. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +225 -145
  138. data/vendor/faiss/faiss/impl/ProductQuantizer.h +29 -10
  139. data/vendor/faiss/faiss/impl/Quantizer.h +43 -0
  140. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +192 -36
  141. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +40 -20
  142. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -0
  143. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +97 -173
  144. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +18 -18
  145. data/vendor/faiss/faiss/impl/index_read.cpp +240 -9
  146. data/vendor/faiss/faiss/impl/index_write.cpp +237 -5
  147. data/vendor/faiss/faiss/impl/kmeans1d.cpp +6 -4
  148. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +56 -16
  149. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +25 -8
  150. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +66 -25
  151. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +75 -27
  152. data/vendor/faiss/faiss/index_factory.cpp +196 -7
  153. data/vendor/faiss/faiss/index_io.h +5 -0
  154. data/vendor/faiss/faiss/invlists/DirectMap.cpp +1 -0
  155. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +4 -1
  156. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +2 -1
  157. data/vendor/faiss/faiss/python/python_callbacks.cpp +27 -0
  158. data/vendor/faiss/faiss/python/python_callbacks.h +15 -0
  159. data/vendor/faiss/faiss/utils/Heap.h +31 -15
  160. data/vendor/faiss/faiss/utils/distances.cpp +380 -56
  161. data/vendor/faiss/faiss/utils/distances.h +113 -15
  162. data/vendor/faiss/faiss/utils/distances_simd.cpp +726 -6
  163. data/vendor/faiss/faiss/utils/extra_distances.cpp +12 -7
  164. data/vendor/faiss/faiss/utils/extra_distances.h +3 -1
  165. data/vendor/faiss/faiss/utils/fp16-fp16c.h +21 -0
  166. data/vendor/faiss/faiss/utils/fp16-inl.h +101 -0
  167. data/vendor/faiss/faiss/utils/fp16.h +11 -0
  168. data/vendor/faiss/faiss/utils/hamming-inl.h +54 -0
  169. data/vendor/faiss/faiss/utils/hamming.cpp +0 -48
  170. data/vendor/faiss/faiss/utils/ordered_key_value.h +10 -0
  171. data/vendor/faiss/faiss/utils/quantize_lut.cpp +62 -0
  172. data/vendor/faiss/faiss/utils/quantize_lut.h +20 -0
  173. data/vendor/faiss/faiss/utils/random.cpp +53 -0
  174. data/vendor/faiss/faiss/utils/random.h +5 -0
  175. data/vendor/faiss/faiss/utils/simdlib_avx2.h +4 -0
  176. data/vendor/faiss/faiss/utils/simdlib_emulated.h +6 -1
  177. data/vendor/faiss/faiss/utils/simdlib_neon.h +7 -2
  178. metadata +37 -3
@@ -9,6 +9,7 @@
9
9
 
10
10
  #include <memory>
11
11
 
12
+ #include <faiss/IndexIVFFastScan.h>
12
13
  #include <faiss/IndexIVFPQ.h>
13
14
  #include <faiss/impl/ProductQuantizer.h>
14
15
  #include <faiss/utils/AlignedTable.h>
@@ -31,36 +32,20 @@ namespace faiss {
31
32
  * 13: idem, collect results in reservoir
32
33
  */
33
34
 
34
- struct IndexIVFPQFastScan : IndexIVF {
35
- bool by_residual; ///< Encode residual or plain vector?
35
+ struct IndexIVFPQFastScan : IndexIVFFastScan {
36
36
  ProductQuantizer pq; ///< produces the codes
37
37
 
38
- // size of the kernel
39
- int bbs; // set at build time
40
-
41
- // M rounded up to a multiple of 2
42
- size_t M2;
43
-
44
38
  /// precomputed tables management
45
39
  int use_precomputed_table = 0;
46
40
  /// if use_precompute_table size (nlist, pq.M, pq.ksub)
47
41
  AlignedTable<float> precomputed_table;
48
42
 
49
- // search-time implementation
50
- int implem = 0;
51
- // skip some parts of the computation (for timing)
52
- int skip = 0;
53
-
54
- // batching factors at search time (0 = default)
55
- int qbs = 0;
56
- size_t qbs2 = 0;
57
-
58
43
  IndexIVFPQFastScan(
59
44
  Index* quantizer,
60
45
  size_t d,
61
46
  size_t nlist,
62
47
  size_t M,
63
- size_t nbits_per_idx,
48
+ size_t nbits,
64
49
  MetricType metric = METRIC_L2,
65
50
  int bbs = 32);
66
51
 
@@ -69,9 +54,6 @@ struct IndexIVFPQFastScan : IndexIVF {
69
54
  // built from an IndexIVFPQ
70
55
  explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32);
71
56
 
72
- /// orig's inverted lists (for debugging)
73
- InvertedLists* orig_invlists = nullptr;
74
-
75
57
  void train_residual(idx_t n, const float* x) override;
76
58
 
77
59
  /// build precomputed table, possibly updating use_precomputed_table
@@ -86,106 +68,19 @@ struct IndexIVFPQFastScan : IndexIVF {
86
68
  uint8_t* codes,
87
69
  bool include_listno = false) const override;
88
70
 
89
- void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
90
-
91
- void search(
92
- idx_t n,
93
- const float* x,
94
- idx_t k,
95
- float* distances,
96
- idx_t* labels) const override;
97
-
98
71
  // prepare look-up tables
99
72
 
73
+ bool lookup_table_is_3d() const override;
74
+
100
75
  void compute_LUT(
101
76
  size_t n,
102
77
  const float* x,
103
78
  const idx_t* coarse_ids,
104
79
  const float* coarse_dis,
105
80
  AlignedTable<float>& dis_tables,
106
- AlignedTable<float>& biases) const;
107
-
108
- void compute_LUT_uint8(
109
- size_t n,
110
- const float* x,
111
- const idx_t* coarse_ids,
112
- const float* coarse_dis,
113
- AlignedTable<uint8_t>& dis_tables,
114
- AlignedTable<uint16_t>& biases,
115
- float* normalizers) const;
116
-
117
- // internal search funcs
81
+ AlignedTable<float>& biases) const override;
118
82
 
119
- template <bool is_max>
120
- void search_dispatch_implem(
121
- idx_t n,
122
- const float* x,
123
- idx_t k,
124
- float* distances,
125
- idx_t* labels) const;
126
-
127
- template <class C>
128
- void search_implem_1(
129
- idx_t n,
130
- const float* x,
131
- idx_t k,
132
- float* distances,
133
- idx_t* labels) const;
134
-
135
- template <class C>
136
- void search_implem_2(
137
- idx_t n,
138
- const float* x,
139
- idx_t k,
140
- float* distances,
141
- idx_t* labels) const;
142
-
143
- // implem 10 and 12 are not multithreaded internally, so
144
- // export search stats
145
- template <class C>
146
- void search_implem_10(
147
- idx_t n,
148
- const float* x,
149
- idx_t k,
150
- float* distances,
151
- idx_t* labels,
152
- int impl,
153
- size_t* ndis_out,
154
- size_t* nlist_out) const;
155
-
156
- template <class C>
157
- void search_implem_12(
158
- idx_t n,
159
- const float* x,
160
- idx_t k,
161
- float* distances,
162
- idx_t* labels,
163
- int impl,
164
- size_t* ndis_out,
165
- size_t* nlist_out) const;
83
+ void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
166
84
  };
167
85
 
168
- struct IVFFastScanStats {
169
- uint64_t times[10];
170
- uint64_t t_compute_distance_tables, t_round;
171
- uint64_t t_copy_pack, t_scan, t_to_flat;
172
- uint64_t reservoir_times[4];
173
-
174
- double Mcy_at(int i) {
175
- return times[i] / (1000 * 1000.0);
176
- }
177
-
178
- double Mcy_reservoir_at(int i) {
179
- return reservoir_times[i] / (1000 * 1000.0);
180
- }
181
- IVFFastScanStats() {
182
- reset();
183
- }
184
- void reset() {
185
- memset(this, 0, sizeof(*this));
186
- }
187
- };
188
-
189
- FAISS_API extern IVFFastScanStats IVFFastScan_stats;
190
-
191
86
  } // namespace faiss
@@ -201,11 +201,11 @@ void IndexIVFPQR::reconstruct_from_offset(
201
201
  }
202
202
  }
203
203
 
204
- void IndexIVFPQR::merge_from(IndexIVF& other_in, idx_t add_id) {
205
- IndexIVFPQR* other = dynamic_cast<IndexIVFPQR*>(&other_in);
204
+ void IndexIVFPQR::merge_from(Index& otherIndex, idx_t add_id) {
205
+ IndexIVFPQR* other = dynamic_cast<IndexIVFPQR*>(&otherIndex);
206
206
  FAISS_THROW_IF_NOT(other);
207
207
 
208
- IndexIVF::merge_from(other_in, add_id);
208
+ IndexIVF::merge_from(otherIndex, add_id);
209
209
 
210
210
  refine_codes.insert(
211
211
  refine_codes.end(),
@@ -51,7 +51,7 @@ struct IndexIVFPQR : IndexIVFPQ {
51
51
  void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
52
52
  const override;
53
53
 
54
- void merge_from(IndexIVF& other, idx_t add_id) override;
54
+ void merge_from(Index& otherIndex, idx_t add_id) override;
55
55
 
56
56
  void search_preassigned(
57
57
  idx_t n,
@@ -293,7 +293,9 @@ struct IVFScanner : InvertedListScanner {
293
293
  } // anonymous namespace
294
294
 
295
295
  InvertedListScanner* IndexIVFSpectralHash::get_InvertedListScanner(
296
- bool store_pairs) const {
296
+ bool store_pairs,
297
+ const IDSelector* sel) const {
298
+ FAISS_THROW_IF_NOT(!sel);
297
299
  switch (code_size) {
298
300
  #define HANDLE_CODE_SIZE(cs) \
299
301
  case cs: \
@@ -70,7 +70,8 @@ struct IndexIVFSpectralHash : IndexIVF {
70
70
  bool include_listnos = false) const override;
71
71
 
72
72
  InvertedListScanner* get_InvertedListScanner(
73
- bool store_pairs) const override;
73
+ bool store_pairs,
74
+ const IDSelector* sel) const override;
74
75
 
75
76
  /** replace the vector transform for an empty (and possibly untrained) index
76
77
  */
@@ -103,9 +103,11 @@ void IndexLSH::search(
103
103
  const float* x,
104
104
  idx_t k,
105
105
  float* distances,
106
- idx_t* labels) const {
106
+ idx_t* labels,
107
+ const SearchParameters* params) const {
108
+ FAISS_THROW_IF_NOT_MSG(
109
+ !params, "search params not supported for this index");
107
110
  FAISS_THROW_IF_NOT(k > 0);
108
-
109
111
  FAISS_THROW_IF_NOT(is_trained);
110
112
  const float* xt = apply_preprocess(n, x);
111
113
  ScopeDeleter<float> del(xt == x ? nullptr : xt);
@@ -49,7 +49,8 @@ struct IndexLSH : IndexFlatCodes {
49
49
  const float* x,
50
50
  idx_t k,
51
51
  float* distances,
52
- idx_t* labels) const override;
52
+ idx_t* labels,
53
+ const SearchParameters* params = nullptr) const override;
53
54
 
54
55
  /// transfer the thresholds to a pre-processing stage (and unset
55
56
  /// train_thresholds)
@@ -118,7 +118,13 @@ void IndexLattice::add(idx_t, const float*) {
118
118
  FAISS_THROW_MSG("not implemented");
119
119
  }
120
120
 
121
- void IndexLattice::search(idx_t, const float*, idx_t, float*, idx_t*) const {
121
+ void IndexLattice::search(
122
+ idx_t,
123
+ const float*,
124
+ idx_t,
125
+ float*,
126
+ idx_t*,
127
+ const SearchParameters*) const {
122
128
  FAISS_THROW_MSG("not implemented");
123
129
  }
124
130
 
@@ -54,7 +54,9 @@ struct IndexLattice : Index {
54
54
  const float* x,
55
55
  idx_t k,
56
56
  float* distances,
57
- idx_t* labels) const override;
57
+ idx_t* labels,
58
+ const SearchParameters* params = nullptr) const override;
59
+
58
60
  void reset() override;
59
61
  };
60
62
 
@@ -135,9 +135,10 @@ void IndexNNDescent::search(
135
135
  const float* x,
136
136
  idx_t k,
137
137
  float* distances,
138
- idx_t* labels) const
139
-
140
- {
138
+ idx_t* labels,
139
+ const SearchParameters* params) const {
140
+ FAISS_THROW_IF_NOT_MSG(
141
+ !params, "search params not supported for this index");
141
142
  FAISS_THROW_IF_NOT_MSG(
142
143
  storage,
143
144
  "Please use IndexNNDescentFlat (or variants) "
@@ -53,7 +53,8 @@ struct IndexNNDescent : Index {
53
53
  const float* x,
54
54
  idx_t k,
55
55
  float* distances,
56
- idx_t* labels) const override;
56
+ idx_t* labels,
57
+ const SearchParameters* params = nullptr) const override;
57
58
 
58
59
  void reconstruct(idx_t key, float* recons) const override;
59
60
 
@@ -78,9 +78,10 @@ void IndexNSG::search(
78
78
  const float* x,
79
79
  idx_t k,
80
80
  float* distances,
81
- idx_t* labels) const
82
-
83
- {
81
+ idx_t* labels,
82
+ const SearchParameters* params) const {
83
+ FAISS_THROW_IF_NOT_MSG(
84
+ !params, "search params not supported for this index");
84
85
  FAISS_THROW_IF_NOT_MSG(
85
86
  storage,
86
87
  "Please use IndexNSGFlat (or variants) instead of IndexNSG directly");
@@ -298,4 +299,37 @@ IndexNSGFlat::IndexNSGFlat(int d, int R, MetricType metric)
298
299
  is_trained = true;
299
300
  }
300
301
 
302
+ /**************************************************************
303
+ * IndexNSGPQ implementation
304
+ **************************************************************/
305
+
306
+ IndexNSGPQ::IndexNSGPQ() {}
307
+
308
+ IndexNSGPQ::IndexNSGPQ(int d, int pq_m, int M)
309
+ : IndexNSG(new IndexPQ(d, pq_m, 8), M) {
310
+ own_fields = true;
311
+ is_trained = false;
312
+ }
313
+
314
+ void IndexNSGPQ::train(idx_t n, const float* x) {
315
+ IndexNSG::train(n, x);
316
+ (dynamic_cast<IndexPQ*>(storage))->pq.compute_sdc_table();
317
+ }
318
+
319
+ /**************************************************************
320
+ * IndexNSGSQ implementation
321
+ **************************************************************/
322
+
323
+ IndexNSGSQ::IndexNSGSQ(
324
+ int d,
325
+ ScalarQuantizer::QuantizerType qtype,
326
+ int M,
327
+ MetricType metric)
328
+ : IndexNSG(new IndexScalarQuantizer(d, qtype, metric), M) {
329
+ is_trained = false;
330
+ own_fields = true;
331
+ }
332
+
333
+ IndexNSGSQ::IndexNSGSQ() {}
334
+
301
335
  } // namespace faiss
@@ -13,6 +13,8 @@
13
13
 
14
14
  #include <faiss/IndexFlat.h>
15
15
  #include <faiss/IndexNNDescent.h>
16
+ #include <faiss/IndexPQ.h>
17
+ #include <faiss/IndexScalarQuantizer.h>
16
18
  #include <faiss/impl/NSG.h>
17
19
  #include <faiss/utils/utils.h>
18
20
 
@@ -64,7 +66,8 @@ struct IndexNSG : Index {
64
66
  const float* x,
65
67
  idx_t k,
66
68
  float* distances,
67
- idx_t* labels) const override;
69
+ idx_t* labels,
70
+ const SearchParameters* params = nullptr) const override;
68
71
 
69
72
  void reconstruct(idx_t key, float* recons) const override;
70
73
 
@@ -82,4 +85,25 @@ struct IndexNSGFlat : IndexNSG {
82
85
  IndexNSGFlat(int d, int R, MetricType metric = METRIC_L2);
83
86
  };
84
87
 
88
+ /** PQ index topped with with a NSG structure to access elements
89
+ * more efficiently.
90
+ */
91
+ struct IndexNSGPQ : IndexNSG {
92
+ IndexNSGPQ();
93
+ IndexNSGPQ(int d, int pq_m, int M);
94
+ void train(idx_t n, const float* x) override;
95
+ };
96
+
97
+ /** SQ index topped with with a NSG structure to access elements
98
+ * more efficiently.
99
+ */
100
+ struct IndexNSGSQ : IndexNSG {
101
+ IndexNSGSQ();
102
+ IndexNSGSQ(
103
+ int d,
104
+ ScalarQuantizer::QuantizerType qtype,
105
+ int M,
106
+ MetricType metric = METRIC_L2);
107
+ };
108
+
85
109
  } // namespace faiss