diff --git a/deeplake/core/vectorstore/vector_search/indra/search_algorithm.py b/deeplake/core/vectorstore/vector_search/indra/search_algorithm.py index 85a3156969..133a572fe6 100644 --- a/deeplake/core/vectorstore/vector_search/indra/search_algorithm.py +++ b/deeplake/core/vectorstore/vector_search/indra/search_algorithm.py @@ -68,14 +68,12 @@ def run( return {"data": return_data, "tql": tql_query} return return_data - @staticmethod + @abstractmethod def _collect_return_data( - view: Union[DeepLakeQueryDataset, DeepLakeDataset], + self, + view: DeepLakeDataset, ): - return_data = {} - for tensor in view.tensors: - return_data[tensor] = utils.parse_tensor_return(view[tensor]) - return return_data + pass @staticmethod def _create_tql_string( @@ -138,14 +136,30 @@ def _get_indra_dataset(self): indra_dataset = dataset_to_libdeeplake(self.deeplake_dataset) return indra_dataset + def _collect_return_data( + self, + view: DeepLakeDataset, + ): + return_data = {} + for tensor in view.tensors: + return_data[tensor] = utils.parse_tensor_return(view[tensor]) + return return_data + class SearchManaged(SearchBasic): def _get_view(self, tql_query, runtime: Optional[Dict] = None): - view, _ = self.deeplake_dataset.query( + view, data = self.deeplake_dataset.query( tql_query, runtime=runtime, return_data=True ) + self.data = data return view + def _collect_return_data( + self, + view: DeepLakeDataset, + ): + return self.data + def search( query_embedding: np.ndarray,