google / xarray-beam

Distributed Xarray with Apache Beam
https://xarray-beam.readthedocs.io
Apache License 2.0
125 stars 7 forks source link

Support missing chunks in ConsolidateChunks #54

Open shoyer opened 1 year ago

shoyer commented 1 year ago

This currently raises an error, but isn't always a programming error:

import xarray
import xarray_beam as xbeam
import numpy as np

ds = xarray.Dataset({"foo": (('x', 'y'), np.zeros((4, 4)))})
key = xbeam.Key({'x': 0, 'y': 0})

([(key, ds)] | xbeam.SplitChunks({'x': 2, 'y': 2}))[:-1] | xbeam.ConsolidateChunks({'x': -1, 'y': -1})
# ValueError: some expected chunks are missing for vars=None shape: [2, 2]
# len(inputs): 3 [while running 'ConsolidateChunks/Consolidate']

It would be nice if we supported missing chunks by filling in values with NaN, like xarray.merge.