add max_bin support

dmlc · Aug 2, 2022 · 2096997 · 2096997
1 parent a1766c0
commit 2096997
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/python-package/xgboost/spark/core.py b/python-package/xgboost/spark/core.py
@@ -611,6 +611,13 @@ def _train_booster(pandas_df_iter):
                 gpu_id = context.partitionId() if is_local else _get_gpu_id(context)
                 booster_params["gpu_id"] = gpu_id
 
+                # max_bin is needed for qdm
+                if (
+                    features_cols_names is not None
+                    and booster_params.get("max_bin", None) is not None
+                ):
+                    dmatrix_kwargs["max_bin"] = booster_params["max_bin"]
+
             _rabit_args = ""
             if context.partitionId() == 0:
                 _rabit_args = str(_get_rabit_args(context, num_workers))