add arguments to st_regionstats and make it into a unary op

Author: tswast

bigframes/bigquery/_operations/geo.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import json
-from typing import Mapping, Union
+from typing import Mapping, Optional, Union
 
 import shapely  # type: ignore
 
@@ -681,25 +681,46 @@ def st_length(
 
 def st_regionstats(
     geography: Union[bigframes.series.Series, bigframes.geopandas.GeoSeries],
-    raster: bigframes.series.Series,
-    band: str,
-    options: Mapping[str, Union[str, int, float]] = {},
+    raster_id: str,
+    band: Optional[str] = None,
+    include: Optional[str] = None,
+    options: Optional[Mapping[str, Union[str, int, float]]] = None,
 ) -> bigframes.dataframe.DataFrame:
-    """Computes statistics for a raster band within a given geography.
+    """Returns statistics summarizing the pixel values of the raster image
+    referenced by raster_id that intersect with geography.
 
-    See: https://cloud.google.com/bigquery/docs/reference/standard-sql/geography_functions#st_regionstats
+    The statistics include the count, minimum, maximum, sum, standard
+    deviation, mean, and area of the valid pixels of the raster band named
+    band_name. Google Earth Engine computes the results of the function call.
 
-    .. warning::
-        This function requires the Earth Engine API to be enabled.
+    See: https://cloud.google.com/bigquery/docs/reference/standard-sql/geography_functions#st_regionstats
 
     Args:
         geography (bigframes.series.Series | bigframes.geopandas.GeoSeries):
-            A series of geography objects.
-        raster (bigframes.series.Series):
-            A series of raster URIs. This can be a Google Cloud Storage URI,
-            or an Earth Engine asset ID.
-        band (str):
-            The name of the raster band to compute statistics for.
+            A series of geography objects to intersect with the raster image.
+        raster_id (str):
+            A string that identifies a raster image. The following formats are
+            supported. A URI from an image table provided by Google Earth Engine
+            in BigQuery sharing (formerly Analytics Hub). A URI for a readable
+            GeoTIFF raster file. A Google Earth Engine asset path that
+            references public catalog data or project-owned assets with read
+            access.
+        band (Optional[str]):
+            A string in one of the following formats:
+            A single band within the raster image specified by raster_id. A
+            formula to compute a value from the available bands in the raster
+            image. The formula uses the Google Earth Engine image expression
+            syntax. Bands can be referenced by their name, band_name, in
+            expressions. If you don't specify a band, the first band of the
+            image is used.
+        include (Optional[str]):
+            An optional string formula that uses the Google Earth Engine image
+            expression syntax to compute a pixel weight. The formula should
+            return values from 0 to 1. Values outside this range are set to the
+            nearest limit, either 0 or 1. A value of 0 means that the pixel is
+            invalid and it's excluded from analysis. A positive value means that
+            a pixel is valid. Values between 0 and 1 represent proportional
+            weights for calculations, such as weighted means.
         options (Mapping[str, Union[str, int, float]], optional):
             A dictionary of options to pass to the function. See the BigQuery
             documentation for a list of available options.
@@ -708,6 +729,11 @@ def st_regionstats(
         bigframes.dataframe.DataFrame:
             A dataframe containing the computed statistics.
     """
-    op = ops.StRegionStatsOp(options=json.dumps(options) if options else None)
-    df = geography._apply_ternary_op(raster, band, op)
+    op = ops.StRegionStatsOp(
+        raster_id=raster_id,
+        band=band,
+        include=include,
+        options=json.dumps(options) if options else None,
+    )
+    df = geography._apply_unary_op(op)
     return df[df.columns[0]].struct.explode()

bigframes/core/compile/ibis_compiler/scalar_op_compiler.py

@@ -288,16 +288,15 @@ def isfinite(arg):
     return arg.isinf().negate() & arg.isnan().negate()
 
 
-@scalar_op_compiler.register_ternary_op(geo_ops.StRegionStatsOp, pass_op=True)
+@scalar_op_compiler.register_unary_op(geo_ops.StRegionStatsOp, pass_op=True)
 def st_regionstats(
     geography: ibis_types.Value,
-    raster: ibis_types.Value,
-    band: ibis_types.Value,
     op: geo_ops.StRegionStatsOp,
 ):
-    args = [geography, raster, band]
+    args = [geography]  # TODO: get band, include, and other properies from op.
     if op.options:
         args.append(bigframes_vendored.ibis.literal(op.options, type="json"))
+    # TODO: We may need a custom ibis op so that we can pass arguments by name instead of position.
     return bigframes_vendored.ibis.remote_function(
         "st_regionstats",
         args,

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -230,14 +230,12 @@ def _add_parentheses(cls, expr: TypedExpr) -> TypedExpr:
 scalar_op_compiler = ScalarOpCompiler()
 
 
-@scalar_op_compiler.register_ternary_op(geo_ops.StRegionStatsOp, pass_op=True)
+@scalar_op_compiler.register_unary_op(geo_ops.StRegionStatsOp, pass_op=True)
 def compile_st_regionstats(
     geography: TypedExpr,
-    raster: TypedExpr,
-    band: TypedExpr,
     op: geo_ops.StRegionStatsOp,
 ):
-    args = [geography.expr, raster.expr, band.expr]
+    args = [geography.expr]  # TODO: get raster, band, include from op.
     if op.options:
         args.append(
             sge.Anonymous(

bigframes/operations/geo_ops.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import dataclasses
-import typing
+from typing import Optional
 
 from bigframes import dtypes
 from bigframes.operations import base_ops
@@ -128,11 +128,14 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 
 
 @dataclasses.dataclass(frozen=True)
-class StRegionStatsOp(base_ops.TernaryOp):
+class StRegionStatsOp(base_ops.UnaryOp):
     """See: https://cloud.google.com/bigquery/docs/reference/standard-sql/geography_functions#st_regionstats"""
 
     name = "st_regionstats"
-    options: typing.Optional[str] = None
+    raster_id: str
+    band: Optional[str]
+    include: Optional[str]
+    options: Optional[str]
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.struct_type(