The scan from Zarr to rows costs as much as a whole dense data loader

Aperta
#311 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 3 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
5/5
Tempo stimato
Più di una settimana
Idoneità per principianti
35/100
Tipo di issue
Funzionalità
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python, pytorch

Direzione di ricerca

Start by comparing the Zarr-to-Arrow scan with the Dask and direct Zarr baselines described in the issue, then read related issues #173 and #276 for context. The proposed work spans scan indexing, mask pushdown, and preserving row order, so first identify which optimization has a defined entry point. Done means a measured improvement against the stated baselines while preserving the reported batch correctness.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

🔥 Claude (torch-sql):

torch-sql is a prototype that runs PyTorch models as SQL on XQL. Its first step into production is to replace Samudra's training data loader with a query through xarray-sql: read a rollout window, normalize it with a join to per-channel statistics, mask land, and hand dense tensors to PyTorch.

A spike measured that on a synthetic store shaped like Samudra's OM4 data: 40 channels, 180 × 360, chunks of one time step, NaN on land, 24 time steps. One rollout window is 5 time steps × 40 channels, 52 MB of float32. The batches are correct (equal to Samudra's own normalization and masking within 1e-5), but they're slow:

Path Seconds per window (warm)
Dask, as Samudra's cpu loader reads it 0.17–0.25
Zarr directly, no Dask 0.16–0.24
xarray-sql on DataFusion 0.40–0.48
the same, dropping land rows 0.36–0.43
xarray-sql on DuckDB 0.49–0.54

These are xarray-sql 0.6.0 and DataFusion 54 on a 12-core x86 machine that other jobs shared, so expect about 30% noise between runs. The time filter prunes partitions well, so the cost is proportional to the window, not the store.

Where the time goes:

  • The scan from Zarr to Arrow rows takes about 240 ms, as long as the baselines' whole pipeline. Each cell becomes a row with three int64 index columns, so the scan writes more index bytes than value bytes.
  • The mask and normalization joins add about 90 ms.
  • Turning rows back into a dense tensor adds about 40 ms, and it needs a scatter because the rows arrive out of order.

Three things would help, roughly in order:

  1. A cheaper scan. Narrower index types, or index columns generated from each chunk's offsets rather than materialized per cell.
  2. A mask as a scan filter. Today, dropping land rows is a join, so land chunks are still read and decoded. Pushing a mask table into the scan would skip chunks that are all land.
  3. A per-partition order fact. Rows that come out ordered by their indices within a partition could be reshaped without a copy (9 µs on 10.5 million rows), instead of scattered (182 ms). Related: #173 and #276.
Lingua principale
Python
Stelle
142
Fork
19
Merge medio
3g 47m
PR unite (30g)
8

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di xqlsystems/xarray-sql

Tutte le issue di xqlsystems/xarray-sql

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.