diff --git a/geoengine/raster.py b/geoengine/raster.py index 3d5199f2..b7793729 100644 --- a/geoengine/raster.py +++ b/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''' diff --git a/geoengine/raster_workflow_rio_writer.py b/geoengine/raster_workflow_rio_writer.py index 7e70a566..e43d700a 100644 --- a/geoengine/raster_workflow_rio_writer.py +++ b/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"