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

Aberta
#293 0 comentários 0 reações 0 responsáveis Ver no GitHub

Mantenedores costumam responder em até 3 dias

Ninguém assumiu esta issue ainda.

Avaliação

Dificuldade
4/5
Tempo estimado
3-5 dias
Facilidade para iniciantes
18/100
Tipo de issue
Bug
Clareza
Razoavelmente clara
Status de atividade
Ativa
Stack de tecnologia
python, sql

Direção de pesquisa

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.

Escrita pelo modelo de indexação a partir do texto da issue.

Descrição

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.

Linguagem predominante
Python
Estrelas
142
Forks
19
Merge médio
3d 47min
PRs com merge (30d)
8

Preparar o ambiente

Primeiros passos

  1. Leia a issue inteira e depois o guia de contribuição do projeto.
  2. Comente na issue dizendo que vai assumir — evita que duas pessoas façam o mesmo trabalho.
  3. Faça um fork do repositório e trabalhe em uma branch.
  4. Abra um pull request que referencie o número da issue.

Mais de xqlsystems/xarray-sql

Todas as issues de xqlsystems/xarray-sql

Issues semelhantes

Mais issues de Python

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.