feat: add bigframes.bigquery.ai.generate_embedding (#2343)

Implement AI.GENERATE_EMBEDDING function in bigframes.bigquery.ai.
---
*PR created automatically by Jules for task
[11924477578091076513](https://jules.google.com/task/11924477578091076513)
started by @tswast*
---------
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: Tim Sweña <swast@google.com>
Co-authored-by: Shenyang Cai <sycai@users.noreply.github.com>

Author: 3 people

bigframes/bigquery/_operations/ai.py

@@ -19,7 +19,7 @@
 from __future__ importannotations
 
 importjson
-fromtypingimportAny, Iterable, List, Literal, Mapping, Tuple, Union
+fromtypingimportAny, Dict, Iterable, List, Literal, Mapping, Optional, Tuple, Union
 
 importpandasaspd
 
@@ -28,6 +28,7 @@
 frombigframesimportseries, session
 frombigframes.coreimportconvert
 frombigframes.core.loggingimportlog_adapter
+importbigframes.core.sql.literals
 frombigframes.mlimportcoreasml_core
 frombigframes.operationsimportai_ops, output_schemas
 
@@ -388,6 +389,113 @@ def generate_double(
 returnseries_list[0]._apply_nary_op(operator, series_list[1:])
 
 
+@log_adapter.method_logger(custom_base_name="bigquery_ai")
+defgenerate_embedding(
+model_name: str,
+data: Union[dataframe.DataFrame, series.Series, pd.DataFrame, pd.Series],
+*,
+output_dimensionality: Optional[int] =None,
+task_type: Optional[str] =None,
+start_second: Optional[float] =None,
+end_second: Optional[float] =None,
+interval_seconds: Optional[float] =None,
+trial_id: Optional[int] =None,
+) ->dataframe.DataFrame:
+"""
+ Creates embeddings that describe an entity—for example, a piece of text or an image.
+
+ **Examples:**
+
+ >>> import bigframes.pandas as bpd
+ >>> import bigframes.bigquery as bbq
+ >>> df = bpd.DataFrame({"content": ["apple", "bear", "pear"]})
+ >>> bbq.ai.generate_embedding(
+ ... "project.dataset.model_name",
+ ... df
+ ... ) # doctest: +SKIP
+
+ Args:
+ model_name (str):
+ The name of a remote model from Vertex AI, such as the
+ multimodalembedding@001 model.
+ 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_dimensionality (int, optional):
+ An INT64 value that specifies the number of dimensions to use when
+ generating embeddings. For example, if you specify 256 AS
+ output_dimensionality, then the embedding output column contains a
+ 256-dimensional embedding for each input value. To find the
+ supported range of output dimensions, read about the available
+ `Google text embedding models <https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#google-models>`_.
+ task_type (str, optional):
+ A STRING literal that specifies the intended downstream application to
+ help the model produce better quality embeddings. For a list of
+ supported task types and how to choose which one to use, see `Choose an
+ embeddings task type <http://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types>`_.
+ start_second (float, optional):
+ The second in the video at which to start the embedding. The default value is 0.
+ end_second (float, optional):
+ The second in the video at which to end the embedding. The default value is 120.
+ interval_seconds (float, optional):
+ The interval to use when creating embeddings. The default value is 16.
+ trial_id (int, optional):
+ An INT64 value that identifies the hyperparameter tuning trial that
+ you want the function to evaluate. The function uses the optimal
+ trial by default. Only specify this argument if you ran
+ hyperparameter tuning when creating the model.
+
+ Returns:
+ bigframes.pandas.DataFrame:
+ A new DataFrame with the generated embeddings. See the `SQL
+ reference for AI.GENERATE_EMBEDDING
+ <https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-generate-embedding#output>`_
+ for details.
+ """
+ifisinstance(data, (pd.DataFrame, pd.Series)):
+data=bpd.read_pandas(data)
+
+ifisinstance(data, series.Series):
+data=data.copy()
+data.name="content"
+data_df=data.to_frame()
+elifisinstance(data, dataframe.DataFrame):
+data_df=data
+else:
+raiseValueError(f"Unsupported data type: {type(data)}")
+
+# We need to get the SQL for the input data to pass as a subquery to the TVF
+source_sql=data_df.sql
+
+struct_fields: Dict[str, bigframes.core.sql.literals.STRUCT_VALUES] = {}
+ifoutput_dimensionalityisnotNone:
+struct_fields["OUTPUT_DIMENSIONALITY"] =output_dimensionality
+iftask_typeisnotNone:
+struct_fields["TASK_TYPE"] =task_type
+ifstart_secondisnotNone:
+struct_fields["START_SECOND"] =start_second
+ifend_secondisnotNone:
+struct_fields["END_SECOND"] =end_second
+ifinterval_secondsisnotNone:
+struct_fields["INTERVAL_SECONDS"] =interval_seconds
+iftrial_idisnotNone:
+struct_fields["TRIAL_ID"] =trial_id
+
+# Construct the TVF query
+query=f"""
+ SELECT *
+ FROM AI.GENERATE_EMBEDDING(
+ MODEL `{model_name}`,
+ ({source_sql}),
+{bigframes.core.sql.literals.struct_literal(struct_fields)})
+ )
+ """
+
+returndata_df._session.read_gbq(query)
+
+
 @log_adapter.method_logger(custom_base_name="bigquery_ai")
 defif_(
 prompt: PROMPT_TYPE,

bigframes/bigquery/ai.py

@@ -22,6 +22,7 @@
 generate,
 generate_bool,
 generate_double,
+generate_embedding,
 generate_int,
 if_,
 score,
@@ -33,6 +34,7 @@
 "generate",
 "generate_bool",
 "generate_double",
+"generate_embedding",
 "generate_int",
 "if_",
 "score",

bigframes/core/sql/literals.py

@@ -0,0 +1,58 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ importannotations
+
+importcollections.abc
+importjson
+fromtypingimportAny, List, Mapping, Union
+
+importbigframes.core.sql
+
+STRUCT_VALUES=Union[
+str, int, float, bool, Mapping[str, str], List[str], Mapping[str, Any]
+]
+STRUCT_TYPE=Mapping[str, STRUCT_VALUES]
+
+
+defstruct_literal(struct_options: STRUCT_TYPE) ->str:
+rendered_options= []
+foroption_name, option_valueinstruct_options.items():
+ifoption_name=="model_params":
+json_str=json.dumps(option_value)
+# Escape single quotes for SQL string literal
+sql_json_str=json_str.replace("'", "''")
+rendered_val=f"JSON'{sql_json_str}'"
+elifisinstance(option_value, collections.abc.Mapping):
+struct_body=", ".join(
+ [
+f"{bigframes.core.sql.simple_literal(v)} AS {k}"
+fork, vinoption_value.items()
+ ]
+ )
+rendered_val=f"STRUCT({struct_body})"
+elifisinstance(option_value, list):
+rendered_val= (
+"["
++", ".join(
+ [bigframes.core.sql.simple_literal(v) forvinoption_value]
+ )
++"]"
+ )
+elifisinstance(option_value, bool):
+rendered_val=str(option_value).lower()
+else:
+rendered_val=bigframes.core.sql.simple_literal(option_value)
+rendered_options.append(f"{rendered_val} AS {option_name}")
+returnf"STRUCT({', '.join(rendered_options)})"

bigframes/core/sql/ml.py

@@ -14,12 +14,11 @@
 
 from __future__ importannotations
 
-importcollections.abc
-importjson
 fromtypingimportAny, Dict, List, Mapping, Optional, Union
 
 importbigframes.core.compile.googlesqlasgooglesql
 importbigframes.core.sql
+importbigframes.core.sql.literals
 
 
 defcreate_model_ddl(
@@ -109,36 +108,7 @@ def _build_struct_sql(
 ) ->str:
 ifnotstruct_options:
 return""
-
-rendered_options= []
-foroption_name, option_valueinstruct_options.items():
-ifoption_name=="model_params":
-json_str=json.dumps(option_value)
-# Escape single quotes for SQL string literal
-sql_json_str=json_str.replace("'", "''")
-rendered_val=f"JSON'{sql_json_str}'"
-elifisinstance(option_value, collections.abc.Mapping):
-struct_body=", ".join(
- [
-f"{bigframes.core.sql.simple_literal(v)} AS {k}"
-fork, vinoption_value.items()
- ]
- )
-rendered_val=f"STRUCT({struct_body})"
-elifisinstance(option_value, list):
-rendered_val= (
-"["
-+", ".join(
- [bigframes.core.sql.simple_literal(v) forvinoption_value]
- )
-+"]"
- )
-elifisinstance(option_value, bool):
-rendered_val=str(option_value).lower()
-else:
-rendered_val=bigframes.core.sql.simple_literal(option_value)
-rendered_options.append(f"{rendered_val} AS {option_name}")
-returnf", STRUCT({', '.join(rendered_options)})"
+returnf", {bigframes.core.sql.literals.struct_literal(struct_options)}"
 
 
 defevaluate(

tests/unit/bigquery/test_ai.py

@@ -0,0 +1,134 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+fromunittestimportmock
+
+importpandasaspd
+importpytest
+
+importbigframes.bigqueryasbbq
+importbigframes.dataframe
+importbigframes.series
+importbigframes.session
+
+
+@pytest.fixture
+defmock_session():
+returnmock.create_autospec(spec=bigframes.session.Session)
+
+
+@pytest.fixture
+defmock_dataframe(mock_session):
+df=mock.create_autospec(spec=bigframes.dataframe.DataFrame)
+df._session=mock_session
+df.sql="SELECT * FROM my_table"
+returndf
+
+
+@pytest.fixture
+defmock_series(mock_session):
+series=mock.create_autospec(spec=bigframes.series.Series)
+series._session=mock_session
+# Mock to_frame to return a mock dataframe
+df=mock.create_autospec(spec=bigframes.dataframe.DataFrame)
+df._session=mock_session
+df.sql="SELECT my_col AS content FROM my_table"
+series.copy.return_value=series
+series.to_frame.return_value=df
+returnseries
+
+
+deftest_generate_embedding_with_dataframe(mock_dataframe, mock_session):
+model_name="project.dataset.model"
+
+bbq.ai.generate_embedding(
+model_name,
+mock_dataframe,
+output_dimensionality=256,
+ )
+
+mock_session.read_gbq.assert_called_once()
+query=mock_session.read_gbq.call_args[0][0]
+
+# Normalize whitespace for comparison
+query=" ".join(query.split())
+
+expected_part_1="SELECT * FROM AI.GENERATE_EMBEDDING("
+expected_part_2=f"MODEL `{model_name}`,"
+expected_part_3="(SELECT * FROM my_table),"
+expected_part_4="STRUCT(256 AS OUTPUT_DIMENSIONALITY)"
+
+assertexpected_part_1inquery
+assertexpected_part_2inquery
+assertexpected_part_3inquery
+assertexpected_part_4inquery
+
+
+deftest_generate_embedding_with_series(mock_series, mock_session):
+model_name="project.dataset.model"
+
+bbq.ai.generate_embedding(
+model_name, mock_series, start_second=0.0, end_second=10.0, interval_seconds=5.0
+ )
+
+mock_session.read_gbq.assert_called_once()
+query=mock_session.read_gbq.call_args[0][0]
+query=" ".join(query.split())
+
+assertf"MODEL `{model_name}`"inquery
+assert"(SELECT my_col AS content FROM my_table)"inquery
+assert (
+"STRUCT(0.0 AS START_SECOND, 10.0 AS END_SECOND, 5.0 AS INTERVAL_SECONDS)"
+inquery
+ )
+
+
+deftest_generate_embedding_defaults(mock_dataframe, mock_session):
+model_name="project.dataset.model"
+
+bbq.ai.generate_embedding(
+model_name,
+mock_dataframe,
+ )
+
+mock_session.read_gbq.assert_called_once()
+query=mock_session.read_gbq.call_args[0][0]
+query=" ".join(query.split())
+
+assertf"MODEL `{model_name}`"inquery
+assert"STRUCT()"inquery
+
+
+@mock.patch("bigframes.pandas.read_pandas")
+deftest_generate_embedding_with_pandas_dataframe(
+read_pandas_mock, mock_dataframe, mock_session
+):
+# This tests that pandas input path works and calls read_pandas
+model_name="project.dataset.model"
+
+# Mock return value of read_pandas to be a BigFrames DataFrame
+read_pandas_mock.return_value=mock_dataframe
+
+pandas_df=pd.DataFrame({"content": ["test"]})
+
+bbq.ai.generate_embedding(
+model_name,
+pandas_df,
+ )
+
+read_pandas_mock.assert_called_once()
+# Check that read_pandas was called with something (the pandas df)
+assertread_pandas_mock.call_args[0][0] ispandas_df
+
+mock_session.read_gbq.assert_called_once()