feat: add bigquery.ai.generate_table function (#2453)

Thank you for opening a Pull Request! Before submitting your PR, there
are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
Fixes #<issue_number_goes_here> 🦕

Author: GarrettWu

bigframes/bigquery/_operations/ai.py

@@ -601,6 +601,101 @@ def generate_text(
 returnsession.read_gbq_query(query)
 
 
+@log_adapter.method_logger(custom_base_name="bigquery_ai")
+defgenerate_table(
+model: Union[bigframes.ml.base.BaseEstimator, str, pd.Series],
+data: Union[dataframe.DataFrame, series.Series, pd.DataFrame, pd.Series],
+*,
+output_schema: str,
+temperature: Optional[float] =None,
+top_p: Optional[float] =None,
+max_output_tokens: Optional[int] =None,
+stop_sequences: Optional[List[str]] =None,
+request_type: Optional[str] =None,
+) ->dataframe.DataFrame:
+"""
+ Generates a table using a BigQuery ML model.
+
+ See the `AI.GENERATE_TABLE function syntax
+ <https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-generate-table>`_
+ for additional reference.
+
+ **Examples:**
+
+ >>> import bigframes.pandas as bpd
+ >>> import bigframes.bigquery as bbq
+ >>> # The user is responsible for constructing a DataFrame that contains
+ >>> # the necessary columns for the model's prompt. For example, a
+ >>> # DataFrame with a 'prompt' column for text classification.
+ >>> df = bpd.DataFrame({'prompt': ["some text to classify"]})
+ >>> result = bbq.ai.generate_table(
+ ... "project.dataset.model_name",
+ ... data=df,
+ ... output_schema="category STRING"
+ ... ) # doctest: +SKIP
+
+ Args:
+ model (bigframes.ml.base.BaseEstimator or str):
+ The model to use for table generation.
+ data (bigframes.pandas.DataFrame or bigframes.pandas.Series):
+ The data to generate embeddings for. If a Series is provided, it is
+ treated as the 'content' column. If a DataFrame is provided, it
+ must contain a 'content' column, or you must rename the column you
+ wish to embed to 'content'.
+ output_schema (str):
+ A string defining the output schema (e.g., "col1 STRING, col2 INT64").
+ temperature (float, optional):
+ A FLOAT64 value that is used for sampling promiscuity. The value
+ must be in the range ``[0.0, 1.0]``.
+ top_p (float, optional):
+ A FLOAT64 value that changes how the model selects tokens for
+ output.
+ max_output_tokens (int, optional):
+ An INT64 value that sets the maximum number of tokens in the
+ generated table.
+ stop_sequences (List[str], optional):
+ An ARRAY<STRING> value that contains the stop sequences for the model.
+ request_type (str, optional):
+ A STRING value that contains the request type for the model.
+
+ Returns:
+ bigframes.pandas.DataFrame:
+ The generated table.
+ """
+data=_to_dataframe(data, series_rename="prompt")
+model_name, session=bq_utils.get_model_name_and_session(model, data)
+table_sql=bq_utils.to_sql(data)
+
+struct_fields_bq: Dict[str, bigframes.core.sql.literals.STRUCT_VALUES] = {
+"output_schema": output_schema
+ }
+iftemperatureisnotNone:
+struct_fields_bq["temperature"] =temperature
+iftop_pisnotNone:
+struct_fields_bq["top_p"] =top_p
+ifmax_output_tokensisnotNone:
+struct_fields_bq["max_output_tokens"] =max_output_tokens
+ifstop_sequencesisnotNone:
+struct_fields_bq["stop_sequences"] =stop_sequences
+ifrequest_typeisnotNone:
+struct_fields_bq["request_type"] =request_type
+
+struct_sql=bigframes.core.sql.literals.struct_literal(struct_fields_bq)
+query=f"""
+ SELECT *
+ FROM AI.GENERATE_TABLE(
+ MODEL `{model_name}`,
+ ({table_sql}),
+{struct_sql}
+ )
+ """
+
+ifsessionisNone:
+returnbpd.read_gbq_query(query)
+else:
+returnsession.read_gbq_query(query)
+
+
 @log_adapter.method_logger(custom_base_name="bigquery_ai")
 defif_(
 prompt: PROMPT_TYPE,

bigframes/bigquery/ai.py

@@ -24,6 +24,7 @@
 generate_double,
 generate_embedding,
 generate_int,
+generate_table,
 generate_text,
 if_,
 score,
@@ -37,6 +38,7 @@
 "generate_double",
 "generate_embedding",
 "generate_int",
+"generate_table",
 "generate_text",
 "if_",
 "score",

tests/system/large/bigquery/test_ai.py

@@ -94,3 +94,20 @@ def test_generate_text_with_options(text_model):
 
 # It basically asserts that the results are still returned.
 assertlen(result) ==2
+
+
+deftest_generate_table(text_model):
+df=bpd.DataFrame(
+ {"prompt": ["Generate a table of 2 programming languages and their creators."]}
+ )
+
+result=ai.generate_table(
+text_model,
+df,
+output_schema="language STRING, creator STRING",
+ )
+
+assert"language"inresult.columns
+assert"creator"inresult.columns
+# The model may not always return the exact number of rows requested.
+assertlen(result) >0

tests/unit/bigquery/test_ai.py

@@ -220,6 +220,55 @@ def test_generate_text_defaults(mock_dataframe, mock_session):
 assert"STRUCT()"inquery
 
 
+deftest_generate_table_with_dataframe(mock_dataframe, mock_session):
+model_name="project.dataset.model"
+
+bbq.ai.generate_table(
+model_name,
+mock_dataframe,
+output_schema="col1 STRING, col2 INT64",
+ )
+
+mock_session.read_gbq_query.assert_called_once()
+query=mock_session.read_gbq_query.call_args[0][0]
+
+# Normalize whitespace for comparison
+query=" ".join(query.split())
+
+expected_part_1="SELECT * FROM AI.GENERATE_TABLE("
+expected_part_2=f"MODEL `{model_name}`,"
+expected_part_3="(SELECT * FROM my_table),"
+expected_part_4="STRUCT('col1 STRING, col2 INT64' AS output_schema)"
+
+assertexpected_part_1inquery
+assertexpected_part_2inquery
+assertexpected_part_3inquery
+assertexpected_part_4inquery
+
+
+deftest_generate_table_with_options(mock_dataframe, mock_session):
+model_name="project.dataset.model"
+
+bbq.ai.generate_table(
+model_name,
+mock_dataframe,
+output_schema="col1 STRING",
+temperature=0.5,
+max_output_tokens=100,
+ )
+
+mock_session.read_gbq_query.assert_called_once()
+query=mock_session.read_gbq_query.call_args[0][0]
+query=" ".join(query.split())
+
+assertf"MODEL `{model_name}`"inquery
+assert"(SELECT * FROM my_table)"inquery
+assert (
+"STRUCT('col1 STRING' AS output_schema, 0.5 AS temperature, 100 AS max_output_tokens)"
+inquery
+ )
+
+
 @mock.patch("bigframes.pandas.read_pandas")
 deftest_generate_text_with_pandas_dataframe(
 read_pandas_mock, mock_dataframe, mock_session