Merge pull request #622 from dcs4cop/forman-617-levels_on_s3

Introduced parameter "base_dataset_id" for writing multi-level datasets

Author: forman

CHANGES.md

@@ -2,6 +2,12 @@
 
 ### Enhancements
 
+* Introduced parameter `base_dataset_id` for writing multi-level 
+  datasets with the "file", "s3", and "memory" data stores. 
+  If given, the base dataset will be linked only with the 
+  value of `base_dataset_id`, instead of being copied as-is.
+  This can save large amounts of storage space. (#617)
+
 ### Fixes
 
 * Fixed `FsDataAccessor.write_data()` implementations, 

test/core/store/fs/test_registry.py

@@ -1,7 +1,8 @@
 import os.path
 import unittest
+import warnings
 from abc import ABC, abstractmethod
-from typing import Type, Union
+from typing import Any, Dict, Optional, Type, Union
 
 import fsspec
 import xarray as xr
@@ -38,7 +39,7 @@ def prepare_fs(cls, fs: fsspec.AbstractFileSystem, root: str):
         # print(f'{fs.protocol}: making root {root}')
         fs.mkdirs(root)
 
-        # Write a text file into each subdirectory so
+        # Write a text file into each subdirectory, so
         # we also test that store.get_data_ids() scans
         # recursively.
         dir_path = root
@@ -54,6 +55,7 @@ def prepare_fs(cls, fs: fsspec.AbstractFileSystem, root: str):
     def test_mldataset_levels(self):
         data_store = self.create_data_store()
         self.assertMultiLevelDatasetFormatSupported(data_store)
+        self.assertMultiLevelDatasetFormatWithLinkSupported(data_store)
 
     def test_dataset_zarr(self):
         data_store = self.create_data_store()
@@ -73,6 +75,47 @@ def assertMultiLevelDatasetFormatSupported(self,
                                     MultiLevelDataset,
                                     MultiLevelDatasetDescriptor)
 
+        # Test that use_saved_levels works
+        self.assertDatasetSupported(data_store,
+                                    '.levels',
+                                    'mldataset',
+                                    MultiLevelDataset,
+                                    MultiLevelDatasetDescriptor,
+                                    write_params=dict(
+                                        use_saved_levels=True,
+                                    ))
+
+    def assertMultiLevelDatasetFormatWithLinkSupported(
+            self,
+            data_store: MutableDataStore
+    ):
+        base_dataset = self.new_cube_data()
+        base_dataset_id = f'{DATA_PATH}/base-ds.zarr'
+        data_store.write_data(base_dataset, base_dataset_id)
+
+        # Test that base_dataset_id works
+        self.assertDatasetSupported(data_store,
+                                    '.levels',
+                                    'mldataset',
+                                    MultiLevelDataset,
+                                    MultiLevelDatasetDescriptor,
+                                    write_params=dict(
+                                        base_dataset_id=base_dataset_id,
+                                    ))
+
+        # Test that base_dataset_id + use_saved_levels works
+        self.assertDatasetSupported(data_store,
+                                    '.levels',
+                                    'mldataset',
+                                    MultiLevelDataset,
+                                    MultiLevelDatasetDescriptor,
+                                    write_params=dict(
+                                        base_dataset_id=base_dataset_id,
+                                        use_saved_levels=True,
+                                    ))
+
+        data_store.delete_data(base_dataset_id)
+
     def assertDatasetFormatSupported(self,
                                      data_store: MutableDataStore,
                                      filename_ext: str):
@@ -89,8 +132,12 @@ def assertDatasetSupported(
             expected_data_type_alias: str,
             expected_type: Union[Type[xr.Dataset],
                                  Type[MultiLevelDataset]],
-            expected_descriptor_type: Union[Type[DatasetDescriptor],
-                                            Type[MultiLevelDatasetDescriptor]]
+            expected_descriptor_type: Union[
+                Type[DatasetDescriptor],
+                Type[MultiLevelDatasetDescriptor]
+            ],
+            write_params: Optional[Dict[str, Any]] = None,
+            open_params: Optional[Dict[str, Any]] = None,
     ):
         """
         Call all DataStore operations to ensure data of type
@@ -102,10 +149,15 @@ def assertDatasetSupported(
         :param expected_data_type_alias: The expected data type alias.
         :param expected_type: The expected data type.
         :param expected_descriptor_type: The expected data descriptor type.
+        :param write_params: Optional write parameters
+        :param open_params: Optional open parameters
         """
 
         data_id = f'{DATA_PATH}/ds{filename_ext}'
 
+        write_params = write_params or {}
+        open_params = open_params or {}
+
         self.assertIsInstance(data_store, MutableDataStore)
 
         self.assertEqual({'dataset', 'mldataset', 'geodataframe'},
@@ -114,36 +166,44 @@ def assertDatasetSupported(
         with self.assertRaises(DataStoreError):
             data_store.get_data_types_for_data(data_id)
         self.assertEqual(False, data_store.has_data(data_id))
-        self.assertEqual([], list(data_store.get_data_ids()))
+        self.assertNotIn(data_id, set(data_store.get_data_ids()))
 
-        data = new_cube(variables=dict(A=8, B=9))
-        written_data_id = data_store.write_data(data, data_id)
+        data = self.new_cube_data()
+        written_data_id = data_store.write_data(data, data_id, **write_params)
         self.assertEqual(data_id, written_data_id)
+
         self.assertEqual({expected_data_type_alias},
                          set(data_store.get_data_types_for_data(data_id)))
         self.assertEqual(True, data_store.has_data(data_id))
-        self.assertEqual([data_id], list(data_store.get_data_ids()))
+        self.assertIn(data_id, set(data_store.get_data_ids()))
 
-        data_descriptors = list(data_store.search_data())
+        data_descriptors = list(data_store.search_data(
+            data_type=expected_type)
+        )
         self.assertEqual(1, len(data_descriptors))
         self.assertIsInstance(data_descriptors[0], DataDescriptor)
         self.assertIsInstance(data_descriptors[0], expected_descriptor_type)
 
-        data = data_store.open_data(data_id)
+        data = data_store.open_data(data_id, **open_params)
         self.assertIsInstance(data, expected_type)
 
         try:
             data_store.delete_data(data_id)
-        except PermissionError:  # Typically occurs on win32 due to fsspec
+        except PermissionError as e:  # May occur on win32 due to fsspec
+            warnings.warn(f'{e}')
             return
         with self.assertRaises(DataStoreError):
             data_store.get_data_types_for_data(data_id)
         self.assertEqual(False, data_store.has_data(data_id))
-        self.assertEqual([], list(data_store.get_data_ids()))
+        self.assertNotIn(data_id, set(data_store.get_data_ids()))
 
+    @staticmethod
+    def new_cube_data():
+        cube = new_cube(variables=dict(A=8.5, B=9.5))
+        return cube.chunk(dict(time=1, lat=90, lon=180))
 
-class FileFsDataStoresTest(FsDataStoresTestMixin, unittest.TestCase):
 
+class FileFsDataStoresTest(FsDataStoresTestMixin, unittest.TestCase):
     def create_data_store(self) -> FsDataStore:
         root = os.path.join(new_temp_dir(prefix='xcube'), ROOT_DIR)
         self.prepare_fs(fsspec.filesystem('file'), root)

xcube/core/store/fs/accessor.py

@@ -21,7 +21,7 @@
 
 import copy
 from abc import ABC, abstractmethod
-from typing import Dict, Any, Tuple
+from typing import Dict, Any, Tuple, Optional
 
 import fsspec
 
@@ -52,6 +52,7 @@
 PROTOCOL_PARAM_NAME = 'protocol'
 STORAGE_OPTIONS_PARAM_NAME = 'storage_options'
 FS_PARAM_NAME = 'fs'
+ROOT_PARAM_NAME = 'root'
 
 
 class FsAccessor:
@@ -74,7 +75,9 @@ def get_storage_options_schema(cls) -> JsonObjectSchema:
 
     @classmethod
     def load_fs(cls, params: Dict[str, Any]) \
-            -> Tuple[fsspec.AbstractFileSystem, Dict[str, Any]]:
+            -> Tuple[fsspec.AbstractFileSystem,
+                     Optional[str],
+                     Dict[str, Any]]:
         """
         Load a filesystem instance from *params*.
 
@@ -85,17 +88,21 @@ def load_fs(cls, params: Dict[str, Any]) \
 
         :param params: Parameters passed to a filesystem
             data store, opener, or writer call.
-        :return: A tuple comprising the filesystem
+        :return: A tuple comprising the filesystem, an optional root path,
             and the modified *params*.
         """
         params = dict(params)
 
-        # Filesystem data stores pass "fs" kwarg to
+        # Filesystem data-stores pass "fs" and "root" kwargs to
         # data opener and writer calls.
         fs = params.pop(FS_PARAM_NAME, None)
+        root = params.pop(ROOT_PARAM_NAME, None)
         if fs is not None:
             assert_instance(fs, fsspec.AbstractFileSystem, name=FS_PARAM_NAME)
-            return fs, params
+        if root is not None:
+            assert_instance(root, str, name=ROOT_PARAM_NAME)
+        if fs:
+            return fs, root, params
 
         protocol = cls.get_protocol()
         if protocol == 'abstract':
@@ -117,9 +124,13 @@ def load_fs(cls, params: Dict[str, Any]) \
                  if storage_options else False)
 
         try:
-            return fsspec.filesystem(protocol,
-                                     use_listings_cache=use_listings_cache,
-                                     **(storage_options or {})), params
+            return (
+                fsspec.filesystem(protocol,
+                                  use_listings_cache=use_listings_cache,
+                                  **(storage_options or {})),
+                root,
+                params
+            )
         except (ValueError, ImportError):
             raise DataStoreError(f"Cannot instantiate"
                                  f" filesystem {protocol!r}")
@@ -189,5 +200,5 @@ def get_delete_data_params_schema(self, data_id: str = None) \
     def delete_data(self,
                     data_id: str,
                     **delete_params):
-        fs, delete_params = self.load_fs(delete_params)
+        fs, _, delete_params = self.load_fs(delete_params)
         fs.delete(data_id, **delete_params)

xcube/core/store/fs/impl/dataset.py

@@ -20,7 +20,7 @@
 # SOFTWARE.
 
 from abc import ABC
-from typing import Tuple
+from typing import Tuple, Optional
 
 import xarray as xr
 import zarr
@@ -121,8 +121,10 @@
             additional_properties=True,
         ),
         consolidated=JsonBooleanSchema(
-            description='If True, apply Zarr’s consolidate_metadata()'
-                        ' function to the store after writing.'
+            description='If True (the default), consolidate all metadata'
+                        ' files ("**/.zarray", "**/.zattrs")'
+                        ' into a single top-level file ".zmetadata"',
+            default=True,
         ),
         append_dim=JsonStringSchema(
             description='If set, the dimension on which the'
@@ -160,9 +162,11 @@ def get_open_data_params_schema(self, data_id: str = None) \
             ZARR_OPEN_DATA_PARAMS_SCHEMA
         )
 
-    def open_data(self, data_id: str, **open_params) -> xr.Dataset:
+    def open_data(self,
+                  data_id: str,
+                  **open_params) -> xr.Dataset:
         assert_instance(data_id, str, name='data_id')
-        fs, open_params = self.load_fs(open_params)
+        fs, root, open_params = self.load_fs(open_params)
         zarr_store = fs.get_mapper(data_id)
         cache_size = open_params.pop('cache_size', None)
         if isinstance(cache_size, int) and cache_size > 0:
@@ -194,7 +198,7 @@ def write_data(self,
                    **write_params) -> str:
         assert_instance(data, xr.Dataset, name='data')
         assert_instance(data_id, str, name='data_id')
-        fs, write_params = self.load_fs(write_params)
+        fs, root, write_params = self.load_fs(write_params)
         zarr_store = fs.get_mapper(data_id, create=True)
         log_access = write_params.pop('log_access', None)
         if log_access:
@@ -214,7 +218,7 @@ def write_data(self,
     def delete_data(self,
                     data_id: str,
                     **delete_params):
-        fs, delete_params = self.load_fs(delete_params)
+        fs, root, delete_params = self.load_fs(delete_params)
         delete_params.pop('recursive', None)
         fs.delete(data_id, recursive=True, **delete_params)
 
@@ -253,7 +257,7 @@ def open_data(self,
                   data_id: str,
                   **open_params) -> xr.Dataset:
         assert_instance(data_id, str, name='data_id')
-        fs, open_params = self.load_fs(open_params)
+        fs, root, open_params = self.load_fs(open_params)
 
         # This doesn't yet work as expected with fsspec and netcdf:
         # engine = open_params.pop('engine', 'scipy')
@@ -281,7 +285,7 @@ def write_data(self,
                    **write_params) -> str:
         assert_instance(data, xr.Dataset, name='data')
         assert_instance(data_id, str, name='data_id')
-        fs, write_params = self.load_fs(write_params)
+        fs, root, write_params = self.load_fs(write_params)
         if not replace and fs.exists(data_id):
             raise DataStoreError(f'Data resource {data_id} already exists')
 

xcube/core/store/fs/impl/geodataframe.py

@@ -20,7 +20,7 @@
 # SOFTWARE.
 
 from abc import abstractmethod, ABC
-from typing import Tuple
+from typing import Tuple, Optional
 
 import geopandas as gpd
 import pandas as pd
@@ -60,7 +60,7 @@ def get_open_data_params_schema(self, data_id: str = None) \
     def open_data(self, data_id: str, **open_params) -> gpd.GeoDataFrame:
         # TODO: implement me correctly,
         #  this is not valid for shapefile AND geojson
-        fs, open_params = self.load_fs(open_params)
+        fs, root, open_params = self.load_fs(open_params)
         is_local = is_local_fs(fs)
         if is_local:
             file_path = data_id
@@ -86,7 +86,7 @@ def write_data(self,
         # TODO: implement me correctly,
         #  this is not valid for shapefile AND geojson
         assert_instance(data, (gpd.GeoDataFrame, pd.DataFrame), 'data')
-        fs, write_params = self.load_fs(write_params)
+        fs, root, write_params = self.load_fs(write_params)
         is_local = is_local_fs(fs)
         if is_local:
             file_path = data_id