#!/usr/bin/env python

import json
import sys

def read_json(filename):
    with open(filename, 'r') as f:
        return json.load(f)

def uniquify_fragments(query_plan):
    fragment_inv = []
    for worker in sorted(query_plan.keys()):
        worker_plan = query_plan[worker]
        for fragment in worker_plan:
            flag = False
            for (i,(x,y)) in enumerate(fragment_inv):
                if y == fragment:
                    fragment_inv[i] = (x + [worker], y)
                    flag = True
                    break
            if flag:
                continue
            fragment_inv.append(([worker], fragment))
    return fragment_inv

def json_pretty(obj):
    return json.dumps(obj, sort_keys=True, indent=4, separators=(',', ': '))

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print >> sys.stderr, "Usage: %s <old json file>" % sys.argv[0]
        sys.exit(1)

    myria_json_plan = read_json(sys.argv[1])
    fragments = []
    frags = uniquify_fragments(myria_json_plan['query_plan'])
    for (ws,ops) in frags:
        fragments.append({
            'workers' : ws,
            'operators' : ops
        })
    output = {
            'raw_datalog' : myria_json_plan['raw_datalog'],
            'logical_ra' : myria_json_plan['logical_ra'],
            'fragments' : fragments
    }

    print json_pretty(output)

@@ -20,30 +20,572 @@
 on%0A%0A
-import json%0Aimport sys
+from collections import defaultdict%0Aimport json%0Aimport sys%0A%0A# Fields to delete for certain operators%0A# .. these are fields that are automatically inferred now.%0Adelete_fields = defaultdict(list)%0Adelete_fields%5B'Consumer'%5D = %5B'arg_worker_ids'%5D%0Adelete_fields%5B'CollectConsumer'%5D = %5B'arg_worker_ids'%5D%0Adelete_fields%5B'CollectProducer'%5D = %5B'arg_worker_id'%5D%0Adelete_fields%5B'EOSController'%5D = %5B'arg_worker_ids'%5D%0Adelete_fields%5B'IDBInput'%5D = %5B'arg_controller_worker_id'%5D%0Adelete_fields%5B'ShuffleConsumer'%5D = %5B'arg_worker_ids'%5D%0Adelete_fields%5B'ShuffleProducer'%5D = %5B'arg_worker_ids'%5D
 %0A%0Ade
@@ -1204,16 +1204,257 @@
 nt_inv%0A%0A
+def clean_up(ops):%0A    def _clean_op(op):%0A        for field in delete_fields%5Bop%5B'op_type'%5D%5D:%0A            try:%0A                del op%5Bfield%5D%0A            except:%0A                pass%0A        return op%0A    return %5B_clean_op(op) for op in ops%5D%0A%0A
 def json
@@ -1925,19 +1925,29 @@
 tors' : 
+clean_up(
 ops
+)
 %0A       

delete worker fields where appropriate