From 4528eb4fad33d8410646940128345281143ef60c Mon Sep 17 00:00:00 2001
From: Garrett Wu <garrettwu@google.com>
Date: Thu, 22 Aug 2024 22:59:15 +0000
Subject: [PATCH] fix: NullIndex in ML model.predict error

---
 bigframes/ml/core.py                       |  3 ++-
 tests/system/large/ml/test_linear_model.py | 10 +++++++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py
index d570945f16..b41426f7d7 100644
--- a/bigframes/ml/core.py
+++ b/bigframes/ml/core.py
@@ -77,7 +77,8 @@ def _apply_ml_tvf(
 
         result_sql = apply_sql_tvf(input_sql)
         df = self._session.read_gbq(result_sql, index_col=index_col_ids)
-        df.index.names = index_labels
+        if df._has_index:
+            df.index.names = index_labels
         # Restore column labels
         df.rename(
             columns={
diff --git a/tests/system/large/ml/test_linear_model.py b/tests/system/large/ml/test_linear_model.py
index 2f4c07fa28..f593ac2983 100644
--- a/tests/system/large/ml/test_linear_model.py
+++ b/tests/system/large/ml/test_linear_model.py
@@ -111,7 +111,7 @@ def test_linear_regression_customized_params_fit_score(
     assert reloaded_model.learning_rate == 0.2
 
 
-def test_unordered_mode_regression_configure_fit_score(
+def test_unordered_mode_linear_regression_configure_fit_score_predict(
     unordered_session, penguins_table_id, dataset_id
 ):
     model = bigframes.ml.linear_model.LinearRegression()
@@ -154,6 +154,14 @@ def test_unordered_mode_regression_configure_fit_score(
     assert reloaded_model.max_iterations == 20
     assert reloaded_model.tol == 0.01
 
+    pred = reloaded_model.predict(df)
+    utils.check_pandas_df_schema_and_index(
+        pred,
+        columns=("predicted_body_mass_g",),
+        col_exact=False,
+        index=334,
+    )
+
 
 # TODO(garrettwu): add tests for param warm_start. Requires a trained model.