feat(bigframes): add AI TVFs to the pandas bq accessor#17402
Conversation
There was a problem hiding this comment.
Code Review
This pull request adds generate_embedding, generate_text, and generate_table accessor methods to the BigQuery DataFrame accessor, along with comprehensive unit tests. The feedback suggests removing pd.Series from the type annotations of the model parameter in these new methods, as BigQuery ML models cannot be pandas Series.
| ) | ||
| return self._to_series(result) | ||
|
|
||
| def generate_embedding( |
There was a problem hiding this comment.
Given that this acts only on a single column (called "content"), I think the Series accessor would be a much better fit.
There was a problem hiding this comment.
Good call! Let me move the contents to the series accessor
| ) | ||
| return self._to_dataframe(result) | ||
|
|
||
| def generate_text( |
There was a problem hiding this comment.
Same here. This acts on a single "prompt" column. I think Series accessor would be a better fit.
To avoid conflicts with generated code, please use a "mixin" pattern to avoid conflicts between the generated Series code and these methods: https://adamj.eu/tech/2025/05/01/python-type-hints-mixin-classes/
| ) | ||
| return self._to_dataframe(result) | ||
|
|
||
| def generate_table( |
There was a problem hiding this comment.
Same here. This acts on a single "prompt" column and would be better suited as a series accessor method.
Also updated the tests to fully utilize the mocking framework.
internal issue: b/517233441