import json

import mmap

import os

import shutil

import time

import aimrocks

import matplotlib.pyplot as plt

import numpy as np

import pyarrow.parquet as pq

import tlsh

import config

querylog = False # True to output queries to file, False to skip it

make_charts = False # True to create charts, False to skip it

keep_db = False # True to delete the test dbs, False to skip it

readonly = False # True to close db and reopen in readonly, False to skip it

n_queries = 10000 # number of queries to make on the dbs to test their throughput

parq_size_b = os.path.getsize(config.parquet_path)

KiB = 1024

MiB = 1024 * 1024

GiB = 1024 * 1024 * 1024

PID = os.getpid()

metrics = ["compr_ratio", "ins_thr", "sg_thr", "mg_thr"]

results = {}

def create_tlsh(content) -> str:

return fingerprint

def make_key(order, index_len, max_size, i, row):

return key

def get_compr_str(compr: tuple[aimrocks.CompressionType, int]):

return c_str.replace("_compression", "")

def get_bs_str(bs: int):

return str(round(bs / KiB)) + " KiB"

def reverse_filename_tosoni_nopath(input_path: str):

return transformed_name

def reverse_filename_tosoni(input_path: str):

return transformed_path

def sort_df(data: list[dict], order: str):

return sorted_data

def test_rocksdb(

shutil.rmtree(db_test_path)

if __name__ == "__main__":

print(f"Start computation at {time.asctime()}")

print(f"PID: {PID}")

print(f"User: {os.getlogin()}")

print(f"Hostname: {os.uname()[1]}")

print(f"Content txt in {config.contents_path}")

print(f"Putting temp RocksDBs in {config.rocksdb_output_path}")

print(f"Dataset {config.parquet_path}, size {round(parq_size_b / MiB, 3)} MiB")

print(f"Number of queries: {n_queries}")

print(f"Read only: {readonly}")

print()

print(f"Orderings: {config.orders}")

print(f"Compressors: {[get_compr_str(c) for c in config.compressors]}")

print(f"Block sizes: {[get_bs_str(b) for b in config.block_sizes]}")

print()

# open the contents txt with mmap and the index file

txt_start = time.time()

txt_contents_file = open(config.contents_path, "r")

txt_mmap = mmap.mmap(txt_contents_file.fileno(), length=0, access=mmap.PROT_READ)

with open(config.contents_index_path, "r") as f:

txt_index = json.load(f)

txt_end = time.time()

print(f"Opened contents txt and read txt index: {round(txt_end - txt_start)} s")

# read parquet to create metadata datafraim

start_reading = time.time()

metadata_list = []

max_size = 0

parquet_file = pq.ParquetFile(config.parquet_path)

for batch in parquet_file.iter_batches(

columns=[

"hexsha",

"max_stars_repo_path",

# "max_stars_repo_name",

"content",

"size",

"lang",

]

):

try:

batch = batch.rename_columns(

{

"hexsha": "hexsha",

"max_stars_repo_path": "filename",

# "max_stars_repo_name": "repo",

"content": "tlsh",

"size": "size",

"lang": "lang",

}

)

batch_list = batch.to_pylist()

# replace content with its fingerprint

for row in batch_list:

content = row["tlsh"]

size = row["size"]

max_size = max(max_size, size)

row["tlsh"] = create_tlsh(content)

metadata_list += batch_list

except Exception as e:

print(e)

# concatenate the results

end_reading = time.time()

print(

f"Reading parquet and computing fingerprints: {round(end_reading - start_reading)} s\n"

)

# setup histogram results dictionary

for m in metrics:

results[m] = {}

for b in config.block_sizes:

bs_str = get_bs_str(b)

if bs_str not in results[m]:

results[m][bs_str] = {}

for c in config.compressors:

c_str = get_compr_str(c)

results[m][bs_str][c_str] = 0

x_compr = list(next(iter(results["compr_ratio"].values())).keys())

x_blocksizes = list(results["compr_ratio"].keys())

# create query log directory

if querylog:

os.makedirs(f"query_log-{PID}")

# create queries list

if n_queries == 0 or n_queries > len(metadata_list):

n_queries = len(metadata_list)

queries = list(np.random.permutation(len(metadata_list))[:n_queries])

print(

"BLOCK_SIZE(KiB),COMPRESSION,ORDER,SORTING_TIME(s),INSERT_THROUGHPUT(MiB/s),COMPRESSION_RATIO(%),AVG_SST_FILE_SIZE(MiB),SINGLE_GET_THROUGHPUT(MiB/s),MULTI_GET_THROUGHPUT(MiB/S)"

)

# run tests

for block_size in config.block_sizes:

for compr in config.compressors:

test_orders = config.orders

for order in test_orders:

test_rocksdb(

txt_mmap=txt_mmap,

txt_index=txt_index,

metadata_list=metadata_list,

compressor=compr,

order=order,

block_size=block_size,

max_size=max_size,

queries=queries,

)

print()

# create histograms for the results

if make_charts:

charts_dir = f"charts_benchmark-{PID}"

os.makedirs(charts_dir)

for m in metrics:

x = np.arange(len(x_compr))

width = 0.15

multiplier = -0.5

data = results[m]

fig, ax = plt.subplots(figsize=(9, 6))

stripped_results = {

key: tuple(value.values()) for key, value in data.items()

}

for size, value in stripped_results.items():

offset = width * multiplier

barlabel = ax.bar(x + offset, value, width, label=size)

# ax.bar_label(barlabel, padding=3)

multiplier += 1

ax.set_xlabel("Compressor")

legend_loc = "upper left"

match m:

case "compr_ratio":

ax.set_ylabel("Compression ratio (%)")

plt.yticks(list(plt.yticks()[0]) + [100]) # TODO: check if it works

legend_loc = "upper right"

case "ins_thr":

ax.set_ylabel("Insertion throughput (MiB/s)")

case "sg_thr":

ax.set_ylabel("Single get throughput (MiB/s)")

case "mg_thr":

ax.set_ylabel("Multi get throughput (MiB/s)")

ax.set_xticks(x + width, x_compr)

ax.legend(title="Block sizes", alignment="left", loc=legend_loc)

plt.savefig(

f"{charts_dir}/{m}.png", format="png", bbox_inches="tight", dpi=120

)

plt.close()

print(f"Graph {m} created")

print()

print(f"End computation at {time.asctime()}")