fixes

code/utils/lmdb_utils.py

@@ -0,0 +1,119 @@
+import pickle
+
+import pandas as pd
+
+
+def save_to_lmdb(env, key, dataset):
+    """
+    Saves the given dataset to an LMDB environment with the given key.
+    :param env: LMDB environment
+    :param key: key to save the dataset to
+    :param dataset: tuple of pandas dataframes
+    """
+
+    with env.begin(write=True) as txn:
+        txn.put(key.encode('ascii'), pickle.dumps(dataset))
+
+
+def load_from_lmdb(env, key):
+    """
+    Loads a dataset from an LMDB environment with the given key.
+    :param env: LMDB environment
+    :param key: key of the dataset to load
+    :return: key and tuple of pandas dataframes (input, context, output)
+    """
+    with env.begin(write=False) as txn:
+        try:
+            data = pickle.loads(txn.get(key.encode('ascii')))
+            return data
+        except TypeError:
+            raise KeyError(key)
+
+
+def delete_from_lmdb(env, key):
+    """
+    Deletes a dataset from an LMDB environment with the given key.
+    :param env: LMDB environment
+    :param key: key of the dataset to delete
+    """
+    with env.begin(write=True) as txn:
+        txn.delete(key.encode('ascii'))
+
+
+def clear_lmdb(env):
+    """
+    Clears all datasets from an LMDB environment.
+    :param env: LMDB environment
+    """
+    with env.begin(write=True) as txn:
+        cursor = txn.cursor()
+        for key, value in cursor:
+            txn.delete(key)
+
+
+def lmdb_dataset_generator(env):
+    """
+    Generator function to yield datasets from an LMDB environment.
+    :param env: LMDB environment
+    :return: generator
+    """
+    with env.begin(write=False) as txn:
+        cursor = txn.cursor()
+        for key, value in cursor:
+            data = pickle.loads(value)
+            yield data
+
+
+def lmdb_contains(env, substring) -> bool:
+    """
+    Checks if the given substring is contained in any of the keys of the LMDB environment.
+    :param env: LMDB environment
+    :param substring: substring to search for
+    :return: boolean
+    """
+    with env.begin(write=False) as txn:
+        cursor = txn.cursor()
+        for key, value in cursor:
+            if substring in key.decode('ascii'):
+                return True
+    return False
+
+
+def lmdb_substring_key_search(env, substring):
+    """
+    Searches for keys in the LMDB environment that contain the given substring.
+    :param env: LMDB environment
+    :param substring: substring to search for
+    :return: list of keys
+    """
+    keys = []
+    with env.begin(write=False) as txn:
+        cursor = txn.cursor()
+        for key, value in cursor:
+            if substring in key.decode('ascii'):
+                keys.append(key)
+    return keys
+
+
+def get_lmdb_keys(env, limit: int = None):
+    """
+    Get all keys in the LMDB environment.
+    :param env: LMDB environment
+    :param limit: maximum number of keys to return
+    :return: list of keys
+    """
+    with env.begin(write=False) as txn:
+        with txn.cursor() as cursor:
+            keys = [key.decode("ascii") for key in cursor.iternext(keys=True, values=False)]
+
+    return keys
+
+
+def get_lmdb_keyspace_size(env):
+    """
+    Get the number of keys in the LMDB environment.
+    :param env: LMDB environment
+    :return: number of keys
+    """
+    with env.begin(write=False) as txn:
+        return txn.stat()['entries']