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

未关闭
#293 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 3 天内回复

还没有人认领这个 Issue。

评估

难度
4/5
预计耗时
3-5 天
新手友好度
18/100
Issue 类型
缺陷
描述清晰度
基本清楚
活跃度
活跃
技术栈
python, sql

调研方向

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.

由索引模型根据 Issue 内容生成。

描述

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.

主要语言
Python
星标
142
派生
19
平均合并
3 天 47 分钟
30 天内合并 PR
8

环境准备

  • 没有 Dockerfile 或 Docker Compose 文件
  • 没有 Pull Request 模板
  • 阅读贡献指南

从这里开始

  1. 先读完整个 Issue,再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库,在一个分支上完成修改。
  4. 提交 Pull Request,并在描述里引用这个 Issue 编号。

xqlsystems/xarray-sql 的其他 Issue

查看 xqlsystems/xarray-sql 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。