|
| 1 | +import hnswlib |
| 2 | +""" |
| 3 | + A python wrapper for lazy indexing, preserves the same api as hnswlib.Index but initializes the index only after adding items for the first time with `add_items`. |
| 4 | +""" |
| 5 | +class LazyIndex(hnswlib.Index): |
| 6 | + def __init__(self, space, dim,max_elements=1024, ef_construction=200, M=16): |
| 7 | + super().__init__(space, dim) |
| 8 | + self.init_max_elements=max_elements |
| 9 | + self.init_ef_construction=ef_construction |
| 10 | + self.init_M=M |
| 11 | + def init_index(self, max_elements=0,M=0,ef_construction=0): |
| 12 | + if max_elements>0: |
| 13 | + self.init_max_elements=max_elements |
| 14 | + if ef_construction>0: |
| 15 | + self.init_ef_construction=ef_construction |
| 16 | + if M>0: |
| 17 | + self.init_M=M |
| 18 | + super().init_index(self.init_max_elements, self.init_M, self.init_ef_construction) |
| 19 | + def add_items(self, data, ids=None, num_threads=-1): |
| 20 | + if self.max_elements==0: |
| 21 | + self.init_index() |
| 22 | + return super().add_items(data,ids, num_threads) |
| 23 | + def get_items(self, ids=None): |
| 24 | + if self.max_elements==0: |
| 25 | + return [] |
| 26 | + return super().get_items(ids) |
| 27 | + def knn_query(self, data,k=1, num_threads=-1): |
| 28 | + if self.max_elements==0: |
| 29 | + return [], [] |
| 30 | + return super().knn_query(data, k, num_threads) |
| 31 | + def resize_index(self, size): |
| 32 | + if self.max_elements==0: |
| 33 | + return self.init_index(size) |
| 34 | + else: |
| 35 | + return super().resize_index(size) |
| 36 | + def set_ef(self, ef): |
| 37 | + if self.max_elements==0: |
| 38 | + self.init_ef_construction=ef |
| 39 | + return |
| 40 | + super().set_ef(ef) |
| 41 | + def get_max_elements(self): |
| 42 | + return self.max_elements |
| 43 | + def get_current_count(self): |
| 44 | + return self.element_count |
0 commit comments