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

Aberta
  • Algo errado?
#311 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
5/5
Tempo estimado
Mais de uma semana
Facilidade para iniciantes
35/100
Tipo de issue
Funcionalidade
Clareza
Razoavelmente clara
Status de atividade
Ativa
Stack de tecnologia
python, pytorch
Domínio
data-engineering, performance

Direção de pesquisa

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.

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

Descrição

🔥 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.
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.