From 43d4ae092073b28e4cc53f9998f3f9e2aff7f314 Mon Sep 17 00:00:00 2001
From: Shuowei Li <shuowei@google.com>
Date: Tue, 17 Mar 2026 20:03:38 +0000
Subject: [PATCH 1/2] test(ml): Update arima plus score tests to ignore extra
 BQML columns

---
 tests/system/small/ml/test_forecasting.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/system/small/ml/test_forecasting.py b/tests/system/small/ml/test_forecasting.py
index 134f82e96e9..23487983ee3 100644
--- a/tests/system/small/ml/test_forecasting.py
+++ b/tests/system/small/ml/test_forecasting.py
@@ -493,7 +493,7 @@ def test_arima_plus_score(
             dtype="Float64",
         )
     pd.testing.assert_frame_equal(
-        result,
+        result[expected.columns],
         expected,
         rtol=0.1,
         check_index_type=False,
@@ -594,7 +594,7 @@ def test_arima_plus_score_series(
             dtype="Float64",
         )
     pd.testing.assert_frame_equal(
-        result,
+        result[expected.columns],
         expected,
         rtol=0.1,
         check_index_type=False,

From 04ae725dc14700b4560c30c134683aa0dac2852d Mon Sep 17 00:00:00 2001
From: Shuowei Li <shuowei@google.com>
Date: Wed, 18 Mar 2026 07:17:20 +0000
Subject: [PATCH 2/2] fix: use col_exact=False to tolerate new metrics in BQML
 EVALUATE

---
 tests/system/large/ml/test_forecasting.py | 1 +
 tests/system/large/ml/test_llm.py         | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/tests/system/large/ml/test_forecasting.py b/tests/system/large/ml/test_forecasting.py
index 72a0ee469b5..8500ad9d5f1 100644
--- a/tests/system/large/ml/test_forecasting.py
+++ b/tests/system/large/ml/test_forecasting.py
@@ -88,6 +88,7 @@ def test_arima_plus_model_fit_score(
         result,
         columns=expected_columns,
         index=2 if id_col_name else 1,
+        col_exact=False,
     )
 
     # save, load to ensure configuration was kept
diff --git a/tests/system/large/ml/test_llm.py b/tests/system/large/ml/test_llm.py
index 6e2695b1b53..6e51d14a3af 100644
--- a/tests/system/large/ml/test_llm.py
+++ b/tests/system/large/ml/test_llm.py
@@ -198,6 +198,7 @@ def test_llm_gemini_score(llm_fine_tune_df_default_index, model_name):
             "evaluation_status",
         ],
         index=1,
+        col_exact=False,
     )
 
 
@@ -226,6 +227,7 @@ def test_llm_gemini_pro_score_params(llm_fine_tune_df_default_index, model_name)
             "label",
             "evaluation_status",
         ],
+        col_exact=False,
     )