Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 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
18 changes: 15 additions & 3 deletions geoengine/raster_workflow_rio_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,20 +168,32 @@ 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):
async def query_and_write(self, query: QueryRectangle, skip_empty_times=True):
''' Query the raster workflow and write the tiles to the dataset.'''

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