PySpark XGBoost integration

.github/workflows/main.yml

@@ -141,7 +141,7 @@ jobs:
     - name: Install Python packages
       run: |
         python -m pip install wheel setuptools
-        python -m pip install pylint cpplint numpy scipy scikit-learn
+        python -m pip install pylint cpplint numpy scipy scikit-learn pyspark pandas cloudpickle
     - name: Run lint
       run: |
         make lint

.github/workflows/python_tests.yml

@@ -92,6 +92,7 @@ jobs:
   python-tests-on-macos:
     name: Test XGBoost Python package on ${{ matrix.config.os }}
     runs-on: ${{ matrix.config.os }}
+    timeout-minutes: 90
     strategy:
       matrix:
         config:

python-package/setup.py

@@ -351,7 +351,8 @@ def run(self) -> None:
               'scikit-learn': ['scikit-learn'],
               'dask': ['dask', 'pandas', 'distributed'],
               'datatable': ['datatable'],
-              'plotting': ['graphviz', 'matplotlib']
+              'plotting': ['graphviz', 'matplotlib'],
+              "pyspark": ["pyspark", "scikit-learn", "cloudpickle"],
           },
           maintainer='Hyunsu Cho',
           maintainer_email='chohyu01@cs.washington.edu',

python-package/xgboost/spark/__init__.py

@@ -0,0 +1,22 @@
+# type: ignore
+"""PySpark XGBoost integration interface
+"""
+
+try:
+    import pyspark
+except ImportError as e:
+    raise ImportError("pyspark package needs to be installed to use this module") from e
+
+from .estimator import (
+    SparkXGBClassifier,
+    SparkXGBClassifierModel,
+    SparkXGBRegressor,
+    SparkXGBRegressorModel,
+)
+
+__all__ = [
+    "SparkXGBClassifier",
+    "SparkXGBClassifierModel",
+    "SparkXGBRegressor",
+    "SparkXGBRegressorModel",
+]