def bench()

in benchmarks/tpch/tpch.py [0:0]

• 48 lines of code
• 9 McCabe index (conditional complexity)

def bench(data_path, query_path):
    with open("results.csv", "w") as results:
        # register tables
        start = time.time()
        total_time_millis = 0

        # create context
        # runtime = (
        #     RuntimeConfig()
        #     .with_disk_manager_os()
        #     .with_fair_spill_pool(10000000)
        # )
        # config = (
        #     SessionConfig()
        #     .with_create_default_catalog_and_schema(True)
        #     .with_default_catalog_and_schema("datafusion", "tpch")
        #     .with_information_schema(True)
        # )
        # ctx = SessionContext(config, runtime)

        ctx = SessionContext()
        print("Configuration:\n", ctx)

        # register tables
        with open("create_tables.sql") as f:
            sql = ""
            for line in f.readlines():
                if line.startswith("--"):
                    continue
                sql = sql + line
                if sql.strip().endswith(";"):
                    sql = sql.strip().replace("$PATH", data_path)
                    ctx.sql(sql)
                    sql = ""

        end = time.time()
        time_millis = (end - start) * 1000
        total_time_millis += time_millis
        print("setup,{}".format(round(time_millis, 1)))
        results.write("setup,{}\n".format(round(time_millis, 1)))
        results.flush()

        # run queries
        for query in range(1, 23):
            with open("{}/q{}.sql".format(query_path, query)) as f:
                text = f.read()
                tmp = text.split(";")
                queries = []
                for str in tmp:
                    if len(str.strip()) > 0:
                        queries.append(str.strip())

                try:
                    start = time.time()
                    for sql in queries:
                        print(sql)
                        df = ctx.sql(sql)
                        # result_set = df.collect()
                        df.show()
                    end = time.time()
                    time_millis = (end - start) * 1000
                    total_time_millis += time_millis
                    print("q{},{}".format(query, round(time_millis, 1)))
                    results.write(
                        "q{},{}\n".format(query, round(time_millis, 1))
                    )
                    results.flush()
                except Exception as e:
                    print("query", query, "failed", e)

        print("total,{}".format(round(total_time_millis, 1)))
        results.write("total,{}\n".format(round(total_time_millis, 1)))