DuckDB

"""DuckDB: the BookNest benchmark. Usage: bench_duckdb.py T1|T2|T3"""
import sys, time
import duckdb
D = "/home/dev/v7-l3/ch05/data"
LINES = "order_lines_x10" if sys.argv[1] == "T3" else "order_lines"   # T3: 10x rows
OUT = f"/home/dev/v7-l3/bench/out/duckdb_{sys.argv[1]}.parquet"
t0 = time.perf_counter()
con = duckdb.connect()
con.execute("SET TimeZone = 'UTC'")
if sys.argv[1] in ("T1", "T3"):
    sql = f"""
    SELECT date_trunc('month', l.order_ts) AS month, b.genre, c.country,
           sum(l.qty * l.unit_price::DOUBLE) AS revenue, count(*) AS lines
    FROM '{D}/{LINES}.parquet/*.parquet' l
    JOIN '{D}/books.parquet/*.parquet' b ON l.book_id = b.id
    JOIN '{D}/customers.parquet/*.parquet' c USING (customer_id)
    WHERE l.status = 'delivered'
    GROUP BY ALL ORDER BY ALL"""
else:
    sql = f"""
    SELECT substr(order_ts, 1, 7) AS month, sum(i.qty * i.unit_price) AS revenue
    FROM (SELECT order_ts, unnest(items) AS i
          FROM read_json('{D}/raw/orders.jsonl', format = 'newline_delimited', columns = {{
            order_ts: 'VARCHAR', status: 'VARCHAR',
            items: 'STRUCT(book_id INT, qty INT, unit_price DOUBLE)[]'}})
          WHERE status = 'delivered')
    GROUP BY ALL ORDER BY ALL"""
con.execute(f"COPY ({sql}) TO '{OUT}' (FORMAT parquet)")
rows, revenue = con.execute(f"SELECT count(*), sum(revenue) FROM '{OUT}'").fetchone()
print(f"duckdb {sys.argv[1]} rows={rows} revenue={revenue:.2f} "
      f"query={time.perf_counter() - t0:.2f}s")