From 02a739e63c9b6359295166b8265d3aa7583183ac Mon Sep 17 00:00:00 2001
From: WeichenXu <weichen.xu@databricks.com>
Date: Sun, 27 Nov 2022 14:58:54 +0800
Subject: [PATCH] [pyspark] Make QDM optional based on cuDF check (#8471)

---
 python-package/xgboost/compat.py     | 17 +++++++++++++++++
 python-package/xgboost/spark/core.py | 13 ++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/python-package/xgboost/compat.py b/python-package/xgboost/compat.py
index 275b6621064d..fab734a01361 100644
--- a/python-package/xgboost/compat.py
+++ b/python-package/xgboost/compat.py
@@ -43,6 +43,7 @@ def lazy_isinstance(instance: Any, module: str, name: str) -> bool:
     pandas_concat = None
     PANDAS_INSTALLED = False
 
+
 # sklearn
 try:
     from sklearn.base import BaseEstimator as XGBModelBase
@@ -72,6 +73,22 @@ def lazy_isinstance(instance: Any, module: str, name: str) -> bool:
     XGBStratifiedKFold = None
 
 
+_logger = logging.getLogger(__name__)
+
+
+def is_cudf_available() -> bool:
+    """Check cuDF package available or not"""
+    if importlib.util.find_spec("cudf") is None:
+        return False
+    try:
+        import cudf
+
+        return True
+    except ImportError:
+        _logger.exception("Importing cuDF failed, use DMatrix instead of QDM")
+        return False
+
+
 class XGBoostLabelEncoder(LabelEncoder):
     """Label encoder with JSON serialization methods."""
 
diff --git a/python-package/xgboost/spark/core.py b/python-package/xgboost/spark/core.py
index 1e544c34f228..56f63ac0f277 100644
--- a/python-package/xgboost/spark/core.py
+++ b/python-package/xgboost/spark/core.py
@@ -32,6 +32,7 @@
     ShortType,
 )
 from scipy.special import expit, softmax  # pylint: disable=no-name-in-module
+from xgboost.compat import is_cudf_available
 from xgboost.core import Booster
 from xgboost.training import train as worker_train
 
@@ -755,7 +756,8 @@ 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")
+
+        use_hist = 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
@@ -769,6 +771,15 @@ def _train_booster(pandas_df_iter):
 
             gpu_id = None
 
+            # If cuDF is not installed, then using DMatrix instead of QDM,
+            # because without cuDF, DMatrix performs better than QDM.
+            # Note: Checking `is_cudf_available` in spark worker side because
+            # spark worker might has different python environment with driver side.
+            if use_gpu:
+                use_qdm = use_hist and is_cudf_available()
+            else:
+                use_qdm = use_hist
+
             if use_qdm and (booster_params.get("max_bin", None) is not None):
                 dmatrix_kwargs["max_bin"] = booster_params["max_bin"]