From 0182945ac2876bcbbf9cc69629081d66f0441de5 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 20:54:44 +0000 Subject: [PATCH 1/6] chore(bigframes): remove deprecated ai and semantics accessors on df --- packages/bigframes/bigframes/dataframe.py | 16 - packages/bigframes/bigframes/operations/ai.py | 847 ----------- .../bigframes/operations/semantics.py | 1175 -------------- .../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 ----------------- 7 files changed, 4395 deletions(-) delete mode 100644 packages/bigframes/bigframes/operations/ai.py delete mode 100644 packages/bigframes/bigframes/operations/semantics.py delete mode 100644 packages/bigframes/tests/system/large/operations/__init__.py delete mode 100644 packages/bigframes/tests/system/large/operations/conftest.py delete mode 100644 packages/bigframes/tests/system/large/operations/test_ai.py delete mode 100644 packages/bigframes/tests/system/large/operations/test_semantics.py diff --git a/packages/bigframes/bigframes/dataframe.py b/packages/bigframes/bigframes/dataframe.py index 07ba6aba1c0e..e32bc760c532 100644 --- a/packages/bigframes/bigframes/dataframe.py +++ b/packages/bigframes/bigframes/dataframe.py @@ -5287,19 +5287,3 @@ def _throw_if_null_index(self, opname: str): 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/operations/ai.py b/packages/bigframes/bigframes/operations/ai.py deleted file mode 100644 index bba0bf5a8362..000000000000 --- a/packages/bigframes/bigframes/operations/ai.py +++ /dev/null @@ -1,847 +0,0 @@ -# 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 deleted file mode 100644 index 3cad3258e035..000000000000 --- a/packages/bigframes/bigframes/operations/semantics.py +++ /dev/null @@ -1,1175 +0,0 @@ -# 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/tests/system/large/operations/__init__.py b/packages/bigframes/tests/system/large/operations/__init__.py deleted file mode 100644 index 6d5e14bcf4a0..000000000000 --- a/packages/bigframes/tests/system/large/operations/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# 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 deleted file mode 100644 index 0122c860ca74..000000000000 --- a/packages/bigframes/tests/system/large/operations/conftest.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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 deleted file mode 100644 index cd63025cd388..000000000000 --- a/packages/bigframes/tests/system/large/operations/test_ai.py +++ /dev/null @@ -1,963 +0,0 @@ -# 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 deleted file mode 100644 index 147b80f795aa..000000000000 --- a/packages/bigframes/tests/system/large/operations/test_semantics.py +++ /dev/null @@ -1,1348 +0,0 @@ -# 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 From bcb64b5e20e86f8f313b0ebae86e00ad5c2ca734 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 20:58:07 +0000 Subject: [PATCH 2/6] fix lint --- packages/bigframes/bigframes/dataframe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/bigframes/bigframes/dataframe.py b/packages/bigframes/bigframes/dataframe.py index e32bc760c532..0c23f59228d6 100644 --- a/packages/bigframes/bigframes/dataframe.py +++ b/packages/bigframes/bigframes/dataframe.py @@ -5286,4 +5286,3 @@ 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." ) - From d95240bc93a343c5ebc7e6fac6a5d1c53e8c1232 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 21:13:18 +0000 Subject: [PATCH 3/6] delete more stuffs --- .../bigframes/_config/compute_options.py | 20 -- .../bigframes/_config/experiment_options.py | 43 --- packages/bigframes/bigframes/dataframe.py | 2 - .../tests/system/small/operations/test_ai.py | 280 ------------------ .../system/small/operations/test_semantics.py | 144 --------- 5 files changed, 489 deletions(-) delete mode 100644 packages/bigframes/tests/system/small/operations/test_ai.py delete mode 100644 packages/bigframes/tests/system/small/operations/test_semantics.py diff --git a/packages/bigframes/bigframes/_config/compute_options.py b/packages/bigframes/bigframes/_config/compute_options.py index 027566ae075f..2ef1e5b72132 100644 --- a/packages/bigframes/bigframes/_config/compute_options.py +++ b/packages/bigframes/bigframes/_config/compute_options.py @@ -168,26 +168,6 @@ 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 202b47b738c1..8f70c8952f65 100644 --- a/packages/bigframes/bigframes/_config/experiment_options.py +++ b/packages/bigframes/bigframes/_config/experiment_options.py @@ -25,52 +25,9 @@ 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 0c23f59228d6..8d05ba4ddb06 100644 --- a/packages/bigframes/bigframes/dataframe.py +++ b/packages/bigframes/bigframes/dataframe.py @@ -80,9 +80,7 @@ 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 diff --git a/packages/bigframes/tests/system/small/operations/test_ai.py b/packages/bigframes/tests/system/small/operations/test_ai.py deleted file mode 100644 index 7a34e9b8414f..000000000000 --- a/packages/bigframes/tests/system/small/operations/test_ai.py +++ /dev/null @@ -1,280 +0,0 @@ -# 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 deleted file mode 100644 index d8336328483e..000000000000 --- a/packages/bigframes/tests/system/small/operations/test_semantics.py +++ /dev/null @@ -1,144 +0,0 @@ -# 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 From 7857f0cdc85a6e77ef948c42f36c900b84058baa Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 21:22:39 +0000 Subject: [PATCH 4/6] remove unused tests --- .../unit/_config/test_experiment_options.py | 30 ------------------- 1 file changed, 30 deletions(-) diff --git a/packages/bigframes/tests/unit/_config/test_experiment_options.py b/packages/bigframes/tests/unit/_config/test_experiment_options.py index 0e69dfe36d75..8365b77ba9d7 100644 --- a/packages/bigframes/tests/unit/_config/test_experiment_options.py +++ b/packages/bigframes/tests/unit/_config/test_experiment_options.py @@ -18,36 +18,6 @@ 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(): options = experiment_options.ExperimentOptions() From 832cb8a0120da4049024ef532fe17e1406704087 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 21:25:59 +0000 Subject: [PATCH 5/6] fix lint --- packages/bigframes/tests/unit/_config/test_experiment_options.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/bigframes/tests/unit/_config/test_experiment_options.py b/packages/bigframes/tests/unit/_config/test_experiment_options.py index 8365b77ba9d7..0d66b2156aba 100644 --- a/packages/bigframes/tests/unit/_config/test_experiment_options.py +++ b/packages/bigframes/tests/unit/_config/test_experiment_options.py @@ -15,7 +15,6 @@ import pytest import bigframes._config.experiment_options as experiment_options -import bigframes.exceptions as bfe def test_sql_compiler_default_stable(): From c042bacb2ac6b616734a900b5eb27c94334d3009 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 21:36:57 +0000 Subject: [PATCH 6/6] remove more tests --- .../bigframes/tests/unit/test_dataframe.py | 21 ------------------- 1 file changed, 21 deletions(-) diff --git a/packages/bigframes/tests/unit/test_dataframe.py b/packages/bigframes/tests/unit/test_dataframe.py index 5f1d16bd19ae..d045bf7c3fca 100644 --- a/packages/bigframes/tests/unit/test_dataframe.py +++ b/packages/bigframes/tests/unit/test_dataframe.py @@ -223,24 +223,3 @@ 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