diff --git a/bigsi/graph/index.py b/bigsi/graph/index.py index 2ac2952..03d4a36 100644 --- a/bigsi/graph/index.py +++ b/bigsi/graph/index.py @@ -69,7 +69,7 @@ def __kmers_to_hashes(self, kmers): return d def __batch_get_rows(self, row_indexes): - return dict(zip(row_indexes, self.bitmatrix.get_rows(row_indexes))) + return dict(zip(row_indexes, self.bitmatrix.get_rows(row_indexes, slice=False))) def __bitwise_and_kmers(self, kmer_to_hashes, rows): d = {} diff --git a/bigsi/matrix/bitmatrix.py b/bigsi/matrix/bitmatrix.py index abba573..b387ceb 100644 --- a/bigsi/matrix/bitmatrix.py +++ b/bigsi/matrix/bitmatrix.py @@ -27,9 +27,14 @@ def create(cls, storage, rows, num_rows, num_cols): def get_row(self, row_index): return self.storage.get_bitarray(row_index)[: self.num_cols] - def get_rows(self, row_indexes): + def get_rows(self, row_indexes, slice=True): + ## Only need to slice for merging (it's a lot slower) # Takes advantage of batching in storage engine if available - return (ba[: self.num_cols] for ba in self.storage.get_bitarrays(row_indexes)) + bitarrays=self.storage.get_bitarrays(row_indexes) + if slice: + return (ba[: self.num_cols] for ba in bitarrays) + else: + return bitarrays def set_row(self, row_index, bitarray): return self.storage.set_bitarray(row_index, bitarray)