xarrayでZarrを読み込む、書き込む

結論

xarrayから整備されたZarrデータを読み込む分には問題ないが、xarrayを使いZarrデータを書き込むのは時期尚早で、挑戦者以外は安定版を待つべきである。

執筆時点(2025年2月12日)でZarrはまだまだ開発途上であり、Zarr2までからの破壊的アップデートを含むZarr3が出たばかりである。当然xarrayのZarr対応も途上である。この記事の内容は急速に過去のものになるだろう。

以上の経緯によりZarr3には xarray バージョン v2025.01.1 以降のまだ非安定版のみ限定的に対応している。本記事ではバージョン v2025.01.2 を元にしているがまだわかりやすいバグ等も目立つ。基本的にはまだnetCDFやGRIBで保存した方がいいだろう。

はじめに

本記事はxarrayでZarrを使おうというアーリーアダプターのために日本語での初歩的な導入を記載する。パフォーマンス計測や最適化には主眼を置かない。基本的には公式のユーザーガイドから一部を抜粋するが、個人の感想も含む。

バージョン情報

• Python: 3.12.7
• xarray: v2025.01.2
• zarr: 3.0.2

Zarr とは

Zarr は比較的新しい多次元データの保存フォーマットである。巨大なデータを分割・圧縮して保存することを前提としている点が膨大なデータIOを要求する機械学習と相性が良く、気象データでも普及が始まっている。必要な部分を必要な時に読み込める設計のためネットワーク越しでインタラクティブに解析したいときなどはnetCDF等に比べてかなり高速に処理できる。最近ではECMWFのERA5再解析データがGoogleを通じてZarrフォーマットで公開されており、事前ダウンロードなしで即座に解析をできるようになった。

Zarrの特徴は Directory store という巨大なデータを構造化したディレクトリ内の無数の小ファイルに分割して保存する形式をとっていることである。考え方としては、Google Map等で使われているように、巨大な画像を小さいタイルに分けて保存し、都度必要な範囲のみを読み込んで繋ぎ合わせるのに似ている。数値計算の文脈では並列の各プロセスが同時に別のデータを読み書きできるため巨大でシーケンシャルな1ファイルを読み込むより効率が良くなる。

試しにGoogleが公開しているERA5のデータをクラウドから直接読み込んでみよう。上記ページにあるサンプルのように以下のPythonコードを実行する。

import xarray as xr
era5 = xr.open_zarr(
    "gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2",
    chunks={'time': 48},
    consolidated=True,
)
print(era5)

Output:

<xarray.Dataset>
Dimensions:                                           (time: 552264, latitude: 721,  longitude: 1440, level: 37)
Coordinates:
  * latitude                                          (latitude) float32 90.0...
  * level                                             (level) int64 1 2 ... 1000
  * longitude                                         (longitude) float32 0.0...
  * time                                              (time) datetime64[ns] 1...
Data variables: (12/31)
    10m_u_component_of_wind                           (time, latitude, longitude) float32 dask.array<chunksize=(48, 721, 1440), meta=np.ndarray>
    10m_v_component_of_wind                           (time, latitude, longitude) float32 dask.array<chunksize=(48, 721, 1440), meta=np.ndarray>
    2m_temperature                                    (time, latitude, longitude) float32 dask.array<chunksize=(48, 721, 1440), meta=np.ndarray>
    angle_of_sub_gridscale_orography                  (latitude, longitude) float32 dask.array<chunksize=(721, 1440), meta=np.ndarray>
    anisotropy_of_sub_gridscale_orography             (latitude, longitude) float32 dask.array<chunksize=(721, 1440), meta=np.ndarray>
    geopotential                                      (time, level, latitude, longitude) float32 dask.array<chunksize=(48, 37, 721, 1440), meta=np.ndarray>
    ...                                                ...
    total_precipitation                               (time, latitude, longitude) float32 dask.array<chunksize=(48, 721, 1440), meta=np.ndarray>
    type_of_high_vegetation                           (latitude, longitude) float32 dask.array<chunksize=(721, 1440), meta=np.ndarray>
    type_of_low_vegetation                            (latitude, longitude) float32 dask.array<chunksize=(721, 1440), meta=np.ndarray>
    u_component_of_wind                               (time, level, latitude, longitude) float32 dask.array<chunksize=(48, 37, 721, 1440), meta=np.ndarray>
    v_component_of_wind                               (time, level, latitude, longitude) float32 dask.array<chunksize=(48, 37, 721, 1440), meta=np.ndarray>
    vertical_velocity                                 (time, level, latitude, longitude) float32 dask.array<chunksize=(48, 37, 721, 1440), meta=np.ndarray>

このようにERA5の気圧面データのメタデータが xr.Dataset 形式で得られる。驚異的なのは次元サイズ (time: 552264, latitude: 721, longitude: 1440, level: 37) と変数の数、そしてその読み込みにインターネット越しであるにもかかわらず数十秒しかかからないことだ。筆者の環境では12.9秒だった。このように構造化されたデータのうち必要な情報（ここではメタデータ）のみを読みにいくので最低限の待ち時間でレスポンスを得られる。

このようなメリットの一方で、大量の小ファイルができることはストレージ管理の面からはあまり嬉しいことではなく、意図しないファイルがディレクトリに紛れ込むリスクも持ちたくはない。そのためこのディレクトリをzipファイルにまとめて読み書きできる Zip store という処理形式に拡張しようとしているが、まだまだ開発途上でありxarrayは対応していない。現状はデフォルトの Directory store でのファイルハンドリングが要求される。

xarray がZarrフォーマットへの対応をどう進めるかは 公式ドキュメントに記してある。Zarr形式はnetCDFを主眼に置いて開発されてきたxarrayとはそのままでは整合しないため、xarrayは専用のメタデータを作成することで互換する必要がある。そのため全てのZarrデータが読み込めるわけではないことに留意が必要である。

Zarr を xarray で読み込む

xr.open_zarr()

xarrayはZarrを読み込む関数xr.open_zarr()を用意しており、基本的にはこの関数を使えば良い。

df = xr.open_zarr('./test.zarr')
print(df)

Output:

<xarray.Dataset> Size: 26GB
Dimensions:  (y: 239, x: 217, p: 30, time: 13, to_m: 40)
Coordinates:
    lat      (y, x) float64 415kB dask.array<chunksize=(120, 217), meta=np.ndarray>
    lon      (y, x) float64 415kB dask.array<chunksize=(120, 217), meta=np.ndarray>
  * p        (p) float32 120B 4e+03 5.103e+03 6.409e+03 ... 1.01e+05 1.05e+05
  * time     (time) datetime64[ns] 104B 2022-08-26T06:00:00 ... 2022-08-26T06...
Dimensions without coordinates: y, x, to_m
Data variables:
    q        (time, y, x, p, to_m) float64 6GB dask.array<chunksize=(1, 64, 64, 10, 40), meta=np.ndarray>
    t        (time, y, x, p, to_m) float64 6GB dask.array<chunksize=(1, 64, 64, 10, 40), meta=np.ndarray>
    u        (time, y, x, p, to_m) float64 6GB dask.array<chunksize=(1, 64, 64, 10, 40), meta=np.ndarray>
    v        (time, y, x, p, to_m) float64 6GB dask.array<chunksize=(1, 64, 64, 10, 40), meta=np.ndarray> -1.4773546067619743e-12

計26GBのデータが読み込まれている。｀dask.array<chunksize=(1, 64, 64, 10, 40), meta=np.ndarray>｀の部分はzarrに保存する際にデータが持っていたチャンクの大きさであり、zarrが処理する単位データサイズでもある。

xr.open_dataset()

ドキュメントには公式の方法として記載があるが、個人的にこの方法での読み込みは非推奨である。理由は後述。

xarrayは従来の xr.open_dataset() 関数を用いてzarrを読み込むオプションも提供している。この場合、xr.open_dataset() は自動でzarrエンジンを適用するが engine='zarr'を明示的に指定した方がエンジン判定をスキップできるため早い。余計なwarningも抑制できる。

df = xr.open_dataset('./test.zarr', engine='zarr')
print(df)

<xarray.Dataset> Size: 26GB
Dimensions:  (y: 239, x: 217, p: 30, time: 13, to_m: 40)
Coordinates:
    lat      (y, x) float64 415kB ...
    lon      (y, x) float64 415kB ...
  * p        (p) float32 120B 4e+03 5.103e+03 6.409e+03 ... 1.01e+05 1.05e+05
  * time     (time) datetime64[ns] 104B 2022-08-26T06:00:00 ... 2022-08-26T06...
Dimensions without coordinates: y, x, to_m
Data variables:
    q        (time, y, x, p, to_m) float64 6GB ...
    t        (time, y, x, p, to_m) float64 6GB ...
    u        (time, y, x, p, to_m) float64 6GB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0
    v        (time, y, x, p, to_m) float64 6GB ... -1.4773546067042494e-12

同じデータが読み込まれているが、xr.open_zarr()のアウトプットにはあった｀dask.array<chunksize=(1, 64, 64, 10, 40), meta=np.ndarray>｀が見られない。xr.open_dataset()ではzarrのチャンク構造を自動的に読まないため、従来のnetCDF等と同じように、zarr内の全データを1つの大きなファイルと同じように扱う。そのためdaskを活用した分割処理を行うためには再度自分でチャンク指定を行う必要がある。もしそのチャンク指定がzarrデータのチャンク構造と不整合な場合、大きくパフォーマンスが落ちることになる。

また、読み込みにxr.open_zarr()とxr.open_dataset()を使ったxr.Datasetではそのままでも計算パフォーマンスに差が出る。たとえば一度領域平均を取る計算では両者に大きな差はみられなかったが、一度読み込んだデータを使って2回目以降の計算を行ったところ、xr.open_zarr()を使ったデータの方が10倍以上早かった(おそらく分割キャッシュなどの機構がある？)。以上の理由から特別の理由がない限りはxr.open_zarr()を常用すべきであろう。

Zarr を xarray で書き込む

xr.Dataset.to_zarr()

とりあえず書き込む

xarrayはxr.Dataset¹をzarr形式で書き出すxr.Dataset.to_zarr()メソッドを提供しているためそれを使えばいい。

df.to_zarr(store=f"./test.zarr")

追加で書き込む

ZarrのnetCDFに比較して大きな利点は、データを一部だけ書き込んだり、後から追加することが容易であることだ。例えば、数字が順に大きくなっていくある次元(例：時刻)に沿って順番にデータを追加していくことができる。

この追加書き込みモードにはバグがある。特に時刻型(datetime64など)を持つ次元に沿って追加すると時刻の値がおかしくなる。信頼が必要な解析にはまだ使うべきではない。

df.to_zarr(store=f"./test.zarr", mode="a", append_dim="time")

mode="a"はappendモードを意味している。appnd_dim=オプションでxr.Dataset df に含まれるtime次元にそってデータを積み上げていくよう指示している。注意が必要なのが、もし追加するデータが既存データと同じtimeの値を持っていた場合、既存データは追加データで上書きされることだ。warningを出すオプションがあるかは不明。

特定の領域だけ書き込む

データを追加するだけでなく、あらかじめデータの型だけを作って中身の値だけを後から書き込むこともできる。その場合はスライスしたデータをregionオプションでどこに書き込むかを指定する。

path = f"./test.zarr"
# 値は書き込まずにデータの型だけを作成する
df.to_zarr(store=path, compute=False)


#　x次元でスライスしたデータ（1チャンク分）を書き込む。データ位置は自動判定
df.isel(x=slice(0, 64)).to_zarr(path, region="auto")
#　どの次元に沿って自動判定するかを指定
df.isel(x=slice(64, 128)).to_zarr(path, region={"x": "auto"})
#　どの次元のどの位置に書き込むかを明示的に指定
df.isel(x=slice(128, 192)).to_zarr(path, region={"x": slice(128, 192)})

1. おそらく同様に xr.Dataarray と xr.Datatree も ↩