diff options
Diffstat (limited to 'gn3')
-rw-r--r-- | gn3/db/matrix.py | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/gn3/db/matrix.py b/gn3/db/matrix.py new file mode 100644 index 0000000..d40f3c3 --- /dev/null +++ b/gn3/db/matrix.py @@ -0,0 +1,44 @@ +"""Methods for fetching data from the matrix stored in LMDB""" +from typing import Optional +from dataclasses import dataclass + +import struct +import json +import lmdb + +BLOB_HASH_DIGEST = 32 + + +@dataclass +class Matrix: + """Store sample data and any other relevant metadata""" + + data: list + metadata: dict + + +def get_current_matrix(db_path: str) -> Optional[Matrix]: + """Get the most recent matrix from DB_PATH. This is functionally + equivalent to get_nth_matrix(0, db_path)""" + env = lmdb.open(db_path) + with env.begin(write=False) as txn: + current_hash = txn.get(b"current") or b"" + matrix_hash = txn.get(current_hash + b":matrix") or b"" + row_pointers = txn.get(matrix_hash + b":row-pointers") + nrows = 0 + if matrix_hash: + (nrows,) = struct.unpack("<Q", txn.get(matrix_hash + b":nrows")) + data, metadata = None, None + if row_pointers: + return Matrix( + data=[ + json.loads(txn.get(row_pointers[i: i + 32]).decode()) + for i in range(0, nrows * 32, 32) + ], + metadata=json.loads( + txn.get(matrix_hash + b":metadata") + .rstrip(b"\x00") + .decode() + ), + ) + return None |