Merge branch 'qihui/store_many_vectors' of git://github.com/xieqihui/NearPy into xieqihui-qihui/store_many_vectors

Author: pixelogik

nearpy/engine.py

@@ -96,6 +96,22 @@ def store_vector(self, v, data=None):
  self.storage.store_vector(lshash.hash_name, bucket_key,
  nv, data)
 
+ def store_many_vectors(self, vs, data=None):
+ """
+        Store a batch of vectors.
+        Hashes vector vs and stores them in all matching buckets in the storage.
+        The data argument must be either None or a list of JSON-serializable
+        object. It is stored with the vector and will be returned in search
+        results.
+        """
+ # We will store the normalized vector (used during retrieval)
+ nvs = [unitvec(i) for i in vs]
+ # Store vector in each bucket of all hashes
+ for lshash in self.lshashes:
+ bucket_keys = [lshash.hash_vector(i)[0] for i in vs]
+ self.storage.store_many_vectors(lshash.hash_name, bucket_keys,
+ nvs, data)
+
  def delete_vector(self, data, v=None):
  """
         Deletes vector v and his id (data) in all matching buckets in the storage.

nearpy/storage/storage.py

@@ -30,6 +30,13 @@ def store_vector(self, hash_name, bucket_key, v, data):
         """
  raise NotImplementedError
 
+ def store_many_vectors(self, hash_name, bucket_keys, vs, data):
+ """
+        Store a batch of vectors.
+        Stores vector and JSON-serializable data in bucket with specified key.
+        """
+ raise NotImplementedError
+
  def get_all_bucket_keys(self, hash_name):
  """
         Returns all bucket keys for the given hash as iterable of strings

nearpy/storage/storage_memory.py

@@ -28,7 +28,9 @@
 # THE SOFTWARE.
 
 from future.utils import viewkeys
+from future.builtins import zip
 from nearpy.storage.storage import Storage
+import itertools
 
 
 class MemoryStorage(Storage):
@@ -50,6 +52,16 @@ def store_vector(self, hash_name, bucket_key, v, data):
  self.buckets[hash_name][bucket_key] = []
  self.buckets[hash_name][bucket_key].append((v, data))
 
+ def store_many_vectors(self, hash_name, bucket_keys, vs, data):
+ """
+        Store a batch of vectors.
+        Stores vector and JSON-serializable data in bucket with specified key.
+        """
+ if data is None:
+ data = itertools.repeat(data)
+ for v, k, d in zip(vs, bucket_keys, data):
+ self.store_vector(hash_name, k, v, d)
+
  def get_all_bucket_keys(self, hash_name):
  return viewkeys(self.buckets[hash_name])
 

nearpy/storage/storage_redis.py

@@ -50,6 +50,24 @@ def store_vector(self, hash_name, bucket_key, v, data):
  """
         Stores vector and JSON-serializable data in bucket with specified key.
         """
+ self._add_vector(hash_name, bucket_key, v, data, self.redis_object)
+
+ def store_many_vectors(self, hash_name, bucket_keys, vs, data):
+ """
+        Store a batch of vectors in Redis.
+        Stores vector and JSON-serializable data in bucket with specified key.
+        """
+ with self.redis_object.pipeline() as pipeline:
+ if data is None:
+ data = [None] * len(vs)
+ for bucket_key, data, v in zip(bucket_keys, data, vs):
+ self._add_vector(hash_name, bucket_key, v, data, pipeline)
+ pipeline.execute()
+
+ def _add_vector(self, hash_name, bucket_key, v, data, redis_object):
+ '''
+        Store vector and JSON-serializable data in bucket with specified key.
+        '''
  redis_key = self._format_redis_key(hash_name, bucket_key)
 
  val_dict = {}

tests/storage_tests.py

@@ -27,6 +27,7 @@
 from mockredis import MockRedis as Redis
 
 from future.builtins import range
+from future.builtins import zip
 
 from nearpy.storage import MemoryStorage, RedisStorage
 
@@ -55,6 +56,24 @@ def check_store_vector(self, x):
  self.storage.clean_all_buckets()
  self.assertEqual(self.storage.get_bucket('testHash', bucket_key), [])
 
+ def check_store_many_vectors(self, xs):
+ num_vector = len(xs)
+ bucket_keys = list(map(str,
+ list(range(10000000,
+ 10000000 + num_vector))))
+ x_data = list(range(0, num_vector))
+ self.storage.store_many_vectors('testHash', bucket_keys, xs, x_data)
+ for bucket_key, x, data in zip(bucket_keys, xs, x_data):
+ bucket = self.storage.get_bucket('testHash', bucket_key)
+ self.assertEqual(len(bucket), 1)
+ y, y_data = bucket[0]
+ self.assertEqual(type(y), type(x))
+ self.assertEqual(y.shape, x.shape)
+ self.assertEqual(max(abs(y - x)), 0)
+ self.assertEqual(y_data, data)
+ self.storage.clean_all_buckets()
+ self.assertEqual(self.storage.get_bucket('testHash', bucket_key), [])
+
  def check_get_all_bucket_keys(self):
  x, x_data = numpy.ones(100), "data"
  hash_config = [
@@ -135,5 +154,10 @@ def test_store_zero(self):
  _, data = bucket[0]
  self.assertEqual(data, 0)
 
+ def test_store_many_vectors(self):
+ x = numpy.random.randn(100, 10)
+ self.check_store_many_vectors(x)
+
+
 if __name__ == '__main__':
  unittest.main()