feat: support time_series_id_col in ARIMAPlus (#1282)

* manually port from shuowei-arima-plus branch, now I cannot pass format test

* I have resolve all conflicts after manual porting

* use inherritance for arima plus model, and add sql models

* resolve unexpected indent for docstring

---------

Co-authored-by: Shuowei Li <[email protected]>

Author: Shuowei Li

bigframes/ml/base.py

@@ -165,6 +165,34 @@ def fit(
         return self._fit(X, y)
 
 
+class SupervisedTrainableWithIdColPredictor(SupervisedTrainablePredictor):
+    """Inherits from SupervisedTrainablePredictor,
+    but adds an optional id_col parameter to fit()."""
+
+    def __init__(self):
+        super().__init__()
+        self.id_col = None
+
+    def _fit(
+        self,
+        X: utils.ArrayType,
+        y: utils.ArrayType,
+        transforms=None,
+        id_col: Optional[utils.ArrayType] = None,
+    ):
+        return self
+
+    def fit(
+        self,
+        X: utils.ArrayType,
+        y: utils.ArrayType,
+        transforms=None,
+        id_col: Optional[utils.ArrayType] = None,
+    ):
+        self.id_col = id_col
+        return self._fit(X, y, transforms=transforms, id_col=self.id_col)
+
+
 class TrainableWithEvaluationPredictor(TrainablePredictor):
     """A BigQuery DataFrames ML Model base class that can be used to fit and predict outputs.
 

bigframes/ml/core.py

@@ -181,15 +181,23 @@ def detect_anomalies(
 
     def forecast(self, options: Mapping[str, int | float]) -> bpd.DataFrame:
         sql = self._model_manipulation_sql_generator.ml_forecast(struct_options=options)
-        return self._session.read_gbq(sql, index_col="forecast_timestamp").reset_index()
+        timestamp_col_name = "forecast_timestamp"
+        index_cols = [timestamp_col_name]
+        first_col_name = self._session.read_gbq(sql).columns.values[0]
+        if timestamp_col_name != first_col_name:
+            index_cols.append(first_col_name)
+        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
         )
-        return self._session.read_gbq(
-            sql, index_col="time_series_timestamp"
-        ).reset_index()
+        timestamp_col_name = "time_series_timestamp"
+        index_cols = [timestamp_col_name]
+        first_col_name = self._session.read_gbq(sql).columns.values[0]
+        if timestamp_col_name != first_col_name:
+            index_cols.append(first_col_name)
+        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(
@@ -390,6 +398,7 @@ def create_time_series_model(
         self,
         X_train: bpd.DataFrame,
         y_train: bpd.DataFrame,
+        id_col: Optional[bpd.DataFrame] = None,
         transforms: Optional[Iterable[str]] = None,
         options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> BqmlModel:
@@ -399,13 +408,21 @@ def create_time_series_model(
         assert (
             y_train.columns.size == 1
         ), "Time stamp data input must only contain 1 column."
+        assert id_col is None or (
+            id_col is not None and id_col.columns.size == 1
+        ), "Time series id input is either None or must only contain 1 column."
 
         options = dict(options)
         # Cache dataframes to make sure base table is not a snapshot
         # cached dataframe creates a full copy, never uses snapshot
-        input_data = X_train.join(y_train, how="outer").cache()
+        input_data = X_train.join(y_train, how="outer")
+        if id_col is not None:
+            input_data = input_data.join(id_col, how="outer")
+        input_data = input_data.cache()
         options.update({"TIME_SERIES_TIMESTAMP_COL": X_train.columns.tolist()[0]})
         options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]})
+        if id_col is not None:
+            options.update({"TIME_SERIES_ID_COL": id_col.columns.tolist()[0]})
 
         session = X_train._session
         model_ref = self._create_model_ref(session._anonymous_dataset)

bigframes/ml/forecasting.py

@@ -45,7 +45,7 @@
 
 
 @log_adapter.class_logger
-class ARIMAPlus(base.SupervisedTrainablePredictor):
+class ARIMAPlus(base.SupervisedTrainableWithIdColPredictor):
     """Time Series ARIMA Plus model.
 
     Args:
@@ -183,37 +183,53 @@ def _fit(
         X: utils.ArrayType,
         y: utils.ArrayType,
         transforms: Optional[List[str]] = None,
-    ):
+        id_col: Optional[utils.ArrayType] = None,
+    ) -> ARIMAPlus:
         """Fit the model to training data.
 
         Args:
-            X (bigframes.dataframe.DataFrame or bigframes.series.Series):
-                A dataframe of training timestamp.
-
-            y (bigframes.dataframe.DataFrame or bigframes.series.Series):
+            X (bigframes.dataframe.DataFrame or bigframes.series.Series,
+            or pandas.core.frame.DataFrame or pandas.core.series.Series):
+                A dataframe or series of trainging timestamp.
+            y (bigframes.dataframe.DataFrame, or bigframes.series.Series,
+            or pandas.core.frame.DataFrame, or pandas.core.series.Series):
                 Target values for training.
             transforms (Optional[List[str]], default None):
                 Do not use. Internal param to be deprecated.
                 Use bigframes.ml.pipeline instead.
+            id_col (Optional[bigframes.dataframe.DataFrame]
+            or Optional[bigframes.series.Series]
+            or Optional[pandas.core.frame.DataFrame]
+            or Optional[pandas.core.frame.Series]
+            or None, default None):
+                An optional dataframe or series of training id col.
 
         Returns:
             ARIMAPlus: Fitted estimator.
         """
         X, y = utils.batch_convert_to_dataframe(X, y)
 
         if X.columns.size != 1:
-            raise ValueError(
-                "Time series timestamp input X must only contain 1 column."
-            )
+            raise ValueError("Time series timestamp input X contain at least 1 column.")
         if y.columns.size != 1:
             raise ValueError("Time series data input y must only contain 1 column.")
 
+        if id_col is not None:
+            (id_col,) = utils.batch_convert_to_dataframe(id_col)
+
+            if id_col.columns.size != 1:
+                raise ValueError(
+                    "Time series id input id_col must only contain 1 column."
+                )
+
         self._bqml_model = self._bqml_model_factory.create_time_series_model(
             X,
             y,
+            id_col=id_col,
             transforms=transforms,
             options=self._bqml_options,
         )
+        return self
 
     def predict(
         self, X=None, *, horizon: int = 3, confidence_level: float = 0.95
@@ -237,7 +253,7 @@ def predict(
 
         Returns:
             bigframes.dataframe.DataFrame: The predicted DataFrames. Which
-                contains 2 columns: "forecast_timestamp" and "forecast_value".
+                contains 2 columns: "forecast_timestamp", "id" as optional, and "forecast_value".
         """
         if horizon < 1 or horizon > 1000:
             raise ValueError(f"horizon must be [1, 1000], but is {horizon}.")
@@ -345,6 +361,7 @@ def score(
         self,
         X: utils.ArrayType,
         y: utils.ArrayType,
+        id_col: Optional[utils.ArrayType] = None,
     ) -> bpd.DataFrame:
         """Calculate evaluation metrics of the model.
 
@@ -355,13 +372,22 @@ def score(
             for the outputs relevant to this model type.
 
         Args:
-            X (bigframes.dataframe.DataFrame or bigframes.series.Series or pandas.core.frame.DataFrame or pandas.core.series.Series):
-                A BigQuery DataFrame only contains 1 column as
+            X (bigframes.dataframe.DataFrame or bigframes.series.Series
+            or pandas.core.frame.DataFrame or pandas.core.series.Series):
+                A dataframe or series only contains 1 column as
                 evaluation timestamp. The timestamp must be within the horizon
                 of the model, which by default is 1000 data points.
-            y (bigframes.dataframe.DataFrame or bigframes.series.Series or pandas.core.frame.DataFrame or pandas.core.series.Series):
-                A BigQuery DataFrame only contains 1 column as
+            y (bigframes.dataframe.DataFrame or bigframes.series.Series
+            or pandas.core.frame.DataFrame or pandas.core.series.Series):
+                A dataframe or series only contains 1 column as
                 evaluation numeric values.
+            id_col (Optional[bigframes.dataframe.DataFrame]
+            or Optional[bigframes.series.Series]
+            or Optional[pandas.core.frame.DataFrame]
+            or Optional[pandas.core.series.Series]
+            or None, default None):
+                An optional dataframe or series contains at least 1 column as
+                evaluation id column.
 
         Returns:
             bigframes.dataframe.DataFrame: A DataFrame as evaluation result.
@@ -371,6 +397,10 @@ def score(
         X, y = utils.batch_convert_to_dataframe(X, y, session=self._bqml_model.session)
 
         input_data = X.join(y, how="outer")
+        if id_col is not None:
+            (id_col,) = utils.batch_convert_to_dataframe(id_col)
+            input_data = input_data.join(id_col, how="outer")
+
         return self._bqml_model.evaluate(input_data)
 
     def summary(