Merge branch 'main' into sycai_df_merge

Author: sycai

bigframes/bigquery/__init__.py

@@ -40,6 +40,7 @@
     st_intersection,
     st_isclosed,
     st_length,
+    st_regionstats,
     st_simplify,
 )
 from bigframes.bigquery._operations.json import (
@@ -81,6 +82,7 @@
     st_intersection,
     st_isclosed,
     st_length,
+    st_regionstats,
     st_simplify,
     # json ops
     json_extract,

bigframes/bigquery/_operations/geo.py

@@ -14,11 +14,13 @@
 
 from __future__ import annotations
 
-from typing import Union
+import json
+from typing import Mapping, Optional, Union
 
 import shapely  # type: ignore
 
 from bigframes import operations as ops
+import bigframes.dataframe
 import bigframes.geopandas
 import bigframes.series
 
@@ -677,6 +679,65 @@ def st_length(
     return series
 
 
+def st_regionstats(
+    geography: Union[bigframes.series.Series, bigframes.geopandas.GeoSeries],
+    raster_id: str,
+    band: Optional[str] = None,
+    include: Optional[str] = None,
+    options: Optional[Mapping[str, Union[str, int, float]]] = None,
+) -> bigframes.series.Series:
+    """Returns statistics summarizing the pixel values of the raster image
+    referenced by raster_id that intersect with geography.
+
+    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.
+
+    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 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.
+
+    Returns:
+        bigframes.pandas.Series:
+            A STRUCT Series containing the computed statistics.
+    """
+    op = ops.GeoStRegionStatsOp(
+        raster_id=raster_id,
+        band=band,
+        include=include,
+        options=json.dumps(options) if options else None,
+    )
+    return geography._apply_unary_op(op)
+
+
 def st_simplify(
     geography: "bigframes.series.Series",
     tolerance_meters: float,

bigframes/core/compile/ibis_compiler/operations/geo_ops.py

@@ -16,8 +16,10 @@
 
 from typing import cast
 
+from bigframes_vendored import ibis
 from bigframes_vendored.ibis.expr import types as ibis_types
 import bigframes_vendored.ibis.expr.datatypes as ibis_dtypes
+import bigframes_vendored.ibis.expr.operations.geospatial as ibis_geo
 import bigframes_vendored.ibis.expr.operations.udf as ibis_udf
 
 from bigframes.core.compile.ibis_compiler import scalar_op_compiler
@@ -101,6 +103,35 @@ def geo_st_isclosed_op_impl(x: ibis_types.Value):
     return st_isclosed(x)
 
 
+@register_unary_op(ops.GeoStRegionStatsOp, pass_op=True)
+def geo_st_regionstats_op_impl(
+    geography: ibis_types.Value,
+    op: ops.GeoStRegionStatsOp,
+):
+    if op.band:
+        band = ibis.literal(op.band, type=ibis_dtypes.string())
+    else:
+        band = None
+
+    if op.include:
+        include = ibis.literal(op.include, type=ibis_dtypes.string())
+    else:
+        include = None
+
+    if op.options:
+        options = ibis.literal(op.options, type=ibis_dtypes.json())
+    else:
+        options = None
+
+    return ibis_geo.GeoRegionStats(
+        arg=geography,  # type: ignore
+        raster_id=ibis.literal(op.raster_id, type=ibis_dtypes.string()),  # type: ignore
+        band=band,  # type: ignore
+        include=include,  # type: ignore
+        options=options,  # type: ignore
+    ).to_expr()
+
+
 @register_unary_op(ops.GeoStSimplifyOp, pass_op=True)
 def st_simplify_op_impl(x: ibis_types.Value, op: ops.GeoStSimplifyOp):
     x = cast(ibis_types.GeoSpatialValue, x)

bigframes/core/compile/sqlglot/expressions/generic_ops.py

@@ -24,6 +24,7 @@
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
 register_nary_op = scalar_compiler.scalar_op_compiler.register_nary_op
 register_ternary_op = scalar_compiler.scalar_op_compiler.register_ternary_op
 
@@ -159,6 +160,13 @@ def _(*cases_and_outputs: TypedExpr) -> sge.Expression:
     )
 
 
+@register_binary_op(ops.coalesce_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == right.expr:
+        return left.expr
+    return sge.Coalesce(this=left.expr, expressions=[right.expr])
+
+
 @register_nary_op(ops.RowKey)
 def _(*values: TypedExpr) -> sge.Expression:
     # All inputs into hash must be non-null or resulting hash will be null

bigframes/core/compile/sqlglot/expressions/geo_ops.py

@@ -74,6 +74,32 @@ def _(expr: TypedExpr, op: ops.GeoStLengthOp) -> sge.Expression:
     return sge.func("ST_LENGTH", expr.expr)
 
 
+@register_unary_op(ops.GeoStRegionStatsOp, pass_op=True)
+def _(
+    geography: TypedExpr,
+    op: ops.GeoStRegionStatsOp,
+):
+    args = [geography.expr, sge.convert(op.raster_id)]
+    if op.band:
+        args.append(sge.Kwarg(this="band", expression=sge.convert(op.band)))
+    if op.include:
+        args.append(sge.Kwarg(this="include", expression=sge.convert(op.include)))
+    if op.options:
+        args.append(
+            sge.Kwarg(this="options", expression=sge.JSON(this=sge.convert(op.options)))
+        )
+    return sge.func("ST_REGIONSTATS", *args)
+
+
+@register_unary_op(ops.GeoStSimplifyOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.GeoStSimplifyOp) -> sge.Expression:
+    return sge.func(
+        "ST_SIMPLIFY",
+        expr.expr,
+        sge.convert(op.tolerance_meters),
+    )
+
+
 @register_unary_op(ops.geo_x_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("SAFE.ST_X", expr.expr)

bigframes/operations/__init__.py

@@ -121,6 +121,7 @@
     GeoStBufferOp,
     GeoStDistanceOp,
     GeoStLengthOp,
+    GeoStRegionStatsOp,
     GeoStSimplifyOp,
 )
 from bigframes.operations.json_ops import (
@@ -415,12 +416,13 @@
     "geo_st_geogpoint_op",
     "geo_st_intersection_op",
     "geo_st_isclosed_op",
-    "GeoStBufferOp",
-    "GeoStLengthOp",
-    "GeoStSimplifyOp",
     "geo_x_op",
     "geo_y_op",
+    "GeoStBufferOp",
     "GeoStDistanceOp",
+    "GeoStLengthOp",
+    "GeoStRegionStatsOp",
+    "GeoStSimplifyOp",
     # AI ops
     "AIClassify",
     "AIGenerate",

bigframes/operations/geo_ops.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import dataclasses
+from typing import Optional
 
 from bigframes import dtypes
 from bigframes.operations import base_ops
@@ -135,6 +136,29 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         return dtypes.FLOAT_DTYPE
 
 
+@dataclasses.dataclass(frozen=True)
+class GeoStRegionStatsOp(base_ops.UnaryOp):
+    """See: https://cloud.google.com/bigquery/docs/reference/standard-sql/geography_functions#st_regionstats"""
+
+    name = "geo_st_regionstats"
+    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(
+            [
+                ("min", dtypes.FLOAT_DTYPE),
+                ("max", dtypes.FLOAT_DTYPE),
+                ("sum", dtypes.FLOAT_DTYPE),
+                ("count", dtypes.INT_DTYPE),
+                ("mean", dtypes.FLOAT_DTYPE),
+                ("area", dtypes.FLOAT_DTYPE),
+            ]
+        )
+
+
 @dataclasses.dataclass(frozen=True)
 class GeoStSimplifyOp(base_ops.UnaryOp):
     name = "st_simplify"