Merge branch 'feat-st-regionstats' of https://github.com/googleapis/python-bigquery-dataframes into feat-st-regionstats

Author: gcf-owl-bot[bot]

bigframes/bigquery/_operations/ai.py

@@ -19,12 +19,15 @@
 from __future__ import annotations
 
 import json
-from typing import Any, List, Literal, Mapping, Tuple, Union
+from typing import Any, Iterable, List, Literal, Mapping, Tuple, Union
 
 import pandas as pd
 
-from bigframes import clients, dtypes, series, session
+from bigframes import clients, dataframe, dtypes
+from bigframes import pandas as bpd
+from bigframes import 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[
@@ -548,6 +551,91 @@ 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/dataframe.py

@@ -5328,12 +5328,16 @@ def _throw_if_null_index(self, opname: str):
     @property
     def semantics(self):
         msg = bfe.format_message(
-            "The 'semantics' property will be removed. Please use 'ai' instead."
+            "The 'semantics' property will be removed. Please use 'bigframes.bigquery.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,13 +23,14 @@
 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
@@ -46,9 +47,21 @@
 
 
 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.
     """
-    An interactive, paginated table widget for BigFrames DataFrames.
-    """
+
+    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
+    )
 
     def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         """Initialize the TableWidget.
@@ -61,10 +74,11 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
                 "Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use TableWidget."
             )
 
-        super().__init__()
         self._dataframe = dataframe
 
-        # Initialize attributes that might be needed by observers FIRST
+        super().__init__()
+
+        # 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
@@ -73,9 +87,6 @@ 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
 
@@ -84,12 +95,21 @@ 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.row_count = self._batches.total_rows or 0
+        # 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
 
         # 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."""
@@ -100,11 +120,6 @@ 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.
@@ -171,7 +186,10 @@ 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:
-            self._batch_iter = iter(self._batches)
+            if self._batches is None:
+                self._batch_iter = iter([])
+            else:
+                self._batch_iter = iter(self._batches)
         return self._batch_iter
 
     @property
@@ -181,15 +199,22 @@ 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):
+    def _reset_batches_for_new_page_size(self) -> None:
         """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):
+    def _set_table_html(self) -> None:
         """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
 
@@ -211,13 +236,17 @@ def _set_table_html(self):
         )
 
     @traitlets.observe("page")
-    def _page_changed(self, _change: Dict[str, Any]):
+    def _page_changed(self, _change: Dict[str, Any]) -> None:
         """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]):
+    def _page_size_changed(self, _change: Dict[str, Any]) -> None:
         """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,6 +137,12 @@ 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,9 +45,6 @@
         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