Repository: dcs4cop/xcube
Pull request #275 (2020-02-11 16:45:37)

improved efficiency of unchunk()
Optimization: if `"shape"` and `"chunks"` are equal in `${var}/.zarray`, do not open the zarr array. This should improve efficiency of unchunk().


Patch:
diff --git a/CHANGES.md b/CHANGES.md
index f756934b..19f0b4ec 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -110,6 +110,10 @@
 
 ### Other
 
+* The `unchunk_coords` keyword argument of Python API function 
+  `xcube.core.optimize.optimize_dataset()` can now be a name, or list of names  
+  of the coordinate variable(s) to be consolidated. If boolean ``True`` is used
+  all variables will be consolidated.
 * The `xcube serve` API operations `datasets/` and `datasets/{ds_id}` now also
   return the metadata attributes of a given dataset and it variables in a property
   named `attrs`. For variables we added a new metadata property `htmlRepr` that is
diff --git a/xcube/core/optimize.py b/xcube/core/optimize.py
index 6f3d03ee..a2671b4e 100644
--- a/xcube/core/optimize.py
+++ b/xcube/core/optimize.py
@@ -21,7 +21,7 @@
 
 import os.path
 import shutil
-from typing import Type
+from typing import Type, Union, Sequence
 
 import zarr
 
@@ -31,7 +31,7 @@ from xcube.core.unchunk import unchunk_dataset
 def optimize_dataset(input_path: str,
                      output_path: str = None,
                      in_place: bool = False,
-                     unchunk_coords: bool = False,
+                     unchunk_coords: Union[bool, str, Sequence[str]] = False,
                      exception_type: Type[Exception] = ValueError):
     """
     Optimize a dataset for faster access.
@@ -40,17 +40,21 @@ def optimize_dataset(input_path: str,
     Consolidated cubes open much faster from remote locations, e.g. in object storage,
     because obviously much less HTTP requests are required to fetch initial cube meta
     information. That is, it merges all metadata files into a single top-level JSON file ".zmetadata".
-    If *unchunk_coords* is set, it also removes any chunking of coordinate variables
+
+    If *unchunk_coords* is given, it also removes any chunking of coordinate variables
     so they comprise a single binary data file instead of one file per data chunk.
     The primary usage of this function is to optimize data cubes for cloud object storage.
-    The function currently works only for data cubes using ZARR format.
+    The function currently works only for data cubes using Zarr format.
+    *unchunk_coords* can be a name, or list of names of the coordinate variable(s) to be consolidated.
+    If boolean ``True`` is used, coordinate all variables will be consolidated.
 
     :param input_path: Path to input dataset with ZARR format.
     :param output_path: Path to output dataset with ZARR format. May contain "{input}" template string,
-           which is replaced by the input path's file name without file name extentsion.
+           which is replaced by the input path's file name without file name extension.
     :param in_place: Whether to modify the dataset in place.
            If False, a copy is made and *output_path* must be given.
-    :param unchunk_coords: Whether to also consolidate coordinate chunk files.
+    :param unchunk_coords: The name of a coordinate variable or a list of coordinate variables whose chunks should
+           be consolidated. Pass ``True`` to consolidate chunks of all coordinate variables.
     :param exception_type: Type of exception to be used on value errors.
     """
 
@@ -75,6 +79,12 @@ def optimize_dataset(input_path: str,
         shutil.copytree(input_path, output_path)
 
     if unchunk_coords:
-        unchunk_dataset(output_path, coords_only=True)
+        if isinstance(unchunk_coords, str):
+            var_names = (unchunk_coords,)
+        elif isinstance(unchunk_coords, bool):
+            var_names = None
+        else:
+            var_names = tuple(unchunk_coords)
+        unchunk_dataset(output_path, var_names=var_names, coords_only=True)
 
     zarr.convenience.consolidate_metadata(output_path)
diff --git a/xcube/core/unchunk.py b/xcube/core/unchunk.py
index be5bfe60..81745d24 100644
--- a/xcube/core/unchunk.py
+++ b/xcube/core/unchunk.py
@@ -19,6 +19,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import json
 import os.path
 from typing import List, Sequence
 
@@ -38,7 +39,7 @@ def unchunk_dataset(dataset_path: str, var_names: Sequence[str] = None, coords_o
 
     is_zarr = os.path.isfile(os.path.join(dataset_path, '.zgroup'))
     if not is_zarr:
-        raise ValueError(f'{dataset_path!r} is not a valid ZARR directory')
+        raise ValueError(f'{dataset_path!r} is not a valid Zarr directory')
 
     with xr.open_zarr(dataset_path) as dataset:
         if var_names is None:
@@ -46,14 +47,14 @@ def unchunk_dataset(dataset_path: str, var_names: Sequence[str] = None, coords_o
                 var_names = list(dataset.coords)
             else:
                 var_names = list(dataset.variables)
-
-        for var_name in var_names:
-            if coords_only:
-                if var_name not in dataset.coords:
-                    raise ValueError(f'variable {var_name!r} is not a coordinate variable in {dataset_path!r}')
-            else:
-                if var_name not in dataset.variables:
-                    raise ValueError(f'variable {var_name!r} is not a variable in {dataset_path!r}')
+        else:
+            for var_name in var_names:
+                if coords_only:
+                    if var_name not in dataset.coords:
+                        raise ValueError(f'variable {var_name!r} is not a coordinate variable in {dataset_path!r}')
+                else:
+                    if var_name not in dataset.variables:
+                        raise ValueError(f'variable {var_name!r} is not a variable in {dataset_path!r}')
 
     _unchunk_vars(dataset_path, var_names)
 
@@ -61,6 +62,15 @@ def unchunk_dataset(dataset_path: str, var_names: Sequence[str] = None, coords_o
 def _unchunk_vars(dataset_path: str, var_names: List[str]):
     for var_name in var_names:
         var_path = os.path.join(dataset_path, var_name)
+
+        # Optimization: if "shape" and "chunks" are equal in ${var}/.zarray, we are done
+        var_array_info_path = os.path.join(var_path, '.zarray')
+        with open(var_array_info_path, 'r') as fp:
+            var_array_info = json.load(fp)
+            if var_array_info.get('shape') == var_array_info.get('chunks'):
+                continue
+
+        # Open array and remove chunks from the data
         var_array = zarr.convenience.open_array(var_path, 'r+')
         if var_array.shape != var_array.chunks:
             # TODO (forman): Fully loading data is inefficient and dangerous for large arrays.


Test patch:
diff --git a/test/cli/test_optimize.py b/test/cli/test_optimize.py
index 6f09ca2c..041c1102 100644
--- a/test/cli/test_optimize.py
+++ b/test/cli/test_optimize.py
@@ -1,55 +1,55 @@
 from test.cli.helpers import CliDataTest
-from test.core.test_optimize import TEST_CUBE, TEST_CUBE_ZARR, TEST_CUBE_FILE_SET, list_file_set
+from test.core.test_optimize import TEST_CUBE, INPUT_CUBE_PATH, INPUT_CUBE_FILE_SET, list_file_set
 from xcube.core.dsio import rimraf
 
-TEST_CUBE_ZARR_OPTIMIZED_DEFAULT = 'test-optimized.zarr'
-TEST_CUBE_ZARR_OPTIMIZED_USER = 'fast-test.zarr'
+OUTPUT_CUBE_OPTIMIZED_DEFAULT_PATH = 'test-optimized.zarr'
+OUTPUT_CUBE_OPTIMIZED_USER_PATH = 'fast-test.zarr'
 
 
 class OptimizeDataTest(CliDataTest):
     def _clear_outputs(self):
-        rimraf(TEST_CUBE_ZARR)
-        rimraf(TEST_CUBE_ZARR_OPTIMIZED_DEFAULT)
-        rimraf(TEST_CUBE_ZARR_OPTIMIZED_USER)
+        rimraf(INPUT_CUBE_PATH)
+        rimraf(OUTPUT_CUBE_OPTIMIZED_DEFAULT_PATH)
+        rimraf(OUTPUT_CUBE_OPTIMIZED_USER_PATH)
 
     def setUp(self):
         self._clear_outputs()
-        TEST_CUBE.to_zarr(TEST_CUBE_ZARR)
+        TEST_CUBE.to_zarr(INPUT_CUBE_PATH)
 
     def tearDown(self):
         self._clear_outputs()
 
     def test_defaults(self):
-        result = self.invoke_cli(['optimize', TEST_CUBE_ZARR])
+        result = self.invoke_cli(['optimize', INPUT_CUBE_PATH])
         self.assertEqual(0, result.exit_code)
 
-        expected_files = set(TEST_CUBE_FILE_SET)
+        expected_files = set(INPUT_CUBE_FILE_SET)
         expected_files.add('.zmetadata')
-        self.assertEqual(expected_files, list_file_set(TEST_CUBE_ZARR_OPTIMIZED_DEFAULT))
+        self.assertEqual(expected_files, list_file_set(OUTPUT_CUBE_OPTIMIZED_DEFAULT_PATH))
 
     def test_user_output(self):
-        result = self.invoke_cli(['optimize', '-o', TEST_CUBE_ZARR_OPTIMIZED_USER, '-C', TEST_CUBE_ZARR])
+        result = self.invoke_cli(['optimize', '-o', OUTPUT_CUBE_OPTIMIZED_USER_PATH, '-C', INPUT_CUBE_PATH])
         self.assertEqual(0, result.exit_code)
 
-        expected_files = set(TEST_CUBE_FILE_SET)
+        expected_files = set(INPUT_CUBE_FILE_SET)
         expected_files.add('.zmetadata')
         expected_files.remove('time/1')
         expected_files.remove('time/2')
         expected_files.remove('time_bnds/1.0')
         expected_files.remove('time_bnds/2.0')
-        self.assertEqual(expected_files, list_file_set(TEST_CUBE_ZARR_OPTIMIZED_USER))
+        self.assertEqual(expected_files, list_file_set(OUTPUT_CUBE_OPTIMIZED_USER_PATH))
 
     def test_in_place(self):
-        result = self.invoke_cli(['optimize', '-IC', TEST_CUBE_ZARR])
+        result = self.invoke_cli(['optimize', '-IC', INPUT_CUBE_PATH])
         self.assertEqual(0, result.exit_code)
 
-        expected_files = set(TEST_CUBE_FILE_SET)
+        expected_files = set(INPUT_CUBE_FILE_SET)
         expected_files.add('.zmetadata')
         expected_files.remove('time/1')
         expected_files.remove('time/2')
         expected_files.remove('time_bnds/1.0')
         expected_files.remove('time_bnds/2.0')
-        self.assertEqual(expected_files, list_file_set(TEST_CUBE_ZARR))
+        self.assertEqual(expected_files, list_file_set(INPUT_CUBE_PATH))
 
 
 class OptimizeTest(CliDataTest):
diff --git a/test/core/test_optimize.py b/test/core/test_optimize.py
index 98d59b29..ab30ded7 100644
--- a/test/core/test_optimize.py
+++ b/test/core/test_optimize.py
@@ -1,20 +1,22 @@
 import os
 import os.path
 import unittest
-from typing import Set
+from typing import Set, Union, Sequence
 
-from xcube.core.new import new_cube
-from xcube.core.chunk import chunk_dataset
 from xcube.constants import FORMAT_NAME_ZARR
+from xcube.core.chunk import chunk_dataset
 from xcube.core.dsio import rimraf
+from xcube.core.new import new_cube
 from xcube.core.optimize import optimize_dataset
 
 TEST_CUBE = chunk_dataset(new_cube(time_periods=3, variables=dict(A=0.5, B=-1.5)),
                           chunk_sizes=dict(time=1, lat=180, lon=360), format_name=FORMAT_NAME_ZARR)
 
-TEST_CUBE_ZARR = 'test.zarr'
+INPUT_CUBE_PATH = 'test.zarr'
+OUTPUT_CUBE_PATH = 'test_opt.zarr'
+OUTPUT_CUBE_PATTERN = '{input}_opt.zarr'
 
-TEST_CUBE_FILE_SET = {
+INPUT_CUBE_FILE_SET = {
     '.zattrs', '.zgroup',
     'A/.zarray', 'A/.zattrs', 'A/0.0.0', 'A/1.0.0', 'A/2.0.0',
     'B/.zarray', 'B/.zattrs', 'B/0.0.0', 'B/1.0.0', 'B/2.0.0',
@@ -28,35 +30,103 @@ TEST_CUBE_FILE_SET = {
 
 
 class OptimizeDatasetTest(unittest.TestCase):
+    def _clear_outputs(self):
+        rimraf(INPUT_CUBE_PATH)
+        rimraf(OUTPUT_CUBE_PATH)
 
     def setUp(self):
-        rimraf(TEST_CUBE_ZARR)
-        TEST_CUBE.to_zarr(TEST_CUBE_ZARR)
+        self._clear_outputs()
+        TEST_CUBE.to_zarr(INPUT_CUBE_PATH)
+        self.assertEqual(INPUT_CUBE_FILE_SET, list_file_set(INPUT_CUBE_PATH))
 
     def tearDown(self):
-        rimraf(TEST_CUBE_ZARR)
-
-    def test_optimize_dataset_in_place(self):
-        self.assertEqual(TEST_CUBE_FILE_SET, list_file_set(TEST_CUBE_ZARR))
+        self._clear_outputs()
 
-        optimize_dataset(TEST_CUBE_ZARR, in_place=True)
+    def test_optimize_dataset(self):
+        self._test_optimize_dataset(unchunk_coords=False, in_place=False,
+                                    expected_output_path=OUTPUT_CUBE_PATH,
+                                    expected_cons_time=False,
+                                    expected_cons_time_bnds=False)
 
-        expected_files = set(TEST_CUBE_FILE_SET)
-        expected_files.add('.zmetadata')
-        self.assertEqual(expected_files, list_file_set(TEST_CUBE_ZARR))
-
-    def test_optimize_dataset_in_place_unchunk_coords(self):
-        self.assertEqual(TEST_CUBE_FILE_SET, list_file_set(TEST_CUBE_ZARR))
-
-        optimize_dataset(TEST_CUBE_ZARR, in_place=True, unchunk_coords=True)
-
-        expected_files = set(TEST_CUBE_FILE_SET)
+    def test_optimize_dataset_in_place(self):
+        self._test_optimize_dataset(unchunk_coords=False, in_place=True,
+                                    expected_output_path=INPUT_CUBE_PATH,
+                                    expected_cons_time=False,
+                                    expected_cons_time_bnds=False)
+
+    def test_optimize_dataset_unchunk_coords_true(self):
+        self._test_optimize_dataset(unchunk_coords=True, in_place=False,
+                                    expected_output_path=OUTPUT_CUBE_PATH,
+                                    expected_cons_time=True,
+                                    expected_cons_time_bnds=True)
+
+    def test_optimize_dataset_unchunk_coords_true_in_place(self):
+        self._test_optimize_dataset(unchunk_coords=True, in_place=True,
+                                    expected_output_path=INPUT_CUBE_PATH,
+                                    expected_cons_time=True,
+                                    expected_cons_time_bnds=True)
+
+    def test_optimize_dataset_unchunk_coords_str(self):
+        self._test_optimize_dataset(unchunk_coords='time', in_place=False,
+                                    expected_output_path=OUTPUT_CUBE_PATH,
+                                    expected_cons_time=True,
+                                    expected_cons_time_bnds=False)
+
+    def test_optimize_dataset_unchunk_coords_str_in_place(self):
+        self._test_optimize_dataset(unchunk_coords='time', in_place=True,
+                                    expected_output_path=INPUT_CUBE_PATH,
+                                    expected_cons_time=True,
+                                    expected_cons_time_bnds=False)
+
+    def test_optimize_dataset_unchunk_coords_str_list(self):
+        self._test_optimize_dataset(unchunk_coords=['time_bnds'], in_place=False,
+                                    expected_output_path=OUTPUT_CUBE_PATH,
+                                    expected_cons_time=False,
+                                    expected_cons_time_bnds=True)
+
+    def test_optimize_dataset_unchunk_coords_str_list_in_place(self):
+        self._test_optimize_dataset(unchunk_coords=['time_bnds'], in_place=True,
+                                    expected_output_path=INPUT_CUBE_PATH,
+                                    expected_cons_time=False,
+                                    expected_cons_time_bnds=True)
+
+    def test_optimize_dataset_unchunk_coords_str_tuple(self):
+        self._test_optimize_dataset(unchunk_coords=('time_bnds', 'time'), in_place=False,
+                                    expected_output_path=OUTPUT_CUBE_PATH,
+                                    expected_cons_time=True,
+                                    expected_cons_time_bnds=True)
+
+    def test_optimize_dataset_unchunk_coords_str_tuple_in_place(self):
+        self._test_optimize_dataset(unchunk_coords=('time_bnds', 'time'), in_place=True,
+                                    expected_output_path=INPUT_CUBE_PATH,
+                                    expected_cons_time=True,
+                                    expected_cons_time_bnds=True)
+
+    def _test_optimize_dataset(self, unchunk_coords: Union[bool, str, Sequence[str]], in_place: bool,
+                               expected_output_path: str, expected_cons_time: bool = False,
+                               expected_cons_time_bnds: bool = False):
+        if not in_place:
+            optimize_dataset(INPUT_CUBE_PATH, output_path=OUTPUT_CUBE_PATTERN,
+                             in_place=in_place, unchunk_coords=unchunk_coords)
+        else:
+            optimize_dataset(INPUT_CUBE_PATH,
+                             in_place=in_place, unchunk_coords=unchunk_coords)
+        self._assert_consolidated(expected_output_path, expected_cons_time, expected_cons_time_bnds)
+
+    def _assert_consolidated(self,
+                             cube_path: str,
+                             cons_time: bool = False,
+                             cons_time_bnds: bool = False):
+        self.assertTrue(os.path.isdir(cube_path))
+        expected_files = set(INPUT_CUBE_FILE_SET)
         expected_files.add('.zmetadata')
-        expected_files.remove('time/1')
-        expected_files.remove('time/2')
-        expected_files.remove('time_bnds/1.0')
-        expected_files.remove('time_bnds/2.0')
-        self.assertEqual(expected_files, list_file_set(TEST_CUBE_ZARR))
+        if cons_time:
+            expected_files.remove('time/1')
+            expected_files.remove('time/2')
+        if cons_time_bnds:
+            expected_files.remove('time_bnds/1.0')
+            expected_files.remove('time_bnds/2.0')
+        self.assertEqual(expected_files, list_file_set(cube_path))
 
     def test_failures(self):
         with self.assertRaises(RuntimeError) as cm:
@@ -64,15 +134,15 @@ class OptimizeDatasetTest(unittest.TestCase):
         self.assertEqual('Input path must point to ZARR dataset directory.', f'{cm.exception}')
 
         with self.assertRaises(RuntimeError) as cm:
-            optimize_dataset(TEST_CUBE_ZARR, exception_type=RuntimeError)
+            optimize_dataset(INPUT_CUBE_PATH, exception_type=RuntimeError)
         self.assertEqual('Output path must be given.', f'{cm.exception}')
 
         with self.assertRaises(RuntimeError) as cm:
-            optimize_dataset(TEST_CUBE_ZARR, output_path=TEST_CUBE_ZARR, exception_type=RuntimeError)
+            optimize_dataset(INPUT_CUBE_PATH, output_path=INPUT_CUBE_PATH, exception_type=RuntimeError)
         self.assertEqual('Output path already exists.', f'{cm.exception}')
 
         with self.assertRaises(RuntimeError) as cm:
-            optimize_dataset(TEST_CUBE_ZARR, output_path='./' + TEST_CUBE_ZARR, exception_type=RuntimeError)
+            optimize_dataset(INPUT_CUBE_PATH, output_path='./' + INPUT_CUBE_PATH, exception_type=RuntimeError)
         self.assertEqual('Output path already exists.', f'{cm.exception}')
 
 
diff --git a/test/core/test_unchunk.py b/test/core/test_unchunk.py
index a5d6da44..b36e4a3b 100644
--- a/test/core/test_unchunk.py
+++ b/test/core/test_unchunk.py
@@ -1,10 +1,10 @@
 import os.path
 import unittest
 
-from xcube.core.new import new_cube
-from xcube.core.chunk import chunk_dataset
 from xcube.constants import FORMAT_NAME_ZARR
+from xcube.core.chunk import chunk_dataset
 from xcube.core.dsio import rimraf
+from xcube.core.new import new_cube
 from xcube.core.unchunk import unchunk_dataset
 
 
@@ -93,7 +93,16 @@ class UnchunkDatasetTest(unittest.TestCase):
     def test_unchunk_invalid_path(self):
         with self.assertRaises(ValueError) as cm:
             unchunk_dataset(self.TEST_ZARR + '.zip')
-        self.assertEqual("'test.zarr.zip' is not a valid ZARR directory", f'{cm.exception}')
+        self.assertEqual("'test.zarr.zip' is not a valid Zarr directory", f'{cm.exception}')
+
+    def test_unchunk_invalid_vars(self):
+        with self.assertRaises(ValueError) as cm:
+            unchunk_dataset(self.TEST_ZARR, var_names=['times'], coords_only=True)
+        self.assertEqual("variable 'times' is not a coordinate variable in 'test.zarr'", f'{cm.exception}')
+
+        with self.assertRaises(ValueError) as cm:
+            unchunk_dataset(self.TEST_ZARR, var_names=['CHL'], coords_only=False)
+        self.assertEqual("variable 'CHL' is not a variable in 'test.zarr'", f'{cm.exception}')
 
     def _assert_cube_files(self,
                            expected_a_files, expected_b_files,