index

#!/usr/bin/env python3

import argparse
import json
import subprocess
import sys
import time
from datetime import timedelta

def robust_parse_command(name, path, stemmer):
    return (("zcat "
             "$(find {1} -type f -name '*.*z' "
             "\( -path '*/disk4/fr94/[0-9]*/*' "
             "-o -path '*/disk4/ft/ft*' "
             "-o -path '*/disk5/fbis/fb*' "
             "-o -path '*/disk5/latimes/la*' \))"
             " | parse_collection -f trectext "
             "-b 10000 --stemmer {2} --content-parser html -o ./{0}.fwd")
            .format(name, path, stemmer))

def nyt_parse_command(name, path, stemmer):
    return (("/root/miniconda3/bin/python /nyt-corpus-reader/read.py "
             "<(find {1}/data -type f)"
             " | parse_collection -f plaintext "
             "-b 10000 --stemmer {2} --content-parser html -o ./{0}.fwd")
            .format(name, path, stemmer))

def wapo_parse_command(name, path, stemmer):
    return (("cat {1}/data/*"
             " | parse_collection -f wapo "
             "-b 10000 --stemmer {2} --content-parser html -o ./{0}.fwd")
            .format(name, path, stemmer))

def cw09_parse_command(name, path, stemmer):
    return (("zcat "
             "$(find {1} -type f -name '*.gz')"
             " | parse_collection -f warc "
             "-b 10000 --stemmer {2} --content-parser html -o ./{0}.fwd")
            .format(name, path, stemmer))

def cw12_parse_command(name, path, stemmer):
    return (("zcat -r {1}/*"
             " | parse_collection -f warc "
             "-b 10000 --stemmer {2} --content-parser html -o ./{0}.fwd")
            .format(name, path, stemmer))

def gov2_parse_command(name, path, stemmer):
    return (("zcat "
             "$(find {1} -type f -name '*.gz')"
             " | parse_collection -f trecweb "
             "-b 10000 --stemmer {2} --content-parser html -o ./{0}.fwd")
            .format(name, path, stemmer))

def parse_command(name, path, stemmer):
    if name.lower() == "robust04":
        return robust_parse_command(name, path, stemmer)
    elif name.lower() == "core17":
        return nyt_parse_command(name, path, stemmer)
    elif name.lower() == "core18":
        return wapo_parse_command(name, path, stemmer)
    elif name.lower() == "cw09b":
        return cw09_parse_command(name, path, stemmer)
    elif name.lower() == "cw12b":
        return cw12_parse_command(name, path, stemmer)
    elif name.lower() == "gov2":
        return gov2_parse_command(name, path, stemmer)
    else:
        print("Unknown collection: " + name)
        sys.exit(1)

# Lambda is the VBMW block param.
# Searched offline, these values will give a mean block size of ~40 elements/block
# Depends on whether index was reordered or not
def get_lambda(name, reordered):
    if reordered:
        if name.lower() == "robust04":
            return 16.5
        elif name.lower() == "core17":
            return 15
        elif name.lower() == "core18":
            return 13.5
        elif name.lower() == "gov2":
            return 17
        elif name.lower() == "cw09b":
            return 22.5
        elif name.lower() == "cw12b":
            return 25.5
        else:
            print("Lambda unknown for bisection ordering and: " + name)
            print("Using default lambda = 11")
            return 11
    else:
        if name.lower() == "cw12b":
            return 26
        else:
            print("Lambda unknown for default order and: " + name)
            print("Using default lambda = 11")
            return 11    
    
# Allow user to skip reordering step
def get_skip_reordering(options):
    if "skip_reordering" in options:
        return True
    else:
        return False
    
# Allow user to input stemmer, validate
def get_stemmer(options):
    valid = ["porter2", "krovetz"]
    stemmer = "porter2"
    if "stemmer" in options and options["stemmer"].lower() in valid:
        stemmer = options["stemmer"].lower()
    else:
        print("Using default stemmer: " + stemmer)
    return stemmer

# Allow user to input many compressions approaches at once
def get_compressor_list(options):
    valid = ["opt", "block_interpolative", "block_simdbp", "block_optpfor", 
             "block_varintg8iu"] 
    if "compressor" in options:
        compressors = options["compressor"].lower()
        compressor_list = compressors.split(",")
        for comp in compressor_list:
            if not comp in valid:
                print("Invalid compressor: " + comp)
                sys.exit(1)
        print("Found list of compressors: " + str(compressor_list))
        return compressor_list
    return ["block_simdbp"]

# Allow user to select static blocks, validate
def get_block_type(options):
    valid = ["fixed", "variable"]
    block_type = "variable"
    block_size = 0
    if "block_type" in options and options["block_type"].lower() in valid:
        block_type = options["block_type"].lower()
    # Ensure block size is passed in
    if block_type == "fixed" and "block_size" in options:
        try:
            block_size = int(options["block_size"])
        except ValueError:
            print("Block size must be an integer.")
            sys.exit(1)
    if block_size < 0:
        print("Block size must be a positive integer")
        sys.exit(1)
    return block_type, block_size

# Decides which method to use for generating the wand data
def wand_command(name, wand_type, block_size, reordered):
    if wand_type == "variable":
        print("Making variable-block wand data, lambda = " + str(get_lambda(name, reordered)))
        return (("create_wand_data "
                           "-c {0}.inv.bp "
                           "-l {1} "
                           "--variable-block "
                           "-o {0}.vbmw").format(name, get_lambda(name, reordered))) 

    elif wand_type == "fixed":
        print("Making fixed-block wand data, block size = " + str(block_size))
        return (("create_wand_data "
                           "-c {0}.inv.bp "
                           "-b {1} "
                           "-o {0}.bmw").format(name, block_size))
    else:
      print("Invalid wand type. Should be 'fixed' or 'variable'")
      sys.exit(1)


parser = argparse.ArgumentParser()
parser.add_argument("--json", type=json.loads,
                    required=True, help="the json input")

args, _ = parser.parse_known_args()

options = args.json["opts"]
skip_reordering = get_skip_reordering(options)
do_reordering = not skip_reordering
stemmer = get_stemmer(options)
compressor_list = get_compressor_list(options)
wand_type, block_size = get_block_type(options)

# Iterate over the [name]=[path] pairs
for collection in args.json["collections"]:
    name, path = collection["name"], collection["path"]

    start_time = time.time()
    subprocess.check_call(parse_command(name, path, stemmer),
                          shell=True, executable='/bin/bash')
    parse_duration = timedelta(seconds=time.time() - start_time)
    print("--- Parsing took {} ---".format(parse_duration))
    start_time = time.time()
    subprocess.check_call(("invert "
                           "-i {0}.fwd "
                           "-o {0}.inv "
                           "-b 400000 "
                           "--term-count $(cat {0}.fwd.terms | wc -l)")
                          .format(name),
                          shell=True)
    invert_duration = timedelta(seconds=time.time() - start_time)
    print("--- Inverting took {} ---".format(invert_duration))
    bp_duration = timedelta(seconds=0)
    if do_reordering:
        start_time = time.time()
        subprocess.check_call(("recursive_graph_bisection "
                               "-c {0}.inv "
                               "-o {0}.inv.bp "
                               "--documents {0}.fwd.doclex "
                               "--reordered-documents {0}.fwd.bp.doclex "
                               "-m 4096").format(name),
                               shell=True)
        bp_duration = timedelta(seconds=time.time() - start_time)
        print("--- Reordering took {} ---".format(bp_duration))
    else:
        print("--- Skipping the reordering phase ----")
        # Link the index files so the rest of the script continues as usual
        subprocess.check_call(("ln -s {0}.inv.docs {0}.inv.bp.docs").format(name), shell=True)
        subprocess.check_call(("ln -s {0}.inv.freqs {0}.inv.bp.freqs").format(name), shell=True)
        subprocess.check_call(("ln -s {0}.inv.sizes {0}.inv.bp.sizes").format(name), shell=True)
        subprocess.check_call(("ln -s {0}.fwd.doclex {0}.fwd.bp.doclex").format(name), shell=True)

    # Generate an index for each provided coded
    start_time = time.time()
    for compressor in compressor_list:
        per_compressor_start_time = time.time()
        subprocess.check_call(("create_freq_index "
                               "-t {1} "
                               "-c {0}.inv.bp "
                               "-o {0}.{1}").format(name, compressor),
                              shell=True)
        compress_duration = timedelta(seconds=time.time() - per_compressor_start_time)
        print("--- Compressing the {0} index took {1} ---".format(compressor, compress_duration))
    total_compress_duration = timedelta(seconds=time.time() - start_time)
    start_time = time.time()
    subprocess.check_call(wand_command(name, wand_type, block_size, do_reordering),
                          shell=True)
    wand_duration = timedelta(seconds=time.time() - start_time)
    print("--- Extracting block data took {} ---".format(wand_duration))
    print("=== Summary ===")
    print("Parse:\t{}".format(parse_duration))
    print("Invert:\t{}".format(invert_duration))
    print("Reorder:\t{}".format(bp_duration))
    print("Compress:\t{}".format(total_compress_duration))
    print("Block data:\t{}".format(wand_duration))
    print("===============")
    print("Total:\t{}".format(
        parse_duration + invert_duration + bp_duration + compress_duration + wand_duration))