Implemented AnnoyIndex serialization to bytes objects in-memory

.gitignore

@@ -1,5 +1,6 @@
 *.egg-info/
 *.egg/
+*.eggs/
 *.so
 *.o
 build/

annoy/__init__.pyi

@@ -10,6 +10,8 @@ class AnnoyIndex:
     def __init__(self, f: int, metric: Literal["angular", "euclidean", "manhattan", "hamming", "dot"]) -> None: ...
     def load(self, fn: str, prefault: bool = ...) -> Literal[True]: ...
     def save(self, fn: str, prefault: bool = ...) -> Literal[True]: ...
+    def serialize(self) -> bytes: ...
+    def deserialize(self, data: bytes, prefault: bool = ...) -> Literal[True]: ...
     @overload
     def get_nns_by_item(self, i: int, n: int, search_k: int = ..., include_distances: Literal[False] = ...) -> list[int]: ...
     @overload

src/annoygomodule.h

@@ -36,6 +36,15 @@ class AnnoyIndex {
   bool load(const char* filename) {
     return ptr->load(filename, true);
   };
+  vector<uint8_t> serialize() {
+      return ptr->serialize();
+  }
+  bool deserialize(vector<uint8_t>* v, bool prefault) {
+      return ptr->deserialize(v, prefault);
+  }
+  bool deserialize(vector<uint8_t>* v) {
+      return ptr->deserialize(v, true);
+  }
   float getDistance(int i, int j) {
     return ptr->get_distance(i, j);
   };

src/annoylib.h

@@ -912,6 +912,8 @@ class AnnoyIndexInterface {
   virtual bool save(const char* filename, bool prefault=false, char** error=NULL) = 0;
   virtual void unload() = 0;
   virtual bool load(const char* filename, bool prefault=false, char** error=NULL) = 0;
+  virtual vector<uint8_t> serialize(char** error=NULL) const = 0;
+  virtual bool deserialize(vector<uint8_t>* bytes, bool prefault=false, char** error=NULL) = 0;
   virtual T get_distance(S i, S j) const = 0;
   virtual void get_nns_by_item(S item, size_t n, int search_k, vector<S>* result, vector<T>* distances) const = 0;
   virtual void get_nns_by_vector(const T* w, size_t n, int search_k, vector<S>* result, vector<T>* distances) const = 0;
@@ -1221,6 +1223,79 @@ template<typename S, typename T, typename Distance, typename Random, class Threa
     return true;
   }
 
+  vector<uint8_t> serialize(char** error=NULL) const {
+    if (!_built) {
+      set_error_from_string(error, "Index cannot be serialized if it hasn't been built");
+      return {};
+    }
+
+    vector<uint8_t> bytes {};
+
+    S n_items = _n_items;
+    S n_nodes = _n_nodes;
+    size_t roots_size = _roots.size();
+    S nodes_size = _nodes_size;
+
+    bytes.insert(bytes.end(), (uint8_t*)&n_items, (uint8_t*)&n_items + sizeof(n_items));
+    bytes.insert(bytes.end(), (uint8_t*)&n_nodes, (uint8_t*)&n_nodes + sizeof(n_nodes));
+    bytes.insert(bytes.end(), (uint8_t*)&roots_size, (uint8_t*)&roots_size + sizeof(roots_size));
+    bytes.insert(bytes.end(), (uint8_t*)&nodes_size, (uint8_t*)&nodes_size + sizeof(nodes_size));
+
+    uint8_t* roots_buffer = (uint8_t*)_roots.data();
+    bytes.insert(bytes.end(), roots_buffer, roots_buffer + _roots.size() * sizeof(S));
+
+    uint8_t* nodes_buffer = (uint8_t*)_nodes;
+    bytes.insert(bytes.end(), nodes_buffer, nodes_buffer + _nodes_size * _s);
+
+    return bytes;
+  }
+
+  bool deserialize(vector<uint8_t>* bytes, bool prefault=false, char** error=NULL) {
+    if (bytes->empty()) {
+      set_error_from_errno(error, "Size of bytes is zero");
+      return false;
+    }
+
+    int flags = MAP_SHARED;
+    if (prefault) {
+#ifdef MAP_POPULATE
+      flags |= MAP_POPULATE;
+#else
+      annoylib_showUpdate("prefault is set to true, but MAP_POPULATE is not defined on this platform");
+#endif
+    }
+
+    uint8_t* bytes_buffer = (uint8_t*)bytes->data();
+
+    _n_items = *(S*)bytes_buffer;
+    bytes_buffer += sizeof(S);
+
+    S n_nodes = *(S*)bytes_buffer;
+    bytes_buffer += sizeof(S);
+
+    size_t roots_size = *(size_t*)bytes_buffer;
+    bytes_buffer += sizeof(size_t);
+
+    S nodes_size = *(S*)bytes_buffer;
+    bytes_buffer += sizeof(S);
+
+    _roots.clear();
+    _roots.resize(roots_size);
+    _roots.assign((S*) bytes_buffer, (S*) bytes_buffer + roots_size);
+    bytes_buffer += roots_size * sizeof(S);
+
+    _allocate_size((S) nodes_size);
+
+    memcpy(_nodes, bytes_buffer, nodes_size * _s);
+
+    _n_nodes = n_nodes;
+    _loaded = true;
+    _built = true;
+
+    if (_verbose) annoylib_showUpdate("found %zu roots with degree %d\n", _roots.size(), _n_items);
+    return true;
+  }
+
   T get_distance(S i, S j) const {
     return D::normalized_distance(D::distance(_get(i), _get(j), _f));
   }

src/annoyluamodule.cc

@@ -164,6 +164,27 @@ class LuaAnnoy {
     return 1;
   }
 
+  static int serialize(lua_State* L) {
+    Impl* self = getAnnoy(L, 1);
+    int nargs = lua_gettop(L);
+    vector<uint8_t> bytes = self->serialize();
+
+    lua_pushlstring(L, (const char*) bytes.data(), bytes.size());
+
+    return 1;
+  }
+
+  static int deserialize(lua_State* L) {
+    Impl* self = getAnnoy(L, 1);
+    int nargs = lua_gettop(L);
+    const char* bytes_buffer = lua_tostring(L, 2);
+    size_t bytes_buffer_size = lua_rawlen(L, 2);
+    vector<uint8_t> bytes(bytes_buffer, bytes_buffer + bytes_buffer_size);
+    self->deserialize(&bytes);
+
+    return 1;
+  }
+
   static int unload(lua_State* L) {
     Impl* self = getAnnoy(L, 1);
     self->unload();
@@ -260,6 +281,8 @@ class LuaAnnoy {
       {"build", &ThisClass::build},
       {"save", &ThisClass::save},
       {"load", &ThisClass::load},
+      {"serialize", &ThisClass::serialize},
+      {"deserialize", &ThisClass::deserialize},
       {"unload", &ThisClass::unload},
       {"get_nns_by_item", &ThisClass::get_nns_by_item},
       {"get_nns_by_vector", &ThisClass::get_nns_by_vector},

src/annoymodule.cc

@@ -16,6 +16,7 @@
 #include "kissrandom.h"
 #include "Python.h"
 #include "structmember.h"
+#include "bytesobject.h"
 #include <exception>
 #if defined(_MSC_VER) && _MSC_VER == 1500
 typedef signed __int32    int32_t;
@@ -96,6 +97,8 @@ class HammingWrapper : public AnnoyIndexInterface<int32_t, float> {
   bool save(const char* filename, bool prefault, char** error) { return _index.save(filename, prefault, error); };
   void unload() { _index.unload(); };
   bool load(const char* filename, bool prefault, char** error) { return _index.load(filename, prefault, error); };
+  vector<uint8_t> serialize(char** error) const { return _index.serialize(error); };
+  bool deserialize(vector<uint8_t>* bytes, bool prefault, char** error) { return _index.deserialize(bytes, prefault, error); };
   float get_distance(int32_t i, int32_t j) const { return _index.get_distance(i, j); };
   void get_nns_by_item(int32_t item, size_t n, int search_k, vector<int32_t>* result, vector<float>* distances) const {
     if (distances) {
@@ -235,6 +238,51 @@ py_an_save(py_annoy *self, PyObject *args, PyObject *kwargs) {
   Py_RETURN_TRUE;
 }
 
+static PyObject *
+py_an_serialize(py_annoy *self, PyObject *args, PyObject *kwargs) {
+  if (!self->ptr) 
+    return NULL;
+
+  vector<uint8_t> bytes = self->ptr->serialize(NULL);
+
+  return PyBytes_FromStringAndSize((const char*)bytes.data(), bytes.size());
+}
+
+static PyObject *
+py_an_deserialize(py_annoy *self, PyObject *args, PyObject *kwargs) {
+  PyObject* bytes_object;
+  char *error;
+  bool prefault = false;
+  if (!self->ptr) 
+    return NULL;
+
+  static char const * kwlist[] = {"bytes", "prefault", NULL};
+
+  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "S|b", (char**)kwlist, &bytes_object, &prefault))
+    return NULL;
+
+  if (bytes_object == NULL) {
+    PyErr_SetString(PyExc_TypeError, "Expected bytes");
+    return NULL;
+  }
+
+  if (!PyBytes_Check(bytes_object)) {
+    PyErr_SetString(PyExc_TypeError, "Expected bytes");
+    return NULL;
+  }
+
+  Py_ssize_t length = PyBytes_Size(bytes_object);
+  uint8_t* raw_bytes = (uint8_t*)PyBytes_AsString(bytes_object);
+  vector<uint8_t> v(raw_bytes, raw_bytes + length);
+
+  if (!self->ptr->deserialize(&v, prefault, &error)) {
+    PyErr_SetString(PyExc_IOError, error);
+    free(error);
+    return NULL;
+  }
+
+  Py_RETURN_TRUE;
+}
 
 PyObject*
 get_nns_to_python(const vector<int32_t>& result, const vector<float>& distances, int include_distances) {
@@ -575,6 +623,8 @@ py_an_set_seed(py_annoy *self, PyObject *args) {
 static PyMethodDef AnnoyMethods[] = {
   {"load",	(PyCFunction)py_an_load, METH_VARARGS | METH_KEYWORDS, "Loads (mmaps) an index from disk."},
   {"save",	(PyCFunction)py_an_save, METH_VARARGS | METH_KEYWORDS, "Saves the index to disk."},
+  {"serialize",  (PyCFunction)py_an_serialize, METH_VARARGS | METH_KEYWORDS, "Serializes the index to bytes."},
+  {"deserialize", (PyCFunction)py_an_deserialize, METH_VARARGS | METH_KEYWORDS, "Deserializes the index from bytes."},
   {"get_nns_by_item",(PyCFunction)py_an_get_nns_by_item, METH_VARARGS | METH_KEYWORDS, "Returns the `n` closest items to item `i`.\n\n:param search_k: the query will inspect up to `search_k` nodes.\n`search_k` gives you a run-time tradeoff between better accuracy and speed.\n`search_k` defaults to `n_trees * n` if not provided.\n\n:param include_distances: If `True`, this function will return a\n2 element tuple of lists. The first list contains the `n` closest items.\nThe second list contains the corresponding distances."},
   {"get_nns_by_vector",(PyCFunction)py_an_get_nns_by_vector, METH_VARARGS | METH_KEYWORDS, "Returns the `n` closest items to vector `vector`.\n\n:param search_k: the query will inspect up to `search_k` nodes.\n`search_k` gives you a run-time tradeoff between better accuracy and speed.\n`search_k` defaults to `n_trees * n` if not provided.\n\n:param include_distances: If `True`, this function will return a\n2 element tuple of lists. The first list contains the `n` closest items.\nThe second list contains the corresponding distances."},
   {"get_item_vector",(PyCFunction)py_an_get_item_vector, METH_VARARGS, "Returns the vector for item `i` that was previously added."},

test/annoy_test.go

@@ -104,6 +104,51 @@ func (suite *AnnoyTestSuite) TestFileHandling() {
      os.Remove("go_test3.ann")
 }
 
+func (suite *AnnoyTestSuite) TestSerialization() {
+     index := annoyindex.NewAnnoyIndexAngular(3)
+     index.AddItem(0, []float32{0, 0, 1})
+     index.AddItem(1, []float32{0, 1, 0})
+     index.AddItem(2, []float32{1, 0, 0})
+     index.Build(10)
+
+     bytes := index.Serialize()
+
+     index2 := annoyindex.NewAnnoyIndexAngular(3)
+
+     success := index2.Deserialize(bytes)
+
+     if !success {
+         assert.Fail(suite.T(), "Failed to deserialize")
+     }
+
+     itemCountIsSame := index.GetNItems() == index2.GetNItems()
+
+     if !itemCountIsSame {
+         assert.Fail(suite.T(), "Item count is not the same")
+     }
+
+     var resultIndex []int
+     var resultIndex2 []int
+
+     itemCount := index.GetNItems()
+
+     index.GetNnsByItem(0, itemCount, -1, &resultIndex)
+     index2.GetNnsByItem(0, itemCount, -1, &resultIndex2)
+
+     itemsAreSame := true
+
+     for index := 0; index < itemCount; index++ {
+         if resultIndex[index] != resultIndex2[index] {
+             itemsAreSame = false
+             break
+         }
+     }
+
+     if !itemsAreSame {
+         assert.Fail(suite.T(), "Items are not the same")
+     }
+}
+
 func (suite *AnnoyTestSuite) TestOnDiskBuild() {
      index := annoyindex.NewAnnoyIndexAngular(3)
      index.OnDiskBuild("go_test.ann");

test/annoy_test.lua

@@ -496,6 +496,30 @@ describe("index test", function()
         assert.same(u, y)
     end)
 
+    it("serialize_deserialize", function()
+        local f = 2
+        local i = AnnoyIndex(f, 'euclidean')
+        i:add_item(0, {2, 2})
+        i:add_item(1, {3, 2})
+        i:add_item(2, {3, 3})
+        i:add_item(3, {4, 4})
+        i:add_item(4, {5, 5})
+        i:build(10)
+
+        local bytes = i:serialize()
+
+        local j = AnnoyIndex(f, 'euclidean')
+
+        j:deserialize(bytes)
+
+        local item_count = 4
+
+        local first_items = i:get_nns_by_item(0, item_count)
+        local second_items = j:get_nns_by_item(0, item_count)
+
+        assert.same(first_items, second_items)
+    end)
+
     it("on_disk_build", function()
         local f = 2
         local i = AnnoyIndex(f, 'euclidean')

test/serialize_test.py

@@ -0,0 +1,63 @@
+import random
+
+from annoy import AnnoyIndex
+
+def test_serialize_index():
+    f = 32
+
+    index = AnnoyIndex(f, 'angular')
+
+    for iteration in range(1000):
+        vector = [random.gauss(0, 1) for _ in range(f)]
+        index.add_item(iteration, vector)
+
+    index.build(10)
+
+    _ = index.serialize()
+
+
+def test_deserialize_index():
+    f = 32
+
+    index = AnnoyIndex(f, 'angular')
+
+    for iteration in range(1000):
+        vector = [random.gauss(0, 1) for _ in range(f)]
+        index.add_item(iteration, vector)
+
+    index.build(10)
+
+    data = index.serialize()
+
+    index2 = AnnoyIndex(f, 'angular')
+
+    index2.deserialize(data)
+
+    index_item_count = index.get_n_items()
+
+    assert index_item_count == index2.get_n_items()
+    assert index.get_n_trees() == index2.get_n_trees()
+    assert index.get_nns_by_item(0, index_item_count) == index2.get_nns_by_item(0, index_item_count)
+
+def test_serialize_after_load():
+    f = 32
+
+    index1 = AnnoyIndex(f, 'angular')
+
+    for iteration in range(1000):
+        vector = [random.gauss(0, 1) for _ in range(f)]
+        index1.add_item(iteration, vector)
+
+    index1.build(10)
+
+    save_path = "test/test.tree"
+    index1.save(save_path)
+
+    index2 = AnnoyIndex(f, 'angular')
+    index2.load(save_path)
+
+    assert index1.serialize() == index2.serialize()
+    assert index1.get_n_items() == index2.get_n_items()
+    assert index1.get_n_trees() == index2.get_n_trees()
+    assert index1.get_nns_by_item(0, index1.get_n_items()) == index2.get_nns_by_item(0, index1.get_n_items())
+