def test_rechunk_not_all_dimensions(self): data = np.random.RandomState(0).randint(2**30, size=(10, 20, 30)) ds = xarray.Dataset({'foo': (('time', 'x', 'y'), data)}) key = xbeam.Key({'x': 0, 'y': 0}) y_split_with_time_key = ([(key.with_offsets(time=0), ds)] | xbeam.SplitChunks({'y': 3})) x_split = [(key, ds)] | xbeam.SplitChunks({'x': 2}) actual = x_split | xbeam.Rechunk( dim_sizes=ds.sizes, source_chunks={ 'x': 2, 'y': -1 }, target_chunks={ 'x': -1, 'y': 3 }, itemsize=8, max_mem=10_000, ) self.assertIdenticalChunks(actual, y_split_with_time_key) with self.assertRaisesRegex( ValueError, 'source_chunks and target_chunks have different keys', ): xbeam.Rechunk( dim_sizes=ds.sizes, source_chunks={'x': 2}, target_chunks={'y': 3}, itemsize=8, max_mem=10_000, )
def test_rechunk_1d(self, size, max_mem, source_chunks, target_chunks): data = np.random.RandomState(0).randint(2**30, size=(size, )) ds = xarray.Dataset({'foo': ('x', data)}) key = xbeam.Key({'x': 0}) inputs = [(key, ds)] | xbeam.SplitChunks({'x': source_chunks}) expected = [(key, ds)] | xbeam.SplitChunks({'x': target_chunks}) actual = inputs | xbeam.Rechunk( dim_sizes=ds.sizes, source_chunks={'x': source_chunks}, target_chunks={'x': target_chunks}, itemsize=1, max_mem=max_mem, ) self.assertIdenticalChunks(actual, expected)
def main(argv): # By passing chunks=None, we use Xarray's lazy-loading instead of Dask. This # result is much less data being passed from the launch script to workers. source_dataset = xarray.open_zarr( INPUT_PATH.value, chunks=None, consolidated=True, ) # This lazy "template" allows us to setup the Zarr outputs before running the # pipeline. We don't really need to supply a template here because the outputs # are small (the template argument in ChunksToZarr is optional), but it makes # the pipeline slightly more efficient. max_month = source_dataset.time.dt.month.max().item() # normally 12 template = (source_dataset.chunk().pipe(xarray.zeros_like).isel( time=0, drop=True).expand_dims(month=np.arange(1, max_month + 1), hour=np.arange(24))) output_chunks = {'hour': 1, 'month': 1} with beam.Pipeline(runner=RUNNER.value, argv=argv) as root: (root | xbeam.DatasetToChunks(source_dataset, {'time': 31}) | xbeam.SplitChunks({'time': 1}) | beam.MapTuple(rekey_chunk_on_month_hour) | xbeam.Mean.PerKey() | xbeam.ChunksToZarr(OUTPUT_PATH.value, template, output_chunks))
def test_consolidate_and_split_only_some_dims(self): chunk_data = np.arange(0, 10).reshape(2, 5) split = [ (xbeam.Key({ 'x': 0, 'y': 0 }), xarray.Dataset({'foo': (('x', 'y'), chunk_data)})), (xbeam.Key({ 'x': 0, 'y': 5 }), xarray.Dataset({'foo': (('x', 'y'), chunk_data + 10)})), ] all_data = np.concatenate([chunk_data, chunk_data + 10], axis=1) consolidated = [ (xbeam.Key({ 'x': 0, 'y': 0 }), xarray.Dataset({'foo': (('x', 'y'), all_data)})), ] with self.subTest('ConsolidateChunks'): actual = split | xbeam.ConsolidateChunks({'y': 10}) self.assertIdenticalChunks(actual, consolidated) with self.subTest('SplitChunks'): actual = consolidated | xbeam.SplitChunks({'y': 5}) self.assertIdenticalChunks(actual, split)
def test_consolidate_and_split_chunks(self): consolidated = [ (xbeam.Key({'x': 0}), xarray.Dataset({'foo': ('x', np.arange(0, 10))})), (xbeam.Key({'x': 10}), xarray.Dataset({'foo': ('x', np.arange(10, 20))})), ] split = [ (xbeam.Key({'x': 0}), xarray.Dataset({'foo': ('x', np.arange(0, 5))})), (xbeam.Key({'x': 5}), xarray.Dataset({'foo': ('x', np.arange(5, 10))})), (xbeam.Key({'x': 10}), xarray.Dataset({'foo': ('x', np.arange(10, 15))})), (xbeam.Key({'x': 15}), xarray.Dataset({'foo': ('x', np.arange(15, 20))})), ] with self.subTest('ConsolidateChunks'): actual = split | xbeam.ConsolidateChunks({'x': 10}) self.assertIdenticalChunks(actual, consolidated) with self.subTest('SplitChunks'): actual = consolidated | xbeam.SplitChunks({'x': 5}) self.assertIdenticalChunks(actual, split)
def test_rechunk_uneven_2d(self): data = np.random.RandomState(0).randint(2**30, size=(100, 100)) ds = xarray.Dataset({'foo': (('x', 'y'), data)}) key = xbeam.Key({'x': 0, 'y': 0}) inputs = [(key, ds)] | xbeam.SplitChunks({'x': 12}) expected = [(key, ds)] | xbeam.SplitChunks({'y': 15}) actual = inputs | xbeam.Rechunk( dim_sizes=ds.sizes, source_chunks={ 'x': 12, 'y': -1 }, target_chunks={ 'x': -1, 'y': 15 }, itemsize=1, max_mem=100 * 100 // 2, # half the full size ) self.assertIdenticalChunks(actual, expected)
def test_rechunk_end_to_end(self): data = np.random.RandomState(0).randint(2**30, size=(10, 20, 30)) ds = xarray.Dataset({'foo': (('time', 'x', 'y'), data)}) key = xbeam.Key({'time': 0, 'x': 0, 'y': 0}) time_split = [(key, ds)] | xbeam.SplitChunks({'time': 1}) space_split = [(key, ds)] | xbeam.SplitChunks({'x': 5, 'y': 5}) with self.subTest('time-to-space'): actual = time_split | xbeam.Rechunk( dim_sizes=ds.sizes, source_chunks={ 'time': 1, 'x': 20, 'y': 30 }, target_chunks={ 'time': 10, 'x': 5, 'y': 5 }, itemsize=8, max_mem=10_000, ) self.assertIdenticalChunks(actual, space_split) with self.subTest('space-to-time'): actual = space_split | xbeam.Rechunk( dim_sizes=ds.sizes, source_chunks={ 'time': 10, 'x': 5, 'y': 5 }, target_chunks={ 'time': 1, 'x': 20, 'y': 30 }, itemsize=8, max_mem=10_000, ) self.assertIdenticalChunks(actual, time_split)
def test_split_uneven_chunks(self): inputs = [ (xbeam.Key({'x': 0}), xarray.Dataset({'foo': ('x', np.arange(0, 5))})), (xbeam.Key({'x': 5}), xarray.Dataset({'foo': ('x', np.arange(5, 10))})), ] expected = [ (xbeam.Key({'x': 0}), xarray.Dataset({'foo': ('x', np.arange(0, 3))})), (xbeam.Key({'x': 3}), xarray.Dataset({'foo': ('x', np.arange(3, 5))})), (xbeam.Key({'x': 5}), xarray.Dataset({'foo': ('x', np.arange(5, 6))})), (xbeam.Key({'x': 6}), xarray.Dataset({'foo': ('x', np.arange(6, 9))})), (xbeam.Key({'x': 9}), xarray.Dataset({'foo': ('x', np.arange(9, 10))})), ] actual = inputs | xbeam.SplitChunks({'x': 3}) self.assertIdenticalChunks(actual, expected)
def test_2d_chunks_to_zarr(self, coords): dataset = xarray.Dataset( {'foo': (('x', 'y'), np.arange(0, 60, 10).reshape(3, 2))}, coords=coords, ) with self.subTest('partial key'): inputs = [(xbeam.Key({'x': 0}), dataset)] temp_dir = self.create_tempdir().full_path inputs | xbeam.ChunksToZarr(temp_dir) result = xarray.open_zarr(temp_dir, consolidated=True) xarray.testing.assert_identical(dataset, result) with self.subTest('split along partial key'): inputs = [(xbeam.Key({'x': 0}), dataset)] temp_dir = self.create_tempdir().full_path inputs | xbeam.SplitChunks({'x': 1}) | xbeam.ChunksToZarr(temp_dir) result = xarray.open_zarr(temp_dir, consolidated=True) xarray.testing.assert_identical(dataset, result) with self.subTest('full key'): inputs = [(xbeam.Key({'x': 0, 'y': 0}), dataset)] temp_dir = self.create_tempdir().full_path inputs | xbeam.ChunksToZarr(temp_dir) result = xarray.open_zarr(temp_dir, consolidated=True) xarray.testing.assert_identical(dataset, result)