DuckDB backend: plan each step once, and skip catalog sync when nothing changed (0.87 s per GPT step, most of it repeated planning) #216

"""Where a DuckDB training step of xarray-sql's GPT demo spends its time.

Run from an xarray-sql checkout of the ml-demos branch, with ddxdb[duckdb]
0.4.3 installed."""
import importlib.util, sys, time
from collections import defaultdict
import numpy as np
import ddxdb.duckdb as dd

sys.path.insert(0, "benchmarks/ml")
spec = importlib.util.spec_from_file_location("gpt", "benchmarks/ml/02_nanogpt.py")
gpt = importlib.util.module_from_spec(spec); spec.loader.exec_module(gpt)
from _hosts import DuckDB

T = defaultdict(float)

def timed(name, f):
    def g(*a, **k):
        t = time.perf_counter()
        try:
            return f(*a, **k)
        finally:
            T[name] += time.perf_counter() - t
    return g

def materialize(self, name, plan):
    # DuckDBBackend.materialize, its statements timed one by one.
    query = dd.to_sql(plan, "duckdb")
    t = time.perf_counter(); self.con.execute(f"EXPLAIN {query}").fetchall()
    T["  of which planning (EXPLAIN of the same SQL)"] += time.perf_counter() - t
    t = time.perf_counter(); self.con.execute(f"CREATE OR REPLACE TEMP TABLE {dd._ident(name)} AS {query}")
    T["materialize: CREATE TEMP TABLE AS"] += time.perf_counter() - t
    t = time.perf_counter()
    columns = self.con.execute(f"DESCRIBE temp.main.{dd._ident(name)}").fetchall()
    cols = ", ".join(f"{dd._ident(c[0])} {c[1]}" for c in columns)
    self._planner.execute(f"CREATE OR REPLACE TEMP TABLE {dd._ident(name)} ({cols})")
    self._written.add(name)
    T["materialize: DESCRIBE, and the planner's copy"] += time.perf_counter() - t

dd.DuckDBBackend.materialize = materialize
for m in ("_sync", "plan_sql", "plan_key", "view_sql", "select_all", "returns_rows", "drop_table"):
    setattr(dd.DuckDBBackend, m, timed(m, getattr(dd.DuckDBBackend, m)))

text = gpt.shakespeare()
init = gpt.init_params(len(set(text)), np.random.default_rng(0))
h = DuckDB()
for n, v in init.items():
    h.load(n, gpt.as_table(n, v.copy()))
h.load("batch", next(gpt.batches(text, 4, np.random.default_rng(1))))
h.execute(gpt.model())
statements = gpt.step_statements(gpt.LEARNING_RATE)
h.query_all(statements)  # ddx builds its program
T.clear()
t = time.perf_counter(); h.query_all(statements); total = time.perf_counter() - t
print(f"step: {total * 1e3:.0f} ms, of which the EXPLAINs added {T['  of which planning (EXPLAIN of the same SQL)'] * 1e3:.0f} ms")
for k, v in sorted(T.items(), key=lambda kv: -kv[1]):
    print(f"  {k:<48} {v * 1e3:6.0f} ms")