Repository: jina-ai/serve
Pull request: fix the bug in query with faiss
resolves #530

Problem statement:
the training function is missing in the Faiss Indexer
**Describe the bug**
the index object initialized by `faiss.index_factory()` should be trained before loading the indexed data. 

**Describe how you solve it**
- implement the `train()` function for `FaissIndexer`

---

<!-- Optional, but really help us locate the problem faster -->

**Environment**
<!-- Run `jina --version-full` and copy paste the output here -->

**Screenshots**
<!-- If applicable, add screenshots to help explain your problem. -->

Hints from issue discussion:
note that this is not unnecessarily true. as not all indexer in faiss need to be pretrained.

Commit message: fix: clean up imports

Patch to review:
diff --git a/jina/executors/indexers/vector/faiss.py b/jina/executors/indexers/vector/faiss.py
index 0aeabcd95fea7..e7b64ac7f8a65 100644
--- a/jina/executors/indexers/vector/faiss.py
+++ b/jina/executors/indexers/vector/faiss.py
@@ -12,35 +12,73 @@ class FaissIndexer(NumpyIndexer):
     """Faiss powered vector indexer
 
     For more information about the Faiss supported parameters and installation problems, please consult:
-        - https://github.com/spotify/annoy
         - https://github.com/facebookresearch/faiss
 
     .. note::
         Faiss package dependency is only required at the query time.
     """
 
-    def __init__(self, index_key: str, *args, **kwargs):
+    def __init__(self, index_key: str, train_filepath: str = None, *args, **kwargs):
         """
         Initialize an Faiss Indexer
 
         :param index_key: index type supported by ``faiss.index_factory``
+        :param train_filepath: the training data file path, e.g ``faiss.tgz``. The data file is expected to be a gzip
+            file in which `numpy.ndarray` is streamed as binary bytes.
+
+
+        .. highlight:: python
+        .. code-block:: python
+            # generate a training file
+            import gzip
+            import numpy as np
+            train_filepath = 'faiss_train.tgz'
+            train_data = np.random.rand(10000, 128)
+            with gzip.open(train_filepath, 'wb', compresslevel=1) as f:
+                f.write(train_data.astype('float32'))
+
+            from jina.executors.indexers.vector.faiss import FaissIndexer
+            indexer = FaissIndexer('PCA64,FLAT', train_filepath)
         """
         super().__init__(*args, **kwargs)
         self.index_key = index_key
+        self.train_filepath = train_filepath
 
     def get_query_handler(self):
         """Load all vectors (in numpy ndarray) into Faiss indexers """
         import faiss
-        data = super().get_query_handler()
-        _index = faiss.index_factory(self.num_dim, self.index_key)
-        _index.add(data)
-        return _index
+        _index_data = super().get_query_handler()
+        if _index_data is None:
+            self.logger.warning('loading indexing data failed.')
+            return None
+        if _index_data.ndim != 2:
+            self.logger.warning('the index data should be 2D tensor, {} != 2'.format(_index_data.ndim))
+            return None
+        self._index = faiss.index_factory(self.num_dim, self.index_key)
+        if not self.is_trained:
+            _train_data = self._load_training_data(self.train_filepath)
+            if _train_data is None:
+                self.logger.warning('loading training data failed.')
+                return None
+            self.train(_train_data)
+        self._index.add(_index_data.astype('float32'))
+        return self._index
 
     def query(self, keys: 'np.ndarray', top_k: int, *args, **kwargs) -> Tuple['np.ndarray', 'np.ndarray']:
         if keys.dtype != np.float32:
             raise ValueError('vectors should be ndarray of float32')
-
         dist, ids = self.query_handler.search(keys, top_k)
-
-        # ids is already a numpy array
         return self.int2ext_key[ids], dist
+
+    def train(self, data: 'np.ndarray', *args, **kwargs):
+        _num_samples, _num_dim = data.shape
+        if not self.num_dim:
+            self.num_dim = _num_dim
+        if self.num_dim != _num_dim:
+            raise ValueError('training data should have the same number of features as the index, {} != {}'.format(
+                self.num_dim, _num_dim))
+        self._index.train(data)
+        self.is_trained = True
+
+    def _load_training_data(self, train_filepath):
+        return self._load_numpy(train_filepath)
diff --git a/jina/executors/indexers/vector/numpy.py b/jina/executors/indexers/vector/numpy.py
index 2e6ce18015513..2b410e66e7b1a 100644
--- a/jina/executors/indexers/vector/numpy.py
+++ b/jina/executors/indexers/vector/numpy.py
@@ -3,6 +3,7 @@
 
 import gzip
 from typing import Tuple, Optional
+from os import path
 
 import numpy as np
 from jina.executors.indexers import BaseVectorIndexer
@@ -53,14 +54,9 @@ def get_query_handler(self) -> Optional['np.ndarray']:
 
         :return: a numpy ndarray of vectors
         """
-        try:
-            if self.num_dim and self.dtype:
-                with gzip.open(self.index_abspath, 'rb') as fp:
-                    vecs = np.frombuffer(fp.read(), dtype=self.dtype).reshape([-1, self.num_dim])
-        except EOFError:
-            self.logger.error(
-                f'{self.index_abspath} is broken/incomplete, perhaps forgot to ".close()" in the last usage?')
-            return None
+        vecs = self._load_numpy(self.index_abspath)
+        if vecs is None:
+            return vecs
 
         if self.key_bytes and self.key_dtype:
             self.int2ext_key = np.frombuffer(self.key_bytes, dtype=self.key_dtype)
@@ -138,6 +134,20 @@ def query(self, keys: np.ndarray, top_k: int, *args, **kwargs) -> Tuple['np.ndar
         dist = np.take_along_axis(dist, idx, axis=1)
         return self.int2ext_key[idx], dist
 
+    def _load_numpy(self, abspath):
+        if not path.exists(abspath):
+            self.logger.warning('numpy data not found: {}'.format(abspath))
+            return None
+        result = None
+        try:
+            if self.num_dim and self.dtype:
+                with gzip.open(abspath, 'rb') as fp:
+                    result = np.frombuffer(fp.read(), dtype=self.dtype).reshape([-1, self.num_dim])
+        except EOFError:
+            self.logger.error(
+                f'{self.index_abspath} is broken/incomplete, perhaps forgot to ".close()" in the last usage?')
+        return result
+
 
 def _ext_arrs(A, B):
     nA, dim = A.shape


Code review comments:
diff_hunk: @@ -12,35 +12,73 @@ class FaissIndexer(NumpyIndexer):
     """Faiss powered vector indexer
 
     For more information about the Faiss supported parameters and installation problems, please consult:
-        - https://github.com/spotify/annoy
         - https://github.com/facebookresearch/faiss
 
     .. note::
         Faiss package dependency is only required at the query time.
     """
 
-    def __init__(self, index_key: str, *args, **kwargs):
+    def __init__(self, index_key: str, train_filepath: str = None, *args, **kwargs):
         """
         Initialize an Faiss Indexer
 
         :param index_key: index type supported by ``faiss.index_factory``
+        :param train_filepath: the training data file path, e.g ``faiss.tgz``. The data file is expected to be a gzip
+            file in which `numpy.ndarray` is streamed as binary bytes.
+
+
+        .. highlight:: python
+        .. code-block:: python
+            # generate a training file
+            import gzip
+            import numpy as np
+            train_filepath = 'faiss_train.tgz'
+            train_data = np.random.rand(10000, 128)
+            with gzip.open(train_filepath, 'wb', compresslevel=1) as f:
+                f.write(train_data.astype('float32'))
+
+            from jina.executors.indexers.vector.faiss import FaissIndexer
+            indexer = FaissIndexer('PCA64,FLAT', train_filepath)
         """
         super().__init__(*args, **kwargs)
         self.index_key = index_key
+        self.train_filepath = train_filepath
 
     def get_query_handler(self):
         """Load all vectors (in numpy ndarray) into Faiss indexers """
         import faiss
-        data = super().get_query_handler()
-        _index = faiss.index_factory(self.num_dim, self.index_key)
-        _index.add(data)
-        return _index
+        _index_data = super().get_query_handler()
+        if _index_data is None:
+            self.logger.warning('loading indexing data failed.')
+            return None
+        if _index_data.ndim != 2:
+            self.logger.warning('the index data should be 2D tensor, {} != 2'.format(_index_data.ndim))
+            return None
+        self._index = faiss.index_factory(self.num_dim, self.index_key)
+        if not self.is_trained:
+            _train_data = self._load_training_data(self.train_filepath)
+            if _train_data is None:
+                self.logger.warning('loading training data failed.')
+                return None
+            self.train(_train_data)
+        self._index.add(_index_data.astype('float32'))
+        return self._index
 
     def query(self, keys: 'np.ndarray', top_k: int, *args, **kwargs) -> Tuple['np.ndarray', 'np.ndarray']:
         if keys.dtype != np.float32:
             raise ValueError('vectors should be ndarray of float32')
-
         dist, ids = self.query_handler.search(keys, top_k)
-
-        # ids is already a numpy array
         return self.int2ext_key[ids], dist
+
+    def train(self, data: 'np.ndarray', *args, **kwargs):
+        _num_samples, _num_dim = data.shape
+        if not self.num_dim:
+            self.num_dim = _num_dim
+        if self.num_dim != _num_dim:
+            raise ValueError('training data should have the same number of features as the index, {} != {}'.format(
+                self.num_dim, _num_dim))
+        self._index.train(data)
+        self.is_trained = True
+
+    def _load_training_data(self, train_filepath):
+        return self._load_numpy(train_filepath)
path: jina/executors/indexers/vector/faiss.py
text: @user1:
is this necessary? aren't they exactly the same?

@author:
Check out the comments below.