Issue 538: Skip individual split files that already exist

VERSION.txt

@@ -1 +1 @@
-0.2.2
+0.2.3

weather_sp/setup.py

@@ -44,7 +44,7 @@
     packages=find_packages(),
     author='Anthromets',
     author_email='anthromets-ecmwf@google.com',
-    version='0.3.10',
+    version='0.3.11',
     url='https://weather-tools.readthedocs.io/en/latest/weather_sp/',
     description='A tool to split weather data files into per-variable files.',
     install_requires=beam_gcp_requirements + base_requirements,

weather_sp/splitter_pipeline/file_splitters.py

@@ -162,17 +162,12 @@ def split_data(self) -> None:
         if not self.output_info.split_dims():
             raise ValueError('No splitting specified in template.')
 
-        if self.should_skip():
-            metrics.Metrics.counter('file_splitters', 'skipped').inc()
-            self.logger.info('Skipping %s, file already split.',
-                             repr(self.input_path))
-            return
-
         # Here, we keep a map of open file objects (`outputs`). We need these since
         # each output grib file (named `key`) will include multiple `grb` messages
         # each. By writing data to the cache of open file objects, we can keep a
         # minimal amount of data in memory at a time.
         outputs = dict()
+        skipped_keys = set()
         with self._open_grib_locally() as grbs:
             self.logger.info('Splitting & uploading %r...', self.input_path)
             try:
@@ -188,10 +183,20 @@ def split_data(self) -> None:
                                 'Variable not found in grib: %s', dim)
                     key = self.output_info.formatted_output_path(splits)
 
+                    if key in skipped_keys:
+                        del grb
+                        continue
+
                     # Append the current grib message to a set number of output files.
                     # If the target shard doesn't exist, create it.
                     if key not in outputs:
+                        if self.should_skip_file(key):
+                            self.logger.info('Skipping %s, file already split.', repr(key))
+                            skipped_keys.add(key)
+                            del grb
+                            continue
                         outputs[key] = FileSystems.create(key)
+
                     outputs[key].write(grb.tostring())
                     outputs[key].flush()
 
@@ -202,6 +207,12 @@ def split_data(self) -> None:
             finally:
                 for out in outputs.values():
                     out.close()
+
+        if not outputs:
+            metrics.Metrics.counter('file_splitters', 'skipped').inc()
+            self.logger.info('Skipping %s, file already split.',
+                             repr(self.input_path))
+        else:
             self.logger.info('Split %s into %d files',
                              self.input_path, len(outputs))
 
@@ -276,6 +287,7 @@ def split_data(self) -> None:
                 splits = dict(zip(split_dims, line.split(' ')))
                 output_path = self.output_info.formatted_output_path(splits)
                 if self.should_skip_file(output_path):
+                    self.logger.info('Skipping %s, file already split.', repr(output_path))
                     skipped_paths.append(output_path)
                     continue
                 output_paths.append(output_path)
@@ -285,6 +297,7 @@ def split_data(self) -> None:
                                  repr(self.input_path), ', '.join(skipped_paths))
                 return
 
+            output_paths_set = set(output_paths)
             with tempfile.TemporaryDirectory() as tmpdir:
                 self.logger.info('Performing split.')
                 dest = os.path.join(tmpdir, flat_output_template)
@@ -296,9 +309,17 @@ def split_data(self) -> None:
                     subprocess.run([grib_copy_cmd, local_file.name, dest],
                                    check=True)
 
-                files = os.listdir(tmpdir)
-                num_files = len(files)
-                for f in files:
+                # grib_copy generates a file for every message, including ones
+                # that were already split. Drop those so copy_dir won't re-upload
+                # them (issue #538), then flatten the DELIMITER-encoded names back
+                # into their nested directory structure for the recursive copy.
+                num_files = 0
+                for f in os.listdir(tmpdir):
+                    dest_file_path = f'{prefix}{f.replace(delimiter, "/")}'
+                    if dest_file_path not in output_paths_set:
+                        os.remove(os.path.join(tmpdir, f))
+                        continue
+                    num_files += 1
                     if delimiter in f:
                         new_rel_path = f.replace(delimiter, '/').lstrip('/')
                         new_path = os.path.join(tmpdir, new_rel_path)
@@ -322,11 +343,6 @@ def split_data(self) -> None:
             raise ValueError('No splitting specified in template.')
         if any(dim in self._UNSUPPORTED_DIMENSIONS for dim in self.output_info.split_dims()):
             raise ValueError('Unsupported split dimension (lat, lng).')
-        if self.should_skip():
-            metrics.Metrics.counter('file_splitters', 'skipped').inc()
-            self.logger.info('Skipping %s, file already split.',
-                             repr(self.input_path))
-            return
 
         with self._open_dataset_locally() as dataset:
             if any(split not in dataset.dims and split not in ('variable') for split in self.output_info.split_dims()):
@@ -344,13 +360,20 @@ def split_data(self) -> None:
                 iterlists.append(dataset[dim])
             combinations = itertools.product(*iterlists)
             self.logger.info('Splitting & uploading %r...', self.input_path)
+            files_written = 0
             for comb in combinations:
                 selected = comb[0]
                 for da in comb[1:]:
                     for dim in da.coords:
                         selected = selected.sel({dim: getattr(da, dim)})
-                self._write_dataset(selected, filtered_split_dims)
-            self.logger.info('Finished splitting & uploading %r.', self.input_path)
+                if self._write_dataset(selected, filtered_split_dims):
+                    files_written += 1
+            if files_written == 0:
+                metrics.Metrics.counter('file_splitters', 'skipped').inc()
+                self.logger.info('Skipping %s, file already split.',
+                                 repr(self.input_path))
+            else:
+                self.logger.info('Finished splitting & uploading %r.', self.input_path)
 
     @contextmanager
     def _open_dataset_locally(self) -> t.Iterator[xr.Dataset]:
@@ -359,14 +382,19 @@ def _open_dataset_locally(self) -> t.Iterator[xr.Dataset]:
             yield ds
             ds.close()
 
-    def _write_dataset(self, dataset: xr.Dataset, split_dims: t.List[str]) -> None:
-        """Write destination NetCDF file in NETCDF4 format."""
+    def _write_dataset(self, dataset: xr.Dataset, split_dims: t.List[str]) -> bool:
+        """Write destination NetCDF file in NETCDF4 format. Returns True if the file was written."""
         # Here, we need to write the file locally, since only the scipy engine supports file objects or
         # returning bytes. Further, the scipy engine does not support NETCDF4 (which is HDF5 compliant).
         # Storing data in HDF5 is advantageous since it allows opening NetCDF files with buffered readers.
+        output_path = self._get_output_for_dataset(dataset, split_dims)
+        if self.should_skip_file(output_path):
+            self.logger.info('Skipping %s, file already split.', repr(output_path))
+            return False
         with tempfile.NamedTemporaryFile() as tmp:
             dataset.to_netcdf(path=tmp.name, engine='netcdf4', format='NETCDF4')
-            copy(tmp.name, self._get_output_for_dataset(dataset, split_dims))
+            copy(tmp.name, output_path)
+        return True
 
     def _get_output_for_dataset(self, dataset: xr.Dataset, split_dims: t.List[str]) -> str:
         splits = {'variable': list(dataset.data_vars.keys())[0]}

weather_sp/splitter_pipeline/file_splitters_test.py

@@ -191,6 +191,36 @@ def test_skips_existing_split(self, data_dir, grib_splitter):
         assert os.path.exists(f'{data_dir}/split_files/')
         assert splitter.should_skip()
 
+    def test_splits_only_missing_files(self, data_dir, grib_splitter):
+        input_path = f'{data_dir}/era5_sample.grib'
+        split_dir = f'{data_dir}/split_files/'
+        splitter = grib_splitter(
+            input_path,
+            OutFileInfo(
+                f'{data_dir}/split_files/era5_sample',
+                formatting='_{typeOfLevel}_{shortName}',
+                ending='.grib',
+                template_folders=[])
+        )
+        splitter.split_data()
+        assert os.path.exists(split_dir)
+
+        original_mtimes = {
+            f: os.path.getmtime(os.path.join(split_dir, f))
+            for f in os.listdir(split_dir)
+        }
+
+        missing_file = 'era5_sample_isobaricInhPa_z.grib'
+        os.remove(os.path.join(split_dir, missing_file))
+
+        splitter.split_data()
+
+        assert os.path.exists(os.path.join(split_dir, missing_file))
+        for fname, orig_mtime in original_mtimes.items():
+            if fname == missing_file:
+                continue
+            assert os.path.getmtime(os.path.join(split_dir, fname)) == orig_mtime
+
     @patch('weather_sp.splitter_pipeline.file_splitters.FileSplitter.should_skip_file')
     def test_skips_existing_split_with_filter(self, mock_should_skip_file, data_dir):
         input_path = f'{data_dir}/era5_sample.grib'
@@ -304,6 +334,34 @@ def test_skips_existing_split(self, data_dir):
         assert os.path.exists(f'{data_dir}/split_files/')
         assert splitter.should_skip()
 
+    def test_splits_only_missing_files(self, data_dir):
+        input_path = f'{data_dir}/era5_sample.nc'
+        split_dir = f'{data_dir}/split_files/'
+        splitter = NetCdfSplitter(
+            input_path,
+            OutFileInfo(f'{data_dir}/split_files/era5_sample',
+                        formatting='_{time}_{variable}',
+                        ending='.nc',
+                        template_folders=[]))
+        splitter.split_data()
+        assert os.path.exists(split_dir)
+
+        original_mtimes = {
+            f: os.path.getmtime(os.path.join(split_dir, f))
+            for f in os.listdir(split_dir)
+        }
+
+        missing_file = 'era5_sample_2015-01-15T00:00_z.nc'
+        os.remove(os.path.join(split_dir, missing_file))
+
+        splitter.split_data()
+
+        assert os.path.exists(os.path.join(split_dir, missing_file))
+        for fname, orig_mtime in original_mtimes.items():
+            if fname == missing_file:
+                continue
+            assert os.path.getmtime(os.path.join(split_dir, fname)) == orig_mtime
+
     def test_does_not_skip__if_forced(self, data_dir):
         input_path = f'{data_dir}/era5_sample.nc'
         output_base = f'{data_dir}/split_files/era5_sample'