def METHOD_NAME(df_iris):
    ds = df_iris
    ds_train, ds_test = ds.ml.train_test_split(test_size=0.2, verbose=False)
    features = ['sepal_length', 'sepal_width', 'petal_length', 'petal_width']
    booster = vaex.ml.xgboost.XGBoostModel(num_boost_round=10,
                                           params=params_multiclass,
                                           features=features,
                                           target='class_')
    booster.fit(ds_train)
    class_predict = booster.predict(ds_test)
    assert np.all(ds_test.class_.values == class_predict)
    ds_train = booster.transform(ds_train)   # this will add the xgboost_prediction column
    state = ds_train.state_get()
    ds_test.state_set(state)
    assert np.all(ds_test.class_.values == ds_test.xgboost_prediction.values)