| 9 | |
| 10 | |
| 11 | class HNSWSearcher(object): |
| 12 | def __init__(self,table_path,index_path,scale): |
| 13 | tfile = open(table_path,"rb") |
| 14 | tables = pickle.load(tfile) |
| 15 | # For scalability experiments: load a percentage of tables |
| 16 | self.tables = random.sample(tables, int(scale*len(tables))) |
| 17 | print("From %d total data-lake tables, scale down to %d tables" % (len(tables), len(self.tables))) |
| 18 | tfile.close() |
| 19 | self.vec_dim = len(self.tables[1][1][0]) |
| 20 | |
| 21 | index_start_time = time.time() |
| 22 | self.index = hnswlib.Index(space='cosine', dim=self.vec_dim) |
| 23 | self.all_columns, self.col_table_ids = self._preprocess_table_hnsw() |
| 24 | # if not os.path.exists(index_path): |
| 25 | # build index from scratch |
| 26 | # self.index.init_index(max_elements=len(self.all_columns), ef_construction=100, M=16) |
| 27 | self.index.init_index(max_elements=len(self.all_columns), ef_construction=100, M=32) |
| 28 | |
| 29 | self.index.set_ef(10) |
| 30 | self.index.add_items(self.all_columns) |
| 31 | # self.index.save_index(index_path) |
| 32 | print("--- Indexing Time: %s seconds ---" % (time.time() - index_start_time)) |
| 33 | # else: |
| 34 | # # load index |
| 35 | # self.index.load_index(index_path, max_elements = len(self.all_columns)) |
| 36 | |
| 37 | def topk(self, enc, query, K, N=5, threshold=0.6): |
| 38 | # Note: N is the number of columns retrieved from the index |
| 39 | # query是什么 |
| 40 | query_cols = [] |
| 41 | for col in query[1]: |
| 42 | query_cols.append(col) |
| 43 | candidates = self._find_candidates(query_cols, N) |
| 44 | if enc == 'sato': |
| 45 | scores = [] |
| 46 | querySherlock = query[1][:, :1187] |
| 47 | querySato = query[1][0, 1187:] |
| 48 | for table in candidates: |
| 49 | sherlock = table[1][:, :1187] |
| 50 | sato = table[1][0, 1187:] |
| 51 | sScore = self._verify(querySherlock, sherlock, threshold) |
| 52 | sherlockScore = (1/min(len(querySherlock), len(sherlock))) * sScore |
| 53 | satoScore = self._cosine_sim(querySato, sato) |
| 54 | score = sherlockScore + satoScore |
| 55 | scores.append((score, table[0])) |
| 56 | else: # encoder is sherlock |
| 57 | scores = [(self._verify(query[1], table[1], threshold)[0], self._verify(query[1], table[1], threshold)[1], table[0]) for table in candidates] |
| 58 | scores.sort(reverse=True) |
| 59 | scoreLength = len(scores) |
| 60 | return scores[:K], scoreLength |
| 61 | |
| 62 | def _preprocess_table_hnsw(self): |
| 63 | all_columns = [] |
| 64 | col_table_ids = [] |
| 65 | for idx,table in enumerate(self.tables): |
| 66 | for col in table[1]: |
| 67 | all_columns.append(col) |
| 68 | col_table_ids.append(idx) |