Repository: jina-ai/annlite (language: python)

Issue:
fix: some changes



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 07a09bc4..3e165e37 100644
--- a/README.md
+++ b/README.md
@@ -29,63 +29,72 @@ $ git clone https://github.com/jina-ai/pqlite.git \
 ```python
 import random
 import numpy as np
+from jina import Document, DocumentArray
 from pqlite import PQLite
 
 N = 10000 # number of data points
-Nt = 2000
-Nq = 10
+Nq = 10 # number of query data
 D = 128 # dimentionality / number of features
 
-Xt = np.random.random((Nt, D)).astype(np.float32)  # 2,000 128-dim vectors for training
-
 # the column schema: (name:str, dtype:type, create_index: bool)
-pqlite = PQLite(d_vector=D, n_cells=64, n_subvectors=8, columns=[('x', float, True)])
-pqlite.fit(Xt)
+pqlite = PQLite(dim=D, columns=[('x', float, True)], data_path='./data')
 ```
 
 2. Add new data
 
 ```python
 X = np.random.random((N, D)).astype(np.float32)  # 10,000 128-dim vectors to be indexed
-
-tags = [{'x': random.random()} for _ in range(N)]
-pqlite.add(X, ids=list(range(len(X))), doc_tags=tags)
+docs = DocumentArray(
+    [
+        Document(id=f'{i}', embedding=X[i], tags={'x': random.random()})
+        for i in range(N)
+    ]
+)
+pqlite.index(docs)
 ```
 
 3. Search with Filtering
 
 ```python
-query = np.random.random((Nq, D)).astype(np.float32)  # a 128-dim query vector
+Xq = np.random.random((Nq, D)).astype(np.float32)  # a 128-dim query vector
+query = DocumentArray([Document(embedding=Xq[i]) for i in range(Nq)])
 
 # without filtering
-dists, ids = pqlite.search(query, k=5)
+pqlite.search(query, limit=10)
 
 print(f'the result without filtering:')
-for i, (dist, idx) in enumerate(zip(dists, ids)):
-    print(f'query [{i}]: {dist} {idx}')
+for i, q in enumerate(query):
+    print(f'query [{i}]:')
+    for m in q.matches:
+        print(f'\t{m.id} ({m.scores["euclidean"].value})')
 
 # with filtering
 # condition schema: (column_name: str, relation: str, value: any)
 conditions = [('x', '<', 0.3)]
-dists, ids = pqlite.search(query, conditions=conditions, k=5)
-
+pqlite.search(query, conditions=conditions, limit=10)
 print(f'the result with filtering:')
-for i, (dist, idx) in enumerate(zip(dists, ids)):
-    print(f'query [{i}]: {dist} {idx}')
+for i, q in enumerate(query):
+    print(f'query [{i}]:')
+    for m in q.matches:
+        print(f'\t{m.id} {m.scores["euclidean"].value} (x={m.tags["x"]})')
 ```
 4. Update data
 
 ```python
 Xn = np.random.random((10, D)).astype(np.float32)  # 10,000 128-dim vectors to be indexed
-
-tags = [{'x': random.random()} for _ in range(10)]
-pqlite.update(Xn, ids=list(range(len(Xn))), doc_tags=tags)
+docs = DocumentArray(
+    [
+        Document(id=f'{i}', embedding=Xn[i], tags={'x': random.random()})
+        for i in range(10)
+    ]
+)
+pqlite.update(docs)
 ```
 
 5. Delete data
 
 ```python
-pqlite.delete(ids=['1', '2'])
+pqlite.delete(['1', '2'])
 ```
 ## Benchmark
 
diff --git a/bindings/hnsw_bindings.cpp b/bindings/hnsw_bindings.cpp
index cc5379f0..da3594e2 100644
--- a/bindings/hnsw_bindings.cpp
+++ b/bindings/hnsw_bindings.cpp
@@ -590,8 +590,6 @@ class Index {
             else{
                 std::vector<float> norm_array(num_threads*features);
                 ParallelFor(0, rows, num_threads, [&](size_t row, size_t threadId) {
-                                float *data= (float *) items.data(row);
-
                                 size_t start_idx = threadId * dim;
                                 normalize_vector((float *) items.data(row), (norm_array.data()+start_idx));
 
@@ -634,6 +632,126 @@ class Index {
 
     }
 
+
+    py::object knnQuery_with_filter(py::object input, py::object candidate_ids_ = py::none(), size_t k = 1, int num_threads = -1) {
+
+        py::array_t < dist_t, py::array::c_style | py::array::forcecast > items(input);
+        auto buffer = items.request();
+        hnswlib::labeltype *data_numpy_l;
+        dist_t *data_numpy_d;
+
+        if (num_threads <= 0)
+            num_threads = num_threads_default;
+
+        if (buffer.ndim != 2 && buffer.ndim != 1) throw std::runtime_error("data must be a 1d/2d array");
+
+        size_t rows = 1;
+        size_t features = buffer.shape[0];
+
+        if (buffer.ndim == 2) {
+            rows = buffer.shape[0];
+            features = buffer.shape[1];
+        }
+
+        // avoid using threads when the number of searches is small:
+
+        if(rows<=num_threads*4){
+            num_threads=1;
+        }
+
+        // FuseFilter constructing
+        binary_fuse16_t filter(appr_alg->max_elements_);
+
+        if (!candidate_ids_.is_none()) {
+            py::array_t < size_t, py::array::c_style | py::array::forcecast > items(candidate_ids_);
+            auto ids_numpy = items.request();
+
+            if(ids_numpy.ndim==1) {
+                const size_t size = ids_numpy.shape[0];
+                std::vector<uint64_t> big_set;
+                big_set.reserve(size);
+                for (size_t i = 0; i < size; i++) {
+                    big_set[i] = items.data()[i]; // we use contiguous values
+                }
+
+                binary_fuse16_populate(big_set.data(), size, &filter);
+            }
+            else
+                throw std::runtime_error("wrong dimensionality of the filter labels");
+        }
+
+
+        {
+            py::gil_scoped_release l;
+
+            // would like to check the ownership of this data in more detail
+            data_numpy_l = new hnswlib::labeltype[rows * k];
+            data_numpy_d = new dist_t[rows * k];
+
+            if(normalize==false) {
+                ParallelFor(0, rows, num_threads, [&](size_t row, size_t threadId) {
+                                std::priority_queue<std::pair<dist_t, hnswlib::labeltype >> result = appr_alg->searchKnnWithFilter(
+                                        (void *) items.data(row), &filter, k);
+                                if (result.size() != k)
+                                    throw std::runtime_error(
+                                            "Cannot return the results in a contigious 2D array. Probably ef or M is too small");
+                                for (int i = k - 1; i >= 0; i--) {
+                                    auto &result_tuple = result.top();
+                                    data_numpy_d[row * k + i] = result_tuple.first;
+                                    data_numpy_l[row * k + i] = result_tuple.second;
+                                    result.pop();
+                                }
+                            }
+                );
+            }
+            else{
+                std::vector<float> norm_array(num_threads*features);
+                ParallelFor(0, rows, num_threads, [&](size_t row, size_t threadId) {
+                                float *data= (float *) items.data(row);
+
+                                size_t start_idx = threadId * dim;
+                                normalize_vector((float *) items.data(row), (norm_array.data()+start_idx));
+
+                                std::priority_queue<std::pair<dist_t, hnswlib::labeltype >> result = appr_alg->searchKnnWithFilter(
+                                        (void *) (norm_array.data()+start_idx), &filter, k);
+                                if (result.size() != k)
+                                    throw std::runtime_error(
+                                            "Cannot return the results in a contigious 2D array. Probably ef or M is too small");
+                                for (int i = k - 1; i >= 0; i--) {
+                                    auto &result_tuple = result.top();
+                                    data_numpy_d[row * k + i] = result_tuple.first;
+                                    data_numpy_l[row * k + i] = result_tuple.second;
+                                    result.pop();
+                                }
+                            }
+                );
+            }
+        }
+
+        py::capsule free_when_done_l(data_numpy_l, [](void *f) {
+            delete[] f;
+        });
+        py::capsule free_when_done_d(data_numpy_d, [](void *f) {
+            delete[] f;
+        });
+
+
+        return py::make_tuple(
+                py::array_t<hnswlib::labeltype>(
+                        {rows, k}, // shape
+                        {k * sizeof(hnswlib::labeltype),
+                         sizeof(hnswlib::labeltype)}, // C-style contiguous strides for double
+                        data_numpy_l, // the data pointer
+                        free_when_done_l),
+                py::array_t<dist_t>(
+                        {rows, k}, // shape
+                        {k * sizeof(dist_t), sizeof(dist_t)}, // C-style contiguous strides for double
+                        data_numpy_d, // the data pointer
+                        free_when_done_d));
+
+    }
+
+
     void markDeleted(size_t label) {
         appr_alg->markDelete(label);
     }
@@ -664,6 +782,7 @@ PYBIND11_PLUGIN(hnsw_bind) {
         .def(py::init<const std::string &, const int>(), py::arg("space"), py::arg("dim"))
         .def("init_index", &Index<float>::init_new_index, py::arg("max_elements"), py::arg("M")=16, py::arg("ef_construction")=200, py::arg("random_seed")=100)
         .def("knn_query", &Index<float>::knnQuery_return_numpy, py::arg("data"), py::arg("k")=1, py::arg("num_threads")=-1)
+        .def("knn_query_with_filter", &Index<float>::knnQuery_with_filter, py::arg("data"), py::arg("filters") = py::none(), py::arg("k")=1, py::arg("num_threads")=-1)
         .def("add_items", &Index<float>::addItems, py::arg("data"), py::arg("ids") = py::none(), py::arg("num_threads")=-1)
         .def("get_items", &Index<float, float>::getDataReturnList, py::arg("ids") = py::none())
         .def("get_ids_list", &Index<float>::getIdsList)
diff --git a/include/hnswlib/fusefilter.h b/include/hnswlib/fusefilter.h
new file mode 100644
index 00000000..634e769a
--- /dev/null
+++ b/include/hnswlib/fusefilter.h
@@ -0,0 +1,644 @@
+#ifndef BINARYFUSEFILTER_H
+#define BINARYFUSEFILTER_H
+#include <math.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#ifndef XOR_MAX_ITERATIONS
+#define XOR_MAX_ITERATIONS                                                     \
+  100 // probabillity of success should always be > 0.5 so 100 iterations is
+      // highly unlikely
+#endif
+
+/**
+ * We start with a few utilities.
+ ***/
+static inline uint64_t binary_fuse_murmur64(uint64_t h) {
+  h ^= h >> 33;
+  h *= UINT64_C(0xff51afd7ed558ccd);
+  h ^= h >> 33;
+  h *= UINT64_C(0xc4ceb9fe1a85ec53);
+  h ^= h >> 33;
+  return h;
+}
+static inline uint64_t binary_fuse_mix_split(uint64_t key, uint64_t seed) {
+  return binary_fuse_murmur64(key + seed);
+}
+static inline uint64_t binary_fuse_rotl64(uint64_t n, unsigned int c) {
+  return (n << (c & 63)) | (n >> ((-c) & 63));
+}
+static inline uint32_t binary_fuse_reduce(uint32_t hash, uint32_t n) {
+  // http://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
+  return (uint32_t)(((uint64_t)hash * n) >> 32);
+}
+static inline uint64_t binary_fuse8_fingerprint(uint64_t hash) {
+  return hash ^ (hash >> 32);
+}
+
+/**
+ * We need a decent random number generator.
+ **/
+
+// returns random number, modifies the seed
+static inline uint64_t binary_fuse_rng_splitmix64(uint64_t *seed) {
+  uint64_t z = (*seed += UINT64_C(0x9E3779B97F4A7C15));
+  z = (z ^ (z >> 30)) * UINT64_C(0xBF58476D1CE4E5B9);
+  z = (z ^ (z >> 27)) * UINT64_C(0x94D049BB133111EB);
+  return z ^ (z >> 31);
+}
+
+typedef struct binary_fuse8_s {
+  uint64_t Seed;
+  uint32_t SegmentLength;
+  uint32_t SegmentLengthMask;
+  uint32_t SegmentCount;
+  uint32_t SegmentCountLength;
+  uint32_t ArrayLength;
+  uint8_t *Fingerprints;
+} binary_fuse8_t;
+
+#ifdef _MSC_VER
+// Windows programmers who target 32-bit platform may need help:
+uint64_t binary_fuse_mulhi(uint64_t a, uint64_t b) { return __umulh(a, b); }
+#else
+uint64_t binary_fuse_mulhi(uint64_t a, uint64_t b) {
+  return ((__uint128_t)a * b) >> 64;
+}
+#endif
+
+typedef struct binary_hashes_s {
+  uint32_t h0;
+  uint32_t h1;
+  uint32_t h2;
+} binary_hashes_t;
+
+static inline binary_hashes_t binary_fuse8_hash_batch(uint64_t hash,
+                                        const binary_fuse8_t *filter) {
+  uint64_t hi = binary_fuse_mulhi(hash, filter->SegmentCountLength);
+  binary_hashes_t ans;
+  ans.h0 = (uint32_t)hi;
+  ans.h1 = ans.h0 + filter->SegmentLength;
+  ans.h2 = ans.h1 + filter->SegmentLength;
+  ans.h1 ^= (uint32_t)(hash >> 18) & filter->SegmentLengthMask;
+  ans.h2 ^= (uint32_t)(hash)&filter->SegmentLengthMask;
+  return ans;
+}
+static inline uint32_t binary_fuse8_hash(int index, uint64_t hash,
+                                        const binary_fuse8_t *filter) {
+    uint64_t h = binary_fuse_mulhi(hash, filter->SegmentCountLength);
+    h += index * filter->SegmentLength;
+    // keep the lower 36 bits
+    uint64_t hh = hash & ((1UL << 36) - 1);
+    // index 0: right shift by 36; index 1: right shift by 18; index 2: no shift
+    h ^= (size_t)((hh >> (36 - 18 * index)) & filter->SegmentLengthMask);
+    return h;
+}
+
+// Report if the key is in the set, with false positive rate.
+static inline bool binary_fuse8_contain(uint64_t key,
+                                        const binary_fuse8_t *filter) {
+  uint64_t hash = binary_fuse_mix_split(key, filter->Seed);
+  uint8_t f = binary_fuse8_fingerprint(hash);
+  binary_hashes_t hashes = binary_fuse8_hash_batch(hash, filter);
+  f ^= filter->Fingerprints[hashes.h0] ^ filter->Fingerprints[hashes.h1] ^
+       filter->Fingerprints[hashes.h2];
+  return f == 0;
+}
+
+static inline uint32_t binary_fuse_calculate_segment_length(uint32_t arity,
+                                                             uint32_t size) {
+  // These parameters are very sensitive. Replacing 'floor' by 'round' can
+  // substantially affect the construction time.
+  if (arity == 3) {
+    return ((uint32_t)1) << (int)(floor(log((double)(size)) / log(3.33) + 2.25));
+  } else if (arity == 4) {
+    return ((uint32_t)1) << (int)(floor(log((double)(size)) / log(2.91) - 0.5));
+  } else {
+    return 65536;
+  }
+}
+
+double binary_fuse8_max(double a, double b) {
+  if (a < b) {
+    return b;
+  }
+  return a;
+}
+
+static inline double binary_fuse_calculate_size_factor(uint32_t arity,
+                                                        uint32_t size) {
+  if (arity == 3) {
+    return binary_fuse8_max(1.125, 0.875 + 0.25 * log(1000000.0) / log((double)size));
+  } else if (arity == 4) {
+    return binary_fuse8_max(1.075, 0.77 + 0.305 * log(600000.0) / log((double)size));
+  } else {
+    return 2.0;
+  }
+}
+
+// allocate enough capacity for a set containing up to 'size' elements
+// caller is responsible to call binary_fuse8_free(filter)
+// size should be at least 2.
+static inline bool binary_fuse8_allocate(uint32_t size,
+                                         binary_fuse8_t *filter) {
+  uint32_t arity = 3;
+  filter->SegmentLength = size == 0 ? 4 : binary_fuse_calculate_segment_length(arity, size);
+  if (filter->SegmentLength > 262144) {
+    filter->SegmentLength = 262144;
+  }
+  filter->SegmentLengthMask = filter->SegmentLength - 1;
+  double sizeFactor = binary_fuse_calculate_size_factor(arity, size);
+  uint32_t capacity = size <= 1 ? 0 : (uint32_t)(round((double)size * sizeFactor));
+  uint32_t initSegmentCount =
+      (capacity + filter->SegmentLength - 1) / filter->SegmentLength -
+      (arity - 1);
+  filter->ArrayLength = (initSegmentCount + arity - 1) * filter->SegmentLength;
+  filter->SegmentCount =
+      (filter->ArrayLength + filter->SegmentLength - 1) / filter->SegmentLength;
+  if (filter->SegmentCount <= arity - 1) {
+    filter->SegmentCount = 1;
+  } else {
+    filter->SegmentCount = filter->SegmentCount - (arity - 1);
+  }
+  filter->ArrayLength =
+      (filter->SegmentCount + arity - 1) * filter->SegmentLength;
+  filter->SegmentCountLength = filter->SegmentCount * filter->SegmentLength;
+  filter->Fingerprints = (uint8_t*)malloc(filter->ArrayLength);
+  return filter->Fingerprints != NULL;
+}
+
+// report memory usage
+static inline size_t binary_fuse8_size_in_bytes(const binary_fuse8_t *filter) {
+  return filter->ArrayLength * sizeof(uint8_t) + sizeof(binary_fuse8_t);
+}
+
+// release memory
+static inline void binary_fuse8_free(binary_fuse8_t *filter) {
+  free(filter->Fingerprints);
+  filter->Fingerprints = NULL;
+  filter->Seed = 0;
+  filter->SegmentLength = 0;
+  filter->SegmentLengthMask = 0;
+  filter->SegmentCount = 0;
+  filter->SegmentCountLength = 0;
+  filter->ArrayLength = 0;
+}
+
+static inline uint8_t binary_fuse_mod3(uint8_t x) {
+    return x > 2 ? x - 3 : x;
+}
+
+// construct the filter, returns true on success, false on failure.
+// most likely, a failure is due to too high a memory usage
+// size is the number of keys
+// The caller is responsable for calling binary_fuse8_allocate(size,filter)
+// before. The caller is responsible to ensure that there are not too many  duplicated
+// keys. The inner loop will run up to XOR_MAX_ITERATIONS times (default on
+// 100), it should never fail, except if there are many duplicated keys. If it fails,
+// a return value of false is provided.
+//
+bool binary_fuse8_populate(const uint64_t *keys, uint32_t size,
+                           binary_fuse8_t *filter) {
+  uint64_t rng_counter = 0x726b2b9d438b9d4d;
+  filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);
+  uint64_t *reverseOrder = (uint64_t *)calloc((size + 1), sizeof(uint64_t));
+  uint32_t capacity = filter->ArrayLength;
+  uint32_t *alone = (uint32_t *)malloc(capacity * sizeof(uint32_t));
+  uint8_t *t2count = (uint8_t *)calloc(capacity, sizeof(uint8_t));
+  uint8_t *reverseH = (uint8_t *)malloc(size * sizeof(uint8_t));
+  uint64_t *t2hash = (uint64_t *)calloc(capacity, sizeof(uint64_t));
+
+  uint32_t blockBits = 1;
+  while (((uint32_t)1 << blockBits) < filter->SegmentCount) {
+    blockBits += 1;
+  }
+  uint32_t block = ((uint32_t)1 << blockBits);
+  uint32_t *startPos = (uint32_t *)malloc((1 << blockBits) * sizeof(uint32_t));
+  uint32_t h012[5];
+
+  if ((alone == NULL) || (t2count == NULL) || (reverseH == NULL) ||
+      (t2hash == NULL) || (reverseOrder == NULL) || (startPos == NULL)) {
+    free(alone);
+    free(t2count);
+    free(reverseH);
+    free(t2hash);
+    free(reverseOrder);
+    free(startPos);
+    return false;
+  }
+  reverseOrder[size] = 1;
+  for (int loop = 0; true; ++loop) {
+    if (loop + 1 > XOR_MAX_ITERATIONS) {
+      fprintf(stderr, "Too many iterations. Are all your keys unique?");
+      free(alone);
+      free(t2count);
+      free(reverseH);
+      free(t2hash);
+      free(reverseOrder);
+      free(startPos);
+      return false;
+    }
+
+    for (uint32_t i = 0; i < block; i++) {
+      // important : i * size would overflow as a 32-bit number in some
+      // cases.
+      startPos[i] = ((uint64_t)i * size) >> blockBits;
+    }
+
+    uint64_t maskblock = block - 1;
+    for (uint32_t i = 0; i < size; i++) {
+      uint64_t hash = binary_fuse_murmur64(keys[i] + filter->Seed);
+      uint64_t segment_index = hash >> (64 - blockBits);
+      while (reverseOrder[startPos[segment_index]] != 0) {
+        segment_index++;
+        segment_index &= maskblock;
+      }
+      reverseOrder[startPos[segment_index]] = hash;
+      startPos[segment_index]++;
+    }
+    int error = 0;
+    uint32_t duplicates = 0;
+    for (uint32_t i = 0; i < size; i++) {
+      uint64_t hash = reverseOrder[i];
+      uint32_t h0 = binary_fuse8_hash(0, hash, filter);
+      t2count[h0] += 4;
+      t2hash[h0] ^= hash;
+      uint32_t h1= binary_fuse8_hash(1, hash, filter);
+      t2count[h1] += 4;
+      t2count[h1] ^= 1;
+      t2hash[h1] ^= hash;
+      uint32_t h2 = binary_fuse8_hash(2, hash, filter);
+      t2count[h2] += 4;
+      t2hash[h2] ^= hash;
+      t2count[h2] ^= 2;
+      if ((t2hash[h0] & t2hash[h1] & t2hash[h2]) == 0) {
+        if   (((t2hash[h0] == 0) && (t2count[h0] == 8))
+          ||  ((t2hash[h1] == 0) && (t2count[h1] == 8))
+          ||  ((t2hash[h2] == 0) && (t2count[h2] == 8))) {
+					duplicates += 1;
+ 					t2count[h0] -= 4;
+ 					t2hash[h0] ^= hash;
+ 					t2count[h1] -= 4;
+ 					t2count[h1] ^= 1;
+ 					t2hash[h1] ^= hash;
+ 					t2count[h2] -= 4;
+ 					t2count[h2] ^= 2;
+ 					t2hash[h2] ^= hash;
+        }
+      }
+      error = (t2count[h0] < 4) ? 1 : error;
+      error = (t2count[h1] < 4) ? 1 : error;
+      error = (t2count[h2] < 4) ? 1 : error;
+    }
+    if(error) { continue; }
+
+    // End of key addition
+    uint32_t Qsize = 0;
+    // Add sets with one key to the queue.
+    for (uint32_t i = 0; i < capacity; i++) {
+      alone[Qsize] = i;
+      Qsize += ((t2count[i] >> 2) == 1) ? 1 : 0;
+    }
+    uint32_t stacksize = 0;
+    while (Qsize > 0) {
+      Qsize--;
+      uint32_t index = alone[Qsize];
+      if ((t2count[index] >> 2) == 1) {
+        uint64_t hash = t2hash[index];
+
+        //h012[0] = binary_fuse8_hash(0, hash, filter);
+        h012[1] = binary_fuse8_hash(1, hash, filter);
+        h012[2] = binary_fuse8_hash(2, hash, filter);
+        h012[3] = binary_fuse8_hash(0, hash, filter); // == h012[0];
+        h012[4] = h012[1];
+        uint8_t found = t2count[index] & 3;
+        reverseH[stacksize] = found;
+        reverseOrder[stacksize] = hash;
+        stacksize++;
+        uint32_t other_index1 = h012[found + 1];
+        alone[Qsize] = other_index1;
+        Qsize += ((t2count[other_index1] >> 2) == 2 ? 1 : 0);
+
+        t2count[other_index1] -= 4;
+        t2count[other_index1] ^= binary_fuse_mod3(found + 1);
+        t2hash[other_index1] ^= hash;
+
+        uint32_t other_index2 = h012[found + 2];
+        alone[Qsize] = other_index2;
+        Qsize += ((t2count[other_index2] >> 2) == 2 ? 1 : 0);
+        t2count[other_index2] -= 4;
+        t2count[other_index2] ^= binary_fuse_mod3(found + 2);
+        t2hash[other_index2] ^= hash;
+      }
+    }
+    if (stacksize + duplicates == size) {
+      // success
+      size = stacksize;
+      break;
+    }
+    memset(reverseOrder, 0, sizeof(uint64_t[size]));
+    memset(t2count, 0, sizeof(uint8_t[capacity]));
+    memset(t2hash, 0, sizeof(uint64_t[capacity]));
+    filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);
+  }
+
+  for (uint32_t i = size - 1; i < size; i--) {
+    // the hash of the key we insert next
+    uint64_t hash = reverseOrder[i];
+    uint8_t xor2 = binary_fuse8_fingerprint(hash);
+    uint8_t found = reverseH[i];
+    h012[0] = binary_fuse8_hash(0, hash, filter);
+    h012[1] = binary_fuse8_hash(1, hash, filter);
+    h012[2] = binary_fuse8_hash(2, hash, filter);
+    h012[3] = h012[0];
+    h012[4] = h012[1];
+    filter->Fingerprints[h012[found]] = xor2 ^
+                                        filter->Fingerprints[h012[found + 1]] ^
+                                        filter->Fingerprints[h012[found + 2]];
+  }
+  free(alone);
+  free(t2count);
+  free(reverseH);
+  free(t2hash);
+  free(reverseOrder);
+  free(startPos);
+  return true;
+}
+
+//////////////////
+// fuse16
+//////////////////
+
+typedef struct binary_fuse16_s {
+  uint64_t Seed;
+  uint32_t SegmentLength;
+  uint32_t SegmentLengthMask;
+  uint32_t SegmentCount;
+  uint32_t SegmentCountLength;
+  uint32_t ArrayLength;
+  uint16_t *Fingerprints;
+
+  binary_fuse16_s(uint32_t size) {
+      uint32_t arity = 3;
+      SegmentLength = size == 0 ? 4 : binary_fuse_calculate_segment_length(arity, size);
+      if (SegmentLength > 262144) {
+        SegmentLength = 262144;
+      }
+      SegmentLengthMask = SegmentLength - 1;
+      double sizeFactor = size <= 1 ? 0 : binary_fuse_calculate_size_factor(arity, size);
+      uint32_t capacity = (uint32_t)(round((double)size * sizeFactor));
+      uint32_t initSegmentCount =
+          (capacity + SegmentLength - 1) / SegmentLength -
+          (arity - 1);
+      ArrayLength = (initSegmentCount + arity - 1) * SegmentLength;
+      SegmentCount =
+          (ArrayLength + SegmentLength - 1) / SegmentLength;
+      if (SegmentCount <= arity - 1) {
+        SegmentCount = 1;
+      } else {
+        SegmentCount = SegmentCount - (arity - 1);
+      }
+      ArrayLength =
+          (SegmentCount + arity - 1) *SegmentLength;
+      SegmentCountLength = SegmentCount * SegmentLength;
+      Fingerprints = (uint16_t*)malloc(ArrayLength * sizeof(uint16_t));
+  }
+  ~binary_fuse16_s() {
+      free(Fingerprints);
+      Fingerprints = NULL;
+      Seed = 0;
+      SegmentLength = 0;
+      SegmentLengthMask = 0;
+      SegmentCount = 0;
+      SegmentCountLength = 0;
+      ArrayLength = 0;
+  }
+} binary_fuse16_t;
+
+static inline uint64_t binary_fuse16_fingerprint(uint64_t hash) {
+  return hash ^ (hash >> 32);
+}
+
+static inline binary_hashes_t binary_fuse16_hash_batch(uint64_t hash,
+                                        const binary_fuse16_t *filter) {
+  uint64_t hi = binary_fuse_mulhi(hash, filter->SegmentCountLength);
+  binary_hashes_t ans;
+  ans.h0 = (uint32_t)hi;
+  ans.h1 = ans.h0 + filter->SegmentLength;
+  ans.h2 = ans.h1 + filter->SegmentLength;
+  ans.h1 ^= (uint32_t)(hash >> 18) & filter->SegmentLengthMask;
+  ans.h2 ^= (uint32_t)(hash)&filter->SegmentLengthMask;
+  return ans;
+}
+static inline uint32_t binary_fuse16_hash(int index, uint64_t hash,
+                                        const binary_fuse16_t *filter) {
+    uint64_t h = binary_fuse_mulhi(hash, filter->SegmentCountLength);
+    h += index * filter->SegmentLength;
+    // keep the lower 36 bits
+    uint64_t hh = hash & ((1UL << 36) - 1);
+    // index 0: right shift by 36; index 1: right shift by 18; index 2: no shift
+    h ^= (size_t)((hh >> (36 - 18 * index)) & filter->SegmentLengthMask);
+    return h;
+}
+
+// Report if the key is in the set, with false positive rate.
+static inline bool binary_fuse16_contain(uint64_t key,
+                                        const binary_fuse16_t *filter) {
+  uint64_t hash = binary_fuse_mix_split(key, filter->Seed);
+  uint16_t f = binary_fuse16_fingerprint(hash);
+  binary_hashes_t hashes = binary_fuse16_hash_batch(hash, filter);
+  f ^= filter->Fingerprints[hashes.h0] ^ filter->Fingerprints[hashes.h1] ^
+       filter->Fingerprints[hashes.h2];
+  return f == 0;
+}
+
+
+// report memory usage
+static inline size_t binary_fuse16_size_in_bytes(const binary_fuse16_t *filter) {
+  return filter->ArrayLength * sizeof(uint16_t) + sizeof(binary_fuse16_t);
+}
+
+// construct the filter, returns true on success, false on failure.
+// most likely, a failure is due to too high a memory usage
+// size is the number of keys
+// The caller is responsable for calling binary_fuse8_allocate(size,filter)
+// before. The caller is responsible to ensure that there are not too many duplicated
+// keys. The inner loop will run up to XOR_MAX_ITERATIONS times (default on
+// 100), it should never fail, except if there are many duplicated keys. If it fails,
+// a return value of false is provided.
+//
+bool binary_fuse16_populate(const uint64_t *keys, uint32_t size,
+                           binary_fuse16_t *filter) {
+  uint64_t rng_counter = 0x726b2b9d438b9d4d;
+  filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);
+  uint64_t *reverseOrder = (uint64_t *)calloc((size + 1), sizeof(uint64_t));
+  uint32_t capacity = filter->ArrayLength;
+  uint32_t *alone = (uint32_t *)malloc(capacity * sizeof(uint32_t));
+  uint8_t *t2count = (uint8_t *)calloc(capacity, sizeof(uint8_t));
+  uint8_t *reverseH = (uint8_t *)malloc(size * sizeof(uint8_t));
+  uint64_t *t2hash = (uint64_t *)calloc(capacity, sizeof(uint64_t));
+
+  uint32_t blockBits = 1;
+  while (((uint32_t)1 << blockBits) < filter->SegmentCount) {
+    blockBits += 1;
+  }
+  uint32_t block = ((uint32_t)1 << blockBits);
+  uint32_t *startPos = (uint32_t *)malloc((1 << blockBits) * sizeof(uint32_t));
+  uint32_t h012[5];
+
+  if ((alone == NULL) || (t2count == NULL) || (reverseH == NULL) ||
+      (t2hash == NULL) || (reverseOrder == NULL) || (startPos == NULL)) {
+    free(alone);
+    free(t2count);
+    free(reverseH);
+    free(t2hash);
+    free(reverseOrder);
+    free(startPos);
+    return false;
+  }
+  reverseOrder[size] = 1;
+  for (int loop = 0; true; ++loop) {
+    if (loop + 1 > XOR_MAX_ITERATIONS) {
+      fprintf(stderr, "Too many iterations. Are all your keys unique?");
+      free(alone);
+      free(t2count);
+      free(reverseH);
+      free(t2hash);
+      free(reverseOrder);
+      free(startPos);
+      return false;
+    }
+
+    for (uint32_t i = 0; i < block; i++) {
+      // important : i * size would overflow as a 32-bit number in some
+      // cases.
+      startPos[i] = ((uint64_t)i * size) >> blockBits;
+    }
+
+    uint64_t maskblock = block - 1;
+    for (uint32_t i = 0; i < size; i++) {
+      uint64_t hash = binary_fuse_murmur64(keys[i] + filter->Seed);
+      uint64_t segment_index = hash >> (64 - blockBits);
+      while (reverseOrder[startPos[segment_index]] != 0) {
+        segment_index++;
+        segment_index &= maskblock;
+      }
+      reverseOrder[startPos[segment_index]] = hash;
+      startPos[segment_index]++;
+    }
+    int error = 0;
+    uint32_t duplicates = 0;
+    for (uint32_t i = 0; i < size; i++) {
+      uint64_t hash = reverseOrder[i];
+      uint32_t h0 = binary_fuse16_hash(0, hash, filter);
+      t2count[h0] += 4;
+      t2hash[h0] ^= hash;
+      uint32_t h1= binary_fuse16_hash(1, hash, filter);
+      t2count[h1] += 4;
+      t2count[h1] ^= 1;
+      t2hash[h1] ^= hash;
+      uint32_t h2 = binary_fuse16_hash(2, hash, filter);
+      t2count[h2] += 4;
+      t2hash[h2] ^= hash;
+      t2count[h2] ^= 2;
+      if ((t2hash[h0] & t2hash[h1] & t2hash[h2]) == 0) {
+        if   (((t2hash[h0] == 0) && (t2count[h0] == 8))
+          ||  ((t2hash[h1] == 0) && (t2count[h1] == 8))
+          ||  ((t2hash[h2] == 0) && (t2count[h2] == 8))) {
+					duplicates += 1;
+ 					t2count[h0] -= 4;
+ 					t2hash[h0] ^= hash;
+ 					t2count[h1] -= 4;
+ 					t2count[h1] ^= 1;
+ 					t2hash[h1] ^= hash;
+ 					t2count[h2] -= 4;
+ 					t2count[h2] ^= 2;
+ 					t2hash[h2] ^= hash;
+        }
+      }
+      error = (t2count[h0] < 4) ? 1 : error;
+      error = (t2count[h1] < 4) ? 1 : error;
+      error = (t2count[h2] < 4) ? 1 : error;
+    }
+    if(error) { continue; }
+
+    // End of key addition
+    uint32_t Qsize = 0;
+    // Add sets with one key to the queue.
+    for (uint32_t i = 0; i < capacity; i++) {
+      alone[Qsize] = i;
+      Qsize += ((t2count[i] >> 2) == 1) ? 1 : 0;
+    }
+    uint32_t stacksize = 0;
+    while (Qsize > 0) {
+      Qsize--;
+      uint32_t index = alone[Qsize];
+      if ((t2count[index] >> 2) == 1) {
+        uint64_t hash = t2hash[index];
+
+        //h012[0] = binary_fuse16_hash(0, hash, filter);
+        h012[1] = binary_fuse16_hash(1, hash, filter);
+        h012[2] = binary_fuse16_hash(2, hash, filter);
+        h012[3] = binary_fuse16_hash(0, hash, filter); // == h012[0];
+        h012[4] = h012[1];
+        uint8_t found = t2count[index] & 3;
+        reverseH[stacksize] = found;
+        reverseOrder[stacksize] = hash;
+        stacksize++;
+        uint32_t other_index1 = h012[found + 1];
+        alone[Qsize] = other_index1;
+        Qsize += ((t2count[other_index1] >> 2) == 2 ? 1 : 0);
+
+        t2count[other_index1] -= 4;
+        t2count[other_index1] ^= binary_fuse_mod3(found + 1);
+        t2hash[other_index1] ^= hash;
+
+        uint32_t other_index2 = h012[found + 2];
+        alone[Qsize] = other_index2;
+        Qsize += ((t2count[other_index2] >> 2) == 2 ? 1 : 0);
+        t2count[other_index2] -= 4;
+        t2count[other_index2] ^= binary_fuse_mod3(found + 2);
+        t2hash[other_index2] ^= hash;
+      }
+    }
+    if (stacksize + duplicates == size) {
+      // success
+      size = stacksize;
+      break;
+    }
+    memset(reverseOrder, 0, sizeof(uint64_t[size]));
+    memset(t2count, 0, sizeof(uint8_t[capacity]));
+    memset(t2hash, 0, sizeof(uint64_t[capacity]));
+    filter->Seed = binary_fuse_rng_splitmix64(&rng_counter);
+  }
+
+  for (uint32_t i = size - 1; i < size; i--) {
+    // the hash of the key we insert next
+    uint64_t hash = reverseOrder[i];
+    uint16_t xor2 = binary_fuse16_fingerprint(hash);
+    uint8_t found = reverseH[i];
+    h012[0] = binary_fuse16_hash(0, hash, filter);
+    h012[1] = binary_fuse16_hash(1, hash, filter);
+    h012[2] = binary_fuse16_hash(2, hash, filter);
+    h012[3] = h012[0];
+    h012[4] = h012[1];
+    filter->Fingerprints[h012[found]] = xor2 ^
+                                        filter->Fingerprints[h012[found + 1]] ^
+                                        filter->Fingerprints[h012[found + 2]];
+  }
+  free(alone);
+  free(t2count);
+  free(reverseH);
+  free(t2hash);
+  free(reverseOrder);
+  free(startPos);
+  return true;
+}
+
+
+
+
+#endif
\ No newline at end of file
diff --git a/include/hnswlib/hnswalg.h b/include/hnswlib/hnswalg.h
index 67f49ff1..761a93a2 100644
--- a/include/hnswlib/hnswalg.h
+++ b/include/hnswlib/hnswalg.h
@@ -332,6 +332,100 @@ namespace hnswlib {
             return top_candidates;
         }
 
+
+        template <bool has_deletions, bool collect_metrics=false>
+        std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst>
+        searchBaseLayerSTWithFilter(tableint ep_id, const void *data_point, const binary_fuse16_t *filter, size_t ef) const {
+            VisitedList *vl = visited_list_pool_->getFreeVisitedList();
+            vl_type *visited_array = vl->mass;
+            vl_type visited_array_tag = vl->curV;
+
+            std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> top_candidates;
+            std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> candidate_set;
+
+            dist_t lowerBound;
+
+            uint64_t label = getExternalLabel(ep_id);
+            if (binary_fuse16_contain(label, filter)) {
+                dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_);
+                lowerBound = dist;
+                top_candidates.emplace(dist, ep_id);
+                candidate_set.emplace(-dist, ep_id);
+            } else {
+                lowerBound = std::numeric_limits<dist_t>::max();
+                candidate_set.emplace(-lowerBound, ep_id);
+            }
+
+            visited_array[ep_id] = visited_array_tag;
+
+            while (!candidate_set.empty()) {
+
+                std::pair<dist_t, tableint> current_node_pair = candidate_set.top();
+
+                if ((-current_node_pair.first) > lowerBound) {
+                    break;
+                }
+                candidate_set.pop();
+
+                tableint current_node_id = current_node_pair.second;
+                int *data = (int *) get_linklist0(current_node_id);
+                size_t size = getListCount((linklistsizeint*)data);
+//                bool cur_node_deleted = isMarkedDeleted(current_node_id);
+                if(collect_metrics){
+                    metric_hops++;
+                    metric_distance_computations+=size;
+                }
+
+#ifdef USE_SSE
+                _mm_prefetch((char *) (visited_array + *(data + 1)), _MM_HINT_T0);
+                _mm_prefetch((char *) (visited_array + *(data + 1) + 64), _MM_HINT_T0);
+                _mm_prefetch(data_level0_memory_ + (*(data + 1)) * size_data_per_element_ + offsetData_, _MM_HINT_T0);
+                _mm_prefetch((char *) (data + 2), _MM_HINT_T0);
+#endif
+
+                for (size_t j = 1; j <= size; j++) {
+                    int candidate_id = *(data + j);
+//                    if (candidate_id == 0) continue;
+#ifdef USE_SSE
+                    _mm_prefetch((char *) (visited_array + *(data + j + 1)), _MM_HINT_T0);
+                    _mm_prefetch(data_level0_memory_ + (*(data + j + 1)) * size_data_per_element_ + offsetData_,
+                                 _MM_HINT_T0);////////////
+#endif
+                    if (!(visited_array[candidate_id] == visited_array_tag)) {
+
+                        visited_array[candidate_id] = visited_array_tag;
+
+                        char *currObj1 = (getDataByInternalId(candidate_id));
+                        dist_t dist = fstdistfunc_(data_point, currObj1, dist_func_param_);
+
+                        if (top_candidates.size() < ef || lowerBound > dist) {
+                            candidate_set.emplace(-dist, candidate_id);
+#ifdef USE_SSE
+                            _mm_prefetch(data_level0_memory_ + candidate_set.top().second * size_data_per_element_ +
+                                         offsetLevel0_,///////////
+                                         _MM_HINT_T0);////////////////////////
+#endif
+
+                            uint64_t et_label = getExternalLabel(candidate_id);
+
+                            if (binary_fuse16_contain(et_label, filter))
+                                top_candidates.emplace(dist, candidate_id);
+
+                            if (top_candidates.size() > ef)
+                                top_candidates.pop();
+
+                            if (!top_candidates.empty())
+                                lowerBound = top_candidates.top().first;
+                        }
+                    }
+                }
+            }
+
+            visited_list_pool_->releaseVisitedList(vl);
+            return top_candidates;
+        }
+
+
         void getNeighborsByHeuristic2(
                 std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> &top_candidates,
         const size_t M) {
@@ -1153,6 +1247,62 @@ namespace hnswlib {
             return result;
         };
 
+        std::priority_queue<std::pair<dist_t, labeltype >>
+        searchKnnWithFilter(const void *query_data, const binary_fuse16_t *filter, size_t k) const {
+            std::priority_queue<std::pair<dist_t, labeltype >> result;
+            if (cur_element_count == 0) return result;
+
+            tableint currObj = enterpoint_node_;
+            dist_t curdist = fstdistfunc_(query_data, getDataByInternalId(enterpoint_node_), dist_func_param_);
+
+            for (int level = maxlevel_; level > 0; level--) {
+                bool changed = true;
+                while (changed) {
+                    changed = false;
+                    unsigned int *data;
+
+                    data = (unsigned int *) get_linklist(currObj, level);
+                    int size = getListCount(data);
+                    metric_hops++;
+                    metric_distance_computations+=size;
+
+                    tableint *datal = (tableint *) (data + 1);
+                    for (int i = 0; i < size; i++) {
+                        tableint cand = datal[i];
+                        if (cand < 0 || cand > max_elements_)
+                            throw std::runtime_error("cand error");
+                        dist_t d = fstdistfunc_(query_data, getDataByInternalId(cand), dist_func_param_);
+
+                        if (d < curdist) {
+                            curdist = d;
+                            currObj = cand;
+                            changed = true;
+                        }
+                    }
+                }
+            }
+
+            std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> top_candidates;
+            if (has_deletions_) {
+                top_candidates=searchBaseLayerSTWithFilter<true,true>(
+                        currObj, query_data, filter, std::max(ef_, k));
+            }
+            else{
+                top_candidates=searchBaseLayerSTWithFilter<false,true>(
+                        currObj, query_data, filter, std::max(ef_, k));
+            }
+
+            while (top_candidates.size() > k) {
+                top_candidates.pop();
+            }
+            while (top_candidates.size() > 0) {
+                std::pair<dist_t, tableint> rez = top_candidates.top();
+                result.push(std::pair<dist_t, labeltype>(rez.first, getExternalLabel(rez.second)));
+                top_candidates.pop();
+            }
+            return result;
+        };
+
         void checkIntegrity(){
             int connections_checked=0;
             std::vector <int > inbound_connections_num(cur_element_count,0);
diff --git a/include/hnswlib/hnswlib.h b/include/hnswlib/hnswlib.h
index 9409c388..2c667905 100644
--- a/include/hnswlib/hnswlib.h
+++ b/include/hnswlib/hnswlib.h
@@ -105,4 +105,5 @@ namespace hnswlib {
 #include "space_l2.h"
 #include "space_ip.h"
 #include "bruteforce.h"
+#include "fusefilter.h"
 #include "hnswalg.h"
diff --git a/pqlite/container.py b/pqlite/container.py
index 01c78999..3602ed8e 100644
--- a/pqlite/container.py
+++ b/pqlite/container.py
@@ -45,20 +45,9 @@ def __init__(
                 )
                 for _ in range(n_cells)
             ]
-        elif columns is None:
-            self._vec_indexes = [
-                HnswIndex(
-                    dim,
-                    metric=metric,
-                    initial_size=initial_size,
-                    expand_step_size=expand_step_size,
-                    expand_mode=expand_mode,
-                )
-                for _ in range(n_cells)
-            ]
         else:
             self._vec_indexes = [
-                FlatIndex(
+                HnswIndex(
                     dim,
                     metric=metric,
                     initial_size=initial_size,
@@ -67,6 +56,7 @@ def __init__(
                 )
                 for _ in range(n_cells)
             ]
+
         self._doc_stores = [DocStorage(data_path / f'cell_{_}') for _ in range(n_cells)]
 
         self._cell_tables = [
@@ -116,7 +106,6 @@ def ivf_search(
             cell_ids.extend([cell_id] * len(_dists))
             count += len(_dists)
 
-
         cell_ids = np.array(cell_ids, dtype=np.int64)
         if len(dists) != 0:
             dists = np.hstack(dists)
@@ -214,6 +203,7 @@ def update(
             if cell_id == _cell_id:
                 self.vec_index(cell_id).add_with_ids(x.reshape(1, -1), [_offset])
                 self.cell_table(cell_id).undo_delete_by_offset(_offset)
+                self.doc_store(cell_id).update([doc])
 
             elif _cell_id is None:
                 new_data.append(x)
@@ -222,6 +212,7 @@ def update(
             else:
                 # relpace
                 self.cell_table(_cell_id).delete_by_offset(_offset)
+                self.doc_store(_cell_id).delete([doc.id])
 
                 new_data.append(x)
                 new_cells.append(cell_id)
diff --git a/pqlite/core/index/hnsw/index.py b/pqlite/core/index/hnsw/index.py
index 2c3ee430..57708497 100644
--- a/pqlite/core/index/hnsw/index.py
+++ b/pqlite/core/index/hnsw/index.py
@@ -63,17 +63,19 @@ def search(
             _dim == self.dim
         ), f'the query embedding dimension does not match with index dimension: {_dim} vs {self.dim}'
 
-        if indices is not None:
-            raise NotImplementedError(
-                f'the index {self.__class__.__name__} does not pre-filtering now'
-            )
-
         query = query.reshape((-1, self.dim))
 
         ef_search = max(self.ef_search, limit)
         self._index.set_ef(ef_search)
 
-        ids, dists = self._index.knn_query(query, k=limit)
+        if indices is not None:
+            # TODO: add a smart strategy to speed up this case (bruteforce search would be better)
+            if len(indices) < limit:
+                limit = len(indices)
+
+            ids, dists=self._index.knn_query_with_filter(query, filters=indices, k=limit)
+        else:
+            ids, dists = self._index.knn_query(query, k=limit)
         return dists[0], ids[0]
 
     def delete(self, ids: List[int]):


Test patch:
diff --git a/tests/test_index.py b/tests/test_index.py
index 79618e39..2b990650 100644
--- a/tests/test_index.py
+++ b/tests/test_index.py
@@ -82,3 +82,4 @@ def test_index_query_with_filtering(pqlite_with_data):
             query[0].matches[i].scores['euclidean'].value
             <= query[0].matches[i + 1].scores['euclidean'].value
         )
+        assert query[0].matches[i].tags['x'] > 0.6