DuckDB estimates Arrow scans at 1 row, and an MLP layer's aggregate swings 4–6× with the estimate — draft for upstream #293

import duckdb, numpy as np, pyarrow as pa
t = pa.table({"i": np.arange(10**6)})
print(duckdb.connect().sql("EXPLAIN SELECT * FROM t").fetchall()[0][1])
# ARROW_SCAN ... ~1 row

import sys
import time
import duckdb
import numpy as np
import pyarrow as pa

rng = np.random.default_rng(0)
ids, f, nets, inputs = 2048, 512, 64, 186
x = pa.table({"id": np.repeat(np.arange(ids), inputs), "f": np.tile(np.arange(inputs), ids),
              "val": rng.random(ids * inputs)})
# Many 512 x 512 matrices in one long-form table, as a checkpoint's weights.
weights = pa.table({
    "net": np.repeat(np.arange(nets), f * f),
    "f": np.tile(np.repeat(np.arange(f), f), nets),
    "g": np.tile(np.arange(f), f * nets),
    "w": rng.random(nets * f * f),
})
biases = pa.table({"net": np.repeat(np.arange(nets), f), "g": np.tile(np.arange(f), nets), "b": rng.random(nets * f)})
names = pa.table({"net": np.arange(nets), "name": [f"layer_{i}" for i in range(nets)]})

def w(name):
    return f"(SELECT w.f, w.g, w.w FROM nets n JOIN weights w ON w.net = n.net WHERE n.name = '{name}')"

QUERIES = {
    "one layer": f"SELECT x.id, w.g, SUM(x.val * w.w) AS val FROM x JOIN {w('layer_0')} w ON x.f = w.f GROUP BY x.id, w.g",
    "swish": f"""WITH p AS (
  SELECT x.id, w.g AS f, SUM(x.val * w.w) AS val FROM x JOIN {w('layer_0')} w ON x.f = w.f GROUP BY x.id, w.g),
h AS (
  SELECT p.id, p.f, (p.val + b.b) / (1 + exp(-(p.val + b.b))) AS val
  FROM p JOIN (SELECT b.g, b.b FROM nets n JOIN biases b ON b.net = n.net WHERE n.name = 'layer_0') b ON p.f = b.g)
SELECT h.id, w.g, SUM(h.val * w.w) AS val FROM h JOIN {w('layer_1')} w ON h.f = w.f GROUP BY h.id, w.g""",
    "two layers": f"""WITH h AS (
  SELECT x.id, w.g AS f, SUM(x.val * w.w) AS val FROM x JOIN {w('layer_0')} w ON x.f = w.f GROUP BY x.id, w.g)
SELECT h.id, w.g, SUM(h.val * w.w) AS val FROM h JOIN {w('layer_1')} w ON h.f = w.f GROUP BY h.id, w.g""",
}
print("duckdb", duckdb.__version__)
for qname, q in [(k, QUERIES[k]) for k in sys.argv[1:]] or QUERIES.items():
    for source in ("native", "arrow"):
        con = duckdb.connect()
        for name, t in {"x": x, "weights": weights, "biases": biases, "nets": names}.items():
            if source == "arrow" and name != "x":
                con.register(name, t)
            else:
                con.register(f"{name}_arrow", t)
                con.execute(f"CREATE TABLE {name} AS SELECT * FROM {name}_arrow")
        start = time.perf_counter()
        con.sql(f"SELECT count(*), sum(val) FROM ({q})").fetchall()
        print(f"{qname:10} weights {source:6}: {time.perf_counter() - start:5.1f} s", flush=True)