From 0f26abbad2c4af683cd2b34189a9039b82163e28 Mon Sep 17 00:00:00 2001 From: Missing Author Date: Wed, 29 Jul 2026 15:20:00 +0000 Subject: [PATCH] Project import generated by Copybara. GitOrigin-RevId: 34e790815ec4d1db43b3ebaead70e91760d6f60e --- packages/bigframes/.gitignore | 74 + .../bigframes/_config/compute_options.py | 20 + .../bigframes/_config/experiment_options.py | 43 + packages/bigframes/bigframes/dataframe.py | 19 + .../extensions/core/series_accessor.py | 4 +- ...{series_mixins.py => series_tvf_mixins.py} | 95 +- packages/bigframes/bigframes/operations/ai.py | 847 +++++++++++ .../bigframes/operations/semantics.py | 1175 ++++++++++++++ .../templates/core_series_accessor.py.j2 | 4 +- packages/bigframes/setup.py | 10 +- .../bigframes/testing/constraints-3.10.txt | 14 +- .../bigframes/testing/constraints-3.11.txt | 6 +- .../tests/system/large/operations/__init__.py | 13 + .../tests/system/large/operations/conftest.py | 33 + .../tests/system/large/operations/test_ai.py | 963 ++++++++++++ .../system/large/operations/test_semantics.py | 1348 +++++++++++++++++ .../tests/system/small/operations/test_ai.py | 280 ++++ .../system/small/operations/test_semantics.py | 144 ++ .../unit/_config/test_experiment_options.py | 31 + ...es_mixins.py => test_series_tvf_mixins.py} | 136 -- .../bigframes/tests/unit/test_dataframe.py | 21 + 21 files changed, 5044 insertions(+), 236 deletions(-) create mode 100644 packages/bigframes/.gitignore rename packages/bigframes/bigframes/extensions/core/{series_mixins.py => series_tvf_mixins.py} (59%) create mode 100644 packages/bigframes/bigframes/operations/ai.py create mode 100644 packages/bigframes/bigframes/operations/semantics.py create mode 100644 packages/bigframes/tests/system/large/operations/__init__.py create mode 100644 packages/bigframes/tests/system/large/operations/conftest.py create mode 100644 packages/bigframes/tests/system/large/operations/test_ai.py create mode 100644 packages/bigframes/tests/system/large/operations/test_semantics.py create mode 100644 packages/bigframes/tests/system/small/operations/test_ai.py create mode 100644 packages/bigframes/tests/system/small/operations/test_semantics.py rename packages/bigframes/tests/unit/extensions/core/{test_series_mixins.py => test_series_tvf_mixins.py} (64%) diff --git a/packages/bigframes/.gitignore b/packages/bigframes/.gitignore new file mode 100644 index 000000000000..efcdb4857b95 --- /dev/null +++ b/packages/bigframes/.gitignore @@ -0,0 +1,74 @@ +*.py[cod] +*.sw[op] + +# C extensions +*.so + +# Packages +*.egg +*.egg-info +dist +build +eggs +.eggs +parts +bin +var +sdist +develop-eggs +.installed.cfg +lib +lib64 +__pycache__ + +# Installer logs +pip-log.txt + +# Unit test / coverage reports +.coverage +.nox +.cache +.pytest_cache + + +# Mac +.DS_Store + +# JetBrains +.idea + +# VS Code +.vscode + +# emacs +*~ + +# Built documentation +docs/_build +**/docs/_build +docs.metadata + +# Virtual environment +env/ +venv/ + +# Test logs +coverage.xml +*sponge_log.xml + +# System test environment variables. +**/system_tests/local_test_setup +**/system_tests/data + +# Make sure a generated file isn't accidentally committed. +pylintrc +pylintrc.test + +# Ruff cache +.ruff_cache + +# Bazel (created in packages/gapic-generator) +bazel-bin +bazel-gapic-generator +bazel-out +bazel-testlogs diff --git a/packages/bigframes/bigframes/_config/compute_options.py b/packages/bigframes/bigframes/_config/compute_options.py index 2ef1e5b72132..027566ae075f 100644 --- a/packages/bigframes/bigframes/_config/compute_options.py +++ b/packages/bigframes/bigframes/_config/compute_options.py @@ -168,6 +168,26 @@ class ComputeOptions: int | None: Number of rows, if set. """ + semantic_ops_confirmation_threshold: Optional[int] = 0 + """ + Deprecated. + + .. deprecated:: 1.42.0 + Semantic operators are deprecated. Please use the functions in + :mod:`bigframes.bigquery.ai` instead. + + """ + + semantic_ops_threshold_autofail = False + """ + Deprecated. + + .. deprecated:: 1.42.0 + Semantic operators are deprecated. Please use the functions in + :mod:`bigframes.bigquery.ai` instead. + + """ + def assign_extra_query_labels(self, **kwargs: Any) -> None: """ Assigns additional custom labels for query configuration. The method updates the diff --git a/packages/bigframes/bigframes/_config/experiment_options.py b/packages/bigframes/bigframes/_config/experiment_options.py index 8f70c8952f65..202b47b738c1 100644 --- a/packages/bigframes/bigframes/_config/experiment_options.py +++ b/packages/bigframes/bigframes/_config/experiment_options.py @@ -25,9 +25,52 @@ class ExperimentOptions: """ def __init__(self): + self._semantic_operators: bool = False + self._ai_operators: bool = False self._sql_compiler: Literal["legacy", "stable", "experimental"] = "stable" self._enable_python_transpiler: bool = False + @property + def semantic_operators(self) -> bool: + """Deprecated. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.semantic_operators = True # doctest: +SKIP + """ + return self._semantic_operators + + @semantic_operators.setter + def semantic_operators(self, value: bool): + if value is True: + msg = bfe.format_message( + "Semantic operators are deprecated, and will be removed in the future" + ) + warnings.warn(msg, category=FutureWarning) + self._semantic_operators = value + + @property + def ai_operators(self) -> bool: + """If True, allow using the AI operators. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.ai_operators = True # doctest: +SKIP + """ + return self._ai_operators + + @ai_operators.setter + def ai_operators(self, value: bool): + if value is True: + msg = bfe.format_message( + "AI operators are still under experiments, and are subject " + "to change in the future." + ) + warnings.warn(msg, category=bfe.PreviewWarning) + self._ai_operators = value + @property def sql_compiler(self) -> Literal["legacy", "stable", "experimental"]: """Set to 'experimental' to try out the latest in compilation experiments.. diff --git a/packages/bigframes/bigframes/dataframe.py b/packages/bigframes/bigframes/dataframe.py index 8d05ba4ddb06..07ba6aba1c0e 100644 --- a/packages/bigframes/bigframes/dataframe.py +++ b/packages/bigframes/bigframes/dataframe.py @@ -80,7 +80,9 @@ import bigframes.functions import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops +import bigframes.operations.ai import bigframes.operations.plotting as plotting +import bigframes.operations.semantics import bigframes.operations.structs import bigframes.series import bigframes.session._io.bigquery @@ -5284,3 +5286,20 @@ def _throw_if_null_index(self, opname: str): raise bigframes.exceptions.NullIndexError( f"DataFrame cannot perform {opname} as it has no index. Set an index using set_index." ) + + @property + def semantics(self): + msg = bfe.format_message( + "The 'semantics' property will be removed. Please use 'bigframes.bigquery.ai' instead." + ) + warnings.warn(msg, category=FutureWarning) + return bigframes.operations.semantics.Semantics(self) + + @property + def ai(self): + """Returns the accessor for AI operators.""" + msg = bfe.format_message( + "The 'ai' property will be removed. Please use 'bigframes.bigquery.ai' instead." + ) + warnings.warn(msg, category=FutureWarning) + return bigframes.operations.ai.AIAccessor(self) diff --git a/packages/bigframes/bigframes/extensions/core/series_accessor.py b/packages/bigframes/bigframes/extensions/core/series_accessor.py index d7304ad1f7a2..440e6731aba2 100644 --- a/packages/bigframes/bigframes/extensions/core/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/core/series_accessor.py @@ -31,7 +31,7 @@ from bigframes import series, session from bigframes.core import col, sentinels -from bigframes.extensions.core import abstract_series_accessor, series_mixins +from bigframes.extensions.core import abstract_series_accessor, series_tvf_mixins T = TypeVar("T") S = TypeVar("S") @@ -1112,7 +1112,7 @@ def unix_date( return self._to_series(cast(series.Series, result)) -class AiSeriesAccessor(series_mixins.AIMixin[T, S]): +class AiSeriesAccessor(series_tvf_mixins.AITVFMixin[T, S]): """Series accessor for BigQuery ai functions.""" diff --git a/packages/bigframes/bigframes/extensions/core/series_mixins.py b/packages/bigframes/bigframes/extensions/core/series_tvf_mixins.py similarity index 59% rename from packages/bigframes/bigframes/extensions/core/series_mixins.py rename to packages/bigframes/bigframes/extensions/core/series_tvf_mixins.py index 4d1b61ecb0c5..673978bbe46c 100644 --- a/packages/bigframes/bigframes/extensions/core/series_mixins.py +++ b/packages/bigframes/bigframes/extensions/core/series_tvf_mixins.py @@ -14,13 +14,11 @@ from __future__ import annotations -from typing import Any, List, Literal, Mapping, TypeVar +from typing import List, Mapping, TypeVar import pandas as pd -from bigframes import series -from bigframes import session as bf_session -from bigframes.bigquery import ai +from bigframes import session from bigframes.extensions.core import abstract_series_accessor from bigframes.ml import base as ml_base @@ -28,7 +26,7 @@ S = TypeVar("S") -class AIMixin(abstract_series_accessor.AbstractBigQuerySeriesAccessor[T, S]): +class AITVFMixin(abstract_series_accessor.AbstractBigQuerySeriesAccessor[T, S]): def generate_embedding( self, model: ml_base.BaseEstimator | str | pd.Series, @@ -39,7 +37,7 @@ def generate_embedding( end_second: float | None = None, interval_seconds: float | None = None, trial_id: int | None = None, - session: bf_session.Session | None = None, + session: session.Session | None = None, ) -> T: """ Creates embeddings that describe an entity — for example, a piece of text or an image. @@ -47,9 +45,10 @@ def generate_embedding( This is an accessor for :func:`bigframes.bigquery.ai.generate_embedding`. See that function's documentation for detailed parameter descriptions and examples. """ + import bigframes.bigquery.ai bf_series = self._bf_from_series(session) - result = ai.generate_embedding( + result = bigframes.bigquery.ai.generate_embedding( model, bf_series, output_dimensionality=output_dimensionality, @@ -72,7 +71,7 @@ def generate_text( stop_sequences: List[str] | None = None, ground_with_google_search: bool | None = None, request_type: str | None = None, - session: bf_session.Session | None = None, + session: session.Session | None = None, ) -> T: """ Generates text using a BigQuery ML model. @@ -80,8 +79,10 @@ def generate_text( This is an accessor for :func:`bigframes.bigquery.ai.generate_text`. See that function's documentation for detailed parameter descriptions and examples. """ + import bigframes.bigquery.ai + bf_series = self._bf_from_series(session) - result = ai.generate_text( + result = bigframes.bigquery.ai.generate_text( model, bf_series, temperature=temperature, @@ -104,7 +105,7 @@ def generate_table( max_output_tokens: int | None = None, stop_sequences: List[str] | None = None, request_type: str | None = None, - session: bf_session.Session | None = None, + session: session.Session | None = None, ) -> T: """ Generates a table using a BigQuery ML model. @@ -112,8 +113,10 @@ def generate_table( This is an accessor for :func:`bigframes.bigquery.ai.generate_table`. See that function's documentation for detailed parameter descriptions and examples. """ + import bigframes.bigquery.ai + bf_series = self._bf_from_series(session) - result = ai.generate_table( + result = bigframes.bigquery.ai.generate_table( model, bf_series, output_schema=output_schema, @@ -124,73 +127,3 @@ def generate_table( request_type=request_type, ) return self._to_dataframe(result) - - def embed( - self, - *, - endpoint: str | None = None, - model: str | None = None, - task_type: ( - Literal[ - "retrieval_query", - "retrieval_document", - "semantic_similarity", - "classification", - "clustering", - "question_answering", - "fact_verification", - "code_retrieval_query", - ] - | None - ) = None, - title: str | None = None, - model_params: Mapping[Any, Any] | None = None, - connection_id: str | None = None, - session: bf_session.Session | None = None, - ) -> S: - """ - Creates embeddings from text or image data in BigQuery. - - This is an accessor for :func:`bigframes.bigquery.ai.embed`. See that - function's documentation for detailed parameter descriptions and examples. - """ - - bf_series = self._bf_from_series(session) - result = ai.embed( - bf_series, - endpoint=endpoint, - model=model, - task_type=task_type, - title=title, - model_params=model_params, - connection_id=connection_id, - ) - return self._to_series(result) - - def similarity( - self, - other: str | series.Series | pd.Series, - *, - endpoint: str | None = None, - model: str | None = None, - model_params: Mapping[Any, Any] | None = None, - connection_id: str | None = None, - session: bf_session.Session | None = None, - ) -> S: - """ - Returns a FLOAT64 value that represents the cosine similarity between the two inputs. - - This is an accessor for :func:`bigframes.bigquery.ai.similarity`. See that - function's documentation for detailed parameter descriptions and examples. - """ - - bf_series = self._bf_from_series(session) - result = ai.similarity( - bf_series, - other, - endpoint=endpoint, - model=model, - model_params=model_params, - connection_id=connection_id, - ) - return self._to_series(result) diff --git a/packages/bigframes/bigframes/operations/ai.py b/packages/bigframes/bigframes/operations/ai.py new file mode 100644 index 000000000000..bba0bf5a8362 --- /dev/null +++ b/packages/bigframes/bigframes/operations/ai.py @@ -0,0 +1,847 @@ +# 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. + +from __future__ import annotations + +import re +import typing +import warnings +from typing import Dict, Iterable, List, Optional, Sequence, Union + +from bigframes import dtypes, exceptions, options +from bigframes.core import guid +from bigframes.core.logging import log_adapter + + +@log_adapter.class_logger +class AIAccessor: + def __init__(self, df, base_bqml=None) -> None: + import bigframes # Import in the function body to avoid circular imports. + import bigframes.dataframe + from bigframes.ml import core as ml_core + + self._df: bigframes.dataframe.DataFrame = df + self._base_bqml: ml_core.BaseBqml = base_bqml or ml_core.BaseBqml(df._session) + + def filter( + self, + instruction: str, + model, + ground_with_google_search: bool = False, + ): + """ + Filters the DataFrame with the semantics of the user instruction. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.ai_operators = True + >>> bpd.options.compute.ai_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.5-pro") + + >>> df = bpd.DataFrame({"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]}) + >>> df.ai.filter("{city} is the capital of {country}", model) + country city + 1 Germany Berlin + + [1 rows x 2 columns] + + Args: + instruction (str): + An instruction on how to filter the data. This value must contain + column references by name, which should be wrapped in a pair of braces. + For example, if you have a column "food", you can refer to this column + in the instructions like: + "The {food} is healthy." + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by Bigframes ML package. + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.pandas.DataFrame: DataFrame filtered by the instruction. + + Raises: + NotImplementedError: when the AI operator experiment is off. + ValueError: when the instruction refers to a non-existing column, or when no + columns are referred to. + """ + if not options.experiments.ai_operators: + raise NotImplementedError() + + answer_col = "answer" + + output_schema = {answer_col: "bool"} + result = self.map( + instruction, + model, + output_schema, + ground_with_google_search, + ) + + return result[result[answer_col]].drop(answer_col, axis=1) + + def map( + self, + instruction: str, + model, + output_schema: Dict[str, str] | None = None, + ground_with_google_search: bool = False, + ): + """ + Maps the DataFrame with the semantics of the user instruction. The name of the keys in the output_schema parameter carry + semantic meaning, and can be used for information extraction. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.ai_operators = True + >>> bpd.options.compute.ai_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.5-pro") + + >>> df = bpd.DataFrame({"ingredient_1": ["Burger Bun", "Soy Bean"], "ingredient_2": ["Beef Patty", "Bittern"]}) + >>> df.ai.map("What is the food made from {ingredient_1} and {ingredient_2}? One word only.", model=model, output_schema={"food": "string"}) # doctest: +ELLIPSIS + ingredient_1 ingredient_2... + 0 Burger Bun Beef Patty... + 1 Soy Bean Bittern...Tofu + + [2 rows x 3 columns] + + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.ai_operators = True + >>> bpd.options.compute.ai_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.5-pro") + + >>> df = bpd.DataFrame({"text": ["Elmo lives at 123 Sesame Street."]}) + >>> df.ai.map("{text}", model=model, output_schema={"person": "string", "address": "string"}) + text person address + 0 Elmo lives at 123 Sesame Street. Elmo 123 Sesame Street + + [1 rows x 3 columns] + + Args: + instruction (str): + An instruction on how to map the data. This value must contain + column references by name, which should be wrapped in a pair of braces. + For example, if you have a column "food", you can refer to this column + in the instructions like: + "Get the ingredients of {food}." + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by Bigframes ML package. + + output_schema (Dict[str, str] or None, default None): + The schema used to generate structured output as a bigframes DataFrame. The schema is a string key-value pair of :. + Supported types are int64, float64, bool, string, array and struct. If None, generate string result under the column + "ml_generate_text_llm_result". + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.pandas.DataFrame: DataFrame with attached mapping results. + + Raises: + NotImplementedError: when the AI operator experiment is off. + ValueError: when the instruction refers to a non-existing column, or when no + columns are referred to. + """ + if not options.experiments.ai_operators: + raise NotImplementedError() + + import bigframes.dataframe + import bigframes.series + + self._validate_model(model) + columns = self._parse_columns(instruction) + for column in columns: + if column not in self._df.columns: + raise ValueError(f"Column {column} not found.") + + if ground_with_google_search: + msg = exceptions.format_message( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + warnings.warn(msg, category=UserWarning) + + self._confirm_operation(len(self._df)) + + df: bigframes.dataframe.DataFrame = self._df[columns].copy() + has_blob_column = False + for column in columns: + if df[column].dtype == dtypes.OBJ_REF_DTYPE: + # Don't cast ObjectRef columns to string + has_blob_column = True + continue + + if df[column].dtype != dtypes.STRING_DTYPE: + df[column] = df[column].astype(dtypes.STRING_DTYPE) + + user_instruction = self._format_instruction(instruction, columns) + output_instruction = ( + "Based on the provided contenxt, answer the following instruction:" + ) + + if output_schema is None: + output_schema = {"ml_generate_text_llm_result": "string"} + + if has_blob_column: + results = typing.cast( + bigframes.series.Series, + model.predict( + df, + prompt=self._make_multimodal_prompt( + df, columns, user_instruction, output_instruction + ), + temperature=0.0, + ground_with_google_search=ground_with_google_search, + output_schema=output_schema, + ), + ) + else: + results = typing.cast( + bigframes.series.Series, + model.predict( + self._make_text_prompt( + df, columns, user_instruction, output_instruction + ), + temperature=0.0, + ground_with_google_search=ground_with_google_search, + output_schema=output_schema, + ), + ) + + attach_columns = [results[col] for col, _ in output_schema.items()] + + from bigframes.core.reshape.api import concat + + return concat([self._df, *attach_columns], axis=1) + + def classify( + self, + instruction: str, + model, + labels: Sequence[str], + output_column: str = "result", + ground_with_google_search: bool = False, + ): + """ + Classifies the rows of dataframes based on user instruction into the provided labels. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.ai_operators = True + >>> bpd.options.compute.ai_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.5-pro") + + >>> df = bpd.DataFrame({ + ... "feedback_text": [ + ... "The product is amazing, but the shipping was slow.", + ... "I had an issue with my recent bill.", + ... "The user interface is very intuitive." + ... ], + ... }) + >>> df.ai.classify("{feedback_text}", model=model, labels=["Shipping", "Billing", "UI"]) + feedback_text result + 0 The product is amazing, but the shipping was s... Shipping + 1 I had an issue with my recent bill. Billing + 2 The user interface is very intuitive. UI + + [3 rows x 2 columns] + + Args: + instruction (str): + An instruction on how to classify the data. This value must contain + column references by name, which should be wrapped in a pair of braces. + For example, if you have a column "feedback", you can refer to this column + with"{food}". + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by Bigframes ML package. + + labels (Sequence[str]): + A collection of labels (categories). It must contain at least two and at most 20 elements. + Labels are case sensitive. Duplicated labels are not allowed. + + output_column (str, default "result"): + The name of column for the output. + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.pandas.DataFrame: DataFrame with classification result. + + Raises: + NotImplementedError: when the AI operator experiment is off. + ValueError: when the instruction refers to a non-existing column, when no + columns are referred to, or when the count of labels does not meet the + requirement. + """ + if not options.experiments.ai_operators: + raise NotImplementedError() + + if len(labels) < 2 or len(labels) > 20: + raise ValueError( + f"The number of labels should be between 2 and 20 (inclusive), but {len(labels)} labels are provided." + ) + + if len(set(labels)) != len(labels): + raise ValueError("There are duplicate labels.") + + updated_instruction = f"Based on the user instruction {instruction}, you must provide an answer that must exist in the following list of labels: {labels}" + + return self.map( + updated_instruction, + model, + output_schema={output_column: "string"}, + ground_with_google_search=ground_with_google_search, + ) + + def join( + self, + other, + instruction: str, + model, + ground_with_google_search: bool = False, + ): + """ + Joines two dataframes by applying the instruction over each pair of rows from + the left and right table. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.ai_operators = True + >>> bpd.options.compute.ai_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.5-pro") + + >>> cities = bpd.DataFrame({'city': ['Seattle', 'Ottawa', 'Berlin', 'Shanghai', 'New Delhi']}) + >>> continents = bpd.DataFrame({'continent': ['North America', 'Africa', 'Asia']}) + + >>> cities.ai.join(continents, "{city} is in {continent}", model) + city continent + 0 Seattle North America + 1 Ottawa North America + 2 Shanghai Asia + 3 New Delhi Asia + + [4 rows x 2 columns] + + Args: + other (bigframes.pandas.DataFrame): + The other dataframe. + + instruction (str): + An instruction on how left and right rows can be joined. This value must contain + column references by name. which should be wrapped in a pair of braces. + For example: "The {city} belongs to the {country}". + For column names that are shared between two dataframes, you need to add "left." + and "right." prefix for differentiation. This is especially important when you do + self joins. For example: "The {left.employee_name} reports to {right.employee_name}" + For unique column names, this prefix is optional. + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by Bigframes ML package. + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.pandas.DataFrame: The joined dataframe. + + Raises: + ValueError if the amount of data that will be sent for LLM processing is larger than max_rows. + """ + if not options.experiments.ai_operators: + raise NotImplementedError() + + self._validate_model(model) + columns = self._parse_columns(instruction) + + if ground_with_google_search: + msg = exceptions.format_message( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + warnings.warn(msg, category=UserWarning) + + work_estimate = len(self._df) * len(other) + self._confirm_operation(work_estimate) + + left_columns = [] + right_columns = [] + + for col in columns: + if col in self._df.columns and col in other.columns: + raise ValueError(f"Ambiguous column reference: {col}") + + elif col in self._df.columns: + left_columns.append(col) + + elif col in other.columns: + right_columns.append(col) + + elif col.startswith("left."): + original_col_name = col[len("left.") :] + if ( + original_col_name in self._df.columns + and original_col_name in other.columns + ): + left_columns.append(col) + elif original_col_name in self._df.columns: + left_columns.append(col) + instruction = instruction.replace(col, original_col_name) + else: + raise ValueError(f"Column {col} not found") + + elif col.startswith("right."): + original_col_name = col[len("right.") :] + if ( + original_col_name in self._df.columns + and original_col_name in other.columns + ): + right_columns.append(col) + elif original_col_name in other.columns: + right_columns.append(col) + instruction = instruction.replace(col, original_col_name) + else: + raise ValueError(f"Column {col} not found") + + else: + raise ValueError(f"Column {col} not found") + + if not left_columns: + raise ValueError("No left column references.") + + if not right_columns: + raise ValueError("No right column references.") + + # Update column references to be compatible with internal naming scheme. + # That is, "left.col" -> "col_left" and "right.col" -> "col_right" + instruction = re.sub(r"(?>> import bigframes.pandas as bpd + + >>> import bigframes + >>> bigframes.options.experiments.ai_operators = True + >>> bpd.options.compute.ai_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.TextEmbeddingGenerator(model_name="text-embedding-005") + + >>> df = bpd.DataFrame({"creatures": ["salmon", "sea urchin", "frog", "chimpanzee"]}) + >>> df.ai.search("creatures", "monkey", top_k=1, model=model, score_column='distance') + creatures distance + 3 chimpanzee 0.635844 + + [1 rows x 2 columns] + + Args: + search_column: + The name of the column to search from. + query (str): + The search query. + top_k (int): + The number of nearest neighbors to return. + model (TextEmbeddingGenerator): + A TextEmbeddingGenerator provided by Bigframes ML package. + score_column (Optional[str], default None): + The name of the the additional column containning the similarity scores. If None, + this column won't be attached to the result. + + Returns: + DataFrame: the DataFrame with the search result. + + Raises: + ValueError: when the search_column is not found from the the data frame. + TypeError: when the provided model is not TextEmbeddingGenerator. + """ + if not options.experiments.ai_operators: + raise NotImplementedError() + + if search_column not in self._df.columns: + raise ValueError(f"Column `{search_column}` not found") + + self._confirm_operation(len(self._df)) + + import bigframes.ml.llm as llm + + if not isinstance(model, llm.TextEmbeddingGenerator): + raise TypeError(f"Expect a text embedding model, but got: {type(model)}") + + if top_k < 1: + raise ValueError("top_k must be an integer greater than or equal to 1.") + + embedded_df = model.predict(self._df[search_column]) + embedded_table = embedded_df.reset_index().to_gbq() + + import bigframes.pandas as bpd + + embedding_result_column = "ml_generate_embedding_result" + query_df = model.predict(bpd.DataFrame({"query_id": [query]})).rename( + columns={"content": "query_id", embedding_result_column: "embedding"} + ) + + import bigframes.bigquery as bbq + + search_result = ( + bbq.vector_search( + base_table=embedded_table, + column_to_search=embedding_result_column, + query=query_df, + top_k=top_k, + # TODO(tswast): set allow_large_results based on Series size. + # If we expect small results, it could be faster to set + # allow_large_results to False. + allow_large_results=True, + ) + .rename(columns={"content": search_column}) + .set_index("index") + ) + + search_result.index.name = self._df.index.name + + if score_column is not None: + search_result = search_result.rename(columns={"distance": score_column})[ + [search_column, score_column] + ] + else: + search_result = search_result[[search_column]] + + import bigframes.dataframe + + return typing.cast(bigframes.dataframe.DataFrame, search_result) + + def sim_join( + self, + other, + left_on: str, + right_on: str, + model, + top_k: int = 3, + score_column: Optional[str] = None, + max_rows: int = 1000, + ): + """ + Joins two dataframes based on the similarity of the specified columns. + + This method uses BigQuery's VECTOR_SEARCH function to match rows on the left side with the rows that have + nearest embedding vectors on the right. In the worst case scenario, the complexity is around O(M * N * log K). + Therefore, this is a potentially expensive operation. + + ** Examples: ** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.ai_operators = True + >>> bpd.options.compute.ai_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.TextEmbeddingGenerator(model_name="text-embedding-005") + + >>> df1 = bpd.DataFrame({'animal': ['monkey', 'spider']}) + >>> df2 = bpd.DataFrame({'animal': ['scorpion', 'baboon']}) + + >>> res = df1.ai.sim_join(df2, left_on='animal', right_on='animal', model=model, top_k=1) + >>> print("---"); print(res) # doctest: +ELLIPSIS + --- + ... + animal animal_1 + 0 monkey baboon + 1 spider scorpion + + [2 rows x 2 columns] + + Args: + other (DataFrame): + The other data frame to join with. + left_on (str): + The name of the column on left side for the join. + right_on (str): + The name of the column on the right side for the join. + top_k (int, default 3): + The number of nearest neighbors to return. + model (TextEmbeddingGenerator): + A TextEmbeddingGenerator provided by Bigframes ML package. + score_column (Optional[str], default None): + The name of the the additional column containning the similarity scores. If None, + this column won't be attached to the result. + max_rows: + The maximum number of rows allowed to be processed per call. If the result is too large, the method + call will end early with an error. + + Returns: + DataFrame: the data frame with the join result. + + Raises: + ValueError: when the amount of data to be processed exceeds the specified max_rows. + """ + if not options.experiments.ai_operators: + raise NotImplementedError() + + if left_on not in self._df.columns: + raise ValueError(f"Left column {left_on} not found") + if right_on not in self._df.columns: + raise ValueError(f"Right column {right_on} not found") + + import bigframes.ml.llm as llm + + if not isinstance(model, llm.TextEmbeddingGenerator): + raise TypeError(f"Expect a text embedding model, but got: {type(model)}") + + joined_table_rows = len(self._df) * len(other) + if joined_table_rows > max_rows: + raise ValueError( + f"Number of rows that need processing is {joined_table_rows}, which exceeds row limit {max_rows}." + ) + + if top_k < 1: + raise ValueError("top_k must be an integer greater than or equal to 1.") + + work_estimate = len(self._df) * len(other) + self._confirm_operation(work_estimate) + + base_table_embedding_column = guid.generate_guid() + base_table = self._attach_embedding( + other, right_on, base_table_embedding_column, model + ).to_gbq() + query_table = self._attach_embedding(self._df, left_on, "embedding", model) + + import bigframes.bigquery as bbq + + join_result = bbq.vector_search( + base_table=base_table, + column_to_search=base_table_embedding_column, + query=query_table, + top_k=top_k, + ) + + join_result = join_result.drop( + ["embedding", base_table_embedding_column], axis=1 + ) + + if score_column is not None: + join_result = join_result.rename(columns={"distance": score_column}) + else: + del join_result["distance"] + + return join_result + + def forecast( + self, + timestamp_column: str, + data_column: str, + *, + model: str = "TimesFM 2.0", + id_columns: Optional[Iterable[str]] = None, + horizon: int = 10, + confidence_level: float = 0.95, + ): + """ + Forecast time series at future horizon. Using Google Research's open source TimesFM(https://github.com/google-research/timesfm) model. + + .. note:: + + This product or feature is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the + Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is" + and might have limited support. For more information, see the launch stage descriptions + (https://cloud.google.com/products#product-launch-stages). + + Args: + timestamp_column (str): + A str value that specified the name of the time points column. + 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 + data_column (str): + A str value that specifies the name of the data column. The data column contains the data to forecast. + The data column must use one of the following data types: INT64, NUMERIC and FLOAT64 + 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. + id_columns (Iterable[str] or None, default None): + An iterable of str value that specifies the names of one or more ID columns. Each ID identifies a unique time series to forecast. + Specify one or more values for this argument in order to forecast multiple time series using a single query. + The columns that you specify must use one of the following data types: STRING, INT64, ARRAY and ARRAY + horizon (int, default 10): + An int value that specifies the number of time points to forecast. The default value is 10. The valid input range is [1, 10,000]. + 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). + + Returns: + DataFrame: + The forecast dataframe matches that of the BigQuery AI.FORECAST function. + See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-forecast + + Raises: + ValueError: when referring to a non-existing column. + """ + columns = [timestamp_column, data_column] + if id_columns: + columns += id_columns + for column in columns: + if column not in self._df.columns: + raise ValueError(f"Column `{column}` not found") + + options: dict[str, Union[int, float, str, Iterable[str]]] = { + "data_col": data_column, + "timestamp_col": timestamp_column, + "model": model, + "horizon": horizon, + "confidence_level": confidence_level, + } + if id_columns: + options["id_cols"] = id_columns + + return self._base_bqml.ai_forecast(input_data=self._df, options=options) + + @staticmethod + def _attach_embedding(dataframe, source_column: str, embedding_column: str, model): + result_df = dataframe.copy() + embeddings = model.predict(dataframe[source_column])[ + "ml_generate_embedding_result" + ] + result_df[embedding_column] = embeddings + return result_df + + @staticmethod + def _make_multimodal_prompt( + prompt_df, columns, user_instruction: str, output_instruction: str + ): + prompt = [f"{output_instruction}\n{user_instruction}\nContext: "] + for col in columns: + prompt.extend([f"{col} is ", prompt_df[col]]) + + return prompt + + @staticmethod + def _make_text_prompt( + prompt_df, columns, user_instruction: str, output_instruction: str + ): + prompt_df["prompt"] = f"{output_instruction}\n{user_instruction}\nContext: " + + # Combine context from multiple columns. + for col in columns: + prompt_df["prompt"] += f"{col} is `" + prompt_df[col] + "`\n" + + return prompt_df["prompt"] + + @staticmethod + def _parse_columns(instruction: str) -> List[str]: + """Extracts column names enclosed in curly braces from the user instruction. + For example, _parse_columns("{city} is in {continent}") == ["city", "continent"] + """ + columns = re.findall(r"(? str: + """Extracts column names enclosed in curly braces from the user instruction. + For example, `_format_instruction(["city", "continent"], "{city} is in {continent}") + == "city is in continent"` + """ + return instruction.format(**{col: col for col in columns}) + + @staticmethod + def _validate_model(model): + from bigframes.ml.llm import GeminiTextGenerator + + if not isinstance(model, GeminiTextGenerator): + raise TypeError("Model is not GeminiText Generator") + + @staticmethod + def _confirm_operation(row_count: int): + """Raises OperationAbortedError when the confirmation fails""" + import bigframes # Import in the function body to avoid circular imports. + + threshold = bigframes.options.compute.ai_ops_confirmation_threshold + + if threshold is None or row_count <= threshold: + return + + if bigframes.options.compute.ai_ops_threshold_autofail: + raise exceptions.OperationAbortedError( + f"Operation was cancelled because your work estimate is {row_count} rows, which exceeds the threshold {threshold} rows." + ) + + # Separate the prompt out. In IDE such VS Code, leaving prompt in the + # input function makes it less visible to the end user. + print(f"This operation will process about {row_count} rows.") + print( + "You can raise the confirmation threshold by setting `bigframes.options.compute.ai_ops_confirmation_threshold` to a higher value. To completely turn off the confirmation check, set the threshold to `None`." + ) + print("Proceed? [Y/n]") + reply = input().casefold() + if reply not in {"y", "yes", ""}: + raise exceptions.OperationAbortedError("Operation was cancelled.") diff --git a/packages/bigframes/bigframes/operations/semantics.py b/packages/bigframes/bigframes/operations/semantics.py new file mode 100644 index 000000000000..3cad3258e035 --- /dev/null +++ b/packages/bigframes/bigframes/operations/semantics.py @@ -0,0 +1,1175 @@ +# Copyright 2024 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. + + +import re +import typing +import warnings +from typing import List, Optional, cast + +import numpy as np + +from bigframes import dtypes, exceptions +from bigframes.core import guid +from bigframes.core.logging import log_adapter + + +@log_adapter.class_logger +class Semantics: + def __init__(self, df) -> None: + import bigframes # Import in the function body to avoid circular imports. + import bigframes.dataframe + + if not bigframes.options.experiments.semantic_operators: + raise NotImplementedError() + + self._df: bigframes.dataframe.DataFrame = df + + def agg( + self, + instruction: str, + model, + cluster_column: typing.Optional[str] = None, + max_agg_rows: int = 10, + ground_with_google_search: bool = False, + ): + """ + Performs an aggregation over all rows of the table. + + This method recursively aggregates the input data to produce partial answers + in parallel, until a single answer remains. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.semantic_operators = True + >>> bpd.options.compute.semantic_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP + + >>> df = bpd.DataFrame( + ... { + ... "Movies": [ + ... "Titanic", + ... "The Wolf of Wall Street", + ... "Inception", + ... ], + ... "Year": [1997, 2013, 2010], + ... }) + >>> df.semantics.agg( # doctest: +SKIP + ... "Find the first name shared by all actors in {Movies}. One word answer.", + ... model=model, + ... ) + 0 Leonardo + + Name: Movies, dtype: string + + Args: + instruction (str): + An instruction on how to map the data. This value must contain + column references by name enclosed in braces. + For example, to reference a column named "movies", use "{movies}" in the + instruction, like: "Find actor names shared by all {movies}." + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by the Bigframes ML package. + + cluster_column (Optional[str], default None): + If set, aggregates each cluster before performing aggregations across + clusters. Clustering based on semantic similarity can improve accuracy + of the sementic aggregations. + + max_agg_rows (int, default 10): + The maxinum number of rows to be aggregated at a time. + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.dataframe.DataFrame: A new DataFrame with the aggregated answers. + + Raises: + NotImplementedError: when the semantic operator experiment is off. + ValueError: when the instruction refers to a non-existing column, or when + more than one columns are referred to. + """ + import bigframes.bigquery as bbq + import bigframes.dataframe + import bigframes.series + + self._validate_model(model) + columns = self._parse_columns(instruction) + + if max_agg_rows <= 1: + raise ValueError( + f"Invalid value for `max_agg_rows`: {max_agg_rows}." + "It must be greater than 1." + ) + + work_estimate = len(self._df) * int(max_agg_rows / (max_agg_rows - 1)) + self._confirm_operation(work_estimate) + + df: bigframes.dataframe.DataFrame = self._df.copy() + for column in columns: + if column not in self._df.columns: + raise ValueError(f"Column {column} not found.") + + if df[column].dtype != dtypes.STRING_DTYPE: + df[column] = df[column].astype(dtypes.STRING_DTYPE) + + if len(columns) > 1: + raise NotImplementedError( + "Semantic aggregations are limited to a single column." + ) + column = columns[0] + + if ground_with_google_search: + msg = exceptions.format_message( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + warnings.warn(msg, category=UserWarning) + + user_instruction = self._format_instruction(instruction, columns) + + num_cluster = 1 + if cluster_column is not None: + if cluster_column not in df.columns: + raise ValueError(f"Cluster column `{cluster_column}` not found.") + + if df[cluster_column].dtype != dtypes.INT_DTYPE: + raise TypeError( + "Cluster column must be an integer type, not " + f"{type(df[cluster_column])}" + ) + + num_cluster = df[cluster_column].unique().shape[0] + df = df.sort_values(cluster_column) + else: + cluster_column = guid.generate_guid("pid") + df[cluster_column] = 0 + + aggregation_group_id = guid.generate_guid("agg") + group_row_index = guid.generate_guid("gid") + llm_prompt = guid.generate_guid("prompt") + df = ( + df.reset_index(drop=True) + .reset_index() + .rename(columns={"index": aggregation_group_id}) + ) + + output_instruction = ( + "Answer user instructions using the provided context from various sources. " + "Combine all relevant information into a single, concise, well-structured response. " + f"Instruction: {user_instruction}.\n\n" + ) + + while len(df) > 1: + df[group_row_index] = (df[aggregation_group_id] % max_agg_rows + 1).astype( + dtypes.STRING_DTYPE + ) + df[aggregation_group_id] = (df[aggregation_group_id] / max_agg_rows).astype( + dtypes.INT_DTYPE + ) + df[llm_prompt] = "\t\nSource #" + df[group_row_index] + ": " + df[column] + + if len(df) > num_cluster: + # Aggregate within each partition + agg_df = bbq.array_agg( + df.groupby(by=[cluster_column, aggregation_group_id]) + ) + else: + # Aggregate cross partitions + agg_df = bbq.array_agg(df.groupby(by=[aggregation_group_id])) + agg_df[cluster_column] = agg_df[cluster_column].list[0] + + # Skip if the aggregated group only has a single item + single_row_df: bigframes.series.Series = cast( + bigframes.series.Series, + bbq.array_to_string( + agg_df[agg_df[group_row_index].list.len() <= 1][column], + delimiter="", + ), + ) + prompt_s: bigframes.series.Series = cast( + bigframes.series.Series, + bbq.array_to_string( + agg_df[agg_df[group_row_index].list.len() > 1][llm_prompt], + delimiter="", + ), + ) + prompt_s = output_instruction + prompt_s # type:ignore + + # Run model + predict_df = typing.cast( + bigframes.dataframe.DataFrame, + model.predict( + prompt_s, + temperature=0.0, + ground_with_google_search=ground_with_google_search, + ), + ) + agg_df[column] = predict_df["ml_generate_text_llm_result"].combine_first( + single_row_df + ) + + agg_df = agg_df.reset_index() + df = agg_df[[aggregation_group_id, cluster_column, column]] + + return df[column] + + def cluster_by( + self, + column: str, + output_column: str, + model, + n_clusters: int = 5, + ): + """ + Clusters data based on the semantic similarity of text within a specified column. + + This method leverages a language model to generate text embeddings for each value in + the given column. These embeddings capture the semantic meaning of the text. + The data is then grouped into `n` clusters using the k-means clustering algorithm, + which groups data points based on the similarity of their embeddings. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.semantic_operators = True + >>> bpd.options.compute.semantic_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.TextEmbeddingGenerator(model_name="text-embedding-005") + + >>> df = bpd.DataFrame({ + ... "Product": ["Smartphone", "Laptop", "T-shirt", "Jeans"], + ... }) + >>> df.semantics.cluster_by("Product", "Cluster ID", model, n_clusters=2) # doctest: +SKIP + Product Cluster ID + 0 Smartphone 2 + 1 Laptop 2 + 2 T-shirt 1 + 3 Jeans 1 + + [4 rows x 2 columns] + + Args: + column (str): + An column name to perform the similarity clustering. + + output_column (str): + An output column to store the clustering ID. + + model (bigframes.ml.llm.TextEmbeddingGenerator): + A TextEmbeddingGenerator provided by Bigframes ML package. + + n_clusters (int, default 5): + Default 5. Number of clusters to be detected. + + Returns: + bigframes.dataframe.DataFrame: A new DataFrame with the clustering output column. + + Raises: + NotImplementedError: when the semantic operator experiment is off. + ValueError: when the column refers to a non-existing column. + """ + + import bigframes.dataframe + import bigframes.ml.cluster as cluster + import bigframes.ml.llm as llm + + if not isinstance(model, llm.TextEmbeddingGenerator): + raise TypeError(f"Expect a text embedding model, but got: {type(model)}") + + if column not in self._df.columns: + raise ValueError(f"Column {column} not found.") + + if n_clusters <= 1: + raise ValueError( + f"Invalid value for `n_clusters`: {n_clusters}." + "It must be greater than 1." + ) + + self._confirm_operation(len(self._df)) + + df: bigframes.dataframe.DataFrame = self._df.copy() + embeddings_df = model.predict(df[column]) + + cluster_model = cluster.KMeans(n_clusters=n_clusters) + cluster_model.fit(embeddings_df[["ml_generate_embedding_result"]]) + clustered_result = cluster_model.predict(embeddings_df) + df[output_column] = clustered_result["CENTROID_ID"] + return df + + def filter(self, instruction: str, model, ground_with_google_search: bool = False): + """ + Filters the DataFrame with the semantics of the user instruction. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.semantic_operators = True + >>> bpd.options.compute.semantic_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP + + >>> df = bpd.DataFrame({"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]}) + >>> df.semantics.filter("{city} is the capital of {country}", model) # doctest: +SKIP + country city + 1 Germany Berlin + + [1 rows x 2 columns] + + Args: + instruction (str): + An instruction on how to filter the data. This value must contain + column references by name, which should be wrapped in a pair of braces. + For example, if you have a column "food", you can refer to this column + in the instructions like: + "The {food} is healthy." + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by Bigframes ML package. + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.pandas.DataFrame: DataFrame filtered by the instruction. + + Raises: + NotImplementedError: when the semantic operator experiment is off. + ValueError: when the instruction refers to a non-existing column, or when no + columns are referred to. + """ + import bigframes.dataframe + import bigframes.series + + self._validate_model(model) + columns = self._parse_columns(instruction) + for column in columns: + if column not in self._df.columns: + raise ValueError(f"Column {column} not found.") + + if ground_with_google_search: + msg = exceptions.format_message( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + warnings.warn(msg, category=UserWarning) + + self._confirm_operation(len(self._df)) + + df: bigframes.dataframe.DataFrame = self._df[columns].copy() + has_blob_column = False + for column in columns: + if df[column].dtype == dtypes.OBJ_REF_DTYPE: + # Don't cast ObjectRef columns to string + has_blob_column = True + continue + + if df[column].dtype != dtypes.STRING_DTYPE: + df[column] = df[column].astype(dtypes.STRING_DTYPE) + + user_instruction = self._format_instruction(instruction, columns) + output_instruction = "Based on the provided context, reply to the following claim by only True or False:" + + if has_blob_column: + results = typing.cast( + bigframes.dataframe.DataFrame, + model.predict( + df, + prompt=self._make_multimodal_prompt( + df, columns, user_instruction, output_instruction + ), + temperature=0.0, + ground_with_google_search=ground_with_google_search, + ), + ) + else: + results = typing.cast( + bigframes.dataframe.DataFrame, + model.predict( + self._make_text_prompt( + df, columns, user_instruction, output_instruction + ), + temperature=0.0, + ground_with_google_search=ground_with_google_search, + ), + ) + + return self._df[ + results["ml_generate_text_llm_result"].str.lower().str.contains("true") + ] + + def map( + self, + instruction: str, + output_column: str, + model, + ground_with_google_search: bool = False, + ): + """ + Maps the DataFrame with the semantics of the user instruction. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.semantic_operators = True + >>> bpd.options.compute.semantic_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP + + >>> df = bpd.DataFrame({"ingredient_1": ["Burger Bun", "Soy Bean"], "ingredient_2": ["Beef Patty", "Bittern"]}) + >>> df.semantics.map("What is the food made from {ingredient_1} and {ingredient_2}? One word only.", output_column="food", model=model) # doctest: +SKIP + ingredient_1 ingredient_2 food + 0 Burger Bun Beef Patty Burger + + 1 Soy Bean Bittern Tofu + + + [2 rows x 3 columns] + + Args: + instruction (str): + An instruction on how to map the data. This value must contain + column references by name, which should be wrapped in a pair of braces. + For example, if you have a column "food", you can refer to this column + in the instructions like: + "Get the ingredients of {food}." + + output_column (str): + The column name of the mapping result. + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by Bigframes ML package. + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.pandas.DataFrame: DataFrame with attached mapping results. + + Raises: + NotImplementedError: when the semantic operator experiment is off. + ValueError: when the instruction refers to a non-existing column, or when no + columns are referred to. + """ + import bigframes.dataframe + import bigframes.series + + self._validate_model(model) + columns = self._parse_columns(instruction) + for column in columns: + if column not in self._df.columns: + raise ValueError(f"Column {column} not found.") + + if ground_with_google_search: + msg = exceptions.format_message( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + warnings.warn(msg, category=UserWarning) + + self._confirm_operation(len(self._df)) + + df: bigframes.dataframe.DataFrame = self._df[columns].copy() + has_blob_column = False + for column in columns: + if df[column].dtype == dtypes.OBJ_REF_DTYPE: + # Don't cast ObjectRef columns to string + has_blob_column = True + continue + + if df[column].dtype != dtypes.STRING_DTYPE: + df[column] = df[column].astype(dtypes.STRING_DTYPE) + + user_instruction = self._format_instruction(instruction, columns) + output_instruction = ( + "Based on the provided contenxt, answer the following instruction:" + ) + + if has_blob_column: + results = typing.cast( + bigframes.series.Series, + model.predict( + df, + prompt=self._make_multimodal_prompt( + df, columns, user_instruction, output_instruction + ), + temperature=0.0, + ground_with_google_search=ground_with_google_search, + )["ml_generate_text_llm_result"], + ) + else: + results = typing.cast( + bigframes.series.Series, + model.predict( + self._make_text_prompt( + df, columns, user_instruction, output_instruction + ), + temperature=0.0, + ground_with_google_search=ground_with_google_search, + )["ml_generate_text_llm_result"], + ) + + from bigframes.core.reshape.api import concat + + return concat([self._df, results.rename(output_column)], axis=1) + + def join( + self, + other, + instruction: str, + model, + ground_with_google_search: bool = False, + ): + """ + Joines two dataframes by applying the instruction over each pair of rows from + the left and right table. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.semantic_operators = True + >>> bpd.options.compute.semantic_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP + + >>> cities = bpd.DataFrame({'city': ['Seattle', 'Ottawa', 'Berlin', 'Shanghai', 'New Delhi']}) + >>> continents = bpd.DataFrame({'continent': ['North America', 'Africa', 'Asia']}) + + >>> cities.semantics.join(continents, "{city} is in {continent}", model) # doctest: +SKIP + city continent + 0 Seattle North America + 1 Ottawa North America + 2 Shanghai Asia + 3 New Delhi Asia + + [4 rows x 2 columns] + + Args: + other (bigframes.pandas.DataFrame): + The other dataframe. + + instruction (str): + An instruction on how left and right rows can be joined. This value must contain + column references by name. which should be wrapped in a pair of braces. + For example: "The {city} belongs to the {country}". + For column names that are shared between two dataframes, you need to add "left." + and "right." prefix for differentiation. This is especially important when you do + self joins. For example: "The {left.employee_name} reports to {right.employee_name}" + For unique column names, this prefix is optional. + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by Bigframes ML package. + + max_rows (int, default 1000): + The maximum number of rows allowed to be sent to the model per call. If the result is too large, the method + call will end early with an error. + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.pandas.DataFrame: The joined dataframe. + + Raises: + ValueError if the amount of data that will be sent for LLM processing is larger than max_rows. + """ + self._validate_model(model) + columns = self._parse_columns(instruction) + + if ground_with_google_search: + msg = exceptions.format_message( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + warnings.warn(msg, category=UserWarning) + + work_estimate = len(self._df) * len(other) + self._confirm_operation(work_estimate) + + left_columns = [] + right_columns = [] + + for col in columns: + if col in self._df.columns and col in other.columns: + raise ValueError(f"Ambiguous column reference: {col}") + + elif col in self._df.columns: + left_columns.append(col) + + elif col in other.columns: + right_columns.append(col) + + elif col.startswith("left."): + original_col_name = col[len("left.") :] + if ( + original_col_name in self._df.columns + and original_col_name in other.columns + ): + left_columns.append(col) + elif original_col_name in self._df.columns: + left_columns.append(col) + instruction = instruction.replace(col, original_col_name) + else: + raise ValueError(f"Column {col} not found") + + elif col.startswith("right."): + original_col_name = col[len("right.") :] + if ( + original_col_name in self._df.columns + and original_col_name in other.columns + ): + right_columns.append(col) + elif original_col_name in other.columns: + right_columns.append(col) + instruction = instruction.replace(col, original_col_name) + else: + raise ValueError(f"Column {col} not found") + + else: + raise ValueError(f"Column {col} not found") + + if not left_columns: + raise ValueError("No left column references.") + + if not right_columns: + raise ValueError("No right column references.") + + # Update column references to be compatible with internal naming scheme. + # That is, "left.col" -> "col_left" and "right.col" -> "col_right" + instruction = re.sub(r"(?>> import bigframes.pandas as bpd + + >>> import bigframes + >>> bigframes.options.experiments.semantic_operators = True + >>> bpd.options.compute.semantic_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.TextEmbeddingGenerator(model_name="text-embedding-005") # doctest: +SKIP + + >>> df = bpd.DataFrame({"creatures": ["salmon", "sea urchin", "frog", "chimpanzee"]}) + >>> df.semantics.search("creatures", "monkey", top_k=1, model=model, score_column='distance') # doctest: +SKIP + creatures distance + 3 chimpanzee 0.635844 + + [1 rows x 2 columns] + + Args: + search_column: + The name of the column to search from. + query (str): + The search query. + top_k (int): + The number of nearest neighbors to return. + model (TextEmbeddingGenerator): + A TextEmbeddingGenerator provided by Bigframes ML package. + score_column (Optional[str], default None): + The name of the the additional column containning the similarity scores. If None, + this column won't be attached to the result. + + Returns: + DataFrame: the DataFrame with the search result. + + Raises: + ValueError: when the search_column is not found from the the data frame. + TypeError: when the provided model is not TextEmbeddingGenerator. + """ + + if search_column not in self._df.columns: + raise ValueError(f"Column `{search_column}` not found") + + self._confirm_operation(len(self._df)) + + import bigframes.ml.llm as llm + + if not isinstance(model, llm.TextEmbeddingGenerator): + raise TypeError(f"Expect a text embedding model, but got: {type(model)}") + + if top_k < 1: + raise ValueError("top_k must be an integer greater than or equal to 1.") + + embedded_df = model.predict(self._df[search_column]) + embedded_table = embedded_df.reset_index().to_gbq() + + import bigframes.pandas as bpd + + embedding_result_column = "ml_generate_embedding_result" + query_df = model.predict(bpd.DataFrame({"query_id": [query]})).rename( + columns={"content": "query_id", embedding_result_column: "embedding"} + ) + + import bigframes.bigquery as bbq + + search_result = ( + bbq.vector_search( + base_table=embedded_table, + column_to_search=embedding_result_column, + query=query_df, + top_k=top_k, + ) + .rename(columns={"content": search_column}) + .set_index("index") + ) + + search_result.index.name = self._df.index.name + + if score_column is not None: + search_result = search_result.rename(columns={"distance": score_column})[ + [search_column, score_column] + ] + else: + search_result = search_result[[search_column]] + + import bigframes.dataframe + + return typing.cast(bigframes.dataframe.DataFrame, search_result) + + def top_k( + self, + instruction: str, + model, + k: int = 10, + ground_with_google_search: bool = False, + ): + """ + Ranks each tuple and returns the k best according to the instruction. + + This method employs a quick select algorithm to efficiently compare the pivot + with all other items. By leveraging an LLM (Large Language Model), it then + identifies the top 'k' best answers from these comparisons. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.semantic_operators = True + >>> bpd.options.compute.semantic_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP + + >>> df = bpd.DataFrame( + ... { + ... "Animals": ["Dog", "Bird", "Cat", "Horse"], + ... "Sounds": ["Woof", "Chirp", "Meow", "Neigh"], + ... }) + >>> df.semantics.top_k("{Animals} are more popular as pets", model=model, k=2) # doctest: +SKIP + Animals Sounds + 0 Dog Woof + 2 Cat Meow + + [2 rows x 2 columns] + + Args: + instruction (str): + An instruction on how to map the data. This value must contain + column references by name enclosed in braces. + For example, to reference a column named "Animals", use "{Animals}" in the + instruction, like: "{Animals} are more popular as pets" + + model (bigframes.ml.llm.GeminiTextGenerator): + A GeminiTextGenerator provided by the Bigframes ML package. + + k (int, default 10): + The number of rows to return. + + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + + Returns: + bigframes.dataframe.DataFrame: A new DataFrame with the top k rows. + + Raises: + NotImplementedError: when the semantic operator experiment is off. + ValueError: when the instruction refers to a non-existing column, or when no + columns are referred to. + """ + import bigframes.dataframe + import bigframes.series + + self._validate_model(model) + columns = self._parse_columns(instruction) + for column in columns: + if column not in self._df.columns: + raise ValueError(f"Column {column} not found.") + if len(columns) > 1: + raise NotImplementedError("Semantic top K are limited to a single column.") + + if ground_with_google_search: + msg = exceptions.format_message( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + warnings.warn(msg, category=UserWarning) + + work_estimate = int(len(self._df) * (len(self._df) - 1) / 2) + self._confirm_operation(work_estimate) + + df: bigframes.dataframe.DataFrame = self._df[columns].copy() + column = columns[0] + if df[column].dtype != dtypes.STRING_DTYPE: + df[column] = df[column].astype(dtypes.STRING_DTYPE) + + # `index` is reserved for the `reset_index` below. + if column == "index": + raise ValueError( + "Column name 'index' is reserved. Please choose a different name." + ) + + if k < 1: + raise ValueError("k must be an integer greater than or equal to 1.") + + user_instruction = self._format_instruction(instruction, columns) + + n = df.shape[0] + if k >= n: + return df + + # Create a unique index and duplicate it as the "index" column. This workaround + # is needed for the select search algorithm due to unimplemented bigFrame methods. + df = df.reset_index().rename(columns={"index": "old_index"}).reset_index() + + # Initialize a status column to track the selection status of each item. + # - None: Unknown/not yet processed + # - 1.0: Selected as part of the top-k items + # - -1.0: Excluded from the top-k items + status_column = guid.generate_guid("status") + df[status_column] = bigframes.series.Series( + None, dtype=dtypes.FLOAT_DTYPE, session=df._session + ) + + num_selected = 0 + while num_selected < k: + df, num_new_selected = self._topk_partition( + df, + column, + status_column, + user_instruction, + model, + k - num_selected, + ground_with_google_search, + ) + num_selected += num_new_selected + + result_df: bigframes.dataframe.DataFrame = self._df.copy() + return result_df[df.set_index("old_index")[status_column] > 0.0] + + @staticmethod + def _topk_partition( + df, + column: str, + status_column: str, + user_instruction: str, + model, + k: int, + ground_with_google_search: bool, + ): + output_instruction = ( + "Given a question and two documents, choose the document that best answers " + "the question. Respond with 'Document 1' or 'Document 2'. You must choose " + "one, even if neither is ideal. " + ) + + # Random pivot selection for improved average quickselect performance. + pending_df = df[df[status_column].isna()] + pivot_iloc = np.random.randint(0, pending_df.shape[0]) + pivot_index = pending_df.iloc[pivot_iloc]["index"] + pivot_df = pending_df[pending_df["index"] == pivot_index] + + # Build a prompt to compare the pivot item's relevance to other pending items. + prompt_s = pending_df[pending_df["index"] != pivot_index][column] + prompt_s = ( + f"{output_instruction}\n\nQuestion: {user_instruction}\n" + + f"\nDocument 1: {column} " + + pivot_df.iloc[0][column] + + f"\nDocument 2: {column} " + + prompt_s # type:ignore + ) + + import bigframes.dataframe + + predict_df = typing.cast( + bigframes.dataframe.DataFrame, + model.predict( + prompt_s, + temperature=0.0, + ground_with_google_search=ground_with_google_search, + ), + ) + + marks = predict_df["ml_generate_text_llm_result"].str.contains("2") + more_relavant: bigframes.dataframe.DataFrame = df[marks] + less_relavent: bigframes.dataframe.DataFrame = df[~marks] + + num_more_relavant = more_relavant.shape[0] + if k < num_more_relavant: + less_relavent[status_column] = -1.0 + pivot_df[status_column] = -1.0 + df = df.combine_first(less_relavent).combine_first(pivot_df) + return df, 0 + else: # k >= num_more_relavant + more_relavant[status_column] = 1.0 + df = df.combine_first(more_relavant) + if k >= num_more_relavant + 1: + pivot_df[status_column] = 1.0 + df = df.combine_first(pivot_df) + return df, num_more_relavant + 1 + else: + return df, num_more_relavant + + def sim_join( + self, + other, + left_on: str, + right_on: str, + model, + top_k: int = 3, + score_column: Optional[str] = None, + max_rows: int = 1000, + ): + """ + Joins two dataframes based on the similarity of the specified columns. + + This method uses BigQuery's VECTOR_SEARCH function to match rows on the left side with the rows that have + nearest embedding vectors on the right. In the worst case scenario, the complexity is around O(M * N * log K). + Therefore, this is a potentially expensive operation. + + ** Examples: ** + + >>> import bigframes.pandas as bpd + >>> bpd.options.experiments.semantic_operators = True + >>> bpd.options.compute.semantic_ops_confirmation_threshold = 25 + + >>> import bigframes.ml.llm as llm + >>> model = llm.TextEmbeddingGenerator(model_name="text-embedding-005") # doctest: +SKIP + + >>> df1 = bpd.DataFrame({'animal': ['monkey', 'spider']}) + >>> df2 = bpd.DataFrame({'animal': ['scorpion', 'baboon']}) + + >>> df1.semantics.sim_join(df2, left_on='animal', right_on='animal', model=model, top_k=1) # doctest: +SKIP + animal animal_1 + 0 monkey baboon + 1 spider scorpion + + [2 rows x 2 columns] + + Args: + other (DataFrame): + The other data frame to join with. + left_on (str): + The name of the column on left side for the join. + right_on (str): + The name of the column on the right side for the join. + top_k (int, default 3): + The number of nearest neighbors to return. + model (TextEmbeddingGenerator): + A TextEmbeddingGenerator provided by Bigframes ML package. + score_column (Optional[str], default None): + The name of the the additional column containning the similarity scores. If None, + this column won't be attached to the result. + max_rows: + The maximum number of rows allowed to be processed per call. If the result is too large, the method + call will end early with an error. + + Returns: + DataFrame: the data frame with the join result. + + Raises: + ValueError: when the amount of data to be processed exceeds the specified max_rows. + """ + + if left_on not in self._df.columns: + raise ValueError(f"Left column {left_on} not found") + if right_on not in self._df.columns: + raise ValueError(f"Right column {right_on} not found") + + import bigframes.ml.llm as llm + + if not isinstance(model, llm.TextEmbeddingGenerator): + raise TypeError(f"Expect a text embedding model, but got: {type(model)}") + + joined_table_rows = len(self._df) * len(other) + if joined_table_rows > max_rows: + raise ValueError( + f"Number of rows that need processing is {joined_table_rows}, which exceeds row limit {max_rows}." + ) + + if top_k < 1: + raise ValueError("top_k must be an integer greater than or equal to 1.") + + work_estimate = len(self._df) * len(other) + self._confirm_operation(work_estimate) + + base_table_embedding_column = guid.generate_guid() + base_table = self._attach_embedding( + other, right_on, base_table_embedding_column, model + ).to_gbq() + query_table = self._attach_embedding(self._df, left_on, "embedding", model) + + import bigframes.bigquery as bbq + + join_result = bbq.vector_search( + base_table=base_table, + column_to_search=base_table_embedding_column, + query=query_table, + top_k=top_k, + ) + + join_result = join_result.drop( + ["embedding", base_table_embedding_column], axis=1 + ) + + if score_column is not None: + join_result = join_result.rename(columns={"distance": score_column}) + else: + del join_result["distance"] + + return join_result + + @staticmethod + def _attach_embedding(dataframe, source_column: str, embedding_column: str, model): + result_df = dataframe.copy() + embeddings = model.predict(dataframe[source_column])[ + "ml_generate_embedding_result" + ] + result_df[embedding_column] = embeddings + return result_df + + @staticmethod + def _make_multimodal_prompt( + prompt_df, columns, user_instruction: str, output_instruction: str + ): + prompt = [f"{output_instruction}\n{user_instruction}\nContext: "] + for col in columns: + prompt.extend([f"{col} is ", prompt_df[col]]) + + return prompt + + @staticmethod + def _make_text_prompt( + prompt_df, columns, user_instruction: str, output_instruction: str + ): + prompt_df["prompt"] = f"{output_instruction}\n{user_instruction}\nContext: " + + # Combine context from multiple columns. + for col in columns: + prompt_df["prompt"] += f"{col} is `" + prompt_df[col] + "`\n" + + return prompt_df["prompt"] + + @staticmethod + def _parse_columns(instruction: str) -> List[str]: + """Extracts column names enclosed in curly braces from the user instruction. + For example, _parse_columns("{city} is in {continent}") == ["city", "continent"] + """ + columns = re.findall(r"(? str: + """Extracts column names enclosed in curly braces from the user instruction. + For example, `_format_instruction(["city", "continent"], "{city} is in {continent}") + == "city is in continent"` + """ + return instruction.format(**{col: col for col in columns}) + + @staticmethod + def _validate_model(model): + from bigframes.ml.llm import GeminiTextGenerator + + if not isinstance(model, GeminiTextGenerator): + raise TypeError("Model is not GeminiText Generator") + + @staticmethod + def _confirm_operation(row_count: int): + """Raises OperationAbortedError when the confirmation fails""" + import bigframes # Import in the function body to avoid circular imports. + + threshold = bigframes.options.compute.semantic_ops_confirmation_threshold + + if threshold is None or row_count <= threshold: + return + + if bigframes.options.compute.semantic_ops_threshold_autofail: + raise exceptions.OperationAbortedError( + f"Operation was cancelled because your work estimate is {row_count} rows, which exceeds the threshold {threshold} rows." + ) + + # Separate the prompt out. In IDE such VS Code, leaving prompt in the + # input function makes it less visible to the end user. + print(f"This operation will process about {row_count} rows.") + print( + "You can raise the confirmation threshold by setting `bigframes.options.compute.semantic_ops_confirmation_threshold` to a higher value. To completely turn off the confirmation check, set the threshold to `None`." + ) + print("Proceed? [Y/n]") + reply = input().casefold() + if reply not in {"y", "yes", ""}: + raise exceptions.OperationAbortedError("Operation was cancelled.") diff --git a/packages/bigframes/scripts/templates/core_series_accessor.py.j2 b/packages/bigframes/scripts/templates/core_series_accessor.py.j2 index 89decdcbe1e6..ef35d6570cc5 100644 --- a/packages/bigframes/scripts/templates/core_series_accessor.py.j2 +++ b/packages/bigframes/scripts/templates/core_series_accessor.py.j2 @@ -20,7 +20,7 @@ from typing import ( from bigframes import series, session from bigframes.core import col, sentinels -from bigframes.extensions.core import abstract_series_accessor, series_mixins +from bigframes.extensions.core import abstract_series_accessor, series_tvf_mixins T = TypeVar("T") S = TypeVar("S") @@ -28,7 +28,7 @@ S = TypeVar("S") {% for ns in namespaces %} {% if ns.class_name == "AiSeriesAccessor" %} -class {{ ns.class_name }}(series_mixins.AIMixin[T, S]): +class {{ ns.class_name }}(series_tvf_mixins.AITVFMixin[T, S]): {% else %} class {{ ns.class_name }}(abstract_series_accessor.AbstractBigQuerySeriesAccessor[T, S]): {% endif %} diff --git a/packages/bigframes/setup.py b/packages/bigframes/setup.py index e2717fbe5e4b..76b98b88d312 100644 --- a/packages/bigframes/setup.py +++ b/packages/bigframes/setup.py @@ -42,9 +42,9 @@ "google-cloud-bigquery[bqstorage,pandas] >=3.36.0", # 2.30 needed for arrow support. "google-cloud-bigquery-storage >= 2.30.0, < 3.0.0", - "google-cloud-functions >=1.20.2", - "google-cloud-bigquery-connection >=1.18.2", - "google-cloud-resource-manager >=1.14.2", + "google-cloud-functions >=1.12.0", + "google-cloud-bigquery-connection >=1.12.0", + "google-cloud-resource-manager >=1.10.3", "google-cloud-storage >=2.0.0", "google-crc32c >=1.0.0,<2.0.0", "grpc-google-iam-v1 >= 0.14.2", @@ -73,8 +73,8 @@ "tests": [ "freezegun", "pytest-snapshot", - "google-cloud-bigtable >=2.30.0", - "google-cloud-pubsub >=2.29.0", + "google-cloud-bigtable >=2.24.0", + "google-cloud-pubsub >=2.21.4", "tzdata", ], # used for local engine diff --git a/packages/bigframes/testing/constraints-3.10.txt b/packages/bigframes/testing/constraints-3.10.txt index 1dcdd64baa0a..0c76f1dda750 100644 --- a/packages/bigframes/testing/constraints-3.10.txt +++ b/packages/bigframes/testing/constraints-3.10.txt @@ -4,13 +4,13 @@ fsspec==2023.3.0 gcsfs==2023.3.0 geopandas==0.12.2 google-auth==2.15.0 -google-cloud-bigtable==2.30.0 -google-cloud-pubsub==2.29.0 +google-cloud-bigtable==2.24.0 +google-cloud-pubsub==2.21.4 google-cloud-bigquery==3.36.0 -google-cloud-functions==1.20.2 -google-cloud-bigquery-connection==1.18.2 -google-cloud-iam==2.18.2 -google-cloud-resource-manager==1.14.2 +google-cloud-functions==1.12.0 +google-cloud-bigquery-connection==1.12.0 +google-cloud-iam==2.12.1 +google-cloud-resource-manager==1.10.3 google-cloud-storage==2.0.0 grpc-google-iam-v1==0.14.2 numpy==1.24.0 @@ -95,7 +95,7 @@ pluggy==1.6.0 prompt_toolkit==3.0.52 propcache==0.4.1 proto-plus==1.27.1 -protobuf==6.33.5 +protobuf==4.25.8 psygnal==0.15.1 ptyprocess==0.7.0 pure_eval==0.2.3 diff --git a/packages/bigframes/testing/constraints-3.11.txt b/packages/bigframes/testing/constraints-3.11.txt index 17854fda96f8..77645780b7b0 100644 --- a/packages/bigframes/testing/constraints-3.11.txt +++ b/packages/bigframes/testing/constraints-3.11.txt @@ -145,7 +145,7 @@ gitdb==4.0.12 GitPython==3.1.45 glob2==0.7 google==2.0.3 -google-ai-generativelanguage==0.6.17 +google-ai-generativelanguage==0.6.15 google-api-core==2.25.1 google-api-python-client==2.177.0 google-auth==2.38.0 @@ -180,7 +180,7 @@ groovy==0.1.2 grpc-google-iam-v1==0.14.2 grpc-interceptor==0.15.4 grpcio==1.74.0 -grpcio-status==1.72.1 +grpcio-status==1.71.2 grpclib==0.4.8 gspread==6.2.1 gspread-dataframe==4.0.0 @@ -401,7 +401,7 @@ prompt_toolkit==3.0.51 propcache==0.3.2 prophet==1.1.7 proto-plus==1.26.1 -protobuf==6.33.5 +protobuf==5.29.5 psutil==5.9.5 psycopg2==2.9.10 psygnal==0.14.0 diff --git a/packages/bigframes/tests/system/large/operations/__init__.py b/packages/bigframes/tests/system/large/operations/__init__.py new file mode 100644 index 000000000000..6d5e14bcf4a0 --- /dev/null +++ b/packages/bigframes/tests/system/large/operations/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2024 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. diff --git a/packages/bigframes/tests/system/large/operations/conftest.py b/packages/bigframes/tests/system/large/operations/conftest.py new file mode 100644 index 000000000000..0122c860ca74 --- /dev/null +++ b/packages/bigframes/tests/system/large/operations/conftest.py @@ -0,0 +1,33 @@ +# Copyright 2024 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. + +import pytest + +import bigframes.ml.llm as llm + + +@pytest.fixture(scope="session") +def gemini_flash_model(session, bq_connection) -> llm.GeminiTextGenerator: + return llm.GeminiTextGenerator( + session=session, + connection_name=bq_connection, + model_name="gemini-2.5-flash", + ) + + +@pytest.fixture(scope="session") +def text_embedding_generator(session, bq_connection) -> llm.TextEmbeddingGenerator: + return llm.TextEmbeddingGenerator( + session=session, connection_name=bq_connection, model_name="text-embedding-005" + ) diff --git a/packages/bigframes/tests/system/large/operations/test_ai.py b/packages/bigframes/tests/system/large/operations/test_ai.py new file mode 100644 index 000000000000..cd63025cd388 --- /dev/null +++ b/packages/bigframes/tests/system/large/operations/test_ai.py @@ -0,0 +1,963 @@ +# 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. + +from contextlib import nullcontext +from unittest.mock import patch + +import pandas as pd +import pandas.testing +import pytest + +import bigframes +from bigframes import dataframe, exceptions, series + +AI_OP_EXP_OPTION = "experiments.ai_operators" +BLOB_EXP_OPTION = "experiments.blob" +THRESHOLD_OPTION = "compute.ai_ops_confirmation_threshold" + + +def test_filter(session, gemini_flash_model): + df = dataframe.DataFrame( + data={ + "country": ["USA", "Germany"], + "city": ["Seattle", "Berlin"], + "year": [2023, 2024], + }, + session=session, + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = df.ai.filter( + "{city} is the capital of {country} in {year}", gemini_flash_model + ).to_pandas() + + expected_df = pd.DataFrame( + {"country": ["Germany"], "city": ["Berlin"], "year": [2024]}, index=[1] + ) + pandas.testing.assert_frame_equal( + actual_df, expected_df, check_dtype=False, check_index_type=False + ) + + +def test_filter_multi_model(session, gemini_flash_model): + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + BLOB_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + df = session._from_glob_path( + "gs://bigframes-dev-testing/a_multimodal/images/*", name="image" + ) + df["prey"] = series.Series( + ["building", "cross road", "rock", "squirrel", "rabbit"], session=session + ) + result = df.ai.filter( + "The object in {image} feeds on {prey}", + gemini_flash_model, + ).to_pandas() + + assert len(result) <= len(df) + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_filter_with_confirmation(session, gemini_flash_model, reply, monkeypatch): + df = dataframe.DataFrame( + data={ + "country": ["USA", "Germany"], + "city": ["Seattle", "Berlin"], + "year": [2023, 2024], + }, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df.ai.filter("{city} is the capital of {country} in {year}", gemini_flash_model) + + +def test_filter_single_column_reference(session, gemini_flash_model): + df = dataframe.DataFrame( + data={"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]}, + session=session, + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = df.ai.filter( + "{country} is in Europe", gemini_flash_model + ).to_pandas() + + expected_df = pd.DataFrame({"country": ["Germany"], "city": ["Berlin"]}, index=[1]) + pandas.testing.assert_frame_equal( + actual_df, expected_df, check_dtype=False, check_index_type=False + ) + + +@pytest.mark.parametrize( + "instruction", + [ + pytest.param( + "No column reference", + id="zero_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{city} is in the {non_existing_column}", + id="non_existing_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{id}", + id="invalid_type", + marks=pytest.mark.xfail(raises=TypeError), + ), + ], +) +def test_filter_invalid_instruction_raise_error(instruction, gemini_flash_model): + df = dataframe.DataFrame({"id": [1, 2], "city": ["Seattle", "Berlin"]}) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.ai.filter(instruction, gemini_flash_model) + + +def test_filter_invalid_model_raise_error(): + df = dataframe.DataFrame( + {"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]} + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df.ai.filter("{city} is the capital of {country}", None) + + +@pytest.mark.parametrize( + ("output_schema", "output_col"), + [ + pytest.param(None, "ml_generate_text_llm_result", id="default_schema"), + pytest.param({"food": "string"}, "food", id="non_default_schema"), + ], +) +def test_map(session, gemini_flash_model, output_schema, output_col): + df = dataframe.DataFrame( + data={ + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + "gluten-free": [True, True], + }, + session=session, + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = df.ai.map( + "What is the {gluten-free} food made from {ingredient_1} and {ingredient_2}? One word only.", + gemini_flash_model, + output_schema=output_schema, + ).to_pandas() + # Result sanitation + actual_df[output_col] = actual_df[output_col].str.strip().str.lower() + + expected_df = pd.DataFrame( + { + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + "gluten-free": [True, True], + output_col: ["burger", "tofu"], + } + ) + pandas.testing.assert_frame_equal( + actual_df, + expected_df, + check_dtype=False, + check_index_type=False, + check_column_type=False, + ) + + +def test_map_multimodal(session, gemini_flash_model): + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + BLOB_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + df = session._from_glob_path( + "gs://bigframes-dev-testing/a_multimodal/images/*", name="image" + ) + df["scenario"] = series.Series( + ["building", "cross road", "tree", "squirrel", "rabbit"], session=session + ) + result = df.ai.map( + "What is the object in {image} combined with {scenario}? One word only.", + gemini_flash_model, + output_schema={"object": "string"}, + ).to_pandas() + + assert len(result) == len(df) + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_map_with_confirmation(session, gemini_flash_model, reply, monkeypatch): + df = dataframe.DataFrame( + data={ + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + "gluten-free": [True, True], + }, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df.ai.map( + "What is the {gluten-free} food made from {ingredient_1} and {ingredient_2}? One word only.", + gemini_flash_model, + ) + + +@pytest.mark.parametrize( + "instruction", + [ + pytest.param( + "No column reference", + id="zero_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "What is the food made from {ingredient_1} and {non_existing_column}?}", + id="non_existing_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{id}", + id="invalid_type", + marks=pytest.mark.xfail(raises=TypeError), + ), + ], +) +def test_map_invalid_instruction_raise_error(instruction, gemini_flash_model): + df = dataframe.DataFrame( + data={ + "id": [1, 2], + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + } + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.ai.map(instruction, gemini_flash_model, output_schema={"food": "string"}) + + +def test_map_invalid_model_raise_error(): + df = dataframe.DataFrame( + data={ + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + }, + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df.ai.map( + "What is the food made from {ingredient_1} and {ingredient_2}? One word only.", + None, + ) + + +def test_classify(gemini_flash_model, session): + df = dataframe.DataFrame(data={"creature": ["dog", "rose"]}, session=session) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_result = df.ai.classify( + "{creature}", + gemini_flash_model, + labels=["animal", "plant"], + output_column="result", + ).to_pandas() + + expected_result = pd.DataFrame( + { + "creature": ["dog", "rose"], + "result": ["animal", "plant"], + } + ) + pandas.testing.assert_frame_equal( + actual_result, expected_result, check_index_type=False, check_dtype=False + ) + + +@pytest.mark.parametrize( + "instruction", + [ + pytest.param("{city} is in {country}", id="no_dataframe_reference"), + pytest.param("{left.city} is in {country}", id="has_left_dataframe_reference"), + pytest.param( + "{city} is in {right.country}", + id="has_right_dataframe_reference", + ), + pytest.param( + "{left.city} is in {right.country}", id="has_both_dataframe_references" + ), + ], +) +def test_join(instruction, session, gemini_flash_model): + cities = dataframe.DataFrame( + data={ + "city": ["Seattle", "Berlin"], + }, + session=session, + ) + countries = dataframe.DataFrame( + data={"country": ["USA", "UK", "Germany"]}, + session=session, + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = cities.ai.join( + countries, + instruction, + gemini_flash_model, + ).to_pandas() + + expected_df = pd.DataFrame( + { + "city": ["Seattle", "Berlin"], + "country": ["USA", "Germany"], + } + ) + pandas.testing.assert_frame_equal( + actual_df, + expected_df, + check_dtype=False, + check_index_type=False, + check_column_type=False, + ) + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_join_with_confirmation(session, gemini_flash_model, reply, monkeypatch): + cities = dataframe.DataFrame( + data={ + "city": ["Seattle", "Berlin"], + }, + session=session, + ) + countries = dataframe.DataFrame( + data={"country": ["USA", "UK", "Germany"]}, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + cities.ai.join( + countries, + "{city} is in {country}", + gemini_flash_model, + ) + + +def test_self_join(session, gemini_flash_model): + animals = dataframe.DataFrame( + data={ + "animal": ["ant", "elephant"], + }, + session=session, + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = animals.ai.join( + animals, + "{left.animal} is heavier than {right.animal}", + gemini_flash_model, + ).to_pandas() + + expected_df = pd.DataFrame( + { + "animal_left": ["elephant"], + "animal_right": ["ant"], + } + ) + pandas.testing.assert_frame_equal( + actual_df, + expected_df, + check_dtype=False, + check_index_type=False, + check_column_type=False, + ) + + +@pytest.mark.parametrize( + ("instruction", "error_pattern"), + [ + ("No column reference", "No column references"), + pytest.param( + "{city} is in {continent}", r"Column .+ not found", id="non_existing_column" + ), + pytest.param( + "{city} is in {country}", + r"Ambiguous column reference: .+", + id="ambiguous_column", + ), + pytest.param( + "{right.city} is in {country}", r"Column .+ not found", id="wrong_prefix" + ), + pytest.param( + "{city} is in {right.continent}", + r"Column .+ not found", + id="prefix_on_non_existing_column", + ), + ], +) +def test_join_invalid_instruction_raise_error( + instruction, error_pattern, gemini_flash_model +): + df1 = dataframe.DataFrame( + {"city": ["Seattle", "Berlin"], "country": ["USA", "Germany"]} + ) + df2 = dataframe.DataFrame( + { + "country": ["USA", "UK", "Germany"], + "region": ["North America", "Europe", "Europe"], + } + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError, match=error_pattern), + ): + df1.ai.join(df2, instruction, gemini_flash_model) + + +def test_join_invalid_model_raise_error(): + cities = dataframe.DataFrame({"city": ["Seattle", "Berlin"]}) + countries = dataframe.DataFrame({"country": ["USA", "UK", "Germany"]}) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + cities.ai.join(countries, "{city} is in {country}", None) + + +@pytest.mark.parametrize( + "score_column", + [ + pytest.param(None, id="no_score_column"), + pytest.param("distance", id="has_score_column"), + ], +) +def test_search(session, text_embedding_generator, score_column): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_result = df.ai.search( + "creatures", + "monkey", + top_k=2, + model=text_embedding_generator, + score_column=score_column, + ).to_pandas() + + expected_result = pd.Series( + ["baboons", "chimpanzee"], index=[2, 4], name="creatures" + ) + pandas.testing.assert_series_equal( + actual_result["creatures"], + expected_result, + check_dtype=False, + check_index_type=False, + ) + + if score_column is None: + assert len(actual_result.columns) == 1 + else: + assert score_column in actual_result.columns + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_search_with_confirmation( + session, text_embedding_generator, reply, monkeypatch +): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df.ai.search( + "creatures", + "monkey", + top_k=2, + model=text_embedding_generator, + ) + + +def test_search_invalid_column_raises_error(session, text_embedding_generator): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.ai.search("whatever", "monkey", top_k=2, model=text_embedding_generator) + + +def test_search_invalid_model_raises_error(session): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df.ai.search("creatures", "monkey", top_k=2, model=None) + + +def test_search_invalid_top_k_raises_error(session, text_embedding_generator): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.ai.search("creatures", "monkey", top_k=0, model=text_embedding_generator) + + +@pytest.mark.parametrize( + "score_column", + [ + pytest.param(None, id="no_score_column"), + pytest.param("distance", id="has_score_column"), + ], +) +def test_sim_join(session, text_embedding_generator, score_column): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_result = df1.ai.sim_join( + df2, + left_on="creatures", + right_on="creatures", + model=text_embedding_generator, + top_k=1, + score_column=score_column, + ).to_pandas() + + expected_result = pd.DataFrame( + {"creatures": ["salmon", "cat"], "creatures_1": ["tuna", "dog"]} + ) + pandas.testing.assert_frame_equal( + actual_result[["creatures", "creatures_1"]], + expected_result, + check_dtype=False, + check_index_type=False, + ) + + if score_column is None: + assert len(actual_result.columns) == 2 + else: + assert score_column in actual_result.columns + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_sim_join_with_confirmation( + session, text_embedding_generator, reply, monkeypatch +): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df1.ai.sim_join( + df2, + left_on="creatures", + right_on="creatures", + model=text_embedding_generator, + top_k=1, + ) + + +@pytest.mark.parametrize( + ("left_on", "right_on"), + [ + pytest.param("whatever", "creatures", id="incorrect_left_column"), + pytest.param("creatures", "whatever", id="incorrect_right_column"), + ], +) +def test_sim_join_invalid_column_raises_error( + session, text_embedding_generator, left_on, right_on +): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df1.ai.sim_join( + df2, left_on=left_on, right_on=right_on, model=text_embedding_generator + ) + + +def test_sim_join_invalid_model_raises_error(session): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df1.ai.sim_join(df2, left_on="creatures", right_on="creatures", model=None) + + +def test_sim_join_invalid_top_k_raises_error(session, text_embedding_generator): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df1.ai.sim_join( + df2, + left_on="creatures", + right_on="creatures", + top_k=0, + model=text_embedding_generator, + ) + + +def test_sim_join_data_too_large_raises_error(session, text_embedding_generator): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df1.ai.sim_join( + df2, + left_on="creatures", + right_on="creatures", + model=text_embedding_generator, + max_rows=1, + ) + + +@patch("builtins.input", return_value="") +def test_confirm_operation__below_threshold_do_not_confirm(mock_input): + df = dataframe.DataFrame({}) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 3, + ): + df.ai._confirm_operation(1) + + mock_input.assert_not_called() + + +@patch("builtins.input", return_value="") +def test_confirm_operation__threshold_is_none_do_not_confirm(mock_input): + df = dataframe.DataFrame({}) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + None, + ): + df.ai._confirm_operation(100) + + mock_input.assert_not_called() + + +@patch("builtins.input", return_value="") +def test_confirm_operation__threshold_autofail_do_not_confirm(mock_input): + df = dataframe.DataFrame({}) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 1, + "compute.ai_ops_threshold_autofail", + True, + ), + pytest.raises(exceptions.OperationAbortedError), + ): + df.ai._confirm_operation(100) + + mock_input.assert_not_called() + + +@pytest.mark.parametrize( + ("reply", "expectation"), + [ + ("y", nullcontext()), + ("yes", nullcontext()), + ("", nullcontext()), + ("n", pytest.raises(exceptions.OperationAbortedError)), + ("something", pytest.raises(exceptions.OperationAbortedError)), + ], +) +def test_confirm_operation__above_threshold_confirm(reply, expectation, monkeypatch): + monkeypatch.setattr("builtins.input", lambda: reply) + df = dataframe.DataFrame({}) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 3, + ), + expectation as e, + ): + assert df.ai._confirm_operation(4) == e diff --git a/packages/bigframes/tests/system/large/operations/test_semantics.py b/packages/bigframes/tests/system/large/operations/test_semantics.py new file mode 100644 index 000000000000..147b80f795aa --- /dev/null +++ b/packages/bigframes/tests/system/large/operations/test_semantics.py @@ -0,0 +1,1348 @@ +# Copyright 2024 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 contextlib import nullcontext +from unittest.mock import patch + +import pandas as pd +import pandas.testing +import pytest + +import bigframes +from bigframes import dataframe, dtypes, exceptions, series + +pytest.skip( + "Semantics namespace is deprecated. ", + allow_module_level=True, +) + +SEM_OP_EXP_OPTION = "experiments.semantic_operators" +BLOB_EXP_OPTION = "experiments.blob" +THRESHOLD_OPTION = "compute.semantic_ops_confirmation_threshold" + + +def test_semantics_experiment_off_raise_error(): + df = dataframe.DataFrame( + {"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]} + ) + + with ( + bigframes.option_context(SEM_OP_EXP_OPTION, False), + pytest.raises(NotImplementedError), + ): + df.semantics + + +@pytest.mark.parametrize( + ("max_agg_rows", "cluster_column"), + [ + pytest.param(1, None, id="one", marks=pytest.mark.xfail(raises=ValueError)), + pytest.param(2, None, id="two"), + pytest.param(3, None, id="three"), + pytest.param(4, None, id="four"), + pytest.param(5, "Years", id="two_w_cluster_column"), + pytest.param(6, "Years", id="three_w_cluster_column"), + pytest.param(7, "Years", id="four_w_cluster_column"), + ], +) +def test_agg(session, gemini_flash_model, max_agg_rows, cluster_column): + df = dataframe.DataFrame( + data={ + "Movies": [ + "Titanic", + "The Wolf of Wall Street", + "Killers of the Flower Moon", + "The Revenant", + "Inception", + "Shuttle Island", + "The Great Gatsby", + ], + "Years": [1997, 2013, 2023, 2015, 2010, 2010, 2013], + }, + session=session, + ) + instruction = "Find the shared first name of actors in {Movies}. One word answer." + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + actual_s = df.semantics.agg( + instruction, + model=gemini_flash_model, + max_agg_rows=max_agg_rows, + cluster_column=cluster_column, + ).to_pandas() + + expected_s = pd.Series(["Leonardo\n"], dtype=dtypes.STRING_DTYPE) + expected_s.name = "Movies" + pandas.testing.assert_series_equal(actual_s, expected_s, check_index_type=False) + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_agg_with_confirmation(session, gemini_flash_model, reply, monkeypatch): + df = dataframe.DataFrame( + data={ + "Movies": [ + "Titanic", + "The Wolf of Wall Street", + "Killers of the Flower Moon", + "The Revenant", + "Inception", + "Shuttle Island", + "The Great Gatsby", + ], + "Years": [1997, 2013, 2023, 2015, 2010, 2010, 2013], + }, + session=session, + ) + instruction = "Find the shared first name of actors in {Movies}. One word answer." + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df.semantics.agg( + instruction, + model=gemini_flash_model, + ) + + +def test_agg_w_int_column(session, gemini_flash_model): + df = dataframe.DataFrame( + data={ + "Movies": [ + "Killers of the Flower Moon", + "The Great Gatsby", + "The Wolf of Wall Street", + ], + "Years": [2023, 2013, 2013], + }, + session=session, + ) + instruction = "Find the {Years} Leonardo DiCaprio acted in the most movies. Your answer should be the four-digit year, returned as a string." + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_s = df.semantics.agg( + instruction, + model=gemini_flash_model, + ).to_pandas() + + expected_s = pd.Series(["2013\n"], dtype=dtypes.STRING_DTYPE) + expected_s.name = "Years" + pandas.testing.assert_series_equal(actual_s, expected_s, check_index_type=False) + + +@pytest.mark.parametrize( + "instruction", + [ + pytest.param( + "No column reference", + id="zero_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{Movies} is good", + id="non_existing_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{Movies} is better than {Movies}", + id="two_columns", + marks=pytest.mark.xfail(raises=NotImplementedError), + ), + ], +) +def test_agg_invalid_instruction_raise_error(instruction, gemini_flash_model): + df = dataframe.DataFrame( + data={ + "Movies": [ + "Titanic", + "The Wolf of Wall Street", + "Killers of the Flower Moon", + ], + "Year": [1997, 2013, 2023], + }, + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + df.semantics.agg(instruction, gemini_flash_model) + + +@pytest.mark.parametrize( + "cluster_column", + [ + pytest.param( + "non_existing_column", + id="non_existing_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "Movies", id="non_int_column", marks=pytest.mark.xfail(raises=TypeError) + ), + ], +) +def test_agg_invalid_cluster_column_raise_error(gemini_flash_model, cluster_column): + df = dataframe.DataFrame( + data={ + "Movies": [ + "Titanic", + "The Wolf of Wall Street", + "Killers of the Flower Moon", + "The Revenant", + ], + }, + ) + instruction = "Find the shared first name of actors in {Movies}. One word answer." + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + df.semantics.agg(instruction, gemini_flash_model, cluster_column=cluster_column) + + +@pytest.mark.parametrize( + ("n_clusters"), + [ + pytest.param(1, id="one", marks=pytest.mark.xfail(raises=ValueError)), + pytest.param(2, id="two"), + ], +) +def test_cluster_by(session, text_embedding_generator, n_clusters): + df = dataframe.DataFrame( + ( + { + "Item": [ + "Orange", + "Cantaloupe", + "Watermelon", + "Chicken", + "Duck", + "Hen", + "Rooster", + ] + } + ), + session=session, + ) + output_column = "cluster id" + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + result = df.semantics.cluster_by( + "Item", + output_column, + text_embedding_generator, + n_clusters=n_clusters, + ) + + assert output_column in result + # In rare cases, it's possible to have fewer than K clusters due to randomness. + assert len(result[output_column].unique()) <= n_clusters + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_cluster_by_with_confirmation( + session, text_embedding_generator, reply, monkeypatch +): + df = dataframe.DataFrame( + ( + { + "Item": [ + "Orange", + "Cantaloupe", + "Watermelon", + "Chicken", + "Duck", + "Hen", + "Rooster", + ] + } + ), + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df.semantics.cluster_by( + "Item", + "cluster id", + text_embedding_generator, + n_clusters=2, + ) + + +def test_cluster_by_invalid_column(session, text_embedding_generator): + df = dataframe.DataFrame( + ({"Product": ["Smartphone", "Laptop", "Coffee Maker", "T-shirt", "Jeans"]}), + session=session, + ) + output_column = "cluster id" + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.semantics.cluster_by( + "unknown_column", + output_column, + text_embedding_generator, + n_clusters=3, + ) + + +def test_cluster_by_invalid_model(session, gemini_flash_model): + df = dataframe.DataFrame( + ({"Product": ["Smartphone", "Laptop", "Coffee Maker", "T-shirt", "Jeans"]}), + session=session, + ) + output_column = "cluster id" + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df.semantics.cluster_by( + "Product", + output_column, + gemini_flash_model, + n_clusters=3, + ) + + +def test_filter(session, gemini_flash_model): + df = dataframe.DataFrame( + data={ + "country": ["USA", "Germany"], + "city": ["Seattle", "Berlin"], + "year": [2023, 2024], + }, + session=session, + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = df.semantics.filter( + "{city} is the capital of {country} in {year}", gemini_flash_model + ).to_pandas() + + expected_df = pd.DataFrame( + {"country": ["Germany"], "city": ["Berlin"], "year": [2024]}, index=[1] + ) + pandas.testing.assert_frame_equal( + actual_df, expected_df, check_dtype=False, check_index_type=False + ) + + +def test_filter_multi_model(session, gemini_flash_model): + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + BLOB_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + df = session._from_glob_path( + "gs://bigframes-dev-testing/a_multimodal/images/*", name="image" + ) + df["prey"] = series.Series( + ["building", "cross road", "rock", "squirrel", "rabbit"], session=session + ) + result = df.semantics.filter( + "The object in {image} feeds on {prey}", + gemini_flash_model, + ).to_pandas() + + assert len(result) <= len(df) + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_filter_with_confirmation(session, gemini_flash_model, reply, monkeypatch): + df = dataframe.DataFrame( + data={ + "country": ["USA", "Germany"], + "city": ["Seattle", "Berlin"], + "year": [2023, 2024], + }, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df.semantics.filter( + "{city} is the capital of {country} in {year}", gemini_flash_model + ) + + +def test_filter_single_column_reference(session, gemini_flash_model): + df = dataframe.DataFrame( + data={"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]}, + session=session, + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = df.semantics.filter( + "{country} is in Europe", gemini_flash_model + ).to_pandas() + + expected_df = pd.DataFrame({"country": ["Germany"], "city": ["Berlin"]}, index=[1]) + pandas.testing.assert_frame_equal( + actual_df, expected_df, check_dtype=False, check_index_type=False + ) + + +@pytest.mark.parametrize( + "instruction", + [ + pytest.param( + "No column reference", + id="zero_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{city} is in the {non_existing_column}", + id="non_existing_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{id}", + id="invalid_type", + marks=pytest.mark.xfail(raises=TypeError), + ), + ], +) +def test_filter_invalid_instruction_raise_error(instruction, gemini_flash_model): + df = dataframe.DataFrame({"id": [1, 2], "city": ["Seattle", "Berlin"]}) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.semantics.filter(instruction, gemini_flash_model) + + +def test_filter_invalid_model_raise_error(): + df = dataframe.DataFrame( + {"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]} + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df.semantics.filter("{city} is the capital of {country}", None) + + +def test_map(session, gemini_flash_model): + df = dataframe.DataFrame( + data={ + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + "gluten-free": [True, True], + }, + session=session, + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = df.semantics.map( + "What is the {gluten-free} food made from {ingredient_1} and {ingredient_2}? One word only.", + "food", + gemini_flash_model, + ).to_pandas() + # Result sanitation + actual_df["food"] = actual_df["food"].str.strip().str.lower() + + expected_df = pd.DataFrame( + { + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + "gluten-free": [True, True], + "food": ["burger", "tofu"], + } + ) + pandas.testing.assert_frame_equal( + actual_df, + expected_df, + check_dtype=False, + check_index_type=False, + check_column_type=False, + ) + + +def test_map_multimodal(session, gemini_flash_model): + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + BLOB_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + df = session._from_glob_path( + "gs://bigframes-dev-testing/a_multimodal/images/*", name="image" + ) + df["scenario"] = series.Series( + ["building", "cross road", "tree", "squirrel", "rabbit"], session=session + ) + result = df.semantics.map( + "What is the object in {image} combined with {scenario}? One word only.", + "object", + gemini_flash_model, + ).to_pandas() + + assert len(result) == len(df) + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_map_with_confirmation(session, gemini_flash_model, reply, monkeypatch): + df = dataframe.DataFrame( + data={ + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + "gluten-free": [True, True], + }, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df.semantics.map( + "What is the {gluten-free} food made from {ingredient_1} and {ingredient_2}? One word only.", + "food", + gemini_flash_model, + ) + + +@pytest.mark.parametrize( + "instruction", + [ + pytest.param( + "No column reference", + id="zero_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "What is the food made from {ingredient_1} and {non_existing_column}?}", + id="non_existing_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{id}", + id="invalid_type", + marks=pytest.mark.xfail(raises=TypeError), + ), + ], +) +def test_map_invalid_instruction_raise_error(instruction, gemini_flash_model): + df = dataframe.DataFrame( + data={ + "id": [1, 2], + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + } + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.semantics.map(instruction, "food", gemini_flash_model) + + +def test_map_invalid_model_raise_error(): + df = dataframe.DataFrame( + data={ + "ingredient_1": ["Burger Bun", "Soy Bean"], + "ingredient_2": ["Beef Patty", "Bittern"], + }, + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df.semantics.map( + "What is the food made from {ingredient_1} and {ingredient_2}? One word only.", + "food", + None, + ) + + +@pytest.mark.parametrize( + "instruction", + [ + pytest.param("{city} is in {country}", id="no_dataframe_reference"), + pytest.param("{left.city} is in {country}", id="has_left_dataframe_reference"), + pytest.param( + "{city} is in {right.country}", + id="has_right_dataframe_reference", + ), + pytest.param( + "{left.city} is in {right.country}", id="has_both_dataframe_references" + ), + ], +) +def test_join(instruction, session, gemini_flash_model): + cities = dataframe.DataFrame( + data={ + "city": ["Seattle", "Berlin"], + }, + session=session, + ) + countries = dataframe.DataFrame( + data={"country": ["USA", "UK", "Germany"]}, + session=session, + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = cities.semantics.join( + countries, + instruction, + gemini_flash_model, + ).to_pandas() + + expected_df = pd.DataFrame( + { + "city": ["Seattle", "Berlin"], + "country": ["USA", "Germany"], + } + ) + pandas.testing.assert_frame_equal( + actual_df, + expected_df, + check_dtype=False, + check_index_type=False, + check_column_type=False, + ) + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_join_with_confirmation(session, gemini_flash_model, reply, monkeypatch): + cities = dataframe.DataFrame( + data={ + "city": ["Seattle", "Berlin"], + }, + session=session, + ) + countries = dataframe.DataFrame( + data={"country": ["USA", "UK", "Germany"]}, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + cities.semantics.join( + countries, + "{city} is in {country}", + gemini_flash_model, + ) + + +def test_self_join(session, gemini_flash_model): + animals = dataframe.DataFrame( + data={ + "animal": ["ant", "elephant"], + }, + session=session, + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_df = animals.semantics.join( + animals, + "{left.animal} is heavier than {right.animal}", + gemini_flash_model, + ).to_pandas() + + expected_df = pd.DataFrame( + { + "animal_left": ["elephant"], + "animal_right": ["ant"], + } + ) + pandas.testing.assert_frame_equal( + actual_df, + expected_df, + check_dtype=False, + check_index_type=False, + check_column_type=False, + ) + + +@pytest.mark.parametrize( + ("instruction", "error_pattern"), + [ + ("No column reference", "No column references"), + pytest.param( + "{city} is in {continent}", r"Column .+ not found", id="non_existing_column" + ), + pytest.param( + "{city} is in {country}", + r"Ambiguous column reference: .+", + id="ambiguous_column", + ), + pytest.param( + "{right.city} is in {country}", r"Column .+ not found", id="wrong_prefix" + ), + pytest.param( + "{city} is in {right.continent}", + r"Column .+ not found", + id="prefix_on_non_existing_column", + ), + ], +) +def test_join_invalid_instruction_raise_error( + instruction, error_pattern, gemini_flash_model +): + df1 = dataframe.DataFrame( + {"city": ["Seattle", "Berlin"], "country": ["USA", "Germany"]} + ) + df2 = dataframe.DataFrame( + { + "country": ["USA", "UK", "Germany"], + "region": ["North America", "Europe", "Europe"], + } + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError, match=error_pattern), + ): + df1.semantics.join(df2, instruction, gemini_flash_model) + + +def test_join_invalid_model_raise_error(): + cities = dataframe.DataFrame({"city": ["Seattle", "Berlin"]}) + countries = dataframe.DataFrame({"country": ["USA", "UK", "Germany"]}) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + cities.semantics.join(countries, "{city} is in {country}", None) + + +@pytest.mark.parametrize( + "score_column", + [ + pytest.param(None, id="no_score_column"), + pytest.param("distance", id="has_score_column"), + ], +) +def test_search(session, text_embedding_generator, score_column): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_result = df.semantics.search( + "creatures", + "monkey", + top_k=2, + model=text_embedding_generator, + score_column=score_column, + ).to_pandas() + + expected_result = pd.Series( + ["baboons", "chimpanzee"], index=[2, 4], name="creatures" + ) + pandas.testing.assert_series_equal( + actual_result["creatures"], + expected_result, + check_dtype=False, + check_index_type=False, + ) + + if score_column is None: + assert len(actual_result.columns) == 1 + else: + assert score_column in actual_result.columns + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_search_with_confirmation( + session, text_embedding_generator, reply, monkeypatch +): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df.semantics.search( + "creatures", + "monkey", + top_k=2, + model=text_embedding_generator, + ) + + +def test_search_invalid_column_raises_error(session, text_embedding_generator): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.semantics.search( + "whatever", "monkey", top_k=2, model=text_embedding_generator + ) + + +def test_search_invalid_model_raises_error(session): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df.semantics.search("creatures", "monkey", top_k=2, model=None) + + +def test_search_invalid_top_k_raises_error(session, text_embedding_generator): + df = dataframe.DataFrame( + data={"creatures": ["salmon", "sea urchin", "baboons", "frog", "chimpanzee"]}, + session=session, + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.semantics.search( + "creatures", "monkey", top_k=0, model=text_embedding_generator + ) + + +@pytest.mark.parametrize( + "score_column", + [ + pytest.param(None, id="no_score_column"), + pytest.param("distance", id="has_score_column"), + ], +) +def test_sim_join(session, text_embedding_generator, score_column): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + actual_result = df1.semantics.sim_join( + df2, + left_on="creatures", + right_on="creatures", + model=text_embedding_generator, + top_k=1, + score_column=score_column, + ).to_pandas() + + expected_result = pd.DataFrame( + {"creatures": ["salmon", "cat"], "creatures_1": ["tuna", "dog"]} + ) + pandas.testing.assert_frame_equal( + actual_result[["creatures", "creatures_1"]], + expected_result, + check_dtype=False, + check_index_type=False, + ) + + if score_column is None: + assert len(actual_result.columns) == 2 + else: + assert score_column in actual_result.columns + + +@pytest.mark.parametrize( + ("reply"), + [ + pytest.param("y"), + pytest.param( + "n", marks=pytest.mark.xfail(raises=exceptions.OperationAbortedError) + ), + ], +) +def test_sim_join_with_confirmation( + session, text_embedding_generator, reply, monkeypatch +): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + monkeypatch.setattr("builtins.input", lambda: reply) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 0, + ): + df1.semantics.sim_join( + df2, + left_on="creatures", + right_on="creatures", + model=text_embedding_generator, + top_k=1, + ) + + +@pytest.mark.parametrize( + ("left_on", "right_on"), + [ + pytest.param("whatever", "creatures", id="incorrect_left_column"), + pytest.param("creatures", "whatever", id="incorrect_right_column"), + ], +) +def test_sim_join_invalid_column_raises_error( + session, text_embedding_generator, left_on, right_on +): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df1.semantics.sim_join( + df2, left_on=left_on, right_on=right_on, model=text_embedding_generator + ) + + +def test_sim_join_invalid_model_raises_error(session): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(TypeError), + ): + df1.semantics.sim_join( + df2, left_on="creatures", right_on="creatures", model=None + ) + + +def test_sim_join_invalid_top_k_raises_error(session, text_embedding_generator): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df1.semantics.sim_join( + df2, + left_on="creatures", + right_on="creatures", + top_k=0, + model=text_embedding_generator, + ) + + +def test_sim_join_data_too_large_raises_error(session, text_embedding_generator): + df1 = dataframe.DataFrame( + data={"creatures": ["salmon", "cat"]}, + session=session, + ) + df2 = dataframe.DataFrame( + data={"creatures": ["dog", "tuna"]}, + session=session, + ) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df1.semantics.sim_join( + df2, + left_on="creatures", + right_on="creatures", + model=text_embedding_generator, + max_rows=1, + ) + + +@pytest.mark.parametrize( + "instruction", + [ + pytest.param( + "No column reference", + id="zero_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{Animals}", + id="non_existing_column", + marks=pytest.mark.xfail(raises=ValueError), + ), + pytest.param( + "{Animals} and {Animals}", + id="two_columns", + marks=pytest.mark.xfail(raises=NotImplementedError), + ), + pytest.param( + "{index}", + id="preserved", + marks=pytest.mark.xfail(raises=ValueError), + ), + ], +) +def test_top_k_invalid_instruction_raise_error(instruction, gemini_flash_model): + df = dataframe.DataFrame( + { + "Animals": ["Dog", "Cat", "Bird", "Horse"], + "ID": [1, 2, 3, 4], + "index": ["a", "b", "c", "d"], + } + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ): + df.semantics.top_k(instruction, model=gemini_flash_model, k=2) + + +def test_top_k_invalid_k_raise_error(gemini_flash_model): + df = dataframe.DataFrame({"Animals": ["Dog", "Cat", "Bird", "Horse"]}) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 10, + ), + pytest.raises(ValueError), + ): + df.semantics.top_k( + "{Animals} are more popular as pets", + gemini_flash_model, + k=0, + ) + + +@patch("builtins.input", return_value="") +def test_confirm_operation__below_threshold_do_not_confirm(mock_input): + df = dataframe.DataFrame({}) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 3, + ): + df.semantics._confirm_operation(1) + + mock_input.assert_not_called() + + +@patch("builtins.input", return_value="") +def test_confirm_operation__threshold_is_none_do_not_confirm(mock_input): + df = dataframe.DataFrame({}) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + None, + ): + df.semantics._confirm_operation(100) + + mock_input.assert_not_called() + + +@patch("builtins.input", return_value="") +def test_confirm_operation__threshold_autofail_do_not_confirm(mock_input): + df = dataframe.DataFrame({}) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 1, + "compute.semantic_ops_threshold_autofail", + True, + ), + pytest.raises(exceptions.OperationAbortedError), + ): + df.semantics._confirm_operation(100) + + mock_input.assert_not_called() + + +@pytest.mark.parametrize( + ("reply", "expectation"), + [ + ("y", nullcontext()), + ("yes", nullcontext()), + ("", nullcontext()), + ("n", pytest.raises(exceptions.OperationAbortedError)), + ("something", pytest.raises(exceptions.OperationAbortedError)), + ], +) +def test_confirm_operation__above_threshold_confirm(reply, expectation, monkeypatch): + monkeypatch.setattr("builtins.input", lambda: reply) + df = dataframe.DataFrame({}) + + with ( + bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 3, + ), + expectation as e, + ): + assert df.semantics._confirm_operation(4) == e diff --git a/packages/bigframes/tests/system/small/operations/test_ai.py b/packages/bigframes/tests/system/small/operations/test_ai.py new file mode 100644 index 000000000000..7a34e9b8414f --- /dev/null +++ b/packages/bigframes/tests/system/small/operations/test_ai.py @@ -0,0 +1,280 @@ +# 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. + + +# Note that the tests in this files uses fake models for deterministic results. +# Tests that use real LLM models are under system/large/test_ai.py + +import pandas as pd +import pandas.testing +import pytest + +import bigframes +import bigframes.operations.ai +from bigframes import dataframe, dtypes +from bigframes.ml import llm +from bigframes.testing import utils + +AI_OP_EXP_OPTION = "experiments.ai_operators" +THRESHOLD_OPTION = "compute.ai_ops_confirmation_threshold" +AI_FORECAST_COLUMNS = [ + "forecast_timestamp", + "forecast_value", + "confidence_level", + "prediction_interval_lower_bound", + "prediction_interval_upper_bound", + "ai_forecast_status", +] + + +class FakeGeminiTextGenerator(llm.GeminiTextGenerator): + def __init__(self, prediction): + self.prediction = prediction + + def predict(self, *args, **kwargs): + return self.prediction + + +@pytest.mark.parametrize( + ("func", "kwargs"), + [ + pytest.param( + bigframes.operations.ai.AIAccessor.filter, + {"instruction": None, "model": None}, + id="filter", + ), + pytest.param( + bigframes.operations.ai.AIAccessor.map, + {"instruction": None, "model": None}, + id="map", + ), + pytest.param( + bigframes.operations.ai.AIAccessor.classify, + {"instruction": None, "model": None, "labels": None}, + id="classify", + ), + pytest.param( + bigframes.operations.ai.AIAccessor.join, + {"other": None, "instruction": None, "model": None}, + id="join", + ), + pytest.param( + bigframes.operations.ai.AIAccessor.search, + {"search_column": None, "query": None, "top_k": None, "model": None}, + id="search", + ), + pytest.param( + bigframes.operations.ai.AIAccessor.sim_join, + {"other": None, "left_on": None, "right_on": None, "model": None}, + id="sim_join", + ), + ], +) +def test_experiment_off_raise_error(session, func, kwargs): + df = dataframe.DataFrame( + {"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]}, session=session + ) + + with ( + bigframes.option_context(AI_OP_EXP_OPTION, False), + pytest.raises(NotImplementedError), + ): + func(df.ai, **kwargs) + + +def test_filter(session): + df = dataframe.DataFrame({"col": ["A", "B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame( + { + "answer": [True, False], + "full_response": _create_dummy_full_response(2), + }, + session=session, + ), + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + result = df.ai.filter( + "filter {col}", + model=model, + ).to_pandas() + + pandas.testing.assert_frame_equal( + result, + pd.DataFrame({"col": ["A"]}, dtype=dtypes.STRING_DTYPE), + check_index_type=False, + ) + + +def test_map(session): + df = dataframe.DataFrame({"col": ["A", "B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame( + { + "output": ["true", "false"], + "full_response": _create_dummy_full_response(2), + }, + session=session, + ), + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + result = df.ai.map( + "map {col}", model=model, output_schema={"output": "string"} + ).to_pandas() + + pandas.testing.assert_frame_equal( + result, + pd.DataFrame( + {"col": ["A", "B"], "output": ["true", "false"]}, dtype=dtypes.STRING_DTYPE + ), + check_index_type=False, + ) + + +def test_classify(session): + df = dataframe.DataFrame({"col": ["A", "B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame( + { + "result": ["A", "B"], + "full_response": _create_dummy_full_response(2), + }, + session=session, + ), + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + result = df.ai.classify( + "classify {col}", model=model, labels=["A", "B"] + ).to_pandas() + + pandas.testing.assert_frame_equal( + result, + pd.DataFrame( + {"col": ["A", "B"], "result": ["A", "B"]}, dtype=dtypes.STRING_DTYPE + ), + check_index_type=False, + ) + + +@pytest.mark.parametrize( + "labels", + [ + pytest.param([], id="empty-label"), + pytest.param(["A", "A", "B"], id="duplicate-labels"), + ], +) +def test_classify_invalid_labels_raise_error(session, labels): + df = dataframe.DataFrame({"col": ["A", "B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame( + { + "result": ["A", "B"], + "full_response": _create_dummy_full_response(2), + }, + session=session, + ), + ) + + with ( + bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ), + pytest.raises(ValueError), + ): + df.ai.classify("classify {col}", model=model, labels=labels) + + +def test_join(session): + left_df = dataframe.DataFrame({"col_A": ["A"]}, session=session) + right_df = dataframe.DataFrame({"col_B": ["B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame( + { + "answer": [True], + "full_response": _create_dummy_full_response(1), + }, + session=session, + ), + ) + + with bigframes.option_context( + AI_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + result = left_df.ai.join( + right_df, "join {col_A} and {col_B}", model + ).to_pandas() + + pandas.testing.assert_frame_equal( + result, + pd.DataFrame({"col_A": ["A"], "col_B": ["B"]}, dtype=dtypes.STRING_DTYPE), + check_index_type=False, + ) + + +def test_forecast_default(time_series_df_default_index: dataframe.DataFrame): + df = time_series_df_default_index[time_series_df_default_index["id"] == "1"] + + result = df.ai.forecast(timestamp_column="parsed_date", data_column="total_visits") + + utils.check_pandas_df_schema_and_index( + result, + columns=AI_FORECAST_COLUMNS, + index=10, + ) + + +def test_forecast_w_params(time_series_df_default_index: dataframe.DataFrame): + result = time_series_df_default_index.ai.forecast( + timestamp_column="parsed_date", + data_column="total_visits", + id_columns=["id"], + horizon=20, + confidence_level=0.98, + ) + + utils.check_pandas_df_schema_and_index( + result, + columns=["id"] + AI_FORECAST_COLUMNS, + index=20 * 2, # 20 for each id + ) + + +def _create_dummy_full_response(row_count: int) -> pd.Series: + entry = """{"candidates": [{"avg_logprobs": -0.5}]}""" + + return pd.Series([entry] * row_count) diff --git a/packages/bigframes/tests/system/small/operations/test_semantics.py b/packages/bigframes/tests/system/small/operations/test_semantics.py new file mode 100644 index 000000000000..d8336328483e --- /dev/null +++ b/packages/bigframes/tests/system/small/operations/test_semantics.py @@ -0,0 +1,144 @@ +# 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. + + +# Note that the tests in this files uses fake models for deterministic results. +# Tests that use real LLM models are under system/large/test_semantcs.py + +import pandas as pd +import pandas.testing +import pytest + +import bigframes +from bigframes import dataframe, dtypes +from bigframes.ml import llm + +SEM_OP_EXP_OPTION = "experiments.semantic_operators" +THRESHOLD_OPTION = "compute.semantic_ops_confirmation_threshold" + + +class FakeGeminiTextGenerator(llm.GeminiTextGenerator): + def __init__(self, prediction): + self.prediction = prediction + + def predict(self, *args, **kwargs): + return self.prediction + + +def test_semantics_experiment_off_raise_error(session): + df = dataframe.DataFrame( + {"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]}, session=session + ) + + with ( + bigframes.option_context(SEM_OP_EXP_OPTION, False), + pytest.raises(NotImplementedError), + ): + df.semantics + + +def test_filter(session): + df = dataframe.DataFrame({"col": ["A", "B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame( + {"ml_generate_text_llm_result": ["true", "false"]}, session=session + ), + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + result = df.semantics.filter( + "filter {col}", + model=model, + ).to_pandas() + + pandas.testing.assert_frame_equal( + result, + pd.DataFrame({"col": ["A"]}, dtype=dtypes.STRING_DTYPE), + check_index_type=False, + ) + + +def test_map(session): + df = dataframe.DataFrame({"col": ["A", "B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame( + {"ml_generate_text_llm_result": ["true", "false"]}, session=session + ), + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + result = df.semantics.map( + "map {col}", model=model, output_column="output" + ).to_pandas() + + pandas.testing.assert_frame_equal( + result, + pd.DataFrame( + {"col": ["A", "B"], "output": ["true", "false"]}, dtype=dtypes.STRING_DTYPE + ), + check_index_type=False, + ) + + +def test_join(session): + left_df = dataframe.DataFrame({"col_A": ["A"]}, session=session) + right_df = dataframe.DataFrame({"col_B": ["B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame({"ml_generate_text_llm_result": ["true"]}, session=session), + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + result = left_df.semantics.join( + right_df, "join {col_A} and {col_B}", model + ).to_pandas() + + pandas.testing.assert_frame_equal( + result, + pd.DataFrame({"col_A": ["A"], "col_B": ["B"]}, dtype=dtypes.STRING_DTYPE), + check_index_type=False, + ) + + +def test_top_k(session): + df = dataframe.DataFrame({"col": ["A", "B"]}, session=session) + model = FakeGeminiTextGenerator( + dataframe.DataFrame( + {"ml_generate_text_llm_result": ["Document 1"]}, session=session + ), + ) + + with bigframes.option_context( + SEM_OP_EXP_OPTION, + True, + THRESHOLD_OPTION, + 50, + ): + result = df.semantics.top_k("top k of {col}", model, k=1).to_pandas() + + assert len(result) == 1 diff --git a/packages/bigframes/tests/unit/_config/test_experiment_options.py b/packages/bigframes/tests/unit/_config/test_experiment_options.py index 0d66b2156aba..0e69dfe36d75 100644 --- a/packages/bigframes/tests/unit/_config/test_experiment_options.py +++ b/packages/bigframes/tests/unit/_config/test_experiment_options.py @@ -15,6 +15,37 @@ import pytest import bigframes._config.experiment_options as experiment_options +import bigframes.exceptions as bfe + + +def test_semantic_operators_default_false(): + options = experiment_options.ExperimentOptions() + + assert options.semantic_operators is False + + +def test_semantic_operators_set_true_shows_warning(): + options = experiment_options.ExperimentOptions() + + with pytest.warns(FutureWarning): + options.semantic_operators = True + + assert options.semantic_operators is True + + +def test_ai_operators_default_false(): + options = experiment_options.ExperimentOptions() + + assert options.ai_operators is False + + +def test_ai_operators_set_true_shows_warning(): + options = experiment_options.ExperimentOptions() + + with pytest.warns(bfe.PreviewWarning): + options.ai_operators = True + + assert options.ai_operators is True def test_sql_compiler_default_stable(): diff --git a/packages/bigframes/tests/unit/extensions/core/test_series_mixins.py b/packages/bigframes/tests/unit/extensions/core/test_series_tvf_mixins.py similarity index 64% rename from packages/bigframes/tests/unit/extensions/core/test_series_mixins.py rename to packages/bigframes/tests/unit/extensions/core/test_series_tvf_mixins.py index c6e7e4078b17..9d5a24d2db93 100644 --- a/packages/bigframes/tests/unit/extensions/core/test_series_mixins.py +++ b/packages/bigframes/tests/unit/extensions/core/test_series_tvf_mixins.py @@ -246,139 +246,3 @@ def test_bigframes_ai_generate_table(scalar_types_df: bpd.DataFrame, monkeypatch } result_df.to_pandas.assert_not_called() assert actual_result is result_df - - -def test_ai_embed(monkeypatch): - session = mock.create_autospec(bigframes.session.Session) - bf_series = mock.create_autospec(bpd.Series) - session.read_pandas.return_value = bf_series - - mock_embed = mock.MagicMock() - result_series = mock.create_autospec(bpd.Series) - mock_embed.return_value = result_series - expected_result = mock.create_autospec(pd.Series) - result_series.to_pandas.return_value = expected_result - - monkeypatch.setattr(bigframes.bigquery.ai, "embed", mock_embed) - - series = pd.Series(["hello world"], name="content") - actual_result = series.bigquery.ai.embed( # type: ignore - endpoint="my_endpoint", - model="my_model", - task_type="retrieval_query", - title="my_title", - model_params={"key": "val"}, - connection_id="my_connection", - session=session, - ) - - session.read_pandas.assert_called_once() - mock_embed.assert_called_once_with( - bf_series, - endpoint="my_endpoint", - model="my_model", - task_type="retrieval_query", - title="my_title", - model_params={"key": "val"}, - connection_id="my_connection", - ) - result_series.to_pandas.assert_called_once() - assert actual_result is expected_result - - -def test_bigframes_ai_embed(scalar_types_df: bpd.DataFrame, monkeypatch): - session = mock.create_autospec(bigframes.session.Session) - result_series = mock.create_autospec(bpd.Series) - - mock_embed = mock.MagicMock() - mock_embed.return_value = result_series - - monkeypatch.setattr(bigframes.bigquery.ai, "embed", mock_embed) - - scalar_types_series = scalar_types_df["string_col"] - actual_result = scalar_types_series.bigquery.ai.embed( - endpoint="my_endpoint", - session=session, - ) - - session.read_pandas.assert_not_called() - mock_embed.assert_called_once() - args, kwargs = mock_embed.call_args - assert args[0] is scalar_types_series - assert kwargs == { - "endpoint": "my_endpoint", - "model": None, - "task_type": None, - "title": None, - "model_params": None, - "connection_id": None, - } - result_series.to_pandas.assert_not_called() - assert actual_result is result_series - - -def test_ai_similarity(monkeypatch): - session = mock.create_autospec(bigframes.session.Session) - bf_series = mock.create_autospec(bpd.Series) - session.read_pandas.return_value = bf_series - - mock_similarity = mock.MagicMock() - result_series = mock.create_autospec(bpd.Series) - mock_similarity.return_value = result_series - expected_result = mock.create_autospec(pd.Series) - result_series.to_pandas.return_value = expected_result - - monkeypatch.setattr(bigframes.bigquery.ai, "similarity", mock_similarity) - - series = pd.Series(["apple"], name="content") - actual_result = series.bigquery.ai.similarity( # type: ignore - "banana", - endpoint="my_endpoint", - model="my_model", - model_params={"key": "val"}, - connection_id="my_connection", - session=session, - ) - - session.read_pandas.assert_called_once() - mock_similarity.assert_called_once_with( - bf_series, - "banana", - endpoint="my_endpoint", - model="my_model", - model_params={"key": "val"}, - connection_id="my_connection", - ) - result_series.to_pandas.assert_called_once() - assert actual_result is expected_result - - -def test_bigframes_ai_similarity(scalar_types_df: bpd.DataFrame, monkeypatch): - session = mock.create_autospec(bigframes.session.Session) - result_series = mock.create_autospec(bpd.Series) - - mock_similarity = mock.MagicMock() - mock_similarity.return_value = result_series - - monkeypatch.setattr(bigframes.bigquery.ai, "similarity", mock_similarity) - - scalar_types_series = scalar_types_df["string_col"] - actual_result = scalar_types_series.bigquery.ai.similarity( - "other_text", - endpoint="my_endpoint", - session=session, - ) - - session.read_pandas.assert_not_called() - mock_similarity.assert_called_once() - args, kwargs = mock_similarity.call_args - assert args[0] is scalar_types_series - assert args[1] == "other_text" - assert kwargs == { - "endpoint": "my_endpoint", - "model": None, - "model_params": None, - "connection_id": None, - } - result_series.to_pandas.assert_not_called() - assert actual_result is result_series diff --git a/packages/bigframes/tests/unit/test_dataframe.py b/packages/bigframes/tests/unit/test_dataframe.py index d045bf7c3fca..5f1d16bd19ae 100644 --- a/packages/bigframes/tests/unit/test_dataframe.py +++ b/packages/bigframes/tests/unit/test_dataframe.py @@ -223,3 +223,24 @@ def test_dataframe_drop_columns_returns_new_dataframe(monkeypatch: pytest.Monkey new_dataframe = dataframe.drop(columns=["col1", "col3"]) assert dataframe.columns.to_list() == ["col1", "col2", "col3"] assert new_dataframe.columns.to_list() == ["col2"] + + +def test_dataframe_semantics_property_future_warning( + monkeypatch: pytest.MonkeyPatch, +): + dataframe = mocks.create_dataframe(monkeypatch) + + with ( + bigframes.option_context("experiments.semantic_operators", True), + pytest.warns(FutureWarning), + ): + dataframe.semantics + + +def test_dataframe_ai_property_future_warning( + monkeypatch: pytest.MonkeyPatch, +): + dataframe = mocks.create_dataframe(monkeypatch) + + with pytest.warns(FutureWarning): + dataframe.ai