Merge pull request #210 from geo-engine/skip-empty-times

Skip-empty-times

Author: michaelmattig

geoengine/raster.py

@@ -213,6 +213,12 @@ def spatial_partition(self) -> gety.SpatialPartition2D:
     def spatial_resolution(self) -> gety.SpatialResolution:
         return self.geo_transform.spatial_resolution()
 
+    def is_empty(self) -> bool:
+        ''' Returns true if the tile is empty'''
+        num_pixels = self.size_x * self.size_y
+        num_nulls = self.data.null_count
+        return num_pixels == num_nulls
+
     @staticmethod
     def from_ge_record_batch(record_batch: pa.RecordBatch) -> RasterTile2D:
         '''Create a RasterTile2D from an Arrow record batch recieved from the Geo Engine'''

geoengine/raster_workflow_rio_writer.py

@@ -168,20 +168,37 @@ def __create_new_dataset(self, query: QueryRectangle):
             **self.rio_kwargs
         )
 
+        for i, b in enumerate(self.bands, start=1):
+            b_n = b.name
+            b_m = str(b.measurement)
+            rio_dataset.update_tags(i, band_name=b_n, band_measurement=b_m)
+
         self.current_dataset = rio_dataset
 
-    async def query_and_write(self, query: QueryRectangle):
-        ''' Query the raster workflow and write the tiles to the dataset.'''
+    async def query_and_write(self, query: QueryRectangle, skip_empty_times=True):
+        '''
+        Query the raster workflow and write the resulting tiles to a GDAL dataset per timeslice.
+
+        :param query: The QueryRectangle to write to GDAL dataset(s)
+        :param skip_empty_times: Skip timeslices where all pixels are empty/nodata
+        '''
 
         self.create_tiling_geo_transform_width_height(query)
 
-        assert self.workflow is not None, "The workflow must be set"
+        assert self.bands is not None, "The bands must be set"
+        bands = list(range(0, len(self.bands)))
 
+        assert self.workflow is not None, "The workflow must be set"
         try:
-            async for tile in self.workflow.raster_stream(query):
+            async for tile in self.workflow.raster_stream(query, bands=bands):
                 if self.current_time != tile.time:
                     self.close_current_dataset()
                     self.current_time = tile.time
+
+                if tile.is_empty() and skip_empty_times:
+                    continue
+
+                if self.current_dataset is None:
                     self.__create_new_dataset(query)
 
                 assert self.current_time == tile.time, "The time of the current dataset does not match the tile"