Updated combine_and_write
MODIFY
pangeo_forge/tasks/xarray.py
pangeo_forge/tasks/xarray.py
+from typing import List
+
import fsspec
import xarray as xr
from prefect import task
@task
-def combine_and_write(sources, target, append_dim, concat_dim, first=True):
+def combine_and_write(sources: List[str], target: str, concat_dim: str) -> List[str]:
"""
- Combine one or more source datasets into a single `Xarray.Dataset`, then
- write them to a Zarr store.
+ Write a batch of intermediate files to a combined zarr store.
Parameters
----------
- sources : list of str
- Path or url to the source files.
+ sources : List[str]
+ A list of URLs pointing to the intermediate files.
target : str
- Path or url to the target location of the Zarr store.
- append_dim : str
- Name of the dimension of which datasets should be appended during write.
+ The URL for the target combined store.
concat_dim : str
- Name of the dimension of which datasets should be concatenated during read.
- first : bool
- Flag to indicate if the target dataset should be expected to already exist.
+ The dimension to concatenate along.
Returns
-------
- target_url : str
- Path or url in the form of `{cache_location}/hash({source_url})`.
- """
+ target : str
+ The URL of the written combined Zarr store (same as target).
+
+ Examples
+ --------
+ >>> import pangeo_forge.tasks.xarray
+ >>> import fsspec
+ >>> import xarray as xr
+
+ >>> # Load sample data into `sources`.
+ >>> ds = xr.tutorial.open_dataset('rasm').load()
+ >>> a, b = ds.isel(time=slice(18)), ds.isel(time=slice(18, None))
+ >>> fs = fsspec.get_filesystem_class("memory")()
+
+ >>> sources = [f"{fs.protocol}://{source}" for source in fs.ls("cache")]
+ >>> pangeo_forge.tasks.xarray.combine_and_write.run(sources, "memory://out", concat_dim="time")
+
+ """
double_open_files = [fsspec.open(url).open() for url in sources]
ds = xr.open_mfdataset(double_open_files, combine="nested", concat_dim=concat_dim)
-
# by definition, this should be a contiguous chunk
- ds = ds.chunk({append_dim: len(sources)})
+ ds = ds.chunk({concat_dim: len(sources)})
+ mapper = fsspec.get_mapper(target)
- if first:
+ if not len(mapper):
+ # The first write, .
kwargs = dict(mode="w")
else:
- kwargs = dict(mode="a", append_dim=append_dim)
-
- mapper = fsspec.get_mapper(target)
+ kwargs = dict(mode="a", append_dim=concat_dim)
ds.to_zarr(mapper, **kwargs)
+ return target

MODIFY
pangeo_forge/utils.py
pangeo_forge/utils.py
import itertools
+from typing import Any, List, Tuple
+
+from prefect import task
# https://alexwlchan.net/2018/12/iterating-in-fixed-size-chunks/
@@ -9,3 +12,17 @@ def chunked_iterable(iterable, size):
if not chunk:
break
yield chunk
+
+
+@task
+def chunk(sources: List[Any], size: int) -> List[Tuple[Any, ...]]:
+ """
+ Prefect task to chunk a list of sources into batches.
+
+ Examples
+ --------
+ >>> import pangeo_forge.utils
+ >>> pangeo_forge.utils.chunk.run([1, 2, 3, 4, 5], size=2)
+ [(1, 2), (3, 4), (5,)]
+ """
+ return list(chunked_iterable(sources, size))

ADD
tests/tasks/__init__.py
tests/tasks/__init__.py
ADD
tests/tasks/test_xarray.py
+import fsspec
+import xarray as xr
+from prefect import Flow, Task
+
+import pangeo_forge.tasks.xarray
+
+
+def test_combine_and_write():
+ ds = xr.tutorial.open_dataset("rasm").load()
+ dsets = ds.isel(time=slice(18)), ds.isel(time=slice(18, None))
+ fs = fsspec.get_filesystem_class("memory")()
+
+ for i, dset in enumerate(dsets):
+ as_bytes = dset.to_netcdf()
+
+ with fs.open(f"cache/{i}.nc", "wb") as f:
+ f.write(as_bytes)
+
+ sources = [f"memory://{dset}" for dset in fs.ls("cache")]
+
+ # In a flow context
+
+ target = "memory://target.zarr"
+ with Flow("test") as flow:
+ result = pangeo_forge.tasks.xarray.combine_and_write(sources, target, concat_dim="time")
+ assert isinstance(result, Task)
+ flow.validate()
+
+ result = pangeo_forge.tasks.xarray.combine_and_write.run(sources, target, concat_dim="time")
+ assert result == target
+ result = xr.open_zarr(fs.get_mapper("target.zarr"))
+ xr.testing.assert_equal(ds, result)

ADD
tests/test_utils.py
+import pangeo_forge.utils
+
+
+def test_chunk():
+ result = pangeo_forge.utils.chunk.run([1, 2, 3], 2)
+ assert result == [(1, 2), (3,)]

Updated combine_and_write