feat: Add BigFrames.bigquery.st_regionstats method

This commit introduces the `st_regionstats` method in `bigframes.bigquery`, allowing users to compute statistics for a raster band within a given geography.

Key changes:
- Added `StRegionStatsOp` in `bigframes/operations/geo_ops.py`.
- Implemented compilation logic for the operation in both Ibis and SQLGlot compilers.
- Exposed the `st_regionstats` function in `bigframes/bigquery/_operations/geo.py` and the public API.
- Added a new `_apply_ternary_op` method to `bigframes.series.Series`.
- Included a unit test with a snapshot to verify the generated SQL.
- Added a system test that demonstrates the functionality by converting a complex wildfire risk analysis query from SQL to BigFrames.
- Refactored compiler registries to support `pass_op=True` for ternary operations, enabling access to operator parameters during compilation.

Author: google-labs-jules[bot]

bigframes/bigquery/_operations/ai.py

@@ -19,15 +19,12 @@
 from __future__ import annotations
 
 import json
-from typing import Any, Iterable, List, Literal, Mapping, Tuple, Union
+from typing import Any, List, Literal, Mapping, Tuple, Union
 
 import pandas as pd
 
-from bigframes import clients, dataframe, dtypes
-from bigframes import pandas as bpd
-from bigframes import series, session
+from bigframes import clients, dtypes, series, session
 from bigframes.core import convert, log_adapter
-from bigframes.ml import core as ml_core
 from bigframes.operations import ai_ops, output_schemas
 
 PROMPT_TYPE = Union[
@@ -551,91 +548,6 @@ def score(
     return series_list[0]._apply_nary_op(operator, series_list[1:])
 
 
-@log_adapter.method_logger(custom_base_name="bigquery_ai")
-def forecast(
-    df: dataframe.DataFrame | pd.DataFrame,
-    *,
-    data_col: str,
-    timestamp_col: str,
-    model: str = "TimesFM 2.0",
-    id_cols: Iterable[str] | None = None,
-    horizon: int = 10,
-    confidence_level: float = 0.95,
-    context_window: int | None = None,
-) -> dataframe.DataFrame:
-    """
-    Forecast time series at future horizon. Using Google Research's open source TimesFM(https://github.com/google-research/timesfm) model.
-
-    .. note::
-
-        This product or feature is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the
-        Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is"
-        and might have limited support. For more information, see the launch stage descriptions
-        (https://cloud.google.com/products#product-launch-stages).
-
-    Args:
-        df (DataFrame):
-            The dataframe that contains the data that you want to forecast. It could be either a BigFrames Dataframe or
-            a pandas DataFrame. If it's a pandas DataFrame, the global BigQuery session will be used to load the data.
-        data_col (str):
-            A str value that specifies the name of the data column. The data column contains the data to forecast.
-            The data column must use one of the following data types: INT64, NUMERIC and FLOAT64
-        timestamp_col (str):
-            A str value that specified the name of the time points column.
-            The time points column provides the time points used to generate the forecast.
-            The time points column must use one of the following data types: TIMESTAMP, DATE and DATETIME
-        model (str, default "TimesFM 2.0"):
-            A str value that specifies the name of the model. TimesFM 2.0 is the only supported value, and is the default value.
-        id_cols (Iterable[str], optional):
-            An iterable of str value that specifies the names of one or more ID columns. Each ID identifies a unique time series to forecast.
-            Specify one or more values for this argument in order to forecast multiple time series using a single query.
-            The columns that you specify must use one of the following data types: STRING, INT64, ARRAY<STRING> and ARRAY<INT64>
-        horizon (int, default 10):
-            An int value that specifies the number of time points to forecast. The default value is 10. The valid input range is [1, 10,000].
-        confidence_level (float, default 0.95):
-            A FLOAT64 value that specifies the percentage of the future values that fall in the prediction interval.
-            The default value is 0.95. The valid input range is [0, 1).
-        context_window (int, optional):
-            An int value that specifies the context window length used by BigQuery ML's built-in TimesFM model.
-            The context window length determines how many of the most recent data points from the input time series are use by the model.
-            If you don't specify a value, the AI.FORECAST function automatically chooses the smallest possible context window length to use
-            that is still large enough to cover the number of time series data points in your input data.
-
-    Returns:
-        DataFrame:
-            The forecast dataframe matches that of the BigQuery AI.FORECAST function.
-            See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-forecast
-
-    Raises:
-        ValueError: when any column ID does not exist in the dataframe.
-    """
-
-    if isinstance(df, pd.DataFrame):
-        # Load the pandas DataFrame with global session
-        df = bpd.read_pandas(df)
-
-    columns = [timestamp_col, data_col]
-    if id_cols:
-        columns += id_cols
-    for column in columns:
-        if column not in df.columns:
-            raise ValueError(f"Column `{column}` not found")
-
-    options: dict[str, Union[int, float, str, Iterable[str]]] = {
-        "data_col": data_col,
-        "timestamp_col": timestamp_col,
-        "model": model,
-        "horizon": horizon,
-        "confidence_level": confidence_level,
-    }
-    if id_cols:
-        options["id_cols"] = id_cols
-    if context_window:
-        options["context_window"] = context_window
-
-    return ml_core.BaseBqml(df._session).ai_forecast(input_data=df, options=options)
-
-
 def _separate_context_and_series(
     prompt: PROMPT_TYPE,
 ) -> Tuple[List[str | None], List[series.Series]]:

bigframes/bigquery/_operations/geo.py

@@ -680,7 +680,7 @@ def st_length(
 
 
 def st_regionstats(
-    geography: bigframes.geopandas.GeoSeries,
+    geography: Union[bigframes.series.Series, bigframes.geopandas.GeoSeries],
     raster: bigframes.series.Series,
     band: str,
     options: Mapping[str, Union[str, int, float]] = {},
@@ -693,7 +693,7 @@ def st_regionstats(
         This function requires the Earth Engine API to be enabled.
 
     Args:
-        geography (bigframes.geopandas.GeoSeries):
+        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,

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -205,7 +205,9 @@ def compile_st_regionstats(
                 this="_",
                 expressions=[
                     sge.Identifier(this="OPTIONS"),
-                    sge.Anonymous(this="JSON", expressions=[sge.convert(op.options)]),
+                    sge.Anonymous(
+                        this="JSON", expressions=[sge.convert(op.options)]
+                    ),
                 ],
             )
         )

bigframes/dataframe.py

@@ -5328,16 +5328,12 @@ def _throw_if_null_index(self, opname: str):
     @property
     def semantics(self):
         msg = bfe.format_message(
-            "The 'semantics' property will be removed. Please use 'bigframes.bigquery.ai' instead."
+            "The 'semantics' property will be removed. Please use 'ai' instead."
         )
         warnings.warn(msg, category=FutureWarning)
         return bigframes.operations.semantics.Semantics(self)
 
     @property
     def ai(self):
         """Returns the accessor for AI operators."""
-        msg = bfe.format_message(
-            "The 'ai' property will be removed. Please use 'bigframes.bigquery.ai' instead."
-        )
-        warnings.warn(msg, category=FutureWarning)
         return bigframes.operations.ai.AIAccessor(self)

bigframes/display/anywidget.py

@@ -23,14 +23,13 @@
 import pandas as pd
 
 import bigframes
-from bigframes.core import blocks
 import bigframes.dataframe
 import bigframes.display.html
 
-# anywidget and traitlets are optional dependencies. We don't want the import of
-# this module to fail if they aren't installed, though. Instead, we try to
-# limit the surface that these packages could affect. This makes unit testing
-# easier and ensures we don't accidentally make these required packages.
+# anywidget and traitlets are optional dependencies. We don't want the import of this
+# module to fail if they aren't installed, though. Instead, we try to limit the surface that
+# these packages could affect. This makes unit testing easier and ensures we don't
+# accidentally make these required packages.
 try:
     import anywidget
     import traitlets
@@ -47,21 +46,9 @@
 
 
 class TableWidget(WIDGET_BASE):
-    """An interactive, paginated table widget for BigFrames DataFrames.
-
-    This widget provides a user-friendly way to display and navigate through
-    large BigQuery DataFrames within a Jupyter environment.
     """
-
-    page = traitlets.Int(0).tag(sync=True)
-    page_size = traitlets.Int(0).tag(sync=True)
-    row_count = traitlets.Int(0).tag(sync=True)
-    table_html = traitlets.Unicode().tag(sync=True)
-    _initial_load_complete = traitlets.Bool(False).tag(sync=True)
-    _batches: Optional[blocks.PandasBatches] = None
-    _error_message = traitlets.Unicode(allow_none=True, default_value=None).tag(
-        sync=True
-    )
+    An interactive, paginated table widget for BigFrames DataFrames.
+    """
 
     def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         """Initialize the TableWidget.
@@ -74,11 +61,10 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
                 "Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use TableWidget."
             )
 
-        self._dataframe = dataframe
-
         super().__init__()
+        self._dataframe = dataframe
 
-        # Initialize attributes that might be needed by observers first
+        # Initialize attributes that might be needed by observers FIRST
         self._table_id = str(uuid.uuid4())
         self._all_data_loaded = False
         self._batch_iter: Optional[Iterator[pd.DataFrame]] = None
@@ -87,6 +73,9 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # respect display options for initial page size
         initial_page_size = bigframes.options.display.max_rows
 
+        # Initialize data fetching attributes.
+        self._batches = dataframe._to_pandas_batches(page_size=initial_page_size)
+
         # set traitlets properties that trigger observers
         self.page_size = initial_page_size
 
@@ -95,21 +84,12 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # TODO(b/428238610): Start iterating over the result of `to_pandas_batches()`
         # before we get here so that the count might already be cached.
         # TODO(b/452747934): Allow row_count to be None and check to see if
-        # there are multiple pages and show "page 1 of many" in this case
-        self._reset_batches_for_new_page_size()
-        if self._batches is None or self._batches.total_rows is None:
-            self._error_message = "Could not determine total row count. Data might be unavailable or an error occurred."
-            self.row_count = 0
-        else:
-            self.row_count = self._batches.total_rows
+        # there are multiple pages and show "page 1 of many" in this case.
+        self.row_count = self._batches.total_rows or 0
 
         # get the initial page
         self._set_table_html()
 
-        # Signals to the frontend that the initial data load is complete.
-        # Also used as a guard to prevent observers from firing during initialization.
-        self._initial_load_complete = True
-
     @functools.cached_property
     def _esm(self):
         """Load JavaScript code from external file."""
@@ -120,6 +100,11 @@ def _css(self):
         """Load CSS code from external file."""
         return resources.read_text(bigframes.display, "table_widget.css")
 
+    page = traitlets.Int(0).tag(sync=True)
+    page_size = traitlets.Int(25).tag(sync=True)
+    row_count = traitlets.Int(0).tag(sync=True)
+    table_html = traitlets.Unicode().tag(sync=True)
+
     @traitlets.validate("page")
     def _validate_page(self, proposal: Dict[str, Any]) -> int:
         """Validate and clamp the page number to a valid range.
@@ -186,10 +171,7 @@ def _get_next_batch(self) -> bool:
     def _batch_iterator(self) -> Iterator[pd.DataFrame]:
         """Lazily initializes and returns the batch iterator."""
         if self._batch_iter is None:
-            if self._batches is None:
-                self._batch_iter = iter([])
-            else:
-                self._batch_iter = iter(self._batches)
+            self._batch_iter = iter(self._batches)
         return self._batch_iter
 
     @property
@@ -199,22 +181,15 @@ def _cached_data(self) -> pd.DataFrame:
             return pd.DataFrame(columns=self._dataframe.columns)
         return pd.concat(self._cached_batches, ignore_index=True)
 
-    def _reset_batches_for_new_page_size(self) -> None:
+    def _reset_batches_for_new_page_size(self):
         """Reset the batch iterator when page size changes."""
         self._batches = self._dataframe._to_pandas_batches(page_size=self.page_size)
-
         self._cached_batches = []
         self._batch_iter = None
         self._all_data_loaded = False
 
-    def _set_table_html(self) -> None:
+    def _set_table_html(self):
         """Sets the current html data based on the current page and page size."""
-        if self._error_message:
-            self.table_html = (
-                f"<div class='bigframes-error-message'>{self._error_message}</div>"
-            )
-            return
-
         start = self.page * self.page_size
         end = start + self.page_size
 
@@ -236,17 +211,13 @@ def _set_table_html(self) -> None:
         )
 
     @traitlets.observe("page")
-    def _page_changed(self, _change: Dict[str, Any]) -> None:
+    def _page_changed(self, _change: Dict[str, Any]):
         """Handler for when the page number is changed from the frontend."""
-        if not self._initial_load_complete:
-            return
         self._set_table_html()
 
     @traitlets.observe("page_size")
-    def _page_size_changed(self, _change: Dict[str, Any]) -> None:
+    def _page_size_changed(self, _change: Dict[str, Any]):
         """Handler for when the page size is changed from the frontend."""
-        if not self._initial_load_complete:
-            return
         # Reset the page to 0 when page size changes to avoid invalid page states
         self.page = 0
 

bigframes/display/table_widget.js

@@ -137,12 +137,6 @@ function render({ model, el }) {
 		}
 	});
 	model.on(Event.CHANGE_TABLE_HTML, handleTableHTMLChange);
-	model.on(`change:${ModelProperty.ROW_COUNT}`, updateButtonStates);
-	model.on(`change:_initial_load_complete`, (val) => {
-		if (val) {
-			updateButtonStates();
-		}
-	});
 
 	// Assemble the DOM
 	paginationContainer.appendChild(prevPage);

docs/templates/toc.yml

@@ -45,6 +45,9 @@
         uid: bigframes.operations.plotting.PlotAccessor
       - name: StructAccessor
         uid: bigframes.operations.structs.StructFrameAccessor
+      - name: AI
+        uid: bigframes.operations.ai.AIAccessor
+        status: beta
       name: DataFrame
     - items:
       - name: DataFrameGroupBy