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

Ouverte
#293 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub

Les mainteneurs répondent en général sous 3 jours

Personne n'a encore pris cette issue.

Évaluation

Difficulté
4/5
Temps estimé
3-5 jours
Accessibilité débutants
18/100
Type d'issue
Bug
Clarté
Plutôt claire
Activité
Active
Stack technique
python, sql

Piste de recherche

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.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Description

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.

Langage dominant
Python
Étoiles
142
Forks
19
Merge moyen
3 j 47 min
PR mergées (30 j)
8

Préparer son environnement

Par où commencer

  1. Lisez l'issue en entier, puis le guide de contribution du projet.
  2. Signalez en commentaire que vous la prenez — cela évite que deux personnes fassent le même travail.
  3. Forkez le dépôt et travaillez sur une branche.
  4. Ouvrez une pull request qui référence le numéro de l'issue.

Autres issues de xqlsystems/xarray-sql

Toutes les issues de xqlsystems/xarray-sql

Issues similaires

Plus d'issues Python

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.