# xarray-sql：用SQL查询Xarray数据集的实验性项目

- 来源：Hacker News 热门（buzzing.cc 中文翻译）
- 作者：alxmrs
- 发布时间：2026-07-14 19:51
- AIHOT 分数：67
- AIHOT 链接：https://aihot.virxact.com/items/cmrklgp8s0068bit8cqo3kywo
- 原文链接：https://github.com/xqlsystems/xarray-sql/blob/claude/xarray-sql-mnist-demo/benchmarks/nn.py

## AI 摘要

xarray-sql是一个开源实验项目，允许用户通过SQL查询Xarray数据集。该项目托管在GitHub上，旨在将SQL的便捷查询能力与Xarray的多维数组数据处理相结合，降低非Python用户访问科学数据的门槛。项目目前处于早期开发阶段，代码库已公开。

## 正文

/// script

requires-python = ">=3.12"

dependencies = [

"xarray_sql",

"xarray",

"numpy",

"s3fs",

"zarr<3",

]

[tool.uv.sources]

xarray_sql = { path = "..", editable = true }

///

from future import annotations

from typing import Callable

import numpy as np

import xarray as xr

import datetime

import xarray_sql as xql

SIDE=28# images are 28x28; flatten index is height * SIDE + width

WIDTHS= (

SIDE*SIDE,

) # 784 pixels -> 196 -> 32 tanh -> 10 softmax

N_SAMPLES, TRAIN_FRAC=700, 0.7# total samples; fraction used for training

LR, STEPS, CHUNK=0.5, 60, 250

Drop zero-valued pixels from the (dominant) layer-0 contraction. A background

pixel contributes 0 * weight = 0, so skipping those rows shrinks the join

exactly — the result is identical, and the speedup scales with the fraction

of zeros (a dark background). On dense inputs it is a no-op.

Measured 1.8x on real Fashion-MNIST (50% zero pixels): 2.56 -> 1.45 s/step.

SKIP_ZERO_PIXELS=True

def fashion_mnist():

"""The whole training set, left lazy so SQL streams and samples it.

The real path returns a dask-backed (chunked) Dataset — nothing is pulled

into memory here; from_dataset reads it chunk by chunk on demand, and

the random subsample happens later in SQL. The offline fallback is a small

synthetic set built in memory.

try:

ds=xr.open_dataset(

"s3://carbonplan-share/xbatcher/fashion-mnist-train.zarr",

engine="zarr",

chunks=None,

backend_kwargs={"storage_options": {"anon": True}},

if"channel"in ds.dims:

ds=ds.isel(channel=0, drop=True)

To float64, lazily (no full read). This zarr already stores images

as float in [0, 1]; only integer-encoded sources ([0, 255]) rescale.

images=ds["images"].astype("float64")

if not np.issubdtype(ds["images"].dtype, np.floating):

images=images/255.0

ds=ds.assign(images=images, labels=ds["labels"].astype("int64"))

except Exception:

Offline fallback: a separable synthetic set (per-class template +

noise), so the same pipeline still learns without the network. A pool

larger than N_SAMPLES so the SQL subsample still has something to pick.

rng=np.random.default_rng(0)

n=3*N_SAMPLES

templates=rng.standard_normal((10, SIDE, SIDE))

labels=rng.integers(0, 10, n).astype("int64")

images=templates[labels] +0.6*rng.standard_normal((n, SIDE, SIDE))

ds=xr.Dataset(

"images": (("sample", "height", "width"), images),

"labels": (("sample",), labels),

Integer index coords are the SQL join keys (sample, height, width).

return ds[["images", "labels"]].assign_coords(

sample=np.arange(ds.sizes["sample"]),

height=np.arange(ds.sizes["height"]),

width=np.arange(ds.sizes["width"]),

def build_model_with_table_names(

init_weight: Callable[[int, int], np.ndarray],

init_bias: Callable[[int], np.ndarray],

widths=WIDTHS,

) ->tuple[xr.Dataset, dict[tuple[str, ...], str]]:

"""The network as one Dataset that splits into tables per layer.

Layer i is a weight matrix layer_i (inp_i, out_i) and a separate

bias vector bias_i (out_i,).

weights= {

f"layer_{i}": ((f"inp_{i}", f"out_{i}"), init_weight(inp, out))

for i, (inp, out) in enumerate(zip(widths[:-1], widths[1:]))

biases= {

f"bias_{i}": ((f"out_{i}",), init_bias(out))

for i, out in enumerate(widths[1:])

coords= {}

coords.update(

{f"inp_{i}": np.arange(inp) for i, inp in enumerate(widths[:-1])}

coords.update(

{f"out_{i}": np.arange(out) for i, out in enumerate(widths[1:])}

ds=xr.Dataset({**weights, **biases}, coords=coords)

names: dict[tuple[str, ...], str] = {}

for i in range(len(weights)):

names[(f"inp_{i}", f"out_{i}")] =f"layer{i}"

names[(f"out_{i}",)] =f"bias{i}"

return ds, names

def main():

rng=np.random.default_rng(1)

mnist=fashion_mnist()

ctx=xql.XarrayContext()

One Dataset splits into two tables: pixels (sample, height, width) and

labels (sample). The dim names are the join keys.

ctx.from_dataset(

"mnist",

mnist,

chunks=dict(sample=CHUNK),

table_names={

("sample", "height", "width"): "pixels",

("sample",): "labels",

Draw a random N_SAMPLES subset in SQL (ORDER BY random() LIMIT), carrying

each sample's label and a train/test tag. data is the working label

table: cache() pins the chosen subset so every downstream query sees the

same split without rescanning the source. ORDER BY random() shuffles the

whole label column, so the subset is order-independent even if the on-disk

samples are class-sorted.

data=ctx.sql(f"""

SELECT sample, labels,

CASE WHEN random() <{TRAIN_FRAC} THEN 'train' ELSE 'test' END AS split

FROM mnist.labels

ORDER BY random()

LIMIT {N_SAMPLES}

""").cache()

ctx.register_table("data", data)

Materialise just the sampled images once: a single lazy scan of the full

dataset extracts the ~N_SAMPLES subset into pixels, which the per-step

forward joins instead of rescanning the source 60x. Only the subset lives

in memory; the full set stays lazy.

pixels=ctx.sql("""

SELECT p.sample, p.height, p.width, p.images

FROM mnist.pixels p JOIN data d ON p.sample = d.sample

""").cache()

ctx.register_table("pixels", pixels)

FROM bias b LEFT JOIN gb g

ON b.layer = g.layer AND b.out = g.out

""").cache()

ctx.deregister_table("bias")

ctx.register_table("bias", b)

if step%5==0 or step==STEPS-1:

Train cross-entropy (logits span all samples, so filter to train).

loss=ctx.sql(f"""

WITH m AS (SELECT sample, MAX(z) AS m FROM logits GROUP BY sample),

e AS (SELECT logits.sample, logits.out, exp(logits.z - m.m) AS e

FROM logits JOIN m ON logits.sample = m.sample),

s AS (SELECT sample, SUM(e) AS s FROM e GROUP BY sample)

SELECT -AVG(ln(e.e / s.s)) AS loss

FROM e JOIN s ON e.sample = s.sample

JOIN data y ON y.sample = e.sample

WHERE e.out = y.labels

AND e.sample IN (SELECT sample FROM data WHERE split = 'train')

""").to_pandas()["loss"][0]

Accuracy per split: argmax the shared logits, join the split label.

Both come from the one all-samples forward — no second pass.

acc= (

ctx.sql(f"""

WITH pred AS (

SELECT sample, out,

ROW_NUMBER() OVER (PARTITION BY sample ORDER BY z DESC) AS rk

FROM logits)

SELECT d.split,

AVG(CASE WHEN p.out = d.labels THEN 1.0 ELSE 0.0 END) AS acc

FROM pred p JOIN data d ON d.sample = p.sample

WHERE p.rk = 1

GROUP BY d.split

.to_pandas()

.set_index("split")["acc"]

print(

f"step {step:2d}: loss {loss:.3f} "

f"train_acc {acc['train']:.3f} test_acc {acc['test']:.3f}"

The trained parameters come back out as xarray in the *same shape as the

input model*: one weight variable per layer with its own (inp_i, out_i)

dims, plus one bias variable per layer on (out_i,). Each is read from its

relation by the layer column, so the result is a ragged set of per-layer

matrices and vectors — no dense array padded with NaN.

trained=xr.Dataset(

f"layer_{i}": ctx.sql(

f"SELECT inp AS inp_{i}, out AS out_{i}, val AS layer_{i} "

f"FROM weight WHERE layer = {i}"

).to_dataset(dims=[f"inp_{i}", f"out_{i}"])[f"layer_{i}"]

for i in range(len(WIDTHS) -1)

f"bias_{i}": ctx.sql(

f"SELECT out AS out_{i}, val AS bias_{i} "

f"FROM bias WHERE layer = {i}"

).to_dataset(dims=[f"out_{i}"])[f"bias_{i}"]

for i in range(len(WIDTHS) -1)

print(f"trained {WIDTHS} MLP; weights -> xarray {dict(trained.sizes)}.")

print(trained)

trained.to_zarr(

f"fashion_mnist_mlp_"

f"{datetime.datetime.now().isoformat(timespec='seconds')}.zarr"

if name =="main":

main()

You can’t perform that action at this time.