将新的 Xarray DataArray 添加到现有的 Zarr 存储而不重写整个数据集？

Question

如何在不覆盖整个内容的情况下DataArray向现有内容添加新内容？Dataset新的DataArray与现有的有一些坐标，但也有新的。在我当前的实现中，Dataset被完全覆盖，而不是仅仅添加新的东西。

现有的DataArray是分块的 zarr 支持DirectoryStore（尽管我对于 S3 存储也有同样的问题）。

import numpy as np
import xarray as xr
import zarr

arr1 = xr.DataArray(np.random.randn(2, 3),
                   [('x', ['a', 'b']), ('y', [10, 20, 30])],
                   name='arr1')

ds = arr1.chunk({'x': 1, 'y': 3}).to_dataset()

ds看起来像这样：

<xarray.Dataset>
Dimensions:  (x: 2, y: 3)
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30
Data variables:
    arr1     (x, y) float64 dask.array<shape=(2, 3), chunksize=(1, 3)>

我将其写入目录存储：

store = zarr.DirectoryStore('test.zarr')
z = ds.to_zarr(store, group='arr', mode='w')

看上去不错：

$ ls -l test.zarr/arr
total 0
drwxr-xr-x  6 myuser  mygroup  204 Sep 21 11:03 arr1
drwxr-xr-x  5 myuser  mygroup  170 Sep 21 11:03 x
drwxr-xr-x  5 myuser  mygroup  170 Sep 21 11:03 y

我创建一个DataArray与现有坐标共享一些坐标的新坐标，并将其添加到现有的Dataset. 我将首先阅读现有的内容Dataset，因为这就是我在实践中所做的事情。

ds2 = xr.open_zarr(store, group='arr')
arr2 = xr.DataArray(np.random.randn(2, 3),
                   [('x', arr1.x), ('z', [1, 2, 3])],
                   name='arr2')
ds2['arr2'] = arr2

更新后Dataset看起来不错：

<xarray.Dataset>
Dimensions:  (x: 2, y: 3, z: 3)
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30
  * z        (z) int64 1 2 3
Data variables:
    arr1     (x, y) float64 dask.array<shape=(2, 3), chunksize=(1, 3)>
    arr2     (x, z) float64 0.4728 1.118 0.7275 0.4971 -0.3398 -0.3846

...但如果不完全覆盖我就无法写入它。

# I think I'm "appending" to the group `arr`
z2 = ds2.to_zarr(store, group='arr', mode='a')

这给了我一个ValueError: The only supported options for mode are 'w' and 'w-'.

# I think I'm "creating" the new arr2 array in the arr group
z2 = ds2.to_zarr(store, group='arr', mode='w-')

这给了我ValueError: path 'arr' contains a group。

唯一有效的是z2 = ds2.to_zarr(store, group='arr', mode='w')，但这完全覆盖了该组。

原来的DataArray实际上对我的问题来说相当大，所以我真的不想重写它。 有没有办法只写新的DataArray？

谢谢你！

Answer 1

自从这个问题发布以来已经有一段时间了 - 但也许它仍然是对某人有帮助的相关广告（对我来说是！）

0.16.2的版本xarray引入了关键字region to to_zarr，它允许您写入zarr文件的有限区域。这似乎使您能够向现有数据集添加新变量，而无需完全覆盖它。

我的解决方案在您写入dszarr 并ds2在内存中创建新数据后，然后将其写回之前开始运行。

首先，我将每个 zarr 内容的修改时间保存在字典中，以在第二次写入后检查是否确实没有任何更改：

import os
import glob

mtimes = {}
contents = list(glob.glob("test.zarr/arr/*"))
for c in contents:
    mtimes.update({c: os.path.getmtime(c)})

现在我可以写回新变量。要使用region关键字，我需要删除已经存在的所有变量，并且这两个变量是相同的：

ds2_dropped = ds2.drop(["x", "y", "z", "arr1"])

现在我可以编写新变量并检查修改时间（如果确实没有任何更改）：

ds2_dropped.to_zarr("test.zarr/", mode="a", group='arr', region={"x": slice(0, ds2.x.size), "z": slice(0, ds2.z.size)})

for c in contents:
    assert os.path.getmtime(c) == mtimes[c]

# all good!

如果我们再次从 zarr 加载数据集，我们可以看到新变量已成功添加：

print(xr.open_zarr("test.zarr/", group="arr"))

<xarray.Dataset>
Dimensions:  (x: 2, y: 3, z: 3)
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30
Dimensions without coordinates: z
Data variables:
    arr1     (x, y) float64 dask.array<chunksize=(1, 3), meta=np.ndarray>
    arr2     (x, z) float64 dask.array<chunksize=(2, 3), meta=np.ndarray>

Answer 2

不幸的是，据我所知，目前这是不可能的。实现附加模式to_zarr是为了向维度添加新条目，而不是向已写入的条目添加变量。

@davidbrochart在原始 MR 中为用例编写了一个很好的示例：

import xarray as xr
import pandas as pd

ds0 = xr.Dataset({'temperature': (['time'],  [50, 51, 52])}, coords={'time': pd.date_range('2000-01-01', periods=3)})
ds1 = xr.Dataset({'temperature': (['time'],  [53, 54, 55])}, coords={'time': pd.date_range('2000-01-04', periods=3)})

ds0.to_zarr('temp')
ds1.to_zarr('temp', mode='a', append_dim='time')

ds2 = xr.open_zarr('temp')

您将看到这是和在时间维度上ds2的串联版本。ds0ds1

好消息是，可以选择直接与 zarr 商店互动。如果您查看 xarray使用的实现，您会发现在底层zarr库中实际上可以添加新变量。然而，这并未在 xarray API 中实现。