plot_result.py

#!/usr/bin/env python3

## plot result of YCSB

# example of dir format
# --- result dir
#    \-- pxno
#       \-- workloada.result
#       |-- workloadb.result
#    |-- rocksdb
#       \-- workloada.result
#       |-- workloadb.result

import matplotlib.pyplot as plt
import numpy as np
import sys
import os


def main():
    if len(sys.argv) != 2 or sys.argv[1] == '-h':
        print("Usage: {} [result dir]".format(sys.argv[0]))
        sys.exit(-1)

    result_dir = sys.argv[1]
    result_suffix = ".png"

    entries = os.listdir(result_dir)
    db_names = [f for f in entries if os.path.isdir(os.path.join(result_dir, f))]
    workloads = ['workloada', 'workloadb', 'workloadc', 'workloadd', 'workloadf']

    threads = 0
    results = {}
    for db in db_names:
        path = [os.path.join(result_dir, db, w + ".result") for w in workloads]
        for f in path:
            lines = [line.strip().split() for line in open(f)]
            for line in lines:
                if line[0][0] == '#':
                    continue
                num_threads = int(line[0])
                threads = max(threads, num_threads)
                throughput = float(line[1])
                dir_name, workload = os.path.split(f)
                workload = os.path.splitext(workload)[0]

                if db not in results:
                    results[db] = {}
                if workload not in results[db]:
                    results[db][workload] = {}
                if num_threads not in results[db][workload]:
                    results[db][workload][num_threads] = []
                results[db][workload][num_threads].append(throughput)

    width = 0.3

    fig, ax = plt.subplots()
    db_results = {}
    shift = np.linspace(-len(db_names), len(db_names), len(db_names)) * (width / len(db_names))
    for workload in workloads:
        ax.clear()
        count = 0
        for db in db_names:
            keys = sorted(results[db][workload].keys())
            x = list(keys)
            xa = x + shift[count]
            va = []
            for key in keys:
                va.append(max(results[db][workload][key]))
            bar = ax.bar(xa, va, width, label=db)
            ax.set_ylabel('throughput [KOPS]')
            ax.set_xlabel('threads')
            ax.set_title(workload)
            ax.set_xticks(x)
            ax.legend()
            if workloads not in workloads:
                db_results[workload] = {}
            if db not in db_results[workload]:
                db_results[workload][db] = []
            db_results[workload][db].append(bar)
            count += 1
        fig.tight_layout()
        result_fig = os.path.join(result_dir, workload + result_suffix)
        plt.savefig(result_fig)


if __name__ == '__main__':
    main()