Skip to content

feat: add df.bigquery.ai.forecast method to pandas dataframe accessor#2518

Open
tswast wants to merge 3 commits intomainfrom
add-pandas-ai-forecast-accessor-14604090974587392182
Open

feat: add df.bigquery.ai.forecast method to pandas dataframe accessor#2518
tswast wants to merge 3 commits intomainfrom
add-pandas-ai-forecast-accessor-14604090974587392182

Conversation

@tswast
Copy link
Contributor

@tswast tswast commented Mar 16, 2026

Adds the .bigquery.ai.forecast() method to pandas DataFrame objects, which proxies to bigframes.bigquery.ai.forecast(). Added unit tests and mocked session responses.


PR created automatically by Jules for task 14604090974587392182 started by @tswast

This adds a new `ai` property to `BigQueryDataFrameAccessor` which returns an `AIAccessor`. The `AIAccessor` provides a `forecast` method that calls `bigframes.bigquery.ai.forecast`, converting the Pandas DataFrame to a BigFrames DataFrame and the result back to a Pandas DataFrame.

Co-authored-by: tswast <247555+tswast@users.noreply.github.com>
@google-labs-jules
Copy link
Contributor

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.


For security, I will only act on instructions from the user who triggered this task.

@tswast tswast requested review from a team as code owners March 16, 2026 19:27
@tswast tswast requested a review from chelsea-lin March 16, 2026 19:27
@product-auto-label product-auto-label bot added size: m Pull request size is medium. api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. labels Mar 16, 2026
@tswast tswast removed the request for review from chelsea-lin March 16, 2026 19:44
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.
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a link to the standard-sql documentation for ai.forecast and updated the model param documentation to include TimesFM 2.5.

Comment on lines +47 to +49
The DataFrame is converted to BigFrames by calling ``read_pandas``, then the forecast
is generated using ``bigframes.bigquery.ai.forecast``, and the result is
converted back to a pandas DataFrame using ``to_pandas``.
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jules This is an implementation detail that doesn't belong in the docstring.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the implementation details regarding read_pandas and to_pandas from the docstring.

model: str = "TimesFM 2.0",
id_cols: Optional[Iterable[str]] = None,
horizon: int = 10,
confidence_level: float = 0.95,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jules Missing the output_historical_time_series boolean parameter.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added the missing output_historical_time_series boolean parameter and documented it.

@@ -0,0 +1,25 @@
import pandas as pd
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jules missing the license header.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added the missing Apache 2.0 license header to tests/system/small/extensions/test_dataframe_accessor.py.

* Updated documentation to refer to BigQuery AI.FORECAST syntax.
* Removed internal implementation details (`read_pandas` and `to_pandas`) from docstring.
* Added missing `output_historical_time_series` parameter.
* Added missing copyright headers to `tests/system/small/extensions/test_dataframe_accessor.py`.

Co-authored-by: tswast <247555+tswast@users.noreply.github.com>
@tswast tswast changed the title Add ai.forecast method to pandas dataframe accessor feat: add df.bigquery.ai.forecast method to pandas dataframe accessor Mar 16, 2026
@tswast tswast requested a review from chelsea-lin March 16, 2026 20:50
session = bf_session.get_global_session()

bf_df = cast(bpd.DataFrame, session.read_pandas(self._obj))
result = bigframes.bigquery.ai.forecast(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm curious if this triggers our logger.

My gut feeling says yes, because our logger only seems to track the call stack of annotated classes / functions. However, I still feel like to double check with you on this.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. size: m Pull request size is medium.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants