memgraph/tools/plot_ldbc_latency

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''
Latency Barchart (Based on LDBC JSON output).
'''

import json
import os
import numpy as np
from argparse import ArgumentParser
import string

import matplotlib
# Must set 'Agg' backend before importing pyplot
# This is so the script works on headless machines (without X11)
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.cbook import get_sample_data


SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
COLORS = ['#ff7300', '#008cc2'] # TODO: add more colors!
LDBC_TIME_FACTORS = {
    "SECONDS": 1.0,
    "MILLISECONDS": 1000.0,
    "MICROSECONDS": 1000000.0,
    "NANOSECONDS": 1000000000.0
}
TIME_FACTORS = {
    "s": 1.0,
    "ms": 1000,
    "us": 1000000,
    "ns": 1000000000,
}


def parse_args():
    argp = ArgumentParser(description=__doc__)
    argp.add_argument("--vendor-titles", nargs="+",
                      default=["Memgraph", "Market leader"],
                      help="Vender titles that are going to appear "
                           "on the plot, e.g. legend titles.")
    argp.add_argument("--plot-title", default="",
                      help="Plot title.")
    argp.add_argument("--logo-path", default=None,
                      help="Path to the logo that is going to be presented"
                           " instead of title.")
    argp.add_argument("--results", nargs="+", required=True,
                      help="Path to the folder with result files in format "
                           "{{vendor-reference}}-LDBC-results.json")
    argp.add_argument("--time-unit", choices=("s", "ms", "us", "ns"),
                      default="ms", help="The time unit that should be used.")
    argp.add_argument("--output", default="",
                      help="Save plot to file (instead of displaying it).")
    argp.add_argument("--max-label-width", default=11, type=int,
                      help="Maximum length of the x-axis labels (-1 is unlimited)")
    return argp.parse_args()


def autolabel(ax, rects):
    """
    Attach a text label above each bar displaying its height
    """
    for rect in rects:
        height = rect.get_height()
        # TODO: adjust more vendors
        ax.text(rect.get_x() + rect.get_width()/2., 1.00*height,
                '%d' % int(height),
                ha='center', va='bottom')


def main():
    # Read the arguments.
    args = parse_args()

    # Prepare the datastructure.
    vendors = []
    for i, results_file, vendor_title in zip(range(len(args.results)),
                                             args.results,
                                             args.vendor_titles):
        vendor = {}
        vendor['title'] = vendor_title
        vendor['results_file'] = results_file
        vendor['color'] = COLORS[i]
        vendor['results'] = []
        vendors.append(vendor)
    assert len(vendors) == 2, "The graph is tailored for only 2 vendors."

    # Collect the benchmark data.
    print("LDBC Latency Data")
    for vendor in vendors:
        with open(vendor['results_file']) as results_file:
            results_data = json.load(results_file)
            for query_data in results_data["all_metrics"]:
                mean_runtime = (query_data["run_time"]["mean"] /
                        LDBC_TIME_FACTORS[results_data["unit"]] *
                        TIME_FACTORS[args.time_unit])
                query_name = query_data['name']
                vendor['results'].append((query_name, mean_runtime))

    # Sort results.
    for vendor in vendors:
        vendor['results'].sort(key=lambda item: int("".join(filter(
                lambda x: x in string.digits, item[0]))))

    # Print results.
    for vendor in vendors:
        print("Vendor:", vendor['title'])
        for query_name, latency in vendor['results']:
            print("{} -> {:.3f}{}".format(query_name, latency, args.time_unit))

    # Consistency check.
    all_query_names = [tuple(res[0] for res in vd['results']) for vd in vendors]
    assert len(set(all_query_names)) == 1, \
        "Queries between different vendors are different!"
    query_names = all_query_names[0]

    # Plot.
    ind = np.arange(len(query_names))   # the x locations for the groups
    width = 0.40                        # the width of the bars
    fig, ax = plt.subplots()            # figure setup
    fig.set_size_inches(1920 / 96, 1080 / 96) # set figure size
    ax.set_ylabel('Mean Latency (%s)' % (args.time_unit))  # YAxis title
    ax.set_facecolor('#dcdcdc')         # plot bg color (light gray)
    ax.set_xticks(ind + width / len(vendors))  # TODO: adjust (more vendors)

    def shorten_query_name(query_name):
        # Long query names on the x-axis don't look compelling.
        if query_name.lower().startswith('ldbc'):
            query_name = query_name[4:]
        if len(query_name) > args.max_label_width:
            query_name = query_name[:args.max_label_width] + '\N{HORIZONTAL ELLIPSIS}'
        return query_name
    labels = query_names
    if args.max_label_width == 0:
        labels = ["Q{}".format(i) for i, _ in enumerate(query_names)]
    elif args.max_label_width > 0:
        labels = map(shorten_query_name, query_names)
    ax.set_xticklabels(labels, rotation=30)
    # set only horizontal grid lines
    for line in ax.get_xgridlines():
        line.set_linestyle(' ')
    for line in ax.get_ygridlines():
        line.set_linestyle('--')
    ax.set_axisbelow(True)              # put the grid below all other elements
    plt.grid(True)                      # show grid
    # Set plot title
    ax.set_title(args.plot_title)
    # Draw logo or plot title
    if args.logo_path != None:
        # TODO: improve the logo positioning
        im = plt.imread(get_sample_data(os.path.join(os.getcwd(),
                                                     args.logo_path)))
        plt.gcf().subplots_adjust(top=0.85)
        newax = fig.add_axes([0.46, 0.85, 0.12, 0.15], anchor='N')
        newax.imshow(im)
        newax.axis('off')
    # Draw bars
    for index, vendor in enumerate(vendors):
        latencies = [res[1] for res in vendor['results']]
        rects = ax.bar(ind + index * width, latencies, width,
                       color=vendor['color'])
        vendor['rects'] = rects
        autolabel(ax, rects)
    rects = [vd['rects'][0] for vd in vendors]
    titles = [vd['title'] for vd in vendors]
    ax.legend(rects, titles)           # Draw the legend.

    if args.output == "":
        plt.show()
    else:
        plt.savefig(args.output, dpi=96)



if __name__ == '__main__':
    main()