init

Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

python-package/xgboost/compat.py

@@ -43,6 +43,11 @@ def lazy_isinstance(instance: Any, module: str, name: str) -> bool:
     pandas_concat = None
     PANDAS_INSTALLED = False
 
+
+# cuDF
+CUDF_INSTALLED = importlib.util.find_spec("cudf") is not None
+
+
 # sklearn
 try:
     from sklearn.base import BaseEstimator as XGBModelBase

python-package/xgboost/spark/core.py

@@ -37,6 +37,7 @@
 
 import xgboost
 from xgboost import XGBClassifier, XGBRanker, XGBRegressor
+from xgboost.compat import CUDF_INSTALLED
 
 from .data import (
     _read_csr_matrix_from_unwrapped_spark_vec,
@@ -56,6 +57,7 @@
     HasEnableSparseDataOptim,
     HasFeaturesCols,
     HasQueryIdCol,
+    UseQuantileDMatrix,
 )
 from .utils import (
     CommunicatorContext,
@@ -150,6 +152,7 @@ class _SparkXGBParams(
     HasFeaturesCols,
     HasEnableSparseDataOptim,
     HasQueryIdCol,
+    UseQuantileDMatrix,
 ):
     num_workers = Param(
         Params._dummy(),
@@ -755,7 +758,12 @@ def _fit(self, dataset):
             k: v for k, v in train_call_kwargs_params.items() if v is not None
         }
         dmatrix_kwargs = {k: v for k, v in dmatrix_kwargs.items() if v is not None}
-        use_qdm = booster_params.get("tree_method", None) in ("hist", "gpu_hist")
+
+        if self.isDefined(self.use_quantile_dmatrix):
+            use_qdm = self.getOrDefault(self.use_quantile_dmatrix)
+        else:
+            use_qdm = CUDF_INSTALLED and \
+                      booster_params.get("tree_method", None) in ("hist", "gpu_hist")
 
         def _train_booster(pandas_df_iter):
             """Takes in an RDD partition and outputs a booster for that partition after

python-package/xgboost/spark/data.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix
-from xgboost.compat import concat
+from xgboost.compat import CUDF_INSTALLED, concat
 
 from xgboost import DataIter, DMatrix, QuantileDMatrix
 
@@ -81,7 +81,7 @@ def _fetch(self, data: Optional[Sequence[pd.DataFrame]]) -> Optional[pd.DataFram
         if not data:
             return None
 
-        if self._device_id is not None:
+        if self._device_id is not None and CUDF_INSTALLED:
             import cudf  # pylint: disable=import-error
             import cupy as cp  # pylint: disable=import-error
 

python-package/xgboost/spark/params.py

@@ -85,3 +85,19 @@ class HasQueryIdCol(Params):
         "query id column name",
         typeConverter=TypeConverters.toString,
     )
+
+
+class UseQuantileDMatrix(Params):
+    """
+    Mixin for param use_quantile_dmatrix:
+    """
+    use_quantile_dmatrix = Param(
+        Params._dummy(),
+        "use_quantile_dmatrix",
+        "This stores the boolean config of constructing quantile DMatrix instead of plain "
+        "DMatrix as xgboost training input. By default, if 'tree_method' param is 'hist' or "
+        "'gpu_hist' and 'cuDF' package is installed, the config is ON, otherwise the config "
+        "is OFF. Note that if you do not install 'cuDF' package, turning on this config "
+        "might result in performance degradation.",
+        typeConverter=TypeConverters.toBoolean,
+    )