Добавление нового Xarray DataArray в существующее хранилище Zarr без перезаписи всего набора данных?

Question

Как мне добавить новый DataArray существующему Datasetбез перезаписи всего? НовыйDataArrayимеет общие координаты с существующими, но также имеет новые. В моей текущей реализацииDataset полностью перезаписывается, а не просто добавляется новый материал.

Существующий DataArray представляет собой фрагментированный, поддерживаемый зарр 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:

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ввел регион ключевого слова в to_zarr, который позволяет записывать в ограниченную область zarr-файла. Это, по-видимому, позволяет вам добавить новую переменную в существующий набор данных, не перезаписывая его полностью.

Мое решение появляется после того, как вы написали dsк зарру и создал новый 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 это объединенная версия ds0 а также ds1 по временному измерению.

Хорошая новость заключается в том, что есть возможность напрямую взаимодействовать с магазином zarr. Если вы посмотрите на реализацию, которую использует xarray, вы увидите, что добавление новых переменных на самом деле возможно в базовой библиотеке zarr. Однако это не реализовано в API xarray.