Buffers, IP, Factories & more (#40)

* Buffers, IP, Factories & more

* scope & instance data

* export MetricType enum

* fix d.ts

* fix JSDoc

* fix type of sub class

* feedback changes

* fix factory metric

* remove handlescope

* inline constructor causes stability issues

* fix inline variables warning

* update JSDoc

* remove unreached code and its test case

* remove fromFactory from IndexFlatL2 and IndexFlatIP

* update docs and example

---------

Co-authored-by: ewfian <ewfian@vip.qq.com>

CMakeLists.txt

@@ -32,7 +32,7 @@ set(BUILD_TESTING OFF)
 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
 
-set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_CXX_STANDARD 17)
 set(FAISS_ENABLE_GPU OFF)
 set(FAISS_ENABLE_PYTHON OFF)
 

README.md

@@ -24,7 +24,7 @@ $ npm install faiss-node
 ## Usage
 
 ```javascript
-const { IndexFlatL2 } = require('faiss-node');
+const { IndexFlatL2, Index, IndexFlatIP, MetricType } = require('faiss-node');
 
 const dimension = 2;
 const index = new IndexFlatL2(dimension);
@@ -69,6 +69,21 @@ const removedCount = newIndex.removeIds([0]);
 console.log(removedCount); // 1
 console.log(newIndex.ntotal()); // 3
 console.log(newIndex.search([1, 2], 1)); // { distances: [ 0 ], labels: [ 0 ] }
+
+// IndexFlatIP
+const ipIndex = new IndexFlatIP(2);
+ipIndex.add([1, 0]);
+
+// Serialize an index
+const index_buf = newIndex.toBuffer();
+const deserializedIndex = Index.fromBuffer(index_buf);
+console.log(deserializedIndex.ntotal()); // 3
+
+// Factory index
+const hnswIndex = Index.fromFactory(2, 'HNSW,Flat', MetricType.METRIC_INNER_PRODUCT);
+const x = [1, 0, 0, 1];
+hnswIndex.train(x);
+hnswIndex.add(x);
 ```
 
 ## License

examples/index.js

@@ -1,38 +1,60 @@
-const { IndexFlatL2 } = require('../');
+const { IndexFlatL2, Index, IndexFlatIP, MetricType } = require('../');
 
 const dimension = 2;
 const index = new IndexFlatL2(dimension);
 
-console.log(index.getDimension());
-console.log(index.isTrained());
-console.log(index.ntotal());
+console.log(index.getDimension()); // 2
+console.log(index.isTrained()); // true
+console.log(index.ntotal()); // 0
+
+// inserting data into index.
 index.add([1, 0]);
 index.add([1, 2]);
 index.add([1, 3]);
 index.add([1, 1]);
-console.log(index.ntotal());
+
+console.log(index.ntotal()); // 4
 
 const k = 4;
 const results = index.search([1, 0], k);
-console.log(results.labels);
-console.log(results.distances);
+console.log(results.labels); // [ 0, 3, 1, 2 ]
+console.log(results.distances); // [ 0, 1, 4, 9 ]
 
+// Save index
 const fname = 'faiss.index';
 index.write(fname);
 
+// Load saved index
 const index_loaded = IndexFlatL2.read(fname);
-console.log(index_loaded.getDimension());
-console.log(index_loaded.ntotal());
+console.log(index_loaded.getDimension()); //2
+console.log(index_loaded.ntotal()); //4
 const results1 = index_loaded.search([1, 1], 4);
-console.log(results1.labels);
-console.log(results1.distances);
+console.log(results1.labels); // [ 3, 0, 1, 2 ]
+console.log(results1.distances); // [ 0, 1, 1, 4 ]
 
+// Merge index
 const newIndex = new IndexFlatL2(dimension);
 newIndex.mergeFrom(index);
-console.log(newIndex.ntotal());
+console.log(newIndex.ntotal()); // 4
 
-console.log(newIndex.search([1, 2], 1));
+// Remove items
+console.log(newIndex.search([1, 2], 1)); // { distances: [ 0 ], labels: [ 1 ] }
 const removedCount = newIndex.removeIds([0]);
-console.log(removedCount);
-console.log(newIndex.ntotal());
-console.log(newIndex.search([1, 2], 1));
+console.log(removedCount); // 1
+console.log(newIndex.ntotal()); // 3
+console.log(newIndex.search([1, 2], 1)); // { distances: [ 0 ], labels: [ 0 ] }
+
+// IndexFlatIP
+const ipIndex = new IndexFlatIP(2);
+ipIndex.add([1, 0]);
+
+// Serialize an index
+const index_buf = newIndex.toBuffer();
+const deserializedIndex = Index.fromBuffer(index_buf);
+console.log(deserializedIndex.ntotal()); // 3
+
+// Factory index
+const hnswIndex = Index.fromFactory(2, 'HNSW,Flat', MetricType.METRIC_INNER_PRODUCT);
+const x = [1, 0, 0, 1];
+hnswIndex.train(x);
+hnswIndex.add(x);

lib/index.d.ts

@@ -6,12 +6,29 @@ export interface SearchResult {
     labels: number[]
 }
 
+// See faiss/MetricType.h
+export enum MetricType {
+    METRIC_INNER_PRODUCT = 0, ///< maximum inner product search
+    METRIC_L2 = 1,            ///< squared L2 search
+    METRIC_L1,                ///< L1 (aka cityblock)
+    METRIC_Linf,              ///< infinity distance
+    METRIC_Lp,                ///< L_p distance, p is given by a faiss::Index
+    /// metric_arg
+
+    /// some additional metrics defined in scipy.spatial.distance
+    METRIC_Canberra = 20,
+    METRIC_BrayCurtis,
+    METRIC_JensenShannon,
+    METRIC_Jaccard, ///< defined as: sum_i(min(a_i, b_i)) / sum_i(max(a_i, b_i))
+    ///< where a_i, b_i > 0
+}
+
 /**
- * IndexFlatL2 Index.
+ * Index.
  * Index that stores the full vectors and performs exhaustive search.
  * @param {number} d The dimensionality of index.
  */
-export class IndexFlatL2 {
+export class Index {
     constructor(d: number);
     /**
      * returns the number of verctors currently indexed.
@@ -34,6 +51,12 @@ export class IndexFlatL2 {
      * @param {number[]} x Input matrix, size n * d
      */
     add(x: number[]): void;
+    /** 
+     * Train n vectors of dimension d to the index.
+     * Vectors are implicitly assigned labels ntotal .. ntotal + n - 1
+     * @param {number[]} x Input matrix, size n * d
+     */
+    train(x: number[]): void;
     /** 
      * Query n vectors of dimension d to the index.
      * return at most k vectors. If there are not enough results for a
@@ -48,22 +71,91 @@ export class IndexFlatL2 {
      * Write index to a file.
      * @param {string} fname File path to write.
      */
-    write(fname: string): void
+    write(fname: string): void;
+    /** 
+     * Write index to buffer.
+     */
+    toBuffer(): Buffer;
+    /** 
+     * Read index from a file.
+     * @param {string} fname File path to read.
+     * @return {Index} The index read.
+     */
+    static read(fname: string): Index;
+    /** 
+     * Read index from buffer.
+     * @param {Buffer} src Buffer to create index from.
+     * @return {Index} The index read.
+     */
+    static fromBuffer(src: Buffer): Index;
+    /** 
+     * Construct an index from factory descriptor.
+     * @param {number} dims Buffer to create index from.
+     * @param {string} descriptor Factory descriptor.
+     * @param {MetricType} metric Metric type (defaults to L2).
+     * @return {Index} The index read.
+     */
+    static fromFactory(dims: number, descriptor: string, metric?: MetricType): Index;
+    /**
+     * Merge the current index with another Index instance.
+     * @param {Index} otherIndex The other Index instance to merge from.
+     */
+    mergeFrom(otherIndex: Index): void;
+    /**
+     * Remove IDs from the index.
+     * @param {number[]} ids IDs to read.
+     * @return {number} number of IDs removed.
+     */
+    removeIds(ids: number[]): number
+
+}
+
+/**
+ * IndexFlatL2 Index.
+ * IndexFlatL2 that stores the full vectors and performs `squared L2` search.
+ * @param {number} d The dimensionality of index.
+ */
+export class IndexFlatL2 extends Index {
     /** 
      * Read index from a file.
      * @param {string} fname File path to read.
      * @return {IndexFlatL2} The index read.
      */
     static read(fname: string): IndexFlatL2;
+    /** 
+     * Read index from buffer.
+     * @param {Buffer} src Buffer to create index from.
+     * @return {IndexFlatL2} The index read.
+     */
+    static fromBuffer(src: Buffer): IndexFlatL2;
     /**
      * Merge the current index with another IndexFlatL2 instance.
      * @param {IndexFlatL2} otherIndex The other IndexFlatL2 instance to merge from.
      */
     mergeFrom(otherIndex: IndexFlatL2): void;
+}
+
+/**
+ * IndexFlatIP Index.
+ * Index that stores the full vectors and performs `maximum inner product` search.
+ * @param {number} d The dimensionality of index.
+ */
+export class IndexFlatIP extends Index {
+    /** 
+     * Read index from a file.
+     * @param {string} fname File path to read.
+     * @return {IndexFlatIP} The index read.
+     */
+    static read(fname: string): IndexFlatIP;
+    /** 
+     * Read index from buffer.
+     * @param {Buffer} src Buffer to create index from.
+     * @return {IndexFlatIP} The index read.
+     */
+    static fromBuffer(src: Buffer): IndexFlatIP;
     /**
-     * Remove IDs from the index.
-     * @param {string} ids IDs to read.
-     * @return {IndexFlatL2} number of IDs removed.
+     * Merge the current index with another IndexFlatIP instance.
+     * @param {IndexFlatIP} otherIndex The other IndexFlatIP instance to merge from.
      */
-    removeIds(ids: number[]): number
+    mergeFrom(otherIndex: IndexFlatIP): void;
 }

lib/index.js

@@ -1,2 +1,17 @@
 const faiss = require('bindings')('faiss-node');
+
+faiss.MetricType = void 0;
+var MetricType;
+(function (MetricType) {
+  MetricType[MetricType["METRIC_INNER_PRODUCT"] = 0] = "METRIC_INNER_PRODUCT";
+  MetricType[MetricType["METRIC_L2"] = 1] = "METRIC_L2";
+  MetricType[MetricType["METRIC_L1"] = 2] = "METRIC_L1";
+  MetricType[MetricType["METRIC_Linf"] = 3] = "METRIC_Linf";
+  MetricType[MetricType["METRIC_Lp"] = 4] = "METRIC_Lp";
+  MetricType[MetricType["METRIC_Canberra"] = 20] = "METRIC_Canberra";
+  MetricType[MetricType["METRIC_BrayCurtis"] = 21] = "METRIC_BrayCurtis";
+  MetricType[MetricType["METRIC_JensenShannon"] = 22] = "METRIC_JensenShannon";
+  MetricType[MetricType["METRIC_Jaccard"] = 23] = "METRIC_Jaccard";
+})(MetricType || (faiss.MetricType = MetricType = {}));
+
 module.exports = faiss;