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

Offen
  • Stimmt etwas nicht?
#293 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen

Maintainer antworten meist innerhalb von 3 Tagen

Dieses Issue hat noch niemand übernommen.

Bewertung

Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Anfängerfreundlichkeit
18/100
Issue-Typ
Bug
Klarheit
Größtenteils klar
Aktivitätsstatus
Aktiv
Tech-Stack
python, sql
Bereich
databases, performance

Rechercherichtung

The payload is a draft upstream DuckDB report with a repro script; the local part is xql.register in xarray-sql, which should pass a known row count to DuckDB's Arrow scan if the API allows it. Start by checking whether DuckDB's Python Arrow scan accepts a cardinality, then run python arrow_estimate_repro2.py swish to see the 16.1 s vs 3.7 s split. Done means the estimate is set and the timings are rechecked.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Beschreibung

ml

🌀 Claude: DuckDB estimates every Arrow scan at about 1 row, whatever its size. This includes every table xql.register gives it. DuckDB's time for a matrix product written as a join and a GROUP BY then swings 4–6× with that estimate, in either direction. This is a draft for an upstream DuckDB issue.

Arrow scans are estimated at 1 row

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

In DuckDB 1.5.6 it's the same for a pyarrow.dataset.Dataset, a RecordBatchReader and xql.register(con, "t", ds), and a pyarrow.Table is estimated at 1 row in 1.4.4 and 1.5.5 too. A CREATE TABLE copy of the same rows is estimated exactly.

An MLP layer's aggregate depends on the estimate

This is two linear layers with a bias and swish between them (swish is x · sigmoid(x)), on 2,048 rows × 186 features, with 512 × 512 weights. The weights come from one long-form table of 64 matrices, selected through a join with their names, as a model checkpoint's are. The script needs duckdb, numpy and pyarrow:

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)

With python arrow_estimate_repro2.py swish, the weights either as native tables or registered from Arrow (DuckDB 1.5.6, Intel i7-8700, 6 cores):

Weights Wall time Second layer's HASH_GROUP_BY (537M rows into 1M groups) Its estimated groups
native tables 16.1 s 80 s of CPU 177,289
Arrow 3.7 s 14 s of CPU 0

In GraphCast's grid embedding (#270), with the same shape and DeepMind's weights, it goes the other way. The 537M-row aggregate takes 64 s of CPU with native weight tables, estimated at 483,664 groups, and 187 s through xql.register, estimated at 0. Registering the weights as a plain pyarrow.Table is as slow as xql.register (46.7 s wall against 12.3 s native), so this is DuckDB's Arrow scan and not xarray-sql's lazy reads. With a 5–6 GB memory_limit, the whole embedding with its layer norm runs out of memory through xql.register, and filled 9.6 GB of disk spilling before it failed; on native tables it completes in 163 s. The process's peak RSS was 7–8 GB, above DuckDB's 5 GB memory_limit, in both cases.

So DuckDB's hash aggregate, at these sizes, is sensitive to its estimate of the number of groups, and an Arrow source gives it no estimate.

What to do

  • Upstream:
    • Search DuckDB's issues for Arrow scan cardinality, and for a hash aggregate's sensitivity to its estimate.
    • File this script.
    • The aggregate's time shouldn't depend on its estimate this much, and an Arrow source with a known length could report it.
  • Here: check whether DuckDB's Python Arrow scan can take a row count, which xarray-sql knows exactly from the Dataset's shape. If it can, xql.register should pass it.

Measured with DuckDB 1.5.6, pyarrow 23.0.0 and xarray 2026.2.0.

Vorherrschende Sprache
Python
Sterne
142
Forks
19
Ø Merge
3 T. 47 Min.
Gemergte PRs (30 T.)
8

Entwicklungsumgebung

Erste Schritte

  1. Lesen Sie das ganze Issue und danach den Beitragsleitfaden des Projekts.
  2. Schreiben Sie ins Issue, dass Sie es übernehmen — das erspart doppelte Arbeit.
  3. Forken Sie das Repository und arbeiten Sie in einem Branch.
  4. Öffnen Sie einen Pull Request, der die Issue-Nummer nennt.

Mehr aus xqlsystems/xarray-sql

Alle Issues in xqlsystems/xarray-sql

Ähnliche Issues

Weitere Issues zu Python

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.