"""XGBoost pyspark integration submodule for core code."""

import base64

# pylint: disable=fixme, protected-access, no-member
# pylint: disable=too-many-lines, too-many-branches
import json
import logging
import os
from collections import namedtuple
from dataclasses import asdict
from typing import (
    Any,
    Callable,
    Dict,
    Iterator,
    List,
    Optional,
    Tuple,
    Type,
    Union,
    cast,
)

import numpy as np
import pandas as pd
from pyspark import RDD, SparkConf, SparkContext, cloudpickle
from pyspark.ml import Estimator, Model
from pyspark.ml.functions import array_to_vector, vector_to_array
from pyspark.ml.linalg import VectorUDT
from pyspark.ml.param import Param, Params, TypeConverters
from pyspark.ml.param.shared import (
    HasFeaturesCol,
    HasLabelCol,
    HasPredictionCol,
    HasProbabilityCol,
    HasRawPredictionCol,
    HasValidationIndicatorCol,
    HasWeightCol,
)
from pyspark.ml.util import (
    DefaultParamsReader,
    DefaultParamsWriter,
    MLReadable,
    MLReader,
    MLWritable,
    MLWriter,
)
from pyspark.resource import ResourceProfileBuilder, TaskResourceRequests
from pyspark.sql import Column, DataFrame
from pyspark.sql.functions import col, countDistinct, pandas_udf, rand, struct
from pyspark.sql.types import (
    ArrayType,
    BooleanType,
    DoubleType,
    FloatType,
    IntegerType,
    IntegralType,
    LongType,
    ShortType,
)
from scipy.special import expit, softmax  # pylint: disable=no-name-in-module

from .._typing import ArrayLike
from ..collective import Config
from ..compat import import_cupy, is_cudf_available, is_cupy_available
from ..config import config_context, get_config
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,
    create_dmatrix_from_partitions,
    pred_contribs,
    stack_series,
)
from .params import (
    HasArbitraryParamsDict,
    HasBaseMarginCol,
    HasContribPredictionCol,
    HasEnableSparseDataOptim,
    HasFeaturesCols,
    HasQueryIdCol,
)
from .summary import XGBoostTrainingSummary
from .utils import (
    CommunicatorContext,
    _get_default_params_from_func,
    _get_gpu_id,
    _get_host_ip,
    _get_max_num_concurrent_tasks,
    _get_rabit_args,
    _get_spark_session,
    _is_local,
    _is_standalone_or_localcluster,
    deserialize_booster,
    deserialize_xgb_model,
    get_class_name,
    get_logger,
    get_logger_level,
    serialize_booster,
    use_cuda,
)

# Put pyspark specific params here, they won't be passed to XGBoost.
# like `validationIndicatorCol`, `base_margin_col`
_pyspark_specific_params = [
    "featuresCol",
    "labelCol",
    "weightCol",
    "rawPredictionCol",
    "predictionCol",
    "probabilityCol",
    "validationIndicatorCol",
    "base_margin_col",
    "arbitrary_params_dict",
    "force_repartition",
    "num_workers",
    "feature_names",
    "features_cols",
    "enable_sparse_data_optim",
    "qid_col",
    "repartition_random_shuffle",
    "pred_contrib_col",
    "launch_tracker_on_driver",
    "coll_cfg",
]

_non_booster_params = ["missing", "n_estimators", "feature_types", "feature_weights"]

_pyspark_param_alias_map = {
    "features_col": "featuresCol",
    "label_col": "labelCol",
    "weight_col": "weightCol",
    "raw_prediction_col": "rawPredictionCol",
    "prediction_col": "predictionCol",
    "probability_col": "probabilityCol",
    "validation_indicator_col": "validationIndicatorCol",
}

_inverse_pyspark_param_alias_map = {v: k for k, v in _pyspark_param_alias_map.items()}

_unsupported_xgb_params = [
    "enable_categorical",  # Use feature_types param to specify categorical feature instead
    "n_jobs",  # Do not allow user to set it, will use `spark.task.cpus` value instead.
    "nthread",  # Ditto
]

_unsupported_fit_params = {
    "sample_weight",  # Supported by spark param weightCol
    "eval_set",  # Supported by spark param validation_indicator_col
    "sample_weight_eval_set",  # Supported by spark param weight_col + validation_indicator_col
    "base_margin",  # Supported by spark param base_margin_col
    "base_margin_eval_set",  # Supported by spark param base_margin_col + validation_indicator_col
    "group",  # Use spark param `qid_col` instead
    "qid",  # Use spark param `qid_col` instead
    "eval_group",  # Use spark param `qid_col` instead
    "eval_qid",  # Use spark param `qid_col` instead
}

_unsupported_train_params = {
    "evals",  # Supported by spark param validation_indicator_col
    "evals_result",  # Won't support yet+
}

_unsupported_predict_params = {
    # for classification, we can use rawPrediction as margin
    "output_margin",
    "validate_features",  # TODO
    "base_margin",  # Use pyspark base_margin_col param instead.
}

# TODO: supply hint message for all other unsupported params.
_unsupported_params_hint_message = {
    "enable_categorical": "`xgboost.spark` estimators do not have 'enable_categorical' param, "
    "but you can set `feature_types` param and mark categorical features with 'c' string."
}

# Global prediction names
Pred = namedtuple(
    "Pred", ("prediction", "raw_prediction", "probability", "pred_contrib")
)
pred = Pred("prediction", "rawPrediction", "probability", "predContrib")

_INIT_BOOSTER_SAVE_PATH = "init_booster.json"

_LOG_TAG = "XGBoost-PySpark"


class _SparkXGBParams(
    HasFeaturesCol,
    HasLabelCol,
    HasWeightCol,
    HasPredictionCol,
    HasValidationIndicatorCol,
    HasArbitraryParamsDict,
    HasBaseMarginCol,
    HasFeaturesCols,
    HasEnableSparseDataOptim,
    HasQueryIdCol,
    HasContribPredictionCol,
):
    num_workers = Param(
        Params._dummy(),
        "num_workers",
        "The number of XGBoost workers. Each XGBoost worker corresponds to one spark task.",
        TypeConverters.toInt,
    )
    device = Param(
        Params._dummy(),
        "device",
        (
            "The device type for XGBoost executors. Available options are `cpu`,`cuda`"
            " and `gpu`. Set `device` to `cuda` or `gpu` if the executors are running "
            "on GPU instances. Currently, only one GPU per task is supported."
        ),
        TypeConverters.toString,
    )
    force_repartition = Param(
        Params._dummy(),
        "force_repartition",
        "A boolean variable. Set force_repartition=true if you "
        + "want to force the input dataset to be repartitioned before XGBoost training."
        + "Note: The auto repartitioning judgement is not fully accurate, so it is recommended"
        + "to have force_repartition be True.",
        TypeConverters.toBoolean,
    )
    repartition_random_shuffle = Param(
        Params._dummy(),
        "repartition_random_shuffle",
        "A boolean variable. Set repartition_random_shuffle=true if you want to random shuffle "
        "dataset when repartitioning is required. By default is True.",
        TypeConverters.toBoolean,
    )
    feature_names = Param(
        Params._dummy(),
        "feature_names",
        "A list of str to specify feature names.",
        TypeConverters.toList,
    )
    launch_tracker_on_driver = Param(
        Params._dummy(),
        "launch_tracker_on_driver",
        "A boolean variable. Set launch_tracker_on_driver to true if you want the tracker to be "
        "launched on the driver side; otherwise, it will be launched on the executor side.",
        TypeConverters.toBoolean,
    )
    coll_cfg = Param(
        Params._dummy(),
        "coll_cfg",
        "xgboost.collective.Config. The collective configuration.",
        TypeConverters.identity,
    )

    def set_coll_cfg(self, value: Config) -> "_SparkXGBParams":
        """Set collective configuration"""
        assert isinstance(value, Config)
        self.set(self.coll_cfg, value)
        return self

    def set_device(self, value: str) -> "_SparkXGBParams":
        """Set device, optional value: cpu, cuda, gpu"""
        _check_distributed_params({"device": value})
        assert value in ("cpu", "cuda", "gpu")
        self.set(self.device, value)
        return self

    @classmethod
    def _xgb_cls(cls) -> Type[XGBModel]:
        """
        Subclasses should override this method and
        returns an xgboost.XGBModel subclass
        """
        raise NotImplementedError()

    # Parameters for xgboost.XGBModel()
    @classmethod
    def _get_xgb_params_default(cls) -> Dict[str, Any]:
        """Get the xgboost.sklearn.XGBModel default parameters and filter out some"""
        xgb_model_default = cls._xgb_cls()()
        params_dict = xgb_model_default.get_params()
        filtered_params_dict = {
            k: params_dict[k] for k in params_dict if k not in _unsupported_xgb_params
        }
        filtered_params_dict["n_estimators"] = DEFAULT_N_ESTIMATORS
        return filtered_params_dict

    def _set_xgb_params_default(self) -> None:
        """Set xgboost parameters into spark parameters"""
        filtered_params_dict = self._get_xgb_params_default()
        self._setDefault(**filtered_params_dict)

    def _gen_xgb_params_dict(
        self, gen_xgb_sklearn_estimator_param: bool = False
    ) -> Dict[str, Any]:
        """Generate the xgboost parameters which will be passed into xgboost library"""
        xgb_params = {}
        non_xgb_params = (
            set(_pyspark_specific_params)
            | self._get_fit_params_default().keys()
            | self._get_predict_params_default().keys()
        )
        if not gen_xgb_sklearn_estimator_param:
            non_xgb_params |= set(_non_booster_params)
        for param in self.extractParamMap():
            if param.name not in non_xgb_params:
                xgb_params[param.name] = self.getOrDefault(param)

        arbitrary_params_dict = self.getOrDefault(
            self.getParam("arbitrary_params_dict")
        )
        xgb_params.update(arbitrary_params_dict)
        return xgb_params

    # Parameters for xgboost.XGBModel().fit()
    @classmethod
    def _get_fit_params_default(cls) -> Dict[str, Any]:
        """Get the xgboost.XGBModel().fit() parameters"""
        fit_params = _get_default_params_from_func(
            cls._xgb_cls().fit, _unsupported_fit_params
        )
        return fit_params

    def _set_fit_params_default(self) -> None:
        """Get the xgboost.XGBModel().fit() parameters and set them to spark parameters"""
        filtered_params_dict = self._get_fit_params_default()
        self._setDefault(**filtered_params_dict)

    def _gen_fit_params_dict(self) -> Dict[str, Any]:
        """Generate the fit parameters which will be passed into fit function"""
        fit_params_keys = self._get_fit_params_default().keys()
        fit_params = {}
        for param in self.extractParamMap():
            if param.name in fit_params_keys:
                fit_params[param.name] = self.getOrDefault(param)
        return fit_params

    @classmethod
    def _get_predict_params_default(cls) -> Dict[str, Any]:
        """Get the parameters from xgboost.XGBModel().predict()"""
        predict_params = _get_default_params_from_func(
            cls._xgb_cls().predict, _unsupported_predict_params
        )
        return predict_params

    def _set_predict_params_default(self) -> None:
        """Get the parameters from xgboost.XGBModel().predict() and
        set them into spark parameters"""
        filtered_params_dict = self._get_predict_params_default()
        self._setDefault(**filtered_params_dict)

    def _gen_predict_params_dict(self) -> Dict[str, Any]:
        """Generate predict parameters which will be passed into xgboost.XGBModel().predict()"""
        predict_params_keys = self._get_predict_params_default().keys()
        predict_params = {}
        for param in self.extractParamMap():
            if param.name in predict_params_keys:
                predict_params[param.name] = self.getOrDefault(param)
        return predict_params

    def _validate_gpu_params(
        self, spark_version: str, conf: SparkConf, is_local: bool = False
    ) -> None:
        """Validate the gpu parameters and gpu configurations"""

        if self._run_on_gpu():
            if is_local:
                # Supporting GPU training in Spark local mode is just for debugging
                # purposes, so it's okay for printing the below warning instead of
                # checking the real gpu numbers and raising the exception.
                get_logger(self.__class__.__name__).warning(
                    "You have enabled GPU in spark local mode. Please make sure your"
                    " local node has at least %d GPUs",
                    self.getOrDefault(self.num_workers),
                )
            else:
                executor_gpus = conf.get("spark.executor.resource.gpu.amount")
                if executor_gpus is None:
                    raise ValueError(
                        "The `spark.executor.resource.gpu.amount` is required for training"
                        " on GPU."
                    )
                gpu_per_task = conf.get("spark.task.resource.gpu.amount")
                if gpu_per_task is not None and float(gpu_per_task) > 1.0:
                    get_logger(self.__class__.__name__).warning(
                        "The configuration assigns %s GPUs to each Spark task, but each "
                        "XGBoost training task only utilizes 1 GPU, which will lead to "
                        "unnecessary GPU waste",
                        gpu_per_task,
                    )
                # For 3.5.1+, Spark supports task stage-level scheduling for
                #                          Yarn/K8s/Standalone/Local cluster
                # From 3.4.0 ~ 3.5.0, Spark only supports task stage-level scheduing for
                #                           Standalone/Local cluster
                # For spark below 3.4.0, Task stage-level scheduling is not supported.
                #
                # With stage-level scheduling, spark.task.resource.gpu.amount is not required
                # to be set explicitly. Or else, spark.task.resource.gpu.amount is a must-have and
                # must be set to 1.0
                if spark_version < "3.4.0" or (
                    "3.4.0" <= spark_version < "3.5.1"
                    and not _is_standalone_or_localcluster(conf)
                ):
                    if gpu_per_task is not None:
                        if float(gpu_per_task) < 1.0:
                            raise ValueError(
                                "XGBoost doesn't support GPU fractional configurations. Please set "
                                "`spark.task.resource.gpu.amount=spark.executor.resource.gpu."
                                "amount`. To enable GPU fractional configurations, you can try "
                                "standalone/localcluster with spark 3.4.0+ and"
                                "YARN/K8S with spark 3.5.1+"
                            )
                    else:
                        raise ValueError(
                            "The `spark.task.resource.gpu.amount` is required for training"
                            " on GPU."
                        )

    def _validate_params(self) -> None:
        # pylint: disable=too-many-branches
        init_model = self.getOrDefault("xgb_model")
        if init_model is not None and not isinstance(init_model, Booster):
            raise ValueError(
                "The xgb_model param must be set with a `xgboost.core.Booster` "
                "instance."
            )

        if self.getOrDefault(self.num_workers) < 1:
            raise ValueError(
                f"Number of workers was {self.getOrDefault(self.num_workers)}."
                f"It cannot be less than 1 [Default is 1]"
            )

        tree_method = self.getOrDefault(self.getParam("tree_method"))
        if tree_method == "exact":
            raise ValueError(
                "The `exact` tree method is not supported for distributed systems."
            )

        if self.getOrDefault("objective") is not None:
            if not isinstance(self.getOrDefault("objective"), str):
                raise ValueError("Only string type 'objective' param is allowed.")

        eval_metric = "eval_metric"
        if self.getOrDefault(eval_metric) is not None:
            if not (
                isinstance(self.getOrDefault(eval_metric), str)
                or (
                    isinstance(self.getOrDefault(eval_metric), List)
                    and all(
                        isinstance(metric, str)
                        for metric in self.getOrDefault(eval_metric)
                    )
                )
            ):
                raise ValueError(
                    "Only string type or list of string type 'eval_metric' param is allowed."
                )

        if self.getOrDefault("early_stopping_rounds") is not None:
            if not self._col_is_defined_not_empty(self.validationIndicatorCol):
                raise ValueError(
                    "If 'early_stopping_rounds' param is set, you need to set "
                    "'validation_indicator_col' param as well."
                )

        if self.getOrDefault(self.enable_sparse_data_optim):
            if self.getOrDefault("missing") != 0.0:
                # If DMatrix is constructed from csr / csc matrix, then inactive elements
                # in csr / csc matrix are regarded as missing value, but, in pyspark, we
                # are hard to control elements to be active or inactive in sparse vector column,
                # some spark transformers such as VectorAssembler might compress vectors
                # to be dense or sparse format automatically, and when a spark ML vector object
                # is compressed to sparse vector, then all zero value elements become inactive.
                # So we force setting missing param to be 0 when enable_sparse_data_optim config
                # is True.
                raise ValueError(
                    "If enable_sparse_data_optim is True, missing param != 0 is not supported."
                )
            if self.getOrDefault(self.features_cols):
                raise ValueError(
                    "If enable_sparse_data_optim is True, you cannot set multiple feature columns "
                    "but you should set one feature column with values of "
                    "`pyspark.ml.linalg.Vector` type."
                )

        ss = _get_spark_session()
        sc = ss.sparkContext
        self._validate_gpu_params(ss.version, sc.getConf(), _is_local(sc))

    def _run_on_gpu(self) -> bool:
        """If train or transform on the gpu according to the parameters"""

        return use_cuda(self.getOrDefault(self.device))

    def _col_is_defined_not_empty(self, param: "Param[str]") -> bool:
        return self.isDefined(param) and self.getOrDefault(param) not in (None, "")


def _validate_and_convert_feature_col_as_float_col_list(
    dataset: DataFrame, features_col_names: List[str]
) -> List[Column]:
    """Values in feature columns must be integral types or float/double types"""
    feature_cols = []
    for c in features_col_names:
        if isinstance(dataset.schema[c].dataType, DoubleType):
            feature_cols.append(col(c).cast(FloatType()).alias(c))
        elif isinstance(dataset.schema[c].dataType, (FloatType, IntegralType)):
            feature_cols.append(col(c))
        else:
            raise ValueError(
                "Values in feature columns must be integral types or float/double types."
            )
    return feature_cols


def _validate_and_convert_feature_col_as_array_col(
    dataset: DataFrame, features_col_name: str
) -> Column:
    """It handles
    1. Convert vector type to array type
    2. Cast to Array(Float32)"""
    features_col_datatype = dataset.schema[features_col_name].dataType
    features_col = col(features_col_name)
    if isinstance(features_col_datatype, ArrayType):
        if not isinstance(
            features_col_datatype.elementType,
            (DoubleType, FloatType, LongType, IntegerType, ShortType),
        ):
            raise ValueError(
                "If feature column is array type, its elements must be number type, "
                f"got {features_col_datatype.elementType}."
            )
        features_array_col = features_col.cast(ArrayType(FloatType())).alias(alias.data)
    elif isinstance(features_col_datatype, VectorUDT):
        features_array_col = vector_to_array(features_col, dtype="float32").alias(
            alias.data
        )
    else:
        raise ValueError(
            "feature column must be array type or `pyspark.ml.linalg.Vector` type, "
            "if you want to use multiple numetric columns as features, please use "
            "`pyspark.ml.transform.VectorAssembler` to assemble them into a vector "
            "type column first."
        )
    return features_array_col


def _get_unwrap_udt_fn() -> Callable[[Union[Column, str]], Column]:
    try:
        from pyspark.sql.functions import unwrap_udt

        return unwrap_udt
    except ImportError:
        pass

    try:
        from pyspark.databricks.sql.functions import unwrap_udt as databricks_unwrap_udt

        return databricks_unwrap_udt
    except ImportError as exc:
        raise RuntimeError(
            "Cannot import pyspark `unwrap_udt` function. Please install pyspark>=3.4 "
            "or run on Databricks Runtime."
        ) from exc


def _get_unwrapped_vec_cols(feature_col: Column) -> List[Column]:
    unwrap_udt = _get_unwrap_udt_fn()
    features_unwrapped_vec_col = unwrap_udt(feature_col)

    # After a `pyspark.ml.linalg.VectorUDT` type column being unwrapped, it becomes
    # a pyspark struct type column, the struct fields are:
    #  - `type`: byte
    #  - `size`: int
    #  - `indices`: array<int>
    #  - `values`: array<double>
    # For sparse vector, `type` field is 0, `size` field means vector length,
    # `indices` field is the array of active element indices, `values` field
    # is the array of active element values.
    # For dense vector, `type` field is 1, `size` and `indices` fields are None,
    # `values` field is the array of the vector element values.
    return [
        features_unwrapped_vec_col.type.alias("featureVectorType"),
        features_unwrapped_vec_col.size.alias("featureVectorSize"),
        features_unwrapped_vec_col.indices.alias("featureVectorIndices"),
        # Note: the value field is double array type, cast it to float32 array type
        # for speedup following repartitioning.
        features_unwrapped_vec_col.values.cast(ArrayType(FloatType())).alias(
            "featureVectorValues"
        ),
    ]


FeatureProp = namedtuple(
    "FeatureProp",
    ("enable_sparse_data_optim", "has_validation_col", "features_cols_names"),
)

_MODEL_CHUNK_SIZE = 4096 * 1024


class _SparkXGBEstimator(Estimator, _SparkXGBParams, MLReadable, MLWritable):
    _input_kwargs: Dict[str, Any]

    def __init__(self) -> None:
        super().__init__()
        self._set_xgb_params_default()
        self._set_fit_params_default()
        self._set_predict_params_default()
        # Note: The default value for arbitrary_params_dict must always be empty dict.
        #  For additional settings added into "arbitrary_params_dict" by default,
        #  they are added in `setParams`.
        self._setDefault(
            num_workers=1,
            device="cpu",
            force_repartition=False,
            repartition_random_shuffle=False,
            feature_names=None,
            feature_types=None,
            feature_weights=None,
            arbitrary_params_dict={},
            launch_tracker_on_driver=True,
        )

        self.logger = get_logger(self.__class__.__name__)

    def setParams(self, **kwargs: Any) -> None:
        """
        Set params for the estimator.
        """
        _extra_params = {}
        if "arbitrary_params_dict" in kwargs:
            raise ValueError("Invalid param name: 'arbitrary_params_dict'.")

        for k, v in kwargs.items():
            # We're not allowing user use features_cols directly.
            if k == self.features_cols.name:
                raise ValueError(
                    f"Unsupported param '{k}' please use features_col instead."
                )
            if k in _inverse_pyspark_param_alias_map:
                raise ValueError(
                    f"Please use param name {_inverse_pyspark_param_alias_map[k]} instead."
                )
            if k in _pyspark_param_alias_map:
                if k == _inverse_pyspark_param_alias_map[
                    self.featuresCol.name
                ] and isinstance(v, list):
                    real_k = self.features_cols.name
                    k = real_k
                else:
                    real_k = _pyspark_param_alias_map[k]
                    k = real_k

            if self.hasParam(k):
                if k == "features_col" and isinstance(v, list):
                    self._set(**{"features_cols": v})
                else:
                    self._set(**{str(k): v})
            else:
                if (
                    k in _unsupported_xgb_params
                    or k in _unsupported_fit_params
                    or k in _unsupported_predict_params
                    or k in _unsupported_train_params
                ):
                    err_msg = _unsupported_params_hint_message.get(
                        k, f"Unsupported param '{k}'."
                    )
                    raise ValueError(err_msg)
                _extra_params[k] = v

        _check_distributed_params(kwargs)
        _existing_extra_params = self.getOrDefault(self.arbitrary_params_dict)
        self._set(arbitrary_params_dict={**_existing_extra_params, **_extra_params})

    @classmethod
    def _pyspark_model_cls(cls) -> Type["_SparkXGBModel"]:
        """
        Subclasses should override this method and
        returns a _SparkXGBModel subclass
        """
        raise NotImplementedError()

    def _create_pyspark_model(
        self, xgb_model: XGBModel, training_summary: XGBoostTrainingSummary
    ) -> "_SparkXGBModel":
        return self._pyspark_model_cls()(xgb_model, training_summary)

    def _convert_to_sklearn_model(self, booster: bytearray, config: str) -> XGBModel:
        xgb_sklearn_params = self._gen_xgb_params_dict(
            gen_xgb_sklearn_estimator_param=True
        )
        sklearn_model = self._xgb_cls()(**xgb_sklearn_params)
        sklearn_model.load_model(booster)
        sklearn_model._Booster.load_config(config)
        return sklearn_model

    def _repartition_needed(self, dataset: DataFrame) -> bool:
        """
        We repartition the dataset if the number of workers is not equal to the number of
        partitions."""
        if self.getOrDefault(self.force_repartition):
            return True
        num_workers = self.getOrDefault(self.num_workers)
        num_partitions = dataset.rdd.getNumPartitions()
        return not num_workers == num_partitions

    def _get_distributed_train_params(self, dataset: DataFrame) -> Dict[str, Any]:
        """
        This just gets the configuration params for distributed xgboost
        """
        params = self._gen_xgb_params_dict()
        fit_params = self._gen_fit_params_dict()
        verbose_eval = fit_params.pop("verbose", None)

        params.update(fit_params)
        params["verbose_eval"] = verbose_eval
        classification = self._xgb_cls() == XGBClassifier
        if classification:
            num_classes = int(
                dataset.select(countDistinct(alias.label)).collect()[0][0]
            )
            if num_classes <= 2:
                params["objective"] = "binary:logistic"
            else:
                params["objective"] = "multi:softprob"
                params["num_class"] = num_classes
        else:
            # use user specified objective or default objective.
            # e.g., the default objective for Regressor is 'reg:squarederror'
            params["objective"] = self.getOrDefault("objective")

        # TODO: support "num_parallel_tree" for random forest
        params["num_boost_round"] = self.getOrDefault("n_estimators")

        return params

    @classmethod
    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(
            worker_train, _unsupported_train_params
        )
        booster_params, kwargs_params = {}, {}
        for key, value in train_params.items():
            if key in xgb_train_default_args:
                kwargs_params[key] = value
            else:
                booster_params[key] = value

        booster_params = {
            k: v for k, v in booster_params.items() if k not in _non_booster_params
        }
        return booster_params, kwargs_params

    def _prepare_input_columns_and_feature_prop(
        self, dataset: DataFrame
    ) -> Tuple[List[Column], FeatureProp]:
        label_col = col(self.getOrDefault(self.labelCol)).alias(alias.label)

        select_cols = [label_col]
        features_cols_names = None
        enable_sparse_data_optim = self.getOrDefault(self.enable_sparse_data_optim)
        if enable_sparse_data_optim:
            features_col_name = self.getOrDefault(self.featuresCol)
            features_col_datatype = dataset.schema[features_col_name].dataType
            if not isinstance(features_col_datatype, VectorUDT):
                raise ValueError(
                    "If enable_sparse_data_optim is True, the feature column values must be "
                    "`pyspark.ml.linalg.Vector` type."
                )
            select_cols.extend(_get_unwrapped_vec_cols(col(features_col_name)))
        else:
            if self.getOrDefault(self.features_cols):
                features_cols_names = self.getOrDefault(self.features_cols)
                features_cols = _validate_and_convert_feature_col_as_float_col_list(
                    dataset, features_cols_names
                )
                select_cols.extend(features_cols)
            else:
                features_array_col = _validate_and_convert_feature_col_as_array_col(
                    dataset, self.getOrDefault(self.featuresCol)
                )
                select_cols.append(features_array_col)

        if self._col_is_defined_not_empty(self.weightCol):
            select_cols.append(
                col(self.getOrDefault(self.weightCol)).alias(alias.weight)
            )

        has_validation_col = False
        if self._col_is_defined_not_empty(self.validationIndicatorCol):
            select_cols.append(
                col(self.getOrDefault(self.validationIndicatorCol)).alias(alias.valid)
            )
            # In some cases, see https://issues.apache.org/jira/browse/SPARK-40407,
            # the df.repartition can result in some reducer partitions without data,
            # which will cause exception or hanging issue when creating DMatrix.
            has_validation_col = True

        if self._col_is_defined_not_empty(self.base_margin_col):
            select_cols.append(
                col(self.getOrDefault(self.base_margin_col)).alias(alias.margin)
            )

        if self._col_is_defined_not_empty(self.qid_col):
            select_cols.append(col(self.getOrDefault(self.qid_col)).alias(alias.qid))

        feature_prop = FeatureProp(
            enable_sparse_data_optim, has_validation_col, features_cols_names
        )
        return select_cols, feature_prop

    def _prepare_input(self, dataset: DataFrame) -> Tuple[DataFrame, FeatureProp]:
        """Prepare the input including column pruning, repartition and so on"""

        select_cols, feature_prop = self._prepare_input_columns_and_feature_prop(
            dataset
        )

        dataset = dataset.select(*select_cols)

        num_workers = self.getOrDefault(self.num_workers)
        sc = _get_spark_session().sparkContext
        max_concurrent_tasks = _get_max_num_concurrent_tasks(sc)

        if feature_prop.has_validation_col:
            dtype = dataset.schema[alias.valid].dataType
            if not isinstance(dtype, BooleanType):
                raise TypeError("The validation indicator must be boolean type.")

        if num_workers > max_concurrent_tasks:
            get_logger(self.__class__.__name__).warning(
                "The num_workers %s set for xgboost distributed "
                "training is greater than current max number of concurrent "
                "spark task slots, you need wait until more task slots available "
                "or you need increase spark cluster workers.",
                num_workers,
            )

        if self._repartition_needed(dataset):
            if self._col_is_defined_not_empty(self.qid_col):
                # For ranking problem, we need to try best the put the instances with
                # same group into the same partition
                dataset = dataset.repartitionByRange(num_workers, alias.qid)
            else:
                # If validationIndicatorCol defined, and if user unionise train and validation
                # dataset, users must set force_repartition to true to force repartition.
                # Or else some partitions might contain only train or validation dataset.
                if self.getOrDefault(self.repartition_random_shuffle):
                    # In some cases, spark round-robin repartition might cause data skew
                    # use random shuffle can address it.
                    dataset = dataset.repartition(num_workers, rand(1))
                else:
                    dataset = dataset.repartition(num_workers)

        if self._col_is_defined_not_empty(self.qid_col):
            # XGBoost requires qid to be sorted for each partition
            dataset = dataset.sortWithinPartitions(alias.qid, ascending=True)

        return dataset, feature_prop

    def _get_xgb_parameters(
        self, dataset: DataFrame
    ) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
        train_params = self._get_distributed_train_params(dataset)
        booster_params, train_call_kwargs_params = self._get_xgb_train_call_args(
            train_params
        )
        cpu_per_task = int(
            _get_spark_session().sparkContext.getConf().get("spark.task.cpus", "1")
        )

        dmatrix_kwargs = {
            "nthread": cpu_per_task,
            "feature_types": self.getOrDefault("feature_types"),
            "feature_names": self.getOrDefault("feature_names"),
            "feature_weights": self.getOrDefault("feature_weights"),
            "missing": float(self.getOrDefault("missing")),
        }
        if dmatrix_kwargs["feature_types"] is not None:
            dmatrix_kwargs["enable_categorical"] = True
        booster_params["nthread"] = cpu_per_task

        # Remove the parameters whose value is None
        booster_params = {k: v for k, v in booster_params.items() if v is not None}
        train_call_kwargs_params = {
            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}

        return booster_params, train_call_kwargs_params, dmatrix_kwargs

    def _skip_stage_level_scheduling(self, spark_version: str, conf: SparkConf) -> bool:
        # pylint: disable=too-many-return-statements
        """Check if stage-level scheduling is not needed,
        return true to skip stage-level scheduling"""

        if self._run_on_gpu():
            if spark_version < "3.4.0":
                self.logger.info(
                    "Stage-level scheduling in xgboost requires spark version 3.4.0+"
                )
                return True

            if (
                "3.4.0" <= spark_version < "3.5.1"
                and not _is_standalone_or_localcluster(conf)
            ):
                self.logger.info(
                    "For %s, Stage-level scheduling in xgboost requires spark standalone "
                    "or local-cluster mode",
                    spark_version,
                )
                return True

            executor_cores = conf.get("spark.executor.cores")
            executor_gpus = conf.get("spark.executor.resource.gpu.amount")
            if executor_cores is None or executor_gpus is None:
                self.logger.info(
                    "Stage-level scheduling in xgboost requires spark.executor.cores, "
                    "spark.executor.resource.gpu.amount to be set."
                )
                return True

            if int(executor_cores) == 1:
                # there will be only 1 task running at any time.
                self.logger.info(
                    "Stage-level scheduling in xgboost requires spark.executor.cores > 1 "
                )
                return True

            if int(executor_gpus) > 1:
                # For spark.executor.resource.gpu.amount > 1, we suppose user knows how to configure
                # to make xgboost run successfully.
                #
                self.logger.info(
                    "Stage-level scheduling in xgboost will not work "
                    "when spark.executor.resource.gpu.amount>1"
                )
                return True

            task_gpu_amount = conf.get("spark.task.resource.gpu.amount")

            if task_gpu_amount is None:
                # The ETL tasks will not grab a gpu when spark.task.resource.gpu.amount is not set,
                # but with stage-level scheduling, we can make training task grab the gpu.
                return False

            if float(task_gpu_amount) == float(executor_gpus):
                # spark.executor.resource.gpu.amount=spark.task.resource.gpu.amount "
                # results in only 1 task running at a time, which may cause perf issue.
                return True

            # We can enable stage-level scheduling
            return False

        # CPU training doesn't require stage-level scheduling
        return True

    def _try_stage_level_scheduling(self, rdd: RDD) -> RDD:
        """Try to enable stage-level scheduling"""
        ss = _get_spark_session()
        conf = ss.sparkContext.getConf()
        if _is_local(ss.sparkContext) or self._skip_stage_level_scheduling(
            ss.version, conf
        ):
            return rdd

        # executor_cores will not be None
        executor_cores = conf.get("spark.executor.cores")
        assert executor_cores is not None

        # Spark-rapids is a project to leverage GPUs to accelerate spark SQL.
        # If spark-rapids is enabled, to avoid GPU OOM, we don't allow other
        # ETL gpu tasks running alongside training tasks.
        spark_plugins = ss.conf.get("spark.plugins", " ")
        assert spark_plugins is not None
        spark_rapids_sql_enabled = ss.conf.get("spark.rapids.sql.enabled", "true")
        assert spark_rapids_sql_enabled is not None

        task_cores = (
            int(executor_cores)
            if "com.nvidia.spark.SQLPlugin" in spark_plugins
            and "true" == spark_rapids_sql_enabled.lower()
            else (int(executor_cores) // 2) + 1
        )

        # Each training task requires cpu cores > total executor cores//2 + 1 which can
        # make sure the tasks be sent to different executors.
        #
        # Please note that we can't use GPU to limit the concurrent tasks because of
        # https://issues.apache.org/jira/browse/SPARK-45527.

        task_gpus = 1.0
        treqs = TaskResourceRequests().cpus(task_cores).resource("gpu", task_gpus)
        rp = ResourceProfileBuilder().require(treqs).build

        self.logger.info(
            "XGBoost training tasks require the resource(cores=%s, gpu=%s).",
            task_cores,
            task_gpus,
        )
        return rdd.withResources(rp)

    def _get_tracker_args(self) -> Tuple[bool, Dict[str, Any]]:
        """Start the tracker and return the tracker envs on the driver side"""
        launch_tracker_on_driver = self.getOrDefault(self.launch_tracker_on_driver)
        rabit_args = {}
        if launch_tracker_on_driver:
            conf = Config()
            if self.isDefined(self.coll_cfg):
                conf = self.getOrDefault(self.coll_cfg)
                assert isinstance(conf, Config)

            if conf.tracker_host_ip is None:
                conf.tracker_host_ip = (
                    _get_spark_session().sparkContext.getConf().get("spark.driver.host")
                )
            num_workers = self.getOrDefault(self.num_workers)
            rabit_args.update(_get_rabit_args(conf, num_workers))
        else:
            if self.isDefined(self.coll_cfg):
                conf = self.getOrDefault(self.coll_cfg)
                assert isinstance(conf, Config)
                if conf.tracker_host_ip is not None:
                    raise ValueError(
                        f"You must enable launch_tracker_on_driver to use "
                        f"tracker host: {conf.tracker_host_ip}"
                    )
        return launch_tracker_on_driver, rabit_args

    def _fit(self, dataset: DataFrame) -> "_SparkXGBModel":
        # pylint: disable=too-many-statements, too-many-locals
        self._validate_params()

        dataset, feature_prop = self._prepare_input(dataset)

        (
            booster_params,
            train_call_kwargs_params,
            dmatrix_kwargs,
        ) = self._get_xgb_parameters(dataset)

        run_on_gpu = self._run_on_gpu()

        is_local = _is_local(_get_spark_session().sparkContext)

        num_workers = self.getOrDefault(self.num_workers)

        launch_tracker_on_driver, rabit_args = self._get_tracker_args()
        conf: Optional[Config] = (
            self.getOrDefault(self.coll_cfg) if self.isSet(self.coll_cfg) else None
        )

        log_level = get_logger_level(_LOG_TAG)

        use_rmm = get_config()["use_rmm"]

        def _train_booster(
            pandas_df_iter: Iterator[pd.DataFrame],
        ) -> Iterator[pd.DataFrame]:
            """Takes in an RDD partition and outputs a booster for that partition after
            going through the Rabit Ring protocol

            """
            from pyspark import BarrierTaskContext

            context = BarrierTaskContext.get()
            context.barrier()

            dev_ordinal = None
            use_qdm = _can_use_qdm(
                booster_params.get("tree_method", None),
                booster_params.get("device", None),
            )
            verbosity = booster_params.get("verbosity", 1)
            msg = "Training on CPUs"
            if run_on_gpu:
                dev_ordinal = (
                    context.partitionId() if is_local else _get_gpu_id(context)
                )
                booster_params["device"] = "cuda:" + str(dev_ordinal)
                # 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.
                use_qdm = use_qdm and is_cudf_available()
                msg = (
                    f"Leveraging {booster_params['device']} to train with "
                    f"QDM: {'on' if use_qdm else 'off'}"
                )

            if use_qdm and (booster_params.get("max_bin", None) is not None):
                dmatrix_kwargs["max_bin"] = booster_params["max_bin"]
            _rabit_args = rabit_args
            if context.partitionId() == 0:
                if not launch_tracker_on_driver:
                    _conf = conf if conf is not None else Config()
                    _conf.tracker_host_ip = _get_host_ip(context)
                    _rabit_args = _get_rabit_args(_conf, num_workers)
                get_logger(_LOG_TAG, log_level).info(msg)

            worker_message: Dict[str, Any] = {
                "use_qdm": use_qdm,
            }

            if not launch_tracker_on_driver:
                worker_message["rabit_msg"] = _rabit_args

            messages = context.allGather(message=json.dumps(worker_message))
            if len(set(json.loads(x)["use_qdm"] for x in messages)) != 1:
                raise RuntimeError("The workers' cudf environments are in-consistent ")

            if not launch_tracker_on_driver:
                _rabit_args = json.loads(messages[0])["rabit_msg"]

            evals_result: Dict[str, Any] = {}
            with config_context(
                verbosity=verbosity, use_rmm=use_rmm
            ), 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:
                    dval = [(dtrain, "training")]
                booster = worker_train(
                    params=booster_params,
                    dtrain=dtrain,
                    evals=dval,
                    evals_result=evals_result,
                    **train_call_kwargs_params,
                )
            context.barrier()

            if context.partitionId() == 0:
                yield pd.DataFrame({"data": [json.dumps(dict(evals_result))]})
                config = booster.save_config()
                yield pd.DataFrame({"data": [config]})
                booster_json = booster.save_raw("json").decode("utf-8")

                for offset in range(0, len(booster_json), _MODEL_CHUNK_SIZE):
                    booster_chunk = booster_json[offset : offset + _MODEL_CHUNK_SIZE]
                    yield pd.DataFrame({"data": [booster_chunk]})

        def _run_job() -> Tuple[str, str, str]:
            rdd = (
                dataset.mapInPandas(
                    _train_booster,  # type: ignore
                    schema="data string",
                )
                .rdd.barrier()
                .mapPartitions(lambda x: x)
            )
            rdd_with_resource = self._try_stage_level_scheduling(rdd)
            ret = rdd_with_resource.collect()
            data = [v[0] for v in ret]
            return data[0], data[1], "".join(data[2:])

        get_logger(_LOG_TAG).info(
            "Running xgboost-%s on %s workers with"
            "\n\tbooster params: %s"
            "\n\ttrain_call_kwargs_params: %s"
            "\n\tdmatrix_kwargs: %s",
            _py_version(),
            num_workers,
            booster_params,
            train_call_kwargs_params,
            dmatrix_kwargs,
        )
        (evals_result, config, booster) = _run_job()
        get_logger(_LOG_TAG).info("Finished xgboost training!")

        result_xgb_model = self._convert_to_sklearn_model(
            bytearray(booster, "utf-8"), config
        )
        training_summary = XGBoostTrainingSummary.from_metrics(json.loads(evals_result))
        spark_model = self._create_pyspark_model(result_xgb_model, training_summary)
        # According to pyspark ML convention, the model uid should be the same
        # with estimator uid.
        spark_model._resetUid(self.uid)
        return self._copyValues(spark_model)

    def write(self) -> "SparkXGBWriter":
        """
        Return the writer for saving the estimator.
        """
        return SparkXGBWriter(self)

    @classmethod
    def read(cls) -> "SparkXGBReader":
        """
        Return the reader for loading the estimator.
        """
        return SparkXGBReader(cls)


class _SparkXGBModel(Model, _SparkXGBParams, MLReadable, MLWritable):
    def __init__(
        self,
        xgb_sklearn_model: Optional[XGBModel] = None,
        training_summary: Optional[XGBoostTrainingSummary] = None,
    ) -> None:
        super().__init__()
        self._xgb_sklearn_model = xgb_sklearn_model
        self.training_summary = training_summary

    @classmethod
    def _xgb_cls(cls) -> Type[XGBModel]:
        raise NotImplementedError()

    def get_booster(self) -> Booster:
        """
        Return the `xgboost.core.Booster` instance.
        """
        assert self._xgb_sklearn_model is not None
        return self._xgb_sklearn_model.get_booster()

    def get_feature_importances(
        self, importance_type: str = "weight"
    ) -> Dict[str, Union[float, List[float]]]:
        """Get feature importance of each feature.
        Importance type can be defined as:

        * 'weight': the number of times a feature is used to split the data across all trees.
        * 'gain': the average gain across all splits the feature is used in.
        * 'cover': the average coverage across all splits the feature is used in.
        * 'total_gain': the total gain across all splits the feature is used in.
        * 'total_cover': the total coverage across all splits the feature is used in.

        Parameters
        ----------
        importance_type: str, default 'weight'
            One of the importance types defined above.
        """
        return self.get_booster().get_score(importance_type=importance_type)

    def write(self) -> "SparkXGBModelWriter":
        """
        Return the writer for saving the model.
        """
        return SparkXGBModelWriter(self)

    @classmethod
    def read(cls) -> "SparkXGBModelReader":
        """
        Return the reader for loading the model.
        """
        return SparkXGBModelReader(cls)

    def _get_feature_col(
        self, dataset: DataFrame
    ) -> Tuple[List[Column], Optional[List[str]]]:
        """XGBoost model trained with features_cols parameter can also predict
        vector or array feature type. But first we need to check features_cols
        and then featuresCol
        """
        if self.getOrDefault(self.enable_sparse_data_optim):
            feature_col_names = None
            features_col = _get_unwrapped_vec_cols(
                col(self.getOrDefault(self.featuresCol))
            )
            return features_col, feature_col_names

        feature_col_names = self.getOrDefault(self.features_cols)
        features_col = []
        if feature_col_names and set(feature_col_names).issubset(set(dataset.columns)):
            # The model is trained with features_cols and the predicted dataset
            # also contains all the columns specified by features_cols.
            features_col = _validate_and_convert_feature_col_as_float_col_list(
                dataset, feature_col_names
            )
        else:
            # 1. The model was trained by features_cols, but the dataset doesn't contain
            #       all the columns specified by features_cols, so we need to check if
            #       the dataframe has the featuresCol
            # 2. The model was trained by featuresCol, and the predicted dataset must contain
            #       featuresCol column.
            feature_col_names = None
            features_col.append(
                _validate_and_convert_feature_col_as_array_col(
                    dataset, self.getOrDefault(self.featuresCol)
                )
            )
        return features_col, feature_col_names

    def _get_pred_contrib_col_name(self) -> Optional[str]:
        """Return the pred_contrib_col col name"""
        pred_contrib_col_name = None
        if self._col_is_defined_not_empty(self.pred_contrib_col):
            pred_contrib_col_name = self.getOrDefault(self.pred_contrib_col)

        return pred_contrib_col_name

    def _out_schema(self) -> Tuple[bool, str]:
        """Return the bool to indicate if it's a single prediction, true is single prediction,
        and the returned type of the user-defined function. The value must
        be a DDL-formatted type string."""

        if self._get_pred_contrib_col_name() is not None:
            return False, f"{pred.prediction} double, {pred.pred_contrib} array<double>"

        return True, "double"

    def _get_predict_func(self) -> Callable:
        """Return the true prediction function which will be running on the executor side"""

        predict_params = self._gen_predict_params_dict()
        pred_contrib_col_name = self._get_pred_contrib_col_name()

        def _predict(
            model: XGBModel, X: ArrayLike, base_margin: Optional[ArrayLike]
        ) -> Union[pd.DataFrame, pd.Series]:
            data = {}
            preds = model.predict(
                X,
                base_margin=base_margin,
                validate_features=False,
                **predict_params,
            )
            data[pred.prediction] = pd.Series(preds)

            if pred_contrib_col_name is not None:
                contribs = pred_contribs(model, X, base_margin)
                data[pred.pred_contrib] = pd.Series(list(contribs))
                return pd.DataFrame(data=data)

            return data[pred.prediction]

        return _predict

    def _post_transform(self, dataset: DataFrame, pred_col: Column) -> DataFrame:
        """Post process of transform"""
        prediction_col_name = self.getOrDefault(self.predictionCol)
        single_pred, _ = self._out_schema()

        if single_pred:
            if prediction_col_name:
                dataset = dataset.withColumn(prediction_col_name, pred_col)
        else:
            pred_struct_col = "_prediction_struct"
            dataset = dataset.withColumn(pred_struct_col, pred_col)

            if prediction_col_name:
                dataset = dataset.withColumn(
                    prediction_col_name, getattr(col(pred_struct_col), pred.prediction)
                )

            pred_contrib_col_name = self._get_pred_contrib_col_name()
            if pred_contrib_col_name is not None:
                dataset = dataset.withColumn(
                    pred_contrib_col_name,
                    array_to_vector(getattr(col(pred_struct_col), pred.pred_contrib)),
                )

            dataset = dataset.drop(pred_struct_col)
        return dataset

    def _run_on_gpu(self) -> bool:
        """If gpu is used to do the prediction according to the parameters
        and spark configurations"""

        use_gpu_by_params = super()._run_on_gpu()

        if _is_local(_get_spark_session().sparkContext):
            # if it's local model, no need to check the spark configurations
            return use_gpu_by_params

        gpu_per_task = (
            _get_spark_session()
            .sparkContext.getConf()
            .get("spark.task.resource.gpu.amount")
        )

        # User don't set gpu configurations, just use cpu
        if gpu_per_task is None:
            if use_gpu_by_params:
                get_logger(_LOG_TAG).warning(
                    "Do the prediction on the CPUs since "
                    "no gpu configurations are set"
                )
            return False

        # User already sets the gpu configurations.
        return use_gpu_by_params

    def _transform(self, dataset: DataFrame) -> DataFrame:
        # pylint: disable=too-many-statements, too-many-locals
        # Save xgb_sklearn_model and predict_params to be local variable
        # to avoid the `self` object to be pickled to remote.
        xgb_sklearn_model = self._xgb_sklearn_model

        base_margin_col = None
        if self._col_is_defined_not_empty(self.base_margin_col):
            base_margin_col = col(self.getOrDefault(self.base_margin_col)).alias(
                alias.margin
            )
        has_base_margin = base_margin_col is not None

        features_col, feature_col_names = self._get_feature_col(dataset)
        enable_sparse_data_optim = self.getOrDefault(self.enable_sparse_data_optim)

        predict_func = self._get_predict_func()

        _, schema = self._out_schema()

        is_local = _is_local(_get_spark_session().sparkContext)
        run_on_gpu = self._run_on_gpu()

        log_level = get_logger_level(_LOG_TAG)

        @pandas_udf(schema)  # type: ignore
        def predict_udf(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.Series]:
            assert xgb_sklearn_model is not None
            model = xgb_sklearn_model

            from pyspark import TaskContext

            context = TaskContext.get()
            assert context is not None

            dev_ordinal = -1

            msg = "Do the inference on the CPUs"
            if run_on_gpu:
                if is_cudf_available() and is_cupy_available():
                    if is_local:
                        cp = import_cupy()

                        total_gpus = cp.cuda.runtime.getDeviceCount()
                        if total_gpus > 0:
                            partition_id = context.partitionId()
                            # For transform local mode, default the dev_ordinal to
                            # (partition id) % gpus.
                            dev_ordinal = partition_id % total_gpus
                    else:
                        dev_ordinal = _get_gpu_id(context)

                    if dev_ordinal >= 0:
                        device = "cuda:" + str(dev_ordinal)
                        msg = "Do the inference with device: " + device
                        model.set_params(device=device)
                    else:
                        msg = "Couldn't get the correct gpu id, fallback the inference on the CPUs"
                else:
                    msg = "CUDF or Cupy is unavailable, fallback the inference on the CPUs"

            if context.partitionId() == 0:
                get_logger(_LOG_TAG, log_level).info(msg)

            def to_gpu_if_possible(data: ArrayLike) -> ArrayLike:
                """Move the data to gpu if possible"""
                if dev_ordinal >= 0:
                    import cudf
                    import cupy as cp

                    # We must set the device after import cudf, which will change the device id to 0
                    # See https://github.com/rapidsai/cudf/issues/11386
                    cp.cuda.runtime.setDevice(dev_ordinal)  # pylint: disable=I1101
                    df = cudf.DataFrame(data)
                    del data
                    return df
                return data

            for data in iterator:
                if enable_sparse_data_optim:
                    X = _read_csr_matrix_from_unwrapped_spark_vec(data)
                else:
                    if feature_col_names is not None:
                        tmp: ArrayLike = data[feature_col_names]
                    else:
                        tmp = stack_series(data[alias.data])
                    X = to_gpu_if_possible(tmp)

                if has_base_margin:
                    base_margin = to_gpu_if_possible(data[alias.margin])
                else:
                    base_margin = None

                yield predict_func(model, X, base_margin)

        if has_base_margin:
            assert base_margin_col is not None
            pred_col = predict_udf(struct(*features_col, base_margin_col))
        else:
            pred_col = predict_udf(struct(*features_col))

        return self._post_transform(dataset, pred_col)


class _ClassificationModel(  # pylint: disable=abstract-method
    _SparkXGBModel, HasProbabilityCol, HasRawPredictionCol, HasContribPredictionCol
):
    """
    The model returned by :func:`xgboost.spark.SparkXGBClassifier.fit`

    .. Note:: This API is experimental.
    """

    def _out_schema(self) -> Tuple[bool, str]:
        schema = (
            f"{pred.raw_prediction} array<double>, {pred.prediction} double,"
            f" {pred.probability} array<double>"
        )
        if self._get_pred_contrib_col_name() is not None:
            # We will force setting strict_shape to True when predicting contribs,
            # So, it will also output 3-D shape result.
            schema = f"{schema}, {pred.pred_contrib} array<array<double>>"

        return False, schema

    def _get_predict_func(self) -> Callable:
        predict_params = self._gen_predict_params_dict()
        pred_contrib_col_name = self._get_pred_contrib_col_name()

        def transform_margin(margins: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
            if margins.ndim == 1:
                # binomial case
                classone_probs = expit(margins)
                classzero_probs = 1.0 - classone_probs
                raw_preds = np.vstack((-margins, margins)).transpose()
                class_probs = np.vstack((classzero_probs, classone_probs)).transpose()
            else:
                # multinomial case
                raw_preds = margins
                class_probs = softmax(raw_preds, axis=1)
            return raw_preds, class_probs

        def _predict(
            model: XGBModel, X: ArrayLike, base_margin: Optional[np.ndarray]
        ) -> Union[pd.DataFrame, pd.Series]:
            margins = model.predict(
                X,
                base_margin=base_margin,
                output_margin=True,
                validate_features=False,
                **predict_params,
            )
            raw_preds, class_probs = transform_margin(margins)

            # It seems that they use argmax of class probs,
            # not of margin to get the prediction (Note: scala implementation)
            preds = np.argmax(class_probs, axis=1)
            result: Dict[str, pd.Series] = {
                pred.raw_prediction: pd.Series(list(raw_preds)),
                pred.prediction: pd.Series(preds),
                pred.probability: pd.Series(list(class_probs)),
            }

            if pred_contrib_col_name is not None:
                contribs = pred_contribs(model, X, base_margin, strict_shape=True)
                result[pred.pred_contrib] = pd.Series(list(contribs.tolist()))

            return pd.DataFrame(data=result)

        return _predict

    def _post_transform(self, dataset: DataFrame, pred_col: Column) -> DataFrame:
        pred_struct_col = "_prediction_struct"
        dataset = dataset.withColumn(pred_struct_col, pred_col)

        raw_prediction_col_name = self.getOrDefault(self.rawPredictionCol)
        if raw_prediction_col_name:
            dataset = dataset.withColumn(
                raw_prediction_col_name,
                array_to_vector(getattr(col(pred_struct_col), pred.raw_prediction)),
            )

        prediction_col_name = self.getOrDefault(self.predictionCol)
        if prediction_col_name:
            dataset = dataset.withColumn(
                prediction_col_name, getattr(col(pred_struct_col), pred.prediction)
            )

        probability_col_name = self.getOrDefault(self.probabilityCol)
        if probability_col_name:
            dataset = dataset.withColumn(
                probability_col_name,
                array_to_vector(getattr(col(pred_struct_col), pred.probability)),
            )

        pred_contrib_col_name = self._get_pred_contrib_col_name()
        if pred_contrib_col_name is not None:
            dataset = dataset.withColumn(
                pred_contrib_col_name,
                getattr(col(pred_struct_col), pred.pred_contrib),
            )

        return dataset.drop(pred_struct_col)


class _SparkXGBSharedReadWrite:
    @staticmethod
    def saveMetadata(
        instance: Union[_SparkXGBEstimator, _SparkXGBModel],
        path: str,
        sc: SparkContext,
        logger: logging.Logger,
        extraMetadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Save the metadata of an xgboost.spark._SparkXGBEstimator or
        xgboost.spark._SparkXGBModel.
        """
        instance._validate_params()
        skipParams = ["callbacks", "xgb_model", "coll_cfg"]
        jsonParams = {}
        for p, v in instance._paramMap.items():  # pylint: disable=protected-access
            if p.name not in skipParams:
                jsonParams[p.name] = v

        extraMetadata = extraMetadata or {}
        callbacks = instance.getOrDefault("callbacks")
        if callbacks is not None:
            logger.warning(
                "The callbacks parameter is saved using cloudpickle and it "
                "is not a fully self-contained format. It may fail to load "
                "with different versions of dependencies."
            )
            serialized_callbacks = base64.encodebytes(
                cloudpickle.dumps(callbacks)
            ).decode("ascii")
            extraMetadata["serialized_callbacks"] = serialized_callbacks
        init_booster = instance.getOrDefault("xgb_model")
        if init_booster is not None:
            extraMetadata["init_booster"] = _INIT_BOOSTER_SAVE_PATH

        if instance.isDefined("coll_cfg"):
            conf: Config = instance.getOrDefault("coll_cfg")
            if conf is not None:
                extraMetadata["coll_cfg"] = asdict(conf)

        DefaultParamsWriter.saveMetadata(
            instance, path, sc, extraMetadata=extraMetadata, paramMap=jsonParams
        )
        if init_booster is not None:
            ser_init_booster = serialize_booster(init_booster)
            save_path = os.path.join(path, _INIT_BOOSTER_SAVE_PATH)
            _get_spark_session().createDataFrame(
                [(ser_init_booster,)], ["init_booster"]
            ).write.parquet(save_path)

    @staticmethod
    def loadMetadataAndInstance(
        pyspark_xgb_cls: Union[Type[_SparkXGBEstimator], Type[_SparkXGBModel]],
        path: str,
        sc: SparkContext,
        logger: logging.Logger,
    ) -> Tuple[Dict[str, Any], Union[_SparkXGBEstimator, _SparkXGBModel]]:
        """
        Load the metadata and the instance of an xgboost.spark._SparkXGBEstimator or
        xgboost.spark._SparkXGBModel.

        :return: a tuple of (metadata, instance)
        """
        metadata = DefaultParamsReader.loadMetadata(
            path, sc, expectedClassName=get_class_name(pyspark_xgb_cls)
        )
        pyspark_xgb = pyspark_xgb_cls()
        DefaultParamsReader.getAndSetParams(pyspark_xgb, metadata)

        if "serialized_callbacks" in metadata:
            serialized_callbacks = metadata["serialized_callbacks"]
            try:
                callbacks = cloudpickle.loads(
                    base64.decodebytes(serialized_callbacks.encode("ascii"))
                )
                pyspark_xgb.set(pyspark_xgb.callbacks, callbacks)  # type: ignore
            except Exception as e:  # pylint: disable=W0703
                logger.warning(
                    f"Fails to load the callbacks param due to {e}. Please set the "
                    "callbacks param manually for the loaded estimator."
                )
        if "coll_cfg" in metadata:
            pyspark_xgb.set_coll_cfg(Config(**metadata["coll_cfg"]))

        if "init_booster" in metadata:
            load_path = os.path.join(path, metadata["init_booster"])
            ser_init_booster = (
                _get_spark_session().read.parquet(load_path).collect()[0].init_booster
            )
            init_booster = deserialize_booster(ser_init_booster)
            pyspark_xgb.set(pyspark_xgb.xgb_model, init_booster)  # type: ignore

        pyspark_xgb._resetUid(metadata["uid"])  # pylint: disable=protected-access
        return metadata, pyspark_xgb


class SparkXGBWriter(MLWriter):
    """
    Spark Xgboost estimator writer.
    """

    def __init__(self, instance: "_SparkXGBEstimator") -> None:
        super().__init__()
        self.instance = instance
        self.logger = get_logger(self.__class__.__name__, level="WARN")

    def saveImpl(self, path: str) -> None:
        """
        save model.
        """
        _SparkXGBSharedReadWrite.saveMetadata(self.instance, path, self.sc, self.logger)


class SparkXGBReader(MLReader):
    """
    Spark Xgboost estimator reader.
    """

    def __init__(self, cls: Type["_SparkXGBEstimator"]) -> None:
        super().__init__()
        self.cls = cls
        self.logger = get_logger(self.__class__.__name__, level="WARN")

    def load(self, path: str) -> "_SparkXGBEstimator":
        """
        load model.
        """
        _, pyspark_xgb = _SparkXGBSharedReadWrite.loadMetadataAndInstance(
            self.cls, path, self.sc, self.logger
        )
        return cast("_SparkXGBEstimator", pyspark_xgb)


class SparkXGBModelWriter(MLWriter):
    """
    Spark Xgboost model writer.
    """

    def __init__(self, instance: _SparkXGBModel) -> None:
        super().__init__()
        self.instance = instance
        self.logger = get_logger(self.__class__.__name__, level="WARN")

    def saveImpl(self, path: str) -> None:
        """
        Save metadata and model for a :py:class:`_SparkXGBModel`
        - save metadata to path/metadata
        - save model to path/model.json
        """
        xgb_model = self.instance._xgb_sklearn_model
        assert xgb_model is not None
        _SparkXGBSharedReadWrite.saveMetadata(self.instance, path, self.sc, self.logger)
        model_save_path = os.path.join(path, "model")
        booster = xgb_model.get_booster().save_raw("json").decode("utf-8")
        booster_chunks = []

        for offset in range(0, len(booster), _MODEL_CHUNK_SIZE):
            booster_chunks.append(booster[offset : offset + _MODEL_CHUNK_SIZE])

        _get_spark_session().sparkContext.parallelize(booster_chunks, 1).saveAsTextFile(
            model_save_path
        )


class SparkXGBModelReader(MLReader):
    """
    Spark Xgboost model reader.
    """

    def __init__(self, cls: Type["_SparkXGBModel"]) -> None:
        super().__init__()
        self.cls = cls
        self.logger = get_logger(self.__class__.__name__, level="WARN")

    def load(self, path: str) -> "_SparkXGBModel":
        """
        Load metadata and model for a :py:class:`_SparkXGBModel`

        :return: SparkXGBRegressorModel or SparkXGBClassifierModel instance
        """
        _, py_model = _SparkXGBSharedReadWrite.loadMetadataAndInstance(
            self.cls, path, self.sc, self.logger
        )
        py_model = cast("_SparkXGBModel", py_model)

        xgb_sklearn_params = py_model._gen_xgb_params_dict(
            gen_xgb_sklearn_estimator_param=True
        )
        model_load_path = os.path.join(path, "model")

        ser_xgb_model = "".join(
            _get_spark_session().sparkContext.textFile(model_load_path).collect()
        )

        def create_xgb_model() -> "XGBModel":
            return self.cls._xgb_cls()(**xgb_sklearn_params)

        xgb_model = deserialize_xgb_model(ser_xgb_model, create_xgb_model)
        py_model._xgb_sklearn_model = xgb_model
        return py_model
