From f5cb7fa2e74530207d43997cb80dace3d5786e47 Mon Sep 17 00:00:00 2001
From: "xj.lin" <xiaojun.lin@zilliz.com>
Date: Wed, 27 Mar 2019 11:11:45 +0800
Subject: [PATCH] add id

---
 pyengine/engine/ingestion/build_index.py      | 7 ++++---
 pyengine/engine/ingestion/tests/test_build.py | 6 ++++--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/pyengine/engine/ingestion/build_index.py b/pyengine/engine/ingestion/build_index.py
index 63c431d5..6ea9f43e 100644
--- a/pyengine/engine/ingestion/build_index.py
+++ b/pyengine/engine/ingestion/build_index.py
@@ -36,9 +36,10 @@ class DefaultIndex(Index):
         pass
 
     def build(self, d, vectors, vector_ids, DEVICE=INDEXDEVICES.CPU):
-        index = faiss.IndexFlatL2(d)  # trained
-        index.add(vectors)
-        return index
+        index = faiss.IndexFlatL2(d)
+        index2 = faiss.IndexIDMap(index)
+        index2.add_with_ids(vectors, vector_ids)
+        return index2
 
 
 class LowMemoryIndex(Index):
diff --git a/pyengine/engine/ingestion/tests/test_build.py b/pyengine/engine/ingestion/tests/test_build.py
index f4b92bcb..a4bb070b 100644
--- a/pyengine/engine/ingestion/tests/test_build.py
+++ b/pyengine/engine/ingestion/tests/test_build.py
@@ -16,14 +16,16 @@ class TestBuildIndex(unittest.TestCase):
         nb = 10000
         nq = 100
         _, xb, xq = get_dataset(d, nb, 500, nq)
+        ids = np.arange(xb.shape[0])
 
         # Expected result
         index = faiss.IndexFlatL2(d)
-        index.add(xb)
+        index2 = faiss.IndexIDMap(index)
+        index2.add_with_ids(xb, ids)
         Dref, Iref = index.search(xq, 5)
 
         builder = DefaultIndex()
-        index2 = builder.build(d, xb)
+        index2 = builder.build(d, xb, ids)
         Dnew, Inew = index2.search(xq, 5)
 
         assert np.all(Dnew == Dref) and np.all(Inew == Iref)
-- 
GitLab