Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions geoengine/raster.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'''
Expand Down
25 changes: 21 additions & 4 deletions geoengine/raster_workflow_rio_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down