From acdfda6d2466516c689ac77fec7280db9e220114 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Mon, 1 Jun 2026 23:44:42 +0000 Subject: [PATCH 1/4] test(bigframes): Disable ibis fallback for system tests --- packages/bigframes/tests/system/conftest.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/packages/bigframes/tests/system/conftest.py b/packages/bigframes/tests/system/conftest.py index 1adcb051c470..f6fbdd0c510d 100644 --- a/packages/bigframes/tests/system/conftest.py +++ b/packages/bigframes/tests/system/conftest.py @@ -1325,6 +1325,14 @@ def usa_names_grouped_table( return session.bqclient.get_table(table_id) +@pytest.fixture(scope="session", autouse=True) +def use_sqlglot_compiler(): + original_setting = bigframes.options.experiments.sql_compiler + bigframes.options.experiments.sql_compiler = "experimental" + yield + bigframes.options.experiments.sql_compiler = original_setting + + @pytest.fixture() def restore_sampling_settings(): enable_downsampling = bigframes.options.sampling.enable_downsampling From 495d107c232c58532894e447cc37d988a0ed48ea Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Tue, 2 Jun 2026 01:55:48 +0000 Subject: [PATCH 2/4] fix ai prompt construction in sqlglot --- .../bigframes/core/compile/sqlglot/expressions/ai_ops.py | 3 ++- packages/bigframes/tests/system/small/bigquery/test_ai.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/expressions/ai_ops.py b/packages/bigframes/bigframes/core/compile/sqlglot/expressions/ai_ops.py index 12a6b9859a2a..d092f662f0f8 100644 --- a/packages/bigframes/bigframes/core/compile/sqlglot/expressions/ai_ops.py +++ b/packages/bigframes/bigframes/core/compile/sqlglot/expressions/ai_ops.py @@ -111,7 +111,8 @@ def _construct_prompt( else: prompt.append(sge.Literal.string(elem)) - return sge.Kwarg(this=param_name, expression=sge.Tuple(expressions=prompt)) + # Need Struct rather than tuple syntax, as tuple syntax is ambiguous for single arg + return sge.Kwarg(this=param_name, expression=sge.Struct(expressions=prompt)) def _construct_named_args(op: ops.ScalarOp) -> list[sge.Kwarg]: diff --git a/packages/bigframes/tests/system/small/bigquery/test_ai.py b/packages/bigframes/tests/system/small/bigquery/test_ai.py index f3c94edd1969..ef1eba975724 100644 --- a/packages/bigframes/tests/system/small/bigquery/test_ai.py +++ b/packages/bigframes/tests/system/small/bigquery/test_ai.py @@ -401,7 +401,9 @@ def test_ai_classify_multi_model(session, bq_connection): ) image_runtime = bbq.obj.get_access_url(df["image"], mode="R") - result = bbq.ai.classify(image_runtime, ["photo", "cartoon"]) + result = bbq.ai.classify(image_runtime, ["photo", "cartoon"]).to_pandas( + allow_large_results=True + ) assert len(result) == len(df) assert result.dtype == dtypes.STRING_DTYPE From bd4cf05366a6264dd61704cc64664738f2be6010 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Tue, 2 Jun 2026 17:20:15 +0000 Subject: [PATCH 3/4] update snapshots --- .../snapshots/test_ai_ops/test_ai_classify/None/out.sql | 2 +- .../bigframes-dev.us.bigframes-default-connection/out.sql | 2 +- .../test_ai_classify_multi_with_list_examples/out.sql | 2 +- .../test_ai_ops/test_ai_classify_with_output_mode/out.sql | 2 +- .../snapshots/test_ai_ops/test_ai_classify_with_params/out.sql | 2 +- .../expressions/snapshots/test_ai_ops/test_ai_generate/out.sql | 2 +- .../snapshots/test_ai_ops/test_ai_generate_bool/out.sql | 2 +- .../test_ai_generate_bool_with_connection_id/out.sql | 2 +- .../test_ai_ops/test_ai_generate_bool_with_model_param/out.sql | 2 +- .../snapshots/test_ai_ops/test_ai_generate_double/out.sql | 2 +- .../test_ai_generate_double_with_connection_id/out.sql | 2 +- .../test_ai_generate_double_with_model_param/out.sql | 2 +- .../snapshots/test_ai_ops/test_ai_generate_int/out.sql | 2 +- .../test_ai_ops/test_ai_generate_int_with_connection_id/out.sql | 2 +- .../test_ai_ops/test_ai_generate_int_with_model_param/out.sql | 2 +- .../test_ai_ops/test_ai_generate_with_connection_id/out.sql | 2 +- .../test_ai_ops/test_ai_generate_with_model_param/out.sql | 2 +- .../test_ai_ops/test_ai_generate_with_output_schema/out.sql | 2 +- .../expressions/snapshots/test_ai_ops/test_ai_if/None/out.sql | 2 +- .../bigframes-dev.us.bigframes-default-connection/out.sql | 2 +- .../snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql | 2 +- .../snapshots/test_ai_ops/test_ai_score/None/out.sql | 2 +- .../bigframes-dev.us.bigframes-default-connection/out.sql | 2 +- .../test_ai_score_with_endpoint_and_max_error_ratio/out.sql | 2 +- 24 files changed, 24 insertions(+), 24 deletions(-) diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/None/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/None/out.sql index 6771527318fa..fc29d96cc1aa 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/None/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/None/out.sql @@ -1,3 +1,3 @@ SELECT - AI.CLASSIFY(input => (`string_col`), categories => ['greeting', 'rejection']) AS `result` + AI.CLASSIFY(input => STRUCT(`string_col`), categories => ['greeting', 'rejection']) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql index 63c31d94566d..969b946725bc 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql @@ -1,6 +1,6 @@ SELECT AI.CLASSIFY( - input => (`string_col`), + input => STRUCT(`string_col`), categories => ['greeting', 'rejection'], connection_id => 'bigframes-dev.us.bigframes-default-connection' ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_multi_with_list_examples/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_multi_with_list_examples/out.sql index a4a7f783da97..74078e986064 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_multi_with_list_examples/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_multi_with_list_examples/out.sql @@ -1,6 +1,6 @@ SELECT AI.CLASSIFY( - input => (`string_col`), + input => STRUCT(`string_col`), categories => ['greeting', 'rejection'], examples => [('hi', ['greeting', 'positive']), ('bye', ['rejection', 'negative'])], output_mode => 'multi' diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_output_mode/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_output_mode/out.sql index fb3c6af8b0b0..08d7476d77f4 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_output_mode/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_output_mode/out.sql @@ -1,6 +1,6 @@ SELECT AI.CLASSIFY( - input => (`string_col`), + input => STRUCT(`string_col`), categories => ['greeting', 'rejection'], output_mode => 'multi' ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql index 982b747f8927..30542740a2dc 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql @@ -1,6 +1,6 @@ SELECT AI.CLASSIFY( - input => (`string_col`), + input => STRUCT(`string_col`), categories => ['greeting', 'rejection'], examples => [('hi', 'greeting'), ('bye', 'rejection')], endpoint => 'gemini-2.5-flash', diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql index 9593347238f8..622782fa7d65 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), endpoint => 'gemini-2.5-flash', request_type => 'SHARED' ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql index aebccad12217..a71bce037a5d 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_BOOL( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), endpoint => 'gemini-2.5-flash' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql index 8f501a2cc292..db1ec378aaf9 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_BOOL( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', endpoint => 'gemini-2.5-flash' ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_model_param/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_model_param/out.sql index 985f5bb255d7..76af8833e639 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_model_param/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_model_param/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_BOOL( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), model_params => JSON '{}' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql index 3aed8986e179..1cef75687988 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_DOUBLE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), endpoint => 'gemini-2.5-flash' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql index 19b8c18eec14..d0088721e386 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_DOUBLE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', endpoint => 'gemini-2.5-flash' ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_model_param/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_model_param/out.sql index 854acc386739..2b50e05b7fe9 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_model_param/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_model_param/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_DOUBLE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), model_params => JSON '{}' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int/out.sql index 1ea5d0355cc9..9ef143c8b9e4 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_INT( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), endpoint => 'gemini-2.5-flash' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_connection_id/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_connection_id/out.sql index b99a8e9a207e..3fa3e8cc05e1 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_connection_id/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_connection_id/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_INT( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', endpoint => 'gemini-2.5-flash' ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_model_param/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_model_param/out.sql index fb3c9c001013..18adea8a0622 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_model_param/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_model_param/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_INT( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), model_params => JSON '{}' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_connection_id/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_connection_id/out.sql index b122d97b0617..14604cfc8dfd 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_connection_id/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_connection_id/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', endpoint => 'gemini-2.5-flash' ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_model_param/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_model_param/out.sql index 9d818b8c0cc9..090a42d889f5 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_model_param/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_model_param/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), model_params => JSON '{}' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_output_schema/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_output_schema/out.sql index 44abe7085c4e..31c179e7b01a 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_output_schema/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_output_schema/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), endpoint => 'gemini-2.5-flash', output_schema => 'x INT64, y FLOAT64' ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/None/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/None/out.sql index 7696a12c5893..59cf1c02a355 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/None/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/None/out.sql @@ -1,6 +1,6 @@ SELECT AI.IF( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), optimization_mode => 'MINIMIZE_COST', max_error_ratio => 0.5 ) AS `result` diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/bigframes-dev.us.bigframes-default-connection/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/bigframes-dev.us.bigframes-default-connection/out.sql index dc8707487b54..0f26ab3c6ea6 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/bigframes-dev.us.bigframes-default-connection/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/bigframes-dev.us.bigframes-default-connection/out.sql @@ -1,6 +1,6 @@ SELECT AI.IF( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', optimization_mode => 'MINIMIZE_COST', max_error_ratio => 0.5 diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql index 5074584bd72d..4dd910528a41 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql @@ -1,6 +1,6 @@ SELECT AI.IF( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), endpoint => 'gemini-2.5-flash' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/None/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/None/out.sql index 6a16276734ee..37590eec4f0f 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/None/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/None/out.sql @@ -1,3 +1,3 @@ SELECT - AI.SCORE(prompt => (`string_col`, ' is the same as ', `string_col`)) AS `result` + AI.SCORE(prompt => STRUCT(`string_col`, ' is the same as ', `string_col`)) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/bigframes-dev.us.bigframes-default-connection/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/bigframes-dev.us.bigframes-default-connection/out.sql index 92de7cdcdc65..696c7e9f3183 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/bigframes-dev.us.bigframes-default-connection/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/bigframes-dev.us.bigframes-default-connection/out.sql @@ -1,6 +1,6 @@ SELECT AI.SCORE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score_with_endpoint_and_max_error_ratio/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score_with_endpoint_and_max_error_ratio/out.sql index d65590d0b66d..a802e5a396bf 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score_with_endpoint_and_max_error_ratio/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score_with_endpoint_and_max_error_ratio/out.sql @@ -1,6 +1,6 @@ SELECT AI.SCORE( - prompt => (`string_col`, ' is the same as ', `string_col`), + prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), endpoint => 'gemini-2.5-flash', max_error_ratio => 0.5 ) AS `result` From abeefdaf8e85acd1459fbb27ab0b492cb4c9baee Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Tue, 2 Jun 2026 20:38:53 +0000 Subject: [PATCH 4/4] revert test change --- packages/bigframes/tests/system/small/bigquery/test_ai.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/bigframes/tests/system/small/bigquery/test_ai.py b/packages/bigframes/tests/system/small/bigquery/test_ai.py index ef1eba975724..f3c94edd1969 100644 --- a/packages/bigframes/tests/system/small/bigquery/test_ai.py +++ b/packages/bigframes/tests/system/small/bigquery/test_ai.py @@ -401,9 +401,7 @@ def test_ai_classify_multi_model(session, bq_connection): ) image_runtime = bbq.obj.get_access_url(df["image"], mode="R") - result = bbq.ai.classify(image_runtime, ["photo", "cartoon"]).to_pandas( - allow_large_results=True - ) + result = bbq.ai.classify(image_runtime, ["photo", "cartoon"]) assert len(result) == len(df) assert result.dtype == dtypes.STRING_DTYPE