From 60cf2e24fd72ebe8558d65c8bba093f802aeb05f Mon Sep 17 00:00:00 2001
From: Aaron Silvas <aaronsilvas@gmail.com>
Date: Sun, 24 Sep 2023 15:31:45 -0700
Subject: [PATCH] Labels are 64bit but JS bindings were limited to 32bit

---
 lib/index.d.ts           |  6 +++---
 src/faiss.cc             |  2 +-
 test/IndexFlatL2.test.js | 24 ++++++++++++------------
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/lib/index.d.ts b/lib/index.d.ts
index 0bca395..1d4103c 100644
--- a/lib/index.d.ts
+++ b/lib/index.d.ts
@@ -3,7 +3,7 @@ export interface SearchResult {
     /** The disances of the nearest negihbors found, size n*k. */
     distances: number[],
     /** The labels of the nearest neighbors found, size n*k. */
-    labels: number[]
+    labels: BigInt[]
 }
 
 // See faiss/MetricType.h
@@ -103,10 +103,10 @@ export class Index {
     mergeFrom(otherIndex: Index): void;
     /**
      * Remove IDs from the index.
-     * @param {number[]} ids IDs to read.
+     * @param {BigInt[]} ids IDs to read.
      * @return {number} number of IDs removed.
      */
-    removeIds(ids: number[]): number
+    removeIds(ids: BigInt[]): number
 
 }
 
diff --git a/src/faiss.cc b/src/faiss.cc
index f32b6b7..6fe8915 100644
--- a/src/faiss.cc
+++ b/src/faiss.cc
@@ -301,7 +301,7 @@ class IndexBase : public Napi::ObjectWrap<T>
       idx_t label = I[i];
       float distance = D[i];
       arr_distances[i] = Napi::Number::New(env, distance);
-      arr_labels[i] = Napi::Number::New(env, label);
+      arr_labels[i] = Napi::BigInt::New(env, label);
     }
     delete[] I;
     delete[] D;
diff --git a/test/IndexFlatL2.test.js b/test/IndexFlatL2.test.js
index fc70c53..5252eca 100644
--- a/test/IndexFlatL2.test.js
+++ b/test/IndexFlatL2.test.js
@@ -113,9 +113,9 @@ describe('IndexFlatL2', () => {
         });
 
         it('returns search results', () => {
-            expect(index.search([1, 0], 1)).toMatchObject({ distances: [0], labels: [0] });
-            expect(index.search([1, 0], 4)).toMatchObject({ distances: [0, 1, 4, 9], labels: [0, 3, 1, 2] });
-            expect(index.search([1, 1], 4)).toMatchObject({ distances: [0, 1, 1, 4], labels: [3, 0, 1, 2] });
+            expect(index.search([1, 0], 1)).toMatchObject({ distances: [0], labels: [0n] });
+            expect(index.search([1, 0], 4)).toMatchObject({ distances: [0, 1, 4, 9], labels: [0n, 3n, 1n, 2n] });
+            expect(index.search([1, 1], 4)).toMatchObject({ distances: [0, 1, 1, 4], labels: [3n, 0n, 1n, 2n] });
         });
     });
 
@@ -156,15 +156,15 @@ describe('IndexFlatL2', () => {
         it("returns search results on merged index", () => {
             expect(index2.search([1, 0], 1)).toMatchObject({
                 distances: [0],
-                labels: [0],
+                labels: [0n],
             });
             expect(index2.search([1, 0], 4)).toMatchObject({
                 distances: [0, 1, 4, 9],
-                labels: [0, 3, 1, 2],
+                labels: [0n, 3n, 1n, 2n],
             });
             expect(index2.search([1, 1], 4)).toMatchObject({
                 distances: [0, 1, 1, 4],
-                labels: [3, 0, 1, 2],
+                labels: [3n, 0n, 1n, 2n],
             });
         });
     });
@@ -201,21 +201,21 @@ describe('IndexFlatL2', () => {
         });
 
         it("correctly removed", () => {
-            expect(index.search([1, 1], 1)).toMatchObject({ distances: [0], labels: [1] });
+            expect(index.search([1, 1], 1)).toMatchObject({ distances: [0], labels: [1n] });
             expect(index.removeIds([0])).toBe(1);
-            expect(index.search([1, 1], 1)).toMatchObject({ distances: [0], labels: [0] });
+            expect(index.search([1, 1], 1)).toMatchObject({ distances: [0], labels: [0n] });
         });
 
         it("correctly removed multiple elements", () => {
-            expect(index.search([1, 3], 1)).toMatchObject({ distances: [0], labels: [3] });
+            expect(index.search([1, 3], 1)).toMatchObject({ distances: [0], labels: [3n] });
             expect(index.removeIds([0, 1])).toBe(2);
-            expect(index.search([1, 3], 1)).toMatchObject({ distances: [0], labels: [1] });
+            expect(index.search([1, 3], 1)).toMatchObject({ distances: [0], labels: [1n] });
         });
 
         it("correctly removed partal elements", () => {
-            expect(index.search([1, 3], 1)).toMatchObject({ distances: [0], labels: [3] });
+            expect(index.search([1, 3], 1)).toMatchObject({ distances: [0], labels: [3n] });
             expect(index.removeIds([0, 1, 2, 4, 5])).toBe(3);
-            expect(index.search([1, 3], 1)).toMatchObject({ distances: [0], labels: [0] });
+            expect(index.search([1, 3], 1)).toMatchObject({ distances: [0], labels: [0n] });
         });
     });
 });
\ No newline at end of file