Alex Blank
120 lines
3.2 KiB
Python
120 lines
3.2 KiB
Python
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']
|