From 9c947442479883468646d9e66e89391b4756ef1d Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 15 Sep 2025 20:45:09 +0000 Subject: [PATCH 1/5] fix: deflake ai_gen_bool multimodel test --- tests/system/large/bigquery/test_ai.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/system/large/bigquery/test_ai.py b/tests/system/large/bigquery/test_ai.py index be0216a526..bed2f288b7 100644 --- a/tests/system/large/bigquery/test_ai.py +++ b/tests/system/large/bigquery/test_ai.py @@ -27,9 +27,5 @@ def test_ai_generate_bool_multi_model(session): "result" ) - pandas.testing.assert_series_equal( - result.to_pandas(), - pd.Series([True, True, False, False, False], name="result"), - check_dtype=False, - check_index=False, - ) + # Make sure all elements are not null + assert len(result) == result.count() From c98c9e75494f367ac479e50fa814e0ab0c498016 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 15 Sep 2025 20:47:17 +0000 Subject: [PATCH 2/5] fix lint --- tests/system/large/bigquery/test_ai.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/system/large/bigquery/test_ai.py b/tests/system/large/bigquery/test_ai.py index bed2f288b7..d13edbbe05 100644 --- a/tests/system/large/bigquery/test_ai.py +++ b/tests/system/large/bigquery/test_ai.py @@ -12,9 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pandas as pd -import pandas.testing - import bigframes.bigquery as bbq From 97fa173f114767e9914f9a8af6f05060d36904cf Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 15 Sep 2025 21:35:32 +0000 Subject: [PATCH 3/5] fix doctest too --- bigframes/bigquery/_operations/ai.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bigframes/bigquery/_operations/ai.py b/bigframes/bigquery/_operations/ai.py index d7ea29322d..2acd2004d7 100644 --- a/bigframes/bigquery/_operations/ai.py +++ b/bigframes/bigquery/_operations/ai.py @@ -47,13 +47,13 @@ def generate_bool( ... "col_1": ["apple", "bear", "pear"], ... "col_2": ["fruit", "animal", "animal"] ... }) - >>> bbq.ai_generate_bool((df["col_1"], " is a ", df["col_2"])) + >>> bbq.ai.generate_bool((df["col_1"], " is a ", df["col_2"])) 0 {'result': True, 'full_response': '{"candidate... 1 {'result': True, 'full_response': '{"candidate... 2 {'result': False, 'full_response': '{"candidat... dtype: struct[pyarrow] - >>> bbq.ai_generate_bool((df["col_1"], " is a ", df["col_2"])).struct.field("result") + >>> bbq.ai.generate_bool((df["col_1"], " is a ", df["col_2"])).struct.field("result") 0 True 1 True 2 False @@ -66,7 +66,7 @@ def generate_bool( ... } ... } ... } - >>> bbq.ai_generate_bool( + >>> bbq.ai.generate_bool( ... (df["col_1"], " is a ", df["col_2"]), ... endpoint="gemini-2.5-pro", ... model_params=model_params, From 1ce7a3228ef2510829d36a7898cf697021bdfbbe Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 15 Sep 2025 22:58:52 +0000 Subject: [PATCH 4/5] consolidates tests under system/small --- tests/system/large/bigquery/__init__.py | 13 ------ tests/system/large/bigquery/test_ai.py | 28 ------------ tests/system/small/bigquery/test_ai.py | 60 ++++++++++++++++++------- 3 files changed, 45 insertions(+), 56 deletions(-) delete mode 100644 tests/system/large/bigquery/__init__.py delete mode 100644 tests/system/large/bigquery/test_ai.py diff --git a/tests/system/large/bigquery/__init__.py b/tests/system/large/bigquery/__init__.py deleted file mode 100644 index 0a2669d7a2..0000000000 --- a/tests/system/large/bigquery/__init__.py +++ /dev/null @@ -1,13 +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. diff --git a/tests/system/large/bigquery/test_ai.py b/tests/system/large/bigquery/test_ai.py deleted file mode 100644 index d13edbbe05..0000000000 --- a/tests/system/large/bigquery/test_ai.py +++ /dev/null @@ -1,28 +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. - -import bigframes.bigquery as bbq - - -def test_ai_generate_bool_multi_model(session): - df = session.from_glob_path( - "gs://bigframes-dev-testing/a_multimodel/images/*", name="image" - ) - - result = bbq.ai.generate_bool((df["image"], " contains an animal")).struct.field( - "result" - ) - - # Make sure all elements are not null - assert len(result) == result.count() diff --git a/tests/system/small/bigquery/test_ai.py b/tests/system/small/bigquery/test_ai.py index 01050ade04..443d4c54a3 100644 --- a/tests/system/small/bigquery/test_ai.py +++ b/tests/system/small/bigquery/test_ai.py @@ -15,9 +15,10 @@ import sys import pandas as pd -import pandas.testing +import pyarrow as pa import pytest +from bigframes import series import bigframes.bigquery as bbq import bigframes.pandas as bpd @@ -27,15 +28,17 @@ def test_ai_generate_bool(session): s2 = bpd.Series(["fruit", "tree"], session=session) prompt = (s1, " is a ", s2) - result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-flash").struct.field( - "result" - ) + result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-flash") - pandas.testing.assert_series_equal( - result.to_pandas(), - pd.Series([True, False], name="result"), - check_dtype=False, - check_index=False, + assert _contains_no_nulls(result) + assert result.dtype == pd.ArrowDtype( + pa.struct( + ( + pa.field("result", pa.bool_()), + pa.field("full_response", pa.string()), + pa.field("status", pa.string()), + ) + ) ) @@ -52,11 +55,38 @@ def test_ai_generate_bool_with_model_params(session): result = bbq.ai.generate_bool( prompt, endpoint="gemini-2.5-flash", model_params=model_params - ).struct.field("result") + ) + + assert _contains_no_nulls(result) + assert result.dtype == pd.ArrowDtype( + pa.struct( + ( + pa.field("result", pa.bool_()), + pa.field("full_response", pa.string()), + pa.field("status", pa.string()), + ) + ) + ) + - pandas.testing.assert_series_equal( - result.to_pandas(), - pd.Series([True, False], name="result"), - check_dtype=False, - check_index=False, +def test_ai_generate_bool_multi_model(session): + df = session.from_glob_path( + "gs://bigframes-dev-testing/a_multimodel/images/*", name="image" ) + + result = bbq.ai.generate_bool((df["image"], " contains an animal")) + + assert _contains_no_nulls(result) + assert result.dtype == pd.ArrowDtype( + pa.struct( + ( + pa.field("result", pa.bool_()), + pa.field("full_response", pa.string()), + pa.field("status", pa.string()), + ) + ) + ) + + +def _contains_no_nulls(s: series.Series) -> bool: + return len(s) == s.count() From c9fec4c922ae2381e8e621b838d425d0bd00c82d Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 15 Sep 2025 23:49:43 +0000 Subject: [PATCH 5/5] fix doctest --- bigframes/bigquery/_operations/ai.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/bigframes/bigquery/_operations/ai.py b/bigframes/bigquery/_operations/ai.py index 2acd2004d7..d82023e4b5 100644 --- a/bigframes/bigquery/_operations/ai.py +++ b/bigframes/bigquery/_operations/ai.py @@ -59,23 +59,6 @@ def generate_bool( 2 False Name: result, dtype: boolean - >>> model_params = { - ... "generation_config": { - ... "thinking_config": { - ... "thinking_budget": 0 - ... } - ... } - ... } - >>> bbq.ai.generate_bool( - ... (df["col_1"], " is a ", df["col_2"]), - ... endpoint="gemini-2.5-pro", - ... model_params=model_params, - ... ).struct.field("result") - 0 True - 1 True - 2 False - Name: result, dtype: boolean - Args: prompt (series.Series | List[str|series.Series] | Tuple[str|series.Series, ...]): A mixture of Series and string literals that specifies the prompt to send to the model.