如何在不覆盖整个内容的情况下向现有Dataset
添加新DataArray
? 新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
?
谢谢!
这个问题发布已经有一段时间了 - 但也许它仍然对某人有帮助(对我来说是!
xarray
的第0.16.2
版引入了关键字 region 到to_zarr
,它允许您写入 zarr 文件的有限区域。这似乎使您能够向现有数据集添加新变量,而不会完全覆盖它。
在您ds
写入 zarr 并在内存中创建新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>
不幸的是,据我所知,这目前是不可能的。追加模式下的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
是ds0
和ds1
在时间维度上的串联版本。
好消息是,可以选择直接与 zarr 商店互动。如果你看一下 xarray 使用的实现,你可以看到添加新变量实际上是底层 zarr 库中的一种可能性。但是,这在 xarray API 中没有实现。