Introduced keyword arguments source_ds and source_append_offset

forman · forman · commit 056f2f7606f8 · 2024-10-02T13:20:02.000+02:00
diff --git a/tests/contrib/__init__.py b/tests/contrib/__init__.py
@@ -0,0 +1,3 @@
+# Copyright © 2024 Norman Fomferra and contributors
+# Permissions are hereby granted under the terms of the MIT License:
+# https://opensource.org/licenses/MIT.
diff --git a/tests/contrib/test_levels.py b/tests/contrib/test_levels.py
@@ -9,10 +9,10 @@
 import xarray as xr
 
 from zappend.fsutil import FileObj
-from zappend.levels import write_levels
-from zappend.levels import get_variables_config
-from .helpers import clear_memory_fs
-from .helpers import make_test_dataset
+from zappend.contrib import write_levels
+from zappend.contrib.levels import get_variables_config
+from tests.helpers import clear_memory_fs
+from tests.helpers import make_test_dataset
 
 try:
     # noinspection PyUnresolvedReferences
@@ -71,29 +71,30 @@ def test_variables_given(self):
         )
 
 
+source_path = "memory://source.zarr"
+target_path = "memory://target.levels"
+
+
+# noinspection PyMethodMayBeStatic
 @unittest.skipIf(xcube is None, reason="xcube is not installed")
-class WriteLevelsTest(unittest.TestCase):
+class WriteLevelsArgsTest(unittest.TestCase):
     def setUp(self):
         clear_memory_fs()
-
-    # noinspection PyMethodMayBeStatic
-    def test_args_validation(self):
-        source_path = "memory://source.zarr"
-        target_path = "memory://target.levels"
-
         make_test_dataset(
             uri=source_path,
             dims=("time", "lat", "lon"),
             shape=(3, 1024, 2048),
             chunks=(1, 128, 256),
         )
 
+    def test_target_path_not_given(self):
         with pytest.raises(
             ValueError,
             match="missing 'target_path' argument",
         ):
             write_levels(source_path=source_path)
 
+    def test_target_dir_and_target_path_given(self):
         with pytest.raises(
             ValueError,
             match="either 'target_dir' or 'target_path' can be given, not both",
@@ -104,6 +105,7 @@ def test_args_validation(self):
                 target_dir=target_path,
             )
 
+    def test_config_given(self):
         with pytest.raises(
             TypeError,
             match="write_levels\\(\\) got an unexpected keyword argument 'config'",
@@ -114,6 +116,7 @@ def test_args_validation(self):
                 config={"dry_run": True},
             )
 
+    def test_dry_run_and_use_saved_levels_given(self):
         with pytest.raises(
             FileNotFoundError,
             match="Target parent directory does not exist: /target.levels",
@@ -129,8 +132,60 @@ def test_args_validation(self):
                     use_saved_levels=True,
                 )
 
+    def test_source_path_and_source_ds_not_given(self):
+        with pytest.raises(
+            TypeError,
+            match="'source_ds' argument must be of type 'xarray.Dataset',"
+            " but was 'NoneType'",
+        ):
+            write_levels(
+                target_path=target_path,
+            )
+
+    def test_source_path_not_given_and_link_level_zero_is_true(self):
+        with pytest.raises(
+            ValueError,
+            match="'source_path' argument must be provided"
+            " if 'link_level_zero' is used",
+        ):
+            write_levels(
+                source_ds=xr.Dataset(), target_path=target_path, link_level_zero=True
+            )
+
+    def test_source_append_offset_not_int(self):
+        with pytest.raises(
+            TypeError,
+            match="'source_append_offset' argument must be of type 'int',"
+            " but was 'str'",
+        ):
+            # noinspection PyTypeChecker
+            write_levels(
+                source_ds=xr.open_zarr("memory://source.zarr"),
+                source_append_offset="2",
+                target_path=target_path,
+            )
+
+    def test_source_append_offset_out_of_range(self):
+        with pytest.raises(
+            ValueError,
+            match="'source_append_offset' argument must be >=0 and <3, but was 13",
+        ):
+            # noinspection PyTypeChecker
+            write_levels(
+                source_ds=xr.open_zarr("memory://source.zarr"),
+                source_append_offset=13,
+                target_path=target_path,
+            )
+
+
+@unittest.skipIf(xcube is None, reason="xcube is not installed")
+class WriteLevelsTest(unittest.TestCase):
+    def setUp(self):
+        clear_memory_fs()
+
+    # noinspection PyMethodMayBeStatic
+
     def test_force_new(self):
-        source_path = "memory://source.zarr"
         make_test_dataset(
             uri=source_path,
             dims=("time", "lat", "lon"),
@@ -145,12 +200,13 @@ def test_force_new(self):
         (target_dir / "0.zarr" / ".zgroup").write("{}")
         self.assertTrue(target_dir.exists())
 
-        write_levels(source_path=source_path, target_path=target_dir.uri, force_new=True)
+        write_levels(
+            source_path=source_path, target_path=target_dir.uri, force_new=True
+        )
 
         self.assertTrue(target_dir.exists())
 
     def test_default_x_y_with_crs(self):
-        source_path = "memory://source.zarr"
         make_test_dataset(
             uri=source_path,
             dims=("time", "y", "x"),
@@ -185,7 +241,6 @@ def test_default_x_y_with_crs(self):
         self.assert_level(target_dir.uri + "/3.zarr", 3, has_crs=True)
 
     def test_default_lon_lat_no_crs(self):
-        source_path = "memory://source.zarr"
         make_test_dataset(
             uri=source_path,
             dims=("time", "lat", "lon"),
@@ -219,6 +274,78 @@ def test_default_lon_lat_no_crs(self):
         self.assert_level(target_dir.uri + "/2.zarr", 2, xy_dims=xy_dims)
         self.assert_level(target_dir.uri + "/3.zarr", 3, xy_dims=xy_dims)
 
+    def test_default_lon_lat_no_crs_from_source_ds(self):
+        source_ds = make_test_dataset(
+            uri=source_path,
+            dims=("time", "lat", "lon"),
+            shape=(3, 1024, 2048),
+            chunks=(1, 128, 256),
+        )
+
+        target_dir = FileObj("memory://target.levels")
+        self.assertFalse(target_dir.exists())
+
+        write_levels(source_ds=source_ds, target_path=target_dir.uri)
+
+        self.assertTrue(target_dir.exists())
+
+        levels_file = target_dir.for_path(".zlevels")
+        self.assertTrue(levels_file.exists())
+        levels_info = json.loads(levels_file.read())
+        self.assertEqual(
+            {
+                "version": "1.0",
+                "num_levels": 4,
+                "agg_methods": {"chl": "mean", "tsm": "mean"},
+                "use_saved_levels": False,
+            },
+            levels_info,
+        )
+
+        xy_dims = "lon", "lat"
+        self.assert_level(target_dir.uri + "/0.zarr", 0, xy_dims=xy_dims)
+        self.assert_level(target_dir.uri + "/1.zarr", 1, xy_dims=xy_dims)
+        self.assert_level(target_dir.uri + "/2.zarr", 2, xy_dims=xy_dims)
+        self.assert_level(target_dir.uri + "/3.zarr", 3, xy_dims=xy_dims)
+
+    def test_default_lon_lat_no_crs_from_source_ds_with_offset(self):
+        source_ds = make_test_dataset(
+            uri=source_path,
+            dims=("time", "lat", "lon"),
+            shape=(10, 1024, 2048),
+            chunks=(1, 128, 256),
+        )
+
+        target_dir = FileObj("memory://target.levels")
+        self.assertFalse(target_dir.exists())
+
+        write_levels(
+            source_ds=source_ds,
+            source_append_offset=7,
+            target_path=target_dir.uri,
+        )
+
+        self.assertTrue(target_dir.exists())
+
+        levels_file = target_dir.for_path(".zlevels")
+        self.assertTrue(levels_file.exists())
+        levels_info = json.loads(levels_file.read())
+        self.assertEqual(
+            {
+                "version": "1.0",
+                "num_levels": 4,
+                "agg_methods": {"chl": "mean", "tsm": "mean"},
+                "use_saved_levels": False,
+            },
+            levels_info,
+        )
+
+        xy_dims = "lon", "lat"
+        self.assert_level(target_dir.uri + "/0.zarr", 0, xy_dims=xy_dims)
+        self.assert_level(target_dir.uri + "/1.zarr", 1, xy_dims=xy_dims)
+        self.assert_level(target_dir.uri + "/2.zarr", 2, xy_dims=xy_dims)
+        self.assert_level(target_dir.uri + "/3.zarr", 3, xy_dims=xy_dims)
+
     def test_link_level_zero(self):
         source_dir = FileObj("memory://source.zarr")
         make_test_dataset(
diff --git a/zappend/contrib/__init__.py b/zappend/contrib/__init__.py
@@ -0,0 +1,5 @@
+# Copyright © 2024 Norman Fomferra and contributors
+# Permissions are hereby granted under the terms of the MIT License:
+# https://opensource.org/licenses/MIT.
+
+from .levels import write_levels
diff --git a/zappend/contrib/levels.py b/zappend/contrib/levels.py
@@ -17,8 +17,11 @@
 
 
 def write_levels(
-    source_path: str,
+    *,
+    source_ds: xr.Dataset | None = None,
+    source_path: str | None = None,
     source_storage_options: dict[str, Any] | None = None,
+    source_append_offset: int | None = None,
     target_path: str | None = None,
     num_levels: int | None = None,
     tile_size: tuple[int, int] | None = None,
@@ -28,12 +31,13 @@ def write_levels(
     xy_dim_names: tuple[str, str] | None = None,
     **zappend_config,
 ):
-    """Writes a dataset at `source_path` to `target_path` using the
+    """Writes a dataset given by `source_ds` or `source_path` to `target_path`
+    using the
     [multi-level dataset format](https://xcube.readthedocs.io/en/latest/mldatasets.html)
     as specified by
     [xcube](https://github.com/xcube-dev/xcube).
 
-    It resembles the `store.write_data(cube, "<name>.levels", ...)` method
+    It resembles the `store.write_data(dataset, "<name>.levels", ...)` method
     provided by the xcube filesystem data stores ("file", "s3", "memory", etc.).
     The zappend version may be used for potentially very large datasets in terms
     of dimension sizes or for datasets with very large number of chunks.
@@ -58,9 +62,17 @@ def write_levels(
     Important: This function requires the `xcube` package to be installed.
 
     Args:
+        source_ds: The source dataset.
+            Must be given in case `source_path` is not given.
         source_path: The source dataset path.
+            If `source_ds` is provided and `link_level_zero` is true,
+            then `source_path` must also be provided in order
+            to determine the path of the level zero source.
         source_storage_options: Storage options for the source
             dataset's filesystem.
+        source_append_offset: Optional offset in the append dimension.
+            Only slices with indexes greater or equal the offset are
+            appended.
         target_path: The target multi-level dataset path.
             Filename extension should be `.levels`, by convention.
             If not given, `target_dir` should be passed as part of the
@@ -102,14 +114,15 @@ def write_levels(
     from xcube.core.tilingscheme import get_num_levels
     from xcube.util.fspath import get_fs_path_class
 
+    config = zappend_config.pop("config", None)
+    if config is not None:
+        raise TypeError("write_levels() got an unexpected keyword argument 'config'")
+
     dry_run = zappend_config.pop("dry_run", False)
 
     if dry_run and use_saved_levels:
         warnings.warn(f"'use_saved_levels' argument is not applicable if dry_run=True")
         use_saved_levels = False
-    config = zappend_config.pop("config", None)
-    if config is not None:
-        raise TypeError("write_levels() got an unexpected keyword argument 'config'")
 
     target_dir = zappend_config.pop("target_dir", None)
     if not target_path and not target_dir:
@@ -124,16 +137,49 @@ def write_levels(
         target_path, **target_storage_options
     )
 
-    source_fs, source_root = fsspec.core.url_to_fs(
-        source_path,
-        **(
-            source_storage_options
-            if source_storage_options is not None
-            else target_storage_options
-        ),
-    )
-    source_store = source_fs.get_mapper(root=source_root)
-    source_ds = xr.open_zarr(source_store)
+    force_new = zappend_config.pop("force_new", None)
+
+    if source_path is not None:
+        source_fs, source_root = fsspec.core.url_to_fs(
+            source_path,
+            **(
+                source_storage_options
+                if source_storage_options is not None
+                else target_storage_options
+            ),
+        )
+        if source_ds is None:
+            source_store = source_fs.get_mapper(root=source_root)
+            source_ds = xr.open_zarr(source_store)
+    else:
+        source_root = None
+        if not isinstance(source_ds, xr.Dataset):
+            raise TypeError(
+                f"'source_ds' argument must be of type 'xarray.Dataset',"
+                f" but was {type(source_ds).__name__!r}"
+            )
+        if link_level_zero:
+            raise ValueError(
+                f"'source_path' argument must be provided"
+                f" if 'link_level_zero' is used"
+            )
+
+    append_dim = zappend_config.pop("append_dim", "time")
+    append_coord = source_ds.coords[append_dim]
+
+    if source_append_offset is None:
+        source_append_offset = 0
+    elif not isinstance(source_append_offset, int):
+        raise TypeError(
+            f"'source_append_offset' argument must be of type 'int',"
+            f" but was {type(source_append_offset).__name__!r}"
+        )
+    if not (0 <= source_append_offset < append_coord.size):
+        raise ValueError(
+            f"'source_append_offset' argument"
+            f" must be >=0 and <{append_coord.size},"
+            f" but was {source_append_offset}"
+        )
 
     logger = logging.getLogger("zappend")
 
@@ -157,11 +203,6 @@ def write_levels(
         agg_methods=agg_methods,
     )
 
-    force_new = zappend_config.pop("force_new", None)
-
-    append_dim = zappend_config.pop("append_dim", "time")
-    append_coord = source_ds.coords[append_dim]
-
     variables = get_variables_config(
         source_ds,
         {
@@ -207,9 +248,10 @@ def write_levels(
 
     subsample_dataset_kwargs = dict(xy_dim_names=xy_dim_names, agg_methods=agg_methods)
 
-    num_slices = append_coord.size
+    num_slices = append_coord.size - source_append_offset
     for slice_index in range(num_slices):
-        slice_ds_indexer = {append_dim: slice(slice_index, slice_index + 1)}
+        append_index = source_append_offset + slice_index
+        slice_ds_indexer = {append_dim: slice(append_index, append_index + 1)}
         slice_ds = source_ds.isel(slice_ds_indexer)
 
         for level_index in range(num_levels):

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# Copyright © 2024 Norman Fomferra and contributors`
	`2`	`+# Permissions are hereby granted under the terms of the MIT License:`
	`3`	`+# https://opensource.org/licenses/MIT.`