Handle all-nulls in terracotta writing

src/nuthatch/backends/terracotta.py

@@ -3,7 +3,7 @@
 import terracotta as tc
 import sqlalchemy
 import xarray as xr
-import rioxarray # Must import for .rio to work # noqa: F401
+import rioxarray  # Must import for .rio to work # noqa: F401
 import numpy as np
 from rasterio.io import MemoryFile
 from rasterio.enums import Resampling
@@ -13,6 +13,7 @@
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
+
 def base360_to_base180(lons):
     """Converts a list of longitudes from base 360 to base 180.
 
@@ -41,7 +42,6 @@ def base180_to_base360(lons):
     return np.array(val)
 
 
-
 def is_wrapped(lons):
     """Check if the longitudes are wrapped.
 
@@ -193,20 +193,27 @@ def write(self, ds, upsert=False, primary_keys=None):
         ds = ds.rio.reproject('EPSG:3857', resampling=self.resampling, nodata=np.nan)
         ds.rio.write_crs("epsg:3857", inplace=True)
 
-        # Insert the parameters.
         with self.driver.connect(verify=False):
             if self.time_dim in ds.dims:
                 for t in ds.time:
                     # Select just this time and squeeze the dimension
                     sub_ds = ds.sel(time=t)
                     sub_ds = sub_ds.reset_coords('time', drop=True)
 
-                    # Assume the cache_key is now a folder and make sub tifs under it
+                    # Skip time slices with no valid data - terracotta can't
+                    # compute metadata for an all-nodata raster
+                    if not any(bool(sub_ds[v].notnull().any()) for v in sub_ds.data_vars):
+                        logger.warning(f"Skipping {self.cache_key}_{t.values}: no valid (non-NaN) data.")
+                        continue
+
                     sub_cache_key = self.cache_key + '_' + str(t.values)
                     sub_path = os.path.join(self.path, str(t.values) + '.tif')
-                    sub_override_path = os.path.join(self.override_path,  str(t.values) + '.tif')
-
-                    self.write_individual_raster(self.driver, sub_ds, sub_path, sub_cache_key, sub_override_path)
+                    sub_override_path = os.path.join(self.override_path, str(t.values) + '.tif')
+                    try:
+                        self.write_individual_raster(self.driver, sub_ds, sub_path, sub_cache_key, sub_override_path)
+                    except Exception:
+                        logger.exception(f"Failed to write raster {sub_cache_key}; continuing with remaining times.")
+                        continue
             else:
                 path = os.path.join(self.path, '_.tif')
                 override_path = os.path.join(self.override_path, '_.tif')
@@ -223,7 +230,7 @@ def write_individual_raster(self, driver, ds, path, cache_key, override_path):
                 shutil.copyfileobj(mem_dst, f_out)
 
             driver.insert({'key': cache_key.replace('/', '_')}, mem_dst,
-                         override_path=override_path, skip_metadata=False)
+                          override_path=override_path, skip_metadata=False)
 
             logger.info(f"Inserted {cache_key.replace('/', '_')} into the terracotta database.")
 
@@ -233,19 +240,24 @@ def upsert(self, data, upsert_keys=None):
     def read(self, engine):
         datasets_table = sqlalchemy.Table("datasets", self.driver.meta_store.sqla_metadata,
                                           autoload_with=self.driver.meta_store.sqla_engine)
+        base_key = self.cache_key.replace('/', '_')
         stmt = (
             datasets_table.select()
-            .where(datasets_table.c['key'].like(self.cache_key.replace('/', '_') + '%'))
+            .where(datasets_table.c['key'].like(base_key + '%'))
         )
 
         with self.driver.meta_store.connect() as conn:
             result = conn.execute(stmt).all()
 
-        datasets = [row[0] for row in result]
-
+        # LIKE also matches longer cache keys with the same prefix; filter in Python
+        # because "_" is a wildcard in SQL LIKE.
+        time_slice_prefix = base_key + '_'
         ret = []
-        for dataset in datasets:
-            ret.append({dataset: self.driver.get_metadata({'key' : dataset})})
+        for row in result:
+            key = row[0]
+            if key != base_key and not key.startswith(time_slice_prefix):
+                continue
+            ret.append({key: self.driver.get_metadata({'key': key})})
 
         return ret
 
@@ -260,12 +272,10 @@ def delete(self):
             .where(datasets_table.c['key'].like(self.cache_key.replace('/', '_') + '%'))
         )
 
-
         with self.driver.meta_store.connect() as conn:
             result = conn.execute(stmt).all()
 
         datasets = [row[0] for row in result]
         for dataset in datasets:
             logger.info(f"Deleting datasets {datasets} from terracotta.")
             self.driver.delete({'key': dataset})
-

tests/backends/test_terracotta.py

@@ -24,12 +24,11 @@ def simple_latlon_array():
     lats = np.arange(-90, 90, 0.25)
     lons = np.arange(-180, 180, 0.25)
 
-
     data_shape = (len(lats), len(lons))
     empty_data = np.ones(data_shape) * np.random.randint(1000000)
 
     ds = xr.Dataset(
-        data_vars= {
+        data_vars={
             'data': (['lat', 'lon'], empty_data),
         },
         coords={"lat": lats, "lon": lons},
@@ -45,12 +44,32 @@ def time_latlon_array():
     lons = np.arange(-180, 180, 0.25)
     times = pd.date_range("2001-01-01", "2001-01-05", freq='1d')
 
+    data_shape = (len(lats), len(lons), len(times))
+    empty_data = np.ones(data_shape) * np.random.randint(1000000)
+
+    ds = xr.Dataset(
+        data_vars={
+            'data': (['lat', 'lon', 'time'], empty_data),
+        },
+        coords={"lat": lats, "lon": lons, "time": times},
+    )
+
+    return ds
+
+
+@cache(cache_args=[])
+def time_latlon_array_with_null_chunks():
+    """Generate a timeseries dataset where some time slices are all null."""
+    lats = np.arange(-90, 90, 0.25)
+    lons = np.arange(-180, 180, 0.25)
+    times = pd.date_range("2001-01-01", "2001-01-03", freq='1d')
 
     data_shape = (len(lats), len(lons), len(times))
     empty_data = np.ones(data_shape) * np.random.randint(1000000)
+    empty_data[:, :, :2] = np.nan
 
     ds = xr.Dataset(
-        data_vars= {
+        data_vars={
             'data': (['lat', 'lon', 'time'], empty_data),
         },
         coords={"lat": lats, "lon": lons, "time": times},
@@ -68,8 +87,18 @@ def test_terracotta():
     assert len(ds) == 1
 
     # cache the array
+    time_latlon_array_with_null_chunks(backend='zarr')
+    time_latlon_array_with_null_chunks(backend='zarr', storage_backend='terracotta', cache_mode='overwrite')
+
     ds = time_latlon_array(backend='zarr')
     # Verify that write?
     ds = time_latlon_array(backend='zarr', storage_backend='terracotta', cache_mode='overwrite')
     ds = time_latlon_array(backend='terracotta')
     assert len(ds) == 5
+
+
+def test_terracotta_all_null_time_chunks():
+    time_latlon_array_with_null_chunks(backend='zarr')
+    time_latlon_array_with_null_chunks(backend='zarr', storage_backend='terracotta', cache_mode='overwrite')
+    ds = time_latlon_array_with_null_chunks(backend='terracotta')
+    assert len(ds) == 1