ClickHouse/dbms/scripts/merge_algorithm/drawer.py

from __future__ import print_function

import argparse
import matplotlib.pyplot as plt
import ast

TMP_FILE='tmp.tsv'

def parse_args():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-f', '--file', default='data.tsv')
    cfg = parser.parse_args()
    return cfg

def draw():
    place = dict()
    max_coord = 0
    global_top = 0
    for line in open(TMP_FILE):
        numbers = line.split('\t')
        if len(numbers) <= 2:
            continue
        name = numbers[-2]
        if numbers[0] == '1':
            dx = int(numbers[3])
            max_coord += dx
            place[name] = [1, max_coord, 1, dx]
            max_coord += dx
            plt.plot([max_coord - 2 * dx, max_coord], [1, 1])
    for line in open(TMP_FILE):
        numbers = line.split('\t')
        if len(numbers) <= 2:
            continue
        name = numbers[-2]
        if numbers[0] == '2':
            list = ast.literal_eval(numbers[-1])
            coord = [0,0,0,0]
            for cur_name in list:
                coord[0] = max(place[cur_name][0], coord[0])
                coord[1] += place[cur_name][1] * place[cur_name][2]
                coord[2] += place[cur_name][2]
                coord[3] += place[cur_name][3]
            coord[1] /= coord[2]
            coord[0] += 1
            global_top = max(global_top, coord[0])
            place[name] = coord
            for cur_name in list:
                plt.plot([coord[1], place[cur_name][1]],[coord[0], place[cur_name][0]])
            plt.plot([coord[1] - coord[3], coord[1] + coord[3]], [coord[0], coord[0]])
    plt.plot([0], [global_top + 1])
    plt.plot([0], [-1])
    plt.show()


def convert(input_file):
    print(input_file)
    tmp_file = open(TMP_FILE, "w")
    for line in open(input_file):
        numbers = line.split('\t')
        numbers2 = numbers[-2].split('_')
        if numbers2[-2] == numbers2[-3]:
            numbers2[-2] = str(int(numbers2[-2]) + 1)
            numbers2[-3] = str(int(numbers2[-3]) + 1)
            numbers[-2] = '_'.join(numbers2[1:])
            print('\t'.join(numbers), end='', file=tmp_file)
        else:
            print(line, end='', file=tmp_file)

def main():
    cfg = parse_args()
    convert(cfg.file)
    draw()

if __name__ == '__main__':
    main()