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

未关闭
  • 哪里不对？
#311 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 3 天内回复

还没有人认领这个 Issue。

评估

难度
5/5
预计耗时
一周以上
新手友好度
35/100
Issue 类型
功能
描述清晰度
基本清楚
活跃度
活跃
技术栈
python, pytorch
领域
data-engineering, performance

调研方向

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.

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

描述

🔥 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.
主要语言
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 摘要。