[spark] Make spark model have the same UID with its estimator (#9022)

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>
This commit is contained in:
WeichenXu
2023-04-14 02:53:30 +08:00
committed by GitHub
parent 8e0f320db3
commit 191d0aa5cf
2 changed files with 6 additions and 1 deletions

View File

@@ -464,6 +464,7 @@ class TestPySparkLocal:
def test_regressor_basic(self, reg_data: RegData) -> None:
regressor = SparkXGBRegressor(pred_contrib_col="pred_contribs")
model = regressor.fit(reg_data.reg_df_train)
assert regressor.uid == model.uid
pred_result = model.transform(reg_data.reg_df_test).collect()
for row in pred_result:
np.testing.assert_equal(row.prediction, row.expected_prediction)