import collections


def load_run(path):
  """Loads run into a dict of key: query_id, value: list of candidate doc ids."""

  # We want to preserve the order of runs so we can pair the run file with the
  # TFRecord file.
  run = collections.OrderedDict()
  with open(path) as f:
    for i, line in enumerate(f):
      query_id, doc_title, rank = line.split('\t')
      if query_id not in run:
        run[query_id] = []
      run[query_id].append((doc_title, int(rank)))
      if i % 1000000 == 0:
        print('Loading run {}'.format(i))
  # Sort candidate docs by rank.
  sorted_run = collections.OrderedDict()
  for query_id, doc_titles_ranks in run.items():
    sorted(doc_titles_ranks, key=lambda x: x[1])
    doc_titles = [doc_titles for doc_titles, _ in doc_titles_ranks]
    sorted_run[query_id] = doc_titles

  return sorted_run