Skip to content

Commit 612271d

Browse files
authored
fix: NullIndex in ML model.predict error (#917)
1 parent adfaddc commit 612271d

File tree

2 files changed

+11
-2
lines changed

2 files changed

+11
-2
lines changed

bigframes/ml/core.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,8 @@ def _apply_ml_tvf(
7777

7878
result_sql = apply_sql_tvf(input_sql)
7979
df = self._session.read_gbq(result_sql, index_col=index_col_ids)
80-
df.index.names = index_labels
80+
if df._has_index:
81+
df.index.names = index_labels
8182
# Restore column labels
8283
df.rename(
8384
columns={

tests/system/large/ml/test_linear_model.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ def test_linear_regression_customized_params_fit_score(
111111
assert reloaded_model.learning_rate == 0.2
112112

113113

114-
def test_unordered_mode_regression_configure_fit_score(
114+
def test_unordered_mode_linear_regression_configure_fit_score_predict(
115115
unordered_session, penguins_table_id, dataset_id
116116
):
117117
model = bigframes.ml.linear_model.LinearRegression()
@@ -154,6 +154,14 @@ def test_unordered_mode_regression_configure_fit_score(
154154
assert reloaded_model.max_iterations == 20
155155
assert reloaded_model.tol == 0.01
156156

157+
pred = reloaded_model.predict(df)
158+
utils.check_pandas_df_schema_and_index(
159+
pred,
160+
columns=("predicted_body_mass_g",),
161+
col_exact=False,
162+
index=334,
163+
)
164+
157165

158166
# TODO(garrettwu): add tests for param warm_start. Requires a trained model.
159167

0 commit comments

Comments
 (0)