feat: add DataFrame.ai.forecast() support (#1828)

* feat: add DataFrame.ai.forecast() support

* test

* fix

* constructor

* update

* comments

* fix

Author: GarrettWu

bigframes/ml/core.py

@@ -35,7 +35,17 @@ class BaseBqml:
 
     def __init__(self, session: bigframes.session.Session):
         self._session = session
-        self._base_sql_generator = ml_sql.BaseSqlGenerator()
+        self._sql_generator = ml_sql.BaseSqlGenerator()
+
+    def ai_forecast(
+        self,
+        input_data: bpd.DataFrame,
+        options: Mapping[str, Union[str, int, float, Iterable[str]]],
+    ) -> bpd.DataFrame:
+        result_sql = self._sql_generator.ai_forecast(
+            source_sql=input_data.sql, options=options
+        )
+        return self._session.read_gbq(result_sql)
 
 
 class BqmlModel(BaseBqml):
@@ -55,8 +65,8 @@ def __init__(self, session: bigframes.Session, model: bigquery.Model):
         self._model = model
         model_ref = self._model.reference
         assert model_ref is not None
-        self._model_manipulation_sql_generator = ml_sql.ModelManipulationSqlGenerator(
-            model_ref
+        self._sql_generator: ml_sql.ModelManipulationSqlGenerator = (
+            ml_sql.ModelManipulationSqlGenerator(model_ref)
         )
 
     def _apply_ml_tvf(
@@ -126,30 +136,28 @@ def model(self) -> bigquery.Model:
     def recommend(self, input_data: bpd.DataFrame) -> bpd.DataFrame:
         return self._apply_ml_tvf(
             input_data,
-            self._model_manipulation_sql_generator.ml_recommend,
+            self._sql_generator.ml_recommend,
         )
 
     def predict(self, input_data: bpd.DataFrame) -> bpd.DataFrame:
         return self._apply_ml_tvf(
             input_data,
-            self._model_manipulation_sql_generator.ml_predict,
+            self._sql_generator.ml_predict,
         )
 
     def explain_predict(
         self, input_data: bpd.DataFrame, options: Mapping[str, int | float]
     ) -> bpd.DataFrame:
         return self._apply_ml_tvf(
             input_data,
-            lambda source_sql: self._model_manipulation_sql_generator.ml_explain_predict(
+            lambda source_sql: self._sql_generator.ml_explain_predict(
                 source_sql=source_sql,
                 struct_options=options,
             ),
         )
 
     def global_explain(self, options: Mapping[str, bool]) -> bpd.DataFrame:
-        sql = self._model_manipulation_sql_generator.ml_global_explain(
-            struct_options=options
-        )
+        sql = self._sql_generator.ml_global_explain(struct_options=options)
         return (
             self._session.read_gbq(sql)
             .sort_values(by="attribution", ascending=False)
@@ -159,7 +167,7 @@ def global_explain(self, options: Mapping[str, bool]) -> bpd.DataFrame:
     def transform(self, input_data: bpd.DataFrame) -> bpd.DataFrame:
         return self._apply_ml_tvf(
             input_data,
-            self._model_manipulation_sql_generator.ml_transform,
+            self._sql_generator.ml_transform,
         )
 
     def generate_text(
@@ -170,7 +178,7 @@ def generate_text(
         options["flatten_json_output"] = True
         return self._apply_ml_tvf(
             input_data,
-            lambda source_sql: self._model_manipulation_sql_generator.ml_generate_text(
+            lambda source_sql: self._sql_generator.ml_generate_text(
                 source_sql=source_sql,
                 struct_options=options,
             ),
@@ -186,7 +194,7 @@ def generate_embedding(
         options["flatten_json_output"] = True
         return self._apply_ml_tvf(
             input_data,
-            lambda source_sql: self._model_manipulation_sql_generator.ml_generate_embedding(
+            lambda source_sql: self._sql_generator.ml_generate_embedding(
                 source_sql=source_sql,
                 struct_options=options,
             ),
@@ -201,7 +209,7 @@ def generate_table(
     ) -> bpd.DataFrame:
         return self._apply_ml_tvf(
             input_data,
-            lambda source_sql: self._model_manipulation_sql_generator.ai_generate_table(
+            lambda source_sql: self._sql_generator.ai_generate_table(
                 source_sql=source_sql,
                 struct_options=options,
             ),
@@ -216,14 +224,14 @@ def detect_anomalies(
 
         return self._apply_ml_tvf(
             input_data,
-            lambda source_sql: self._model_manipulation_sql_generator.ml_detect_anomalies(
+            lambda source_sql: self._sql_generator.ml_detect_anomalies(
                 source_sql=source_sql,
                 struct_options=options,
             ),
         )
 
     def forecast(self, options: Mapping[str, int | float]) -> bpd.DataFrame:
-        sql = self._model_manipulation_sql_generator.ml_forecast(struct_options=options)
+        sql = self._sql_generator.ml_forecast(struct_options=options)
         timestamp_col_name = "forecast_timestamp"
         index_cols = [timestamp_col_name]
         first_col_name = self._session.read_gbq(sql).columns.values[0]
@@ -232,9 +240,7 @@ def forecast(self, options: Mapping[str, int | float]) -> bpd.DataFrame:
         return self._session.read_gbq(sql, index_col=index_cols).reset_index()
 
     def explain_forecast(self, options: Mapping[str, int | float]) -> bpd.DataFrame:
-        sql = self._model_manipulation_sql_generator.ml_explain_forecast(
-            struct_options=options
-        )
+        sql = self._sql_generator.ml_explain_forecast(struct_options=options)
         timestamp_col_name = "time_series_timestamp"
         index_cols = [timestamp_col_name]
         first_col_name = self._session.read_gbq(sql).columns.values[0]
@@ -243,7 +249,7 @@ def explain_forecast(self, options: Mapping[str, int | float]) -> bpd.DataFrame:
         return self._session.read_gbq(sql, index_col=index_cols).reset_index()
 
     def evaluate(self, input_data: Optional[bpd.DataFrame] = None):
-        sql = self._model_manipulation_sql_generator.ml_evaluate(
+        sql = self._sql_generator.ml_evaluate(
             input_data.sql if (input_data is not None) else None
         )
 
@@ -254,28 +260,24 @@ def llm_evaluate(
         input_data: bpd.DataFrame,
         task_type: Optional[str] = None,
     ):
-        sql = self._model_manipulation_sql_generator.ml_llm_evaluate(
-            input_data.sql, task_type
-        )
+        sql = self._sql_generator.ml_llm_evaluate(input_data.sql, task_type)
 
         return self._session.read_gbq(sql)
 
     def arima_evaluate(self, show_all_candidate_models: bool = False):
-        sql = self._model_manipulation_sql_generator.ml_arima_evaluate(
-            show_all_candidate_models
-        )
+        sql = self._sql_generator.ml_arima_evaluate(show_all_candidate_models)
 
         return self._session.read_gbq(sql)
 
     def arima_coefficients(self) -> bpd.DataFrame:
-        sql = self._model_manipulation_sql_generator.ml_arima_coefficients()
+        sql = self._sql_generator.ml_arima_coefficients()
 
         return self._session.read_gbq(sql)
 
     def centroids(self) -> bpd.DataFrame:
         assert self._model.model_type == "KMEANS"
 
-        sql = self._model_manipulation_sql_generator.ml_centroids()
+        sql = self._sql_generator.ml_centroids()
 
         return self._session.read_gbq(
             sql, index_col=["centroid_id", "feature"]
@@ -284,7 +286,7 @@ def centroids(self) -> bpd.DataFrame:
     def principal_components(self) -> bpd.DataFrame:
         assert self._model.model_type == "PCA"
 
-        sql = self._model_manipulation_sql_generator.ml_principal_components()
+        sql = self._sql_generator.ml_principal_components()
 
         return self._session.read_gbq(
             sql, index_col=["principal_component_id", "feature"]
@@ -293,7 +295,7 @@ def principal_components(self) -> bpd.DataFrame:
     def principal_component_info(self) -> bpd.DataFrame:
         assert self._model.model_type == "PCA"
 
-        sql = self._model_manipulation_sql_generator.ml_principal_component_info()
+        sql = self._sql_generator.ml_principal_component_info()
 
         return self._session.read_gbq(sql)
 
@@ -319,7 +321,7 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel:
         # truncate as Vertex ID only accepts 63 characters, easily exceeding the limit for temp models.
         # The possibility of conflicts should be low.
         vertex_ai_model_id = vertex_ai_model_id[:63]
-        sql = self._model_manipulation_sql_generator.alter_model(
+        sql = self._sql_generator.alter_model(
             options={"vertex_ai_model_id": vertex_ai_model_id}
         )
         # Register the model and wait it to finish

bigframes/ml/sql.py

@@ -49,6 +49,12 @@ def build_parameters(self, **kwargs: Union[str, int, float, Iterable[str]]) -> s
         param_strs = [f"{k}={self.encode_value(v)}" for k, v in kwargs.items()]
         return "\n" + INDENT_STR + f",\n{INDENT_STR}".join(param_strs)
 
+    def build_named_parameters(
+        self, **kwargs: Union[str, int, float, Iterable[str]]
+    ) -> str:
+        param_strs = [f"{k} => {self.encode_value(v)}" for k, v in kwargs.items()]
+        return "\n" + INDENT_STR + f",\n{INDENT_STR}".join(param_strs)
+
     def build_structs(self, **kwargs: Union[int, float, str, Mapping]) -> str:
         """Encode a dict of values into a formatted STRUCT items for SQL"""
         param_strs = []
@@ -187,6 +193,17 @@ def ml_distance(
         https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-distance"""
         return f"""SELECT *, ML.DISTANCE({sql_utils.identifier(col_x)}, {sql_utils.identifier(col_y)}, '{type}') AS {sql_utils.identifier(name)} FROM ({source_sql})"""
 
+    def ai_forecast(
+        self,
+        source_sql: str,
+        options: Mapping[str, Union[int, float, bool, Iterable[str]]],
+    ):
+        """Encode AI.FORECAST.
+        https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-forecast"""
+        named_parameters_sql = self.build_named_parameters(**options)
+
+        return f"""SELECT * FROM AI.FORECAST(({source_sql}),{named_parameters_sql})"""
+
 
 class ModelCreationSqlGenerator(BaseSqlGenerator):
     """Sql generator for creating a model entity. Model id is the standalone id without project id and dataset id."""

bigframes/operations/ai.py

@@ -16,25 +16,24 @@
 
 import re
 import typing
-from typing import Dict, List, Optional, Sequence
+from typing import Dict, Iterable, List, Optional, Sequence, Union
 import warnings
 
 import numpy as np
 
-from bigframes import dtypes, exceptions
+from bigframes import dtypes, exceptions, options
 from bigframes.core import guid, log_adapter
 
 
 @log_adapter.class_logger
 class AIAccessor:
-    def __init__(self, df) -> None:
+    def __init__(self, df, base_bqml=None) -> None:
         import bigframes  # Import in the function body to avoid circular imports.
         import bigframes.dataframe
-
-        if not bigframes.options.experiments.ai_operators:
-            raise NotImplementedError()
+        from bigframes.ml import core as ml_core
 
         self._df: bigframes.dataframe.DataFrame = df
+        self._base_bqml: ml_core.BaseBqml = base_bqml or ml_core.BaseBqml(df._session)
 
     def filter(
         self,
@@ -89,6 +88,8 @@ def filter(
             ValueError: when the instruction refers to a non-existing column, or when no
                 columns are referred to.
         """
+        if not options.experiments.ai_operators:
+            raise NotImplementedError()
 
         answer_col = "answer"
 
@@ -181,6 +182,9 @@ def map(
             ValueError: when the instruction refers to a non-existing column, or when no
                 columns are referred to.
         """
+        if not options.experiments.ai_operators:
+            raise NotImplementedError()
+
         import bigframes.dataframe
         import bigframes.series
 
@@ -320,6 +324,8 @@ def classify(
                 columns are referred to, or when the count of labels does not meet the
                 requirement.
         """
+        if not options.experiments.ai_operators:
+            raise NotImplementedError()
 
         if len(labels) < 2 or len(labels) > 20:
             raise ValueError(
@@ -401,6 +407,9 @@ def join(
         Raises:
             ValueError if the amount of data that will be sent for LLM processing is larger than max_rows.
         """
+        if not options.experiments.ai_operators:
+            raise NotImplementedError()
+
         self._validate_model(model)
         columns = self._parse_columns(instruction)
 
@@ -525,6 +534,8 @@ def search(
             ValueError: when the search_column is not found from the the data frame.
             TypeError: when the provided model is not TextEmbeddingGenerator.
         """
+        if not options.experiments.ai_operators:
+            raise NotImplementedError()
 
         if search_column not in self._df.columns:
             raise ValueError(f"Column `{search_column}` not found")
@@ -640,6 +651,9 @@ def top_k(
             ValueError: when the instruction refers to a non-existing column, or when no
                 columns are referred to.
         """
+        if not options.experiments.ai_operators:
+            raise NotImplementedError()
+
         import bigframes.dataframe
         import bigframes.series
 
@@ -834,6 +848,8 @@ def sim_join(
         Raises:
             ValueError: when the amount of data to be processed exceeds the specified max_rows.
         """
+        if not options.experiments.ai_operators:
+            raise NotImplementedError()
 
         if left_on not in self._df.columns:
             raise ValueError(f"Left column {left_on} not found")
@@ -883,6 +899,73 @@ def sim_join(
 
         return join_result
 
+    def forecast(
+        self,
+        timestamp_column: str,
+        data_column: str,
+        *,
+        model: str = "TimesFM 2.0",
+        id_columns: Optional[Iterable[str]] = None,
+        horizon: int = 10,
+        confidence_level: float = 0.95,
+    ):
+        """
+        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:
+            timestamp_column (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
+            data_column (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
+            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_columns (Iterable[str] or None, default None):
+                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).
+
+        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 referring to a non-existing column.
+        """
+        columns = [timestamp_column, data_column]
+        if id_columns:
+            columns += id_columns
+        for column in columns:
+            if column not in self._df.columns:
+                raise ValueError(f"Column `{column}` not found")
+
+        options: dict[str, Union[int, float, str, Iterable[str]]] = {
+            "data_col": data_column,
+            "timestamp_col": timestamp_column,
+            "model": model,
+            "horizon": horizon,
+            "confidence_level": confidence_level,
+        }
+        if id_columns:
+            options["id_cols"] = id_columns
+
+        return self._base_bqml.ai_forecast(input_data=self._df, options=options)
+
     @staticmethod
     def _attach_embedding(dataframe, source_column: str, embedding_column: str, model):
         result_df = dataframe.copy()