feat: Add BigFrames.bigquery.st_regionstats method

This commit adds the `BigFrames.bigquery.st_regionstats` method, which allows users to compute statistics for a raster band within a given geography.

The implementation includes:
- A new `StRegionStatsOp` in `bigframes/operations/geo_ops.py`.
- Compiler implementations for both the SQLGlot and Ibis backends.
- A unit test with a SQL snapshot.
- A system test that demonstrates the use of the new function.

This commit also addresses feedback from the code review, including:
- Adding `area` to the output struct of `st_regionstats`.
- Making the `options` parameter a positional argument.
- Adding comments to explain the use of `pass_op`.
- Converting the unit test to a pytest-style function.
- Moving the sample code to a system test.

Author: google-labs-jules[bot]

bigframes/bigquery/_operations/geo.py

@@ -683,7 +683,6 @@ def st_regionstats(
     geography: bigframes.geopandas.GeoSeries,
     raster: bigframes.series.Series,
     band: str,
-    *,
     options: Mapping[str, Union[str, int, float]] = {},
 ) -> bigframes.dataframe.DataFrame:
     """Computes statistics for a raster band within a given geography.

bigframes/core/compile/ibis_compiler/scalar_op_compiler.py

@@ -169,6 +169,9 @@ def register_ternary_op(
         Args:
             op_ref (TernaryOp or TernaryOp type):
                 Class or instance of operator that is implemented by the decorated function.
+            pass_op (bool):
+                Set to true if implementation takes the operator object as the last argument.
+                This is needed for parameterized ops where parameters are part of op object.
         """
         key = typing.cast(str, op_ref.name)
 
@@ -296,5 +299,7 @@ def st_regionstats(
     if op.options:
         args.append(bigframes_vendored.ibis.literal(op.options, type="json"))
     return bigframes_vendored.ibis.remote_function(
-        "st_regionstats", args, output_type="struct<min: float, max: float, sum: float, count: int, mean: float>"  # type: ignore
+        "st_regionstats",
+        args,
+        output_type="struct<min: float, max: float, sum: float, count: int, mean: float, area: float>",  # type: ignore
     )

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -201,11 +201,14 @@ def compile_st_regionstats(
     args = [geography.expr, raster.expr, band.expr]
     if op.options:
         args.append(
-            sge.EQ(
-                this=sge.Identifier(this="OPTIONS"),
-                expression=sge.Anonymous(
-                    this="JSON", expressions=[sge.convert(op.options)]
-                ),
+            sge.Anonymous(
+                this="_",
+                expressions=[
+                    sge.Identifier(this="OPTIONS"),
+                    sge.Anonymous(
+                        this="JSON", expressions=[sge.convert(op.options)]
+                    ),
+                ],
             )
         )
     return sge.func("ST_REGIONSTATS", *args)

bigframes/operations/geo_ops.py

@@ -142,6 +142,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
                 ("sum", dtypes.FLOAT_DTYPE),
                 ("count", dtypes.INT_DTYPE),
                 ("mean", dtypes.FLOAT_DTYPE),
+                ("area", dtypes.FLOAT_DTYPE),
             ]
         )
 

samples/snippets/wildfire_risk.py

@@ -0,0 +1,76 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Inspired by the SQL at https://cloud.google.com/blog/products/data-analytics/a-closer-look-at-earth-engine-in-bigquery
+
+import bigframes.bigquery as bbq
+import bigframes.pandas as bpd
+
+
+def test_wildfire_risk(session):
+    # Step 1: Select inputs from datasets that we've subscribed to
+    wildfire_raster = bpd.read_gbq("wildfire_risk_to_community_v0_mosaic.fire")[
+        "assets.image.href"
+    ]
+    places = bpd.read_gbq("bigquery-public-data.geo_us_census_places.places_colorado")[
+        ["place_id", "place_name", "place_geom"]
+    ]
+    places = places.rename(columns={"place_geom": "geo"})
+
+    # Step 2: Compute the weather forecast using WeatherNext Graph forecast data
+    weather_forecast = bpd.read_gbq("weathernext_graph_forecasts.59572747_4_0")
+    weather_forecast = weather_forecast[
+        weather_forecast["init_time"] == "2025-04-28 00:00:00+00:00"
+    ]
+    weather_forecast = weather_forecast.explode("forecast")
+    wind_speed = (
+        weather_forecast["forecast"]["10m_u_component_of_wind"] ** 2
+        + weather_forecast["forecast"]["10m_v_component_of_wind"] ** 2
+    ) ** 0.5
+    weather_forecast = weather_forecast.assign(wind_speed=wind_speed)
+    weather_forecast = weather_forecast[weather_forecast["forecast"]["hours"] < 24]
+    weather_forecast = weather_forecast.merge(
+        places, how="inner", left_on="geography_polygon", right_on="geo"
+    )
+    weather_forecast = weather_forecast.groupby("place_id").agg(
+        place_name=("place_name", "first"),
+        geo=("geo", "first"),
+        average_wind_speed=("wind_speed", "mean"),
+        maximum_wind_speed=("wind_speed", "max"),
+    )
+
+    # Step 3: Combine with wildfire risk for each community
+    wildfire_risk = weather_forecast.assign(
+        wildfire_likelihood=bbq.st_regionstats(
+            weather_forecast["geo"],
+            wildfire_raster,
+            "BP",
+            options={"scale": 1000},
+        )["mean"],
+        wildfire_consequence=bbq.st_regionstats(
+            weather_forecast["geo"],
+            wildfire_raster,
+            "CRPS",
+            options={"scale": 1000},
+        )["mean"],
+    )
+
+    # Step 4: Compute a simple composite index of relative wildfire risk.
+    relative_risk = (
+        wildfire_risk["wildfire_likelihood"].rank(pct=True)
+        + wildfire_risk["wildfire_consequence"].rank(pct=True)
+        + wildfire_risk["average_wind_speed"].rank(pct=True)
+    ) / 3 * 100
+    wildfire_risk = wildfire_risk.assign(relative_risk=relative_risk)
+    assert wildfire_risk is not None

tests/system/small/test_wildfire_risk.py

@@ -21,7 +21,7 @@ WITH `bfcte_1` AS (
 ), `bfcte_4` AS (
   SELECT
     *,
-    ST_REGIONSTATS(`bfcol_1`, `bfcol_5`, 'band1', OPTIONS = JSON('{"scale": 100}')) AS `bfcol_8`
+    ST_REGIONSTATS(`bfcol_1`, `bfcol_5`, 'band1', _(OPTIONS, JSON('{"scale": 100}'))) AS `bfcol_8`
   FROM `bfcte_3`
 ), `bfcte_5` AS (
   SELECT
@@ -30,13 +30,15 @@ WITH `bfcte_1` AS (
     `bfcol_8`.`max` AS `bfcol_11`,
     `bfcol_8`.`sum` AS `bfcol_12`,
     `bfcol_8`.`count` AS `bfcol_13`,
-    `bfcol_8`.`mean` AS `bfcol_14`
+    `bfcol_8`.`mean` AS `bfcol_14`,
+    `bfcol_8`.`area` AS `bfcol_15`
   FROM `bfcte_4`
 )
 SELECT
   `bfcol_10` AS `min`,
   `bfcol_11` AS `max`,
   `bfcol_12` AS `sum`,
   `bfcol_13` AS `count`,
-  `bfcol_14` AS `mean`
+  `bfcol_14` AS `mean`,
+  `bfcol_15` AS `area`
 FROM `bfcte_5`

tests/unit/core/compile/sqlglot/snapshots/test_geo_compiler/test_st_regionstats/out.sql

@@ -21,9 +21,9 @@
 pytest.importorskip("pytest_snapshot")
 
 
-class TestGeoCompiler:
-    def test_st_regionstats(self, compiler_session, snapshot):
-        geos = gpd.GeoSeries(["POINT(1 1)"], session=compiler_session)
-        rasters = bpd.Series(["raster_uri"], dtype="string", session=compiler_session)
-        df = bbq.st_regionstats(geos, rasters, "band1", options={"scale": 100})
-        snapshot.assert_match(df.sql, "out.sql")
+def test_st_regionstats(compiler_session, snapshot):
+    geos = gpd.GeoSeries(["POINT(1 1)"], session=compiler_session)
+    rasters = bpd.Series(["raster_uri"], dtype="string", session=compiler_session)
+    df = bbq.st_regionstats(geos, rasters, "band1", {"scale": 100})
+    assert "area" in df.columns
+    snapshot.assert_match(df.sql, "out.sql")