Skip to content

Commit

Permalink
added python bindings and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
sllynn committed Nov 12, 2024
1 parent 02c8e42 commit 6eb5010
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 0 deletions.
46 changes: 46 additions & 0 deletions python/mosaic/api/raster.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#######################

__all__ = [
"rst_asformat",
"rst_avg",
"rst_bandmetadata",
"rst_boundingbox",
Expand All @@ -23,6 +24,7 @@
"rst_derivedband",
"rst_dtmfromgeoms",
"rst_filter",
"rst_format",
"rst_frombands",
"rst_fromcontent",
"rst_fromfile",
Expand Down Expand Up @@ -85,6 +87,30 @@
]


def rst_asformat(raster_tile: ColumnOrName, driver: ColumnOrName) -> Column:
"""
Translates the raster to the specified format.
Parameters
----------
raster_tile : Column (RasterTileType)
Mosaic raster tile struct column.
driver : Column (StringType)
The format driver to use.
Returns
-------
Column (RasterTileType)
The updated raster.
"""
return config.mosaic_context.invoke_function(
"rst_asformat",
pyspark_to_java_column(raster_tile),
pyspark_to_java_column(driver),
)


def rst_avg(raster_tile: ColumnOrName) -> Column:
"""
Returns an array containing mean value for each band.
Expand Down Expand Up @@ -356,6 +382,26 @@ def rst_filter(raster_tile: ColumnOrName, kernel_size: Any, operation: Any) -> C
)


def rst_format(raster_tile: ColumnOrName) -> Column:
"""
Returns the format of the raster.
Parameters
----------
raster_tile : Column (RasterTileType)
Mosaic raster tile struct column.
Returns
-------
Column (StringType)
The format of the raster (driver required for reading).
"""
return config.mosaic_context.invoke_function(
"rst_format", pyspark_to_java_column(raster_tile)
)


def rst_frombands(bands: ColumnOrName) -> Column:
"""
Stack an array of bands into a raster tile.
Expand Down
3 changes: 3 additions & 0 deletions python/test/test_raster_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,9 @@ def test_raster_scalar_functions(self):
.withColumn("rst_max", api.rst_max("tile"))
.withColumn("rst_median", api.rst_median("tile"))
.withColumn("rst_min", api.rst_min("tile"))
.withColumn("rst_setsrid", api.rst_setsrid("tile", lit(4326)))
.withColumn("rst_format", api.rst_format("rst_setsrid"))
.withColumn("rst_asformat", api.rst_asformat("rst_setsrid", lit("GRIB")))
.withColumn("rst_frombands", api.rst_frombands(array("tile", "tile")))
.withColumn("rst_georeference", api.rst_georeference("tile"))
.withColumn("rst_getnodata", api.rst_getnodata("tile"))
Expand Down

0 comments on commit 6eb5010

Please sign in to comment.