From 43607734e14cf6772f4e10b7de663f9d199a5955 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Tue, 19 Nov 2024 22:42:53 +0800
Subject: [PATCH] Consistent Python relative imports.

This helps align the code base to use relative import consistently similar to the sklearn
codebase. The testing module still has mixed import types as most of the helper code is in
the init module.
---
 python-package/xgboost/spark/core.py      | 39 +++++++++++------------
 python-package/xgboost/spark/data.py      | 10 +++---
 python-package/xgboost/spark/estimator.py |  3 +-
 3 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/python-package/xgboost/spark/core.py b/python-package/xgboost/spark/core.py
index 3d5618d5d8f4..541c4b6ae99e 100644
--- a/python-package/xgboost/spark/core.py
+++ b/python-package/xgboost/spark/core.py
@@ -59,14 +59,12 @@
 )
 from scipy.special import expit, softmax  # pylint: disable=no-name-in-module
 
-import xgboost
-from xgboost import XGBClassifier
-from xgboost.compat import is_cudf_available, is_cupy_available
-from xgboost.core import Booster, _check_distributed_params
-from xgboost.sklearn import DEFAULT_N_ESTIMATORS, XGBModel, _can_use_qdm
-from xgboost.training import train as worker_train
-
 from .._typing import ArrayLike
+from ..compat import is_cudf_available, is_cupy_available
+from ..config import config_context
+from ..core import Booster, _check_distributed_params, _py_version
+from ..sklearn import DEFAULT_N_ESTIMATORS, XGBClassifier, XGBModel, _can_use_qdm
+from ..training import train as worker_train
 from .data import (
     _read_csr_matrix_from_unwrapped_spark_vec,
     alias,
@@ -765,7 +763,7 @@ def _get_xgb_train_call_args(
         cls, train_params: Dict[str, Any]
     ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
         xgb_train_default_args = _get_default_params_from_func(
-            xgboost.train, _unsupported_train_params
+            worker_train, _unsupported_train_params
         )
         booster_params, kwargs_params = {}, {}
         for key, value in train_params.items():
@@ -1136,17 +1134,18 @@ def _train_booster(
                 _rabit_args = json.loads(messages[0])["rabit_msg"]
 
             evals_result: Dict[str, Any] = {}
-            with CommunicatorContext(context, **_rabit_args):
-                with xgboost.config_context(verbosity=verbosity):
-                    dtrain, dvalid = create_dmatrix_from_partitions(
-                        iterator=pandas_df_iter,
-                        feature_cols=feature_prop.features_cols_names,
-                        dev_ordinal=dev_ordinal,
-                        use_qdm=use_qdm,
-                        kwargs=dmatrix_kwargs,
-                        enable_sparse_data_optim=feature_prop.enable_sparse_data_optim,
-                        has_validation_col=feature_prop.has_validation_col,
-                    )
+            with config_context(verbosity=verbosity), CommunicatorContext(
+                context, **_rabit_args
+            ):
+                dtrain, dvalid = create_dmatrix_from_partitions(
+                    iterator=pandas_df_iter,
+                    feature_cols=feature_prop.features_cols_names,
+                    dev_ordinal=dev_ordinal,
+                    use_qdm=use_qdm,
+                    kwargs=dmatrix_kwargs,
+                    enable_sparse_data_optim=feature_prop.enable_sparse_data_optim,
+                    has_validation_col=feature_prop.has_validation_col,
+                )
                 if dvalid is not None:
                     dval = [(dtrain, "training"), (dvalid, "validation")]
                 else:
@@ -1188,7 +1187,7 @@ def _run_job() -> Tuple[str, str]:
             "\n\tbooster params: %s"
             "\n\ttrain_call_kwargs_params: %s"
             "\n\tdmatrix_kwargs: %s",
-            xgboost._py_version(),
+            _py_version(),
             num_workers,
             booster_params,
             train_call_kwargs_params,
diff --git a/python-package/xgboost/spark/data.py b/python-package/xgboost/spark/data.py
index 0277c4be567a..65bdaf6f378f 100644
--- a/python-package/xgboost/spark/data.py
+++ b/python-package/xgboost/spark/data.py
@@ -7,17 +7,17 @@
 import pandas as pd
 from scipy.sparse import csr_matrix
 
-from xgboost import DataIter, DMatrix, QuantileDMatrix, XGBModel
-from xgboost.compat import concat
-
 from .._typing import ArrayLike
-from .utils import get_logger  # type: ignore
+from ..compat import concat
+from ..core import DataIter, DMatrix, QuantileDMatrix
+from ..sklearn import XGBModel
+from .utils import get_logger
 
 
 def stack_series(series: pd.Series) -> np.ndarray:
     """Stack a series of arrays."""
     array = series.to_numpy(copy=False)
-    array = np.stack(array)
+    array = np.stack(array)  # type: ignore
     return array
 
 
diff --git a/python-package/xgboost/spark/estimator.py b/python-package/xgboost/spark/estimator.py
index 8a4840846ac2..cda18abaa0d5 100644
--- a/python-package/xgboost/spark/estimator.py
+++ b/python-package/xgboost/spark/estimator.py
@@ -12,8 +12,7 @@
 from pyspark.ml.param import Param, Params
 from pyspark.ml.param.shared import HasProbabilityCol, HasRawPredictionCol
 
-from xgboost import XGBClassifier, XGBRanker, XGBRegressor
-
+from ..sklearn import XGBClassifier, XGBRanker, XGBRegressor
 from .core import (  # type: ignore
     _ClassificationModel,
     _SparkXGBEstimator,