fix lint errors

Author: tswast

bigframes/bigquery/_operations/ai.py

@@ -880,6 +880,7 @@ def forecast(
     id_cols: Iterable[str] | None = None,
     horizon: int = 10,
     confidence_level: float = 0.95,
+    output_historical_time_series: bool = False,
     context_window: int | None = None,
 ) -> dataframe.DataFrame:
     """
@@ -914,6 +915,15 @@ def forecast(
         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).
+        output_historical_time_series (bool, default False):
+            A BOOL value that determines whether the input data is returned
+            along with the forecasted data. Set this argument to TRUE to return
+            input data. The default value is FALSE.
+
+            Returning the input data along with the forecasted data lets you
+            compare the historical value of the data column with the forecasted
+            value of the data column, or chart the change in the data column
+            values over time.
         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.
@@ -945,6 +955,7 @@ def forecast(
         "timestamp_col": timestamp_col,
         "model": model,
         "horizon": horizon,
+        "output_historical_time_series": output_historical_time_series,
         "confidence_level": confidence_level,
     }
     if id_cols:

patch_tests.diff

@@ -46,8 +46,9 @@ def to_pandas(series, *, ordered):
 
 
 def test_ai_forecast(snapshot, monkeypatch):
-    import bigframes.session
     import bigframes.bigquery.ai
+    import bigframes.session
+
     session = mock.create_autospec(bigframes.session.Session)
     bf_df = mock.create_autospec(bpd.DataFrame)
     session.read_pandas.return_value = bf_df
@@ -59,6 +60,7 @@ def mock_ai_forecast(df, **kwargs):
         return result_df
 
     import bigframes.bigquery.ai
+
     monkeypatch.setattr(bigframes.bigquery.ai, "forecast", mock_ai_forecast)
 
     df = pd.DataFrame({"date": ["2020-01-01"], "value": [1.0]})