snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/mllib/tree.py

@@ -0,0 +1,888 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import sys
+import random
+
+from pyspark import RDD, since
+from pyspark.mllib.common import callMLlibFunc, inherit_doc, JavaModelWrapper
+from pyspark.mllib.linalg import _convert_to_vector
+from pyspark.mllib.regression import LabeledPoint
+from pyspark.mllib.util import JavaLoader, JavaSaveable
+from typing import Dict, Optional, Tuple, Union, overload, TYPE_CHECKING
+from pyspark.rdd import RDD
+
+if TYPE_CHECKING:
+    from pyspark.mllib._typing import VectorLike
+
+
+__all__ = [
+    "DecisionTreeModel",
+    "DecisionTree",
+    "RandomForestModel",
+    "RandomForest",
+    "GradientBoostedTreesModel",
+    "GradientBoostedTrees",
+]
+
+
+class TreeEnsembleModel(JavaModelWrapper, JavaSaveable):
+    """TreeEnsembleModel
+
+    .. versionadded:: 1.3.0
+    """
+
+    @overload
+    def predict(self, x: "VectorLike") -> float:
+        ...
+
+    @overload
+    def predict(self, x: RDD["VectorLike"]) -> RDD[float]:
+        ...
+
+    def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[float, RDD[float]]:
+        """
+        Predict values for a single data point or an RDD of points using
+        the model trained.
+
+        .. versionadded:: 1.3.0
+
+        Notes
+        -----
+        In Python, predict cannot currently be used within an RDD
+        transformation or action.
+        Call predict directly on the RDD instead.
+        """
+        if isinstance(x, RDD):
+            return self.call("predict", x.map(_convert_to_vector))
+
+        else:
+            return self.call("predict", _convert_to_vector(x))
+
+    @since("1.3.0")
+    def numTrees(self) -> int:
+        """
+        Get number of trees in ensemble.
+        """
+        return self.call("numTrees")
+
+    @since("1.3.0")
+    def totalNumNodes(self) -> int:
+        """
+        Get total number of nodes, summed over all trees in the ensemble.
+        """
+        return self.call("totalNumNodes")
+
+    def __repr__(self) -> str:
+        """Summary of model"""
+        return self._java_model.toString()
+
+    @since("1.3.0")
+    def toDebugString(self) -> str:
+        """Full model"""
+        return self._java_model.toDebugString()
+
+
+class DecisionTreeModel(JavaModelWrapper, JavaSaveable, JavaLoader["DecisionTreeModel"]):
+    """
+    A decision tree model for classification or regression.
+
+    .. versionadded:: 1.1.0
+    """
+
+    @overload
+    def predict(self, x: "VectorLike") -> float:
+        ...
+
+    @overload
+    def predict(self, x: RDD["VectorLike"]) -> RDD[float]:
+        ...
+
+    def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[float, RDD[float]]:
+        """
+        Predict the label of one or more examples.
+
+        .. versionadded:: 1.1.0
+
+        Parameters
+        ----------
+        x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD`
+            Data point (feature vector), or an RDD of data points (feature
+            vectors).
+
+        Notes
+        -----
+        In Python, predict cannot currently be used within an RDD
+        transformation or action.
+        Call predict directly on the RDD instead.
+        """
+        if isinstance(x, RDD):
+            return self.call("predict", x.map(_convert_to_vector))
+
+        else:
+            return self.call("predict", _convert_to_vector(x))
+
+    @since("1.1.0")
+    def numNodes(self) -> int:
+        """Get number of nodes in tree, including leaf nodes."""
+        return self._java_model.numNodes()
+
+    @since("1.1.0")
+    def depth(self) -> int:
+        """
+        Get depth of tree (e.g. depth 0 means 1 leaf node, depth 1
+        means 1 internal node + 2 leaf nodes).
+        """
+        return self._java_model.depth()
+
+    def __repr__(self) -> str:
+        """summary of model."""
+        return self._java_model.toString()
+
+    @since("1.2.0")
+    def toDebugString(self) -> str:
+        """full model."""
+        return self._java_model.toDebugString()
+
+    @classmethod
+    def _java_loader_class(cls) -> str:
+        return "org.apache.spark.mllib.tree.model.DecisionTreeModel"
+
+
+class DecisionTree:
+    """
+    Learning algorithm for a decision tree model for classification or
+    regression.
+
+    .. versionadded:: 1.1.0
+    """
+
+    @classmethod
+    def _train(
+        cls,
+        data: RDD[LabeledPoint],
+        type: str,
+        numClasses: int,
+        features: Dict[int, int],
+        impurity: str = "gini",
+        maxDepth: int = 5,
+        maxBins: int = 32,
+        minInstancesPerNode: int = 1,
+        minInfoGain: float = 0.0,
+    ) -> DecisionTreeModel:
+        first = data.first()
+        assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint"
+        model = callMLlibFunc(
+            "trainDecisionTreeModel",
+            data,
+            type,
+            numClasses,
+            features,
+            impurity,
+            maxDepth,
+            maxBins,
+            minInstancesPerNode,
+            minInfoGain,
+        )
+        return DecisionTreeModel(model)
+
+    @classmethod
+    def trainClassifier(
+        cls,
+        data: RDD[LabeledPoint],
+        numClasses: int,
+        categoricalFeaturesInfo: Dict[int, int],
+        impurity: str = "gini",
+        maxDepth: int = 5,
+        maxBins: int = 32,
+        minInstancesPerNode: int = 1,
+        minInfoGain: float = 0.0,
+    ) -> DecisionTreeModel:
+        """
+        Train a decision tree model for classification.
+
+        .. versionadded:: 1.1.0
+
+        Parameters
+        ----------
+        data :  :py:class:`pyspark.RDD`
+            Training data: RDD of LabeledPoint. Labels should take values
+            {0, 1, ..., numClasses-1}.
+        numClasses : int
+            Number of classes for classification.
+        categoricalFeaturesInfo : dict
+            Map storing arity of categorical features. An entry (n -> k)
+            indicates that feature n is categorical with k categories
+            indexed from 0: {0, 1, ..., k-1}.
+        impurity : str, optional
+            Criterion used for information gain calculation.
+            Supported values: "gini" or "entropy".
+            (default: "gini")
+        maxDepth : int, optional
+            Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1
+            means 1 internal node + 2 leaf nodes).
+            (default: 5)
+        maxBins : int, optional
+            Number of bins used for finding splits at each node.
+            (default: 32)
+        minInstancesPerNode : int, optional
+            Minimum number of instances required at child nodes to create
+            the parent split.
+            (default: 1)
+        minInfoGain : float, optional
+            Minimum info gain required to create a split.
+            (default: 0.0)
+
+        Returns
+        -------
+        :py:class:`DecisionTreeModel`
+
+        Examples
+        --------
+        >>> from numpy import array
+        >>> from pyspark.mllib.regression import LabeledPoint
+        >>> from pyspark.mllib.tree import DecisionTree
+        >>>
+        >>> data = [
+        ...     LabeledPoint(0.0, [0.0]),
+        ...     LabeledPoint(1.0, [1.0]),
+        ...     LabeledPoint(1.0, [2.0]),
+        ...     LabeledPoint(1.0, [3.0])
+        ... ]
+        >>> model = DecisionTree.trainClassifier(sc.parallelize(data), 2, {})
+        >>> print(model)
+        DecisionTreeModel classifier of depth 1 with 3 nodes
+
+        >>> print(model.toDebugString())
+        DecisionTreeModel classifier of depth 1 with 3 nodes
+          If (feature 0 <= 0.5)
+           Predict: 0.0
+          Else (feature 0 > 0.5)
+           Predict: 1.0
+        >>> model.predict(array([1.0]))
+        1.0
+        >>> model.predict(array([0.0]))
+        0.0
+        >>> rdd = sc.parallelize([[1.0], [0.0]])
+        >>> model.predict(rdd).collect()
+        [1.0, 0.0]
+        """
+        return cls._train(
+            data,
+            "classification",
+            numClasses,
+            categoricalFeaturesInfo,
+            impurity,
+            maxDepth,
+            maxBins,
+            minInstancesPerNode,
+            minInfoGain,
+        )
+
+    @classmethod
+    @since("1.1.0")
+    def trainRegressor(
+        cls,
+        data: RDD[LabeledPoint],
+        categoricalFeaturesInfo: Dict[int, int],
+        impurity: str = "variance",
+        maxDepth: int = 5,
+        maxBins: int = 32,
+        minInstancesPerNode: int = 1,
+        minInfoGain: float = 0.0,
+    ) -> DecisionTreeModel:
+        """
+        Train a decision tree model for regression.
+
+        Parameters
+        ----------
+        data : :py:class:`pyspark.RDD`
+            Training data: RDD of LabeledPoint. Labels are real numbers.
+        categoricalFeaturesInfo : dict
+            Map storing arity of categorical features. An entry (n -> k)
+            indicates that feature n is categorical with k categories
+            indexed from 0: {0, 1, ..., k-1}.
+        impurity : str, optional
+            Criterion used for information gain calculation.
+            The only supported value for regression is "variance".
+            (default: "variance")
+        maxDepth : int, optional
+            Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1
+            means 1 internal node + 2 leaf nodes).
+            (default: 5)
+        maxBins : int, optional
+            Number of bins used for finding splits at each node.
+            (default: 32)
+        minInstancesPerNode : int, optional
+            Minimum number of instances required at child nodes to create
+            the parent split.
+            (default: 1)
+        minInfoGain : float, optional
+            Minimum info gain required to create a split.
+            (default: 0.0)
+
+        Returns
+        -------
+        :py:class:`DecisionTreeModel`
+
+        Examples
+        --------
+        >>> from pyspark.mllib.regression import LabeledPoint
+        >>> from pyspark.mllib.tree import DecisionTree
+        >>> from pyspark.mllib.linalg import SparseVector
+        >>>
+        >>> sparse_data = [
+        ...     LabeledPoint(0.0, SparseVector(2, {0: 0.0})),
+        ...     LabeledPoint(1.0, SparseVector(2, {1: 1.0})),
+        ...     LabeledPoint(0.0, SparseVector(2, {0: 0.0})),
+        ...     LabeledPoint(1.0, SparseVector(2, {1: 2.0}))
+        ... ]
+        >>>
+        >>> model = DecisionTree.trainRegressor(sc.parallelize(sparse_data), {})
+        >>> model.predict(SparseVector(2, {1: 1.0}))
+        1.0
+        >>> model.predict(SparseVector(2, {1: 0.0}))
+        0.0
+        >>> rdd = sc.parallelize([[0.0, 1.0], [0.0, 0.0]])
+        >>> model.predict(rdd).collect()
+        [1.0, 0.0]
+        """
+        return cls._train(
+            data,
+            "regression",
+            0,
+            categoricalFeaturesInfo,
+            impurity,
+            maxDepth,
+            maxBins,
+            minInstancesPerNode,
+            minInfoGain,
+        )
+
+
+@inherit_doc
+class RandomForestModel(TreeEnsembleModel, JavaLoader["RandomForestModel"]):
+    """
+    Represents a random forest model.
+
+    .. versionadded:: 1.2.0
+    """
+
+    @classmethod
+    def _java_loader_class(cls) -> str:
+        return "org.apache.spark.mllib.tree.model.RandomForestModel"
+
+
+class RandomForest:
+    """
+    Learning algorithm for a random forest model for classification or
+    regression.
+
+    .. versionadded:: 1.2.0
+    """
+
+    supportedFeatureSubsetStrategies: Tuple[str, ...] = ("auto", "all", "sqrt", "log2", "onethird")
+
+    @classmethod
+    def _train(
+        cls,
+        data: RDD[LabeledPoint],
+        algo: str,
+        numClasses: int,
+        categoricalFeaturesInfo: Dict[int, int],
+        numTrees: int,
+        featureSubsetStrategy: str,
+        impurity: str,
+        maxDepth: int,
+        maxBins: int,
+        seed: Optional[int],
+    ) -> RandomForestModel:
+        first = data.first()
+        assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint"
+        if featureSubsetStrategy not in cls.supportedFeatureSubsetStrategies:
+            raise ValueError("unsupported featureSubsetStrategy: %s" % featureSubsetStrategy)
+        if seed is None:
+            seed = random.randint(0, 1 << 30)
+        model = callMLlibFunc(
+            "trainRandomForestModel",
+            data,
+            algo,
+            numClasses,
+            categoricalFeaturesInfo,
+            numTrees,
+            featureSubsetStrategy,
+            impurity,
+            maxDepth,
+            maxBins,
+            seed,
+        )
+        return RandomForestModel(model)
+
+    @classmethod
+    def trainClassifier(
+        cls,
+        data: RDD[LabeledPoint],
+        numClasses: int,
+        categoricalFeaturesInfo: Dict[int, int],
+        numTrees: int,
+        featureSubsetStrategy: str = "auto",
+        impurity: str = "gini",
+        maxDepth: int = 4,
+        maxBins: int = 32,
+        seed: Optional[int] = None,
+    ) -> RandomForestModel:
+        """
+        Train a random forest model for binary or multiclass
+        classification.
+
+        .. versionadded:: 1.2.0
+
+        Parameters
+        ----------
+        data : :py:class:`pyspark.RDD`
+            Training dataset: RDD of LabeledPoint. Labels should take values
+            {0, 1, ..., numClasses-1}.
+        numClasses : int
+            Number of classes for classification.
+        categoricalFeaturesInfo : dict
+            Map storing arity of categorical features. An entry (n -> k)
+            indicates that feature n is categorical with k categories
+            indexed from 0: {0, 1, ..., k-1}.
+        numTrees : int
+            Number of trees in the random forest.
+        featureSubsetStrategy : str, optional
+            Number of features to consider for splits at each node.
+            Supported values: "auto", "all", "sqrt", "log2", "onethird".
+            If "auto" is set, this parameter is set based on numTrees:
+            if numTrees == 1, set to "all";
+            if numTrees > 1 (forest) set to "sqrt".
+            (default: "auto")
+        impurity : str, optional
+            Criterion used for information gain calculation.
+            Supported values: "gini" or "entropy".
+            (default: "gini")
+        maxDepth : int, optional
+            Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1
+            means 1 internal node + 2 leaf nodes).
+            (default: 4)
+        maxBins : int, optional
+            Maximum number of bins used for splitting features.
+            (default: 32)
+        seed : int, Optional
+            Random seed for bootstrapping and choosing feature subsets.
+            Set as None to generate seed based on system time.
+            (default: None)
+
+        Returns
+        -------
+        :py:class:`RandomForestModel`
+            that can be used for prediction.
+
+        Examples
+        --------
+        >>> from pyspark.mllib.regression import LabeledPoint
+        >>> from pyspark.mllib.tree import RandomForest
+        >>>
+        >>> data = [
+        ...     LabeledPoint(0.0, [0.0]),
+        ...     LabeledPoint(0.0, [1.0]),
+        ...     LabeledPoint(1.0, [2.0]),
+        ...     LabeledPoint(1.0, [3.0])
+        ... ]
+        >>> model = RandomForest.trainClassifier(sc.parallelize(data), 2, {}, 3, seed=42)
+        >>> model.numTrees()
+        3
+        >>> model.totalNumNodes()
+        7
+        >>> print(model)
+        TreeEnsembleModel classifier with 3 trees
+        >>> print(model.toDebugString())
+        TreeEnsembleModel classifier with 3 trees
+          Tree 0:
+            Predict: 1.0
+          Tree 1:
+            If (feature 0 <= 1.5)
+             Predict: 0.0
+            Else (feature 0 > 1.5)
+             Predict: 1.0
+          Tree 2:
+            If (feature 0 <= 1.5)
+             Predict: 0.0
+            Else (feature 0 > 1.5)
+             Predict: 1.0
+        >>> model.predict([2.0])
+        1.0
+        >>> model.predict([0.0])
+        0.0
+        >>> rdd = sc.parallelize([[3.0], [1.0]])
+        >>> model.predict(rdd).collect()
+        [1.0, 0.0]
+        """
+        return cls._train(
+            data,
+            "classification",
+            numClasses,
+            categoricalFeaturesInfo,
+            numTrees,
+            featureSubsetStrategy,
+            impurity,
+            maxDepth,
+            maxBins,
+            seed,
+        )
+
+    @classmethod
+    def trainRegressor(
+        cls,
+        data: RDD[LabeledPoint],
+        categoricalFeaturesInfo: Dict[int, int],
+        numTrees: int,
+        featureSubsetStrategy: str = "auto",
+        impurity: str = "variance",
+        maxDepth: int = 4,
+        maxBins: int = 32,
+        seed: Optional[int] = None,
+    ) -> RandomForestModel:
+        """
+        Train a random forest model for regression.
+
+        .. versionadded:: 1.2.0
+
+        Parameters
+        ----------
+        data : :py:class:`pyspark.RDD`
+            Training dataset: RDD of LabeledPoint. Labels are real numbers.
+        categoricalFeaturesInfo : dict
+            Map storing arity of categorical features. An entry (n -> k)
+            indicates that feature n is categorical with k categories
+            indexed from 0: {0, 1, ..., k-1}.
+        numTrees : int
+            Number of trees in the random forest.
+        featureSubsetStrategy : str, optional
+            Number of features to consider for splits at each node.
+            Supported values: "auto", "all", "sqrt", "log2", "onethird".
+            If "auto" is set, this parameter is set based on numTrees:
+
+            - if numTrees == 1, set to "all";
+            - if numTrees > 1 (forest) set to "onethird" for regression.
+
+            (default: "auto")
+        impurity : str, optional
+            Criterion used for information gain calculation.
+            The only supported value for regression is "variance".
+            (default: "variance")
+        maxDepth : int, optional
+            Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1
+            means 1 internal node + 2 leaf nodes).
+            (default: 4)
+        maxBins : int, optional
+            Maximum number of bins used for splitting features.
+            (default: 32)
+        seed : int, optional
+            Random seed for bootstrapping and choosing feature subsets.
+            Set as None to generate seed based on system time.
+            (default: None)
+
+        Returns
+        -------
+        :py:class:`RandomForestModel`
+            that can be used for prediction.
+
+        Examples
+        --------
+        >>> from pyspark.mllib.regression import LabeledPoint
+        >>> from pyspark.mllib.tree import RandomForest
+        >>> from pyspark.mllib.linalg import SparseVector
+        >>>
+        >>> sparse_data = [
+        ...     LabeledPoint(0.0, SparseVector(2, {0: 1.0})),
+        ...     LabeledPoint(1.0, SparseVector(2, {1: 1.0})),
+        ...     LabeledPoint(0.0, SparseVector(2, {0: 1.0})),
+        ...     LabeledPoint(1.0, SparseVector(2, {1: 2.0}))
+        ... ]
+        >>>
+        >>> model = RandomForest.trainRegressor(sc.parallelize(sparse_data), {}, 2, seed=42)
+        >>> model.numTrees()
+        2
+        >>> model.totalNumNodes()
+        4
+        >>> model.predict(SparseVector(2, {1: 1.0}))
+        1.0
+        >>> model.predict(SparseVector(2, {0: 1.0}))
+        0.5
+        >>> rdd = sc.parallelize([[0.0, 1.0], [1.0, 0.0]])
+        >>> model.predict(rdd).collect()
+        [1.0, 0.5]
+        """
+        return cls._train(
+            data,
+            "regression",
+            0,
+            categoricalFeaturesInfo,
+            numTrees,
+            featureSubsetStrategy,
+            impurity,
+            maxDepth,
+            maxBins,
+            seed,
+        )
+
+
+@inherit_doc
+class GradientBoostedTreesModel(TreeEnsembleModel, JavaLoader["GradientBoostedTreesModel"]):
+    """
+    Represents a gradient-boosted tree model.
+
+    .. versionadded:: 1.3.0
+    """
+
+    @classmethod
+    def _java_loader_class(cls) -> str:
+        return "org.apache.spark.mllib.tree.model.GradientBoostedTreesModel"
+
+
+class GradientBoostedTrees:
+    """
+    Learning algorithm for a gradient boosted trees model for
+    classification or regression.
+
+    .. versionadded:: 1.3.0
+    """
+
+    @classmethod
+    def _train(
+        cls,
+        data: RDD[LabeledPoint],
+        algo: str,
+        categoricalFeaturesInfo: Dict[int, int],
+        loss: str,
+        numIterations: int,
+        learningRate: float,
+        maxDepth: int,
+        maxBins: int,
+    ) -> GradientBoostedTreesModel:
+        first = data.first()
+        assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint"
+        model = callMLlibFunc(
+            "trainGradientBoostedTreesModel",
+            data,
+            algo,
+            categoricalFeaturesInfo,
+            loss,
+            numIterations,
+            learningRate,
+            maxDepth,
+            maxBins,
+        )
+        return GradientBoostedTreesModel(model)
+
+    @classmethod
+    def trainClassifier(
+        cls,
+        data: RDD[LabeledPoint],
+        categoricalFeaturesInfo: Dict[int, int],
+        loss: str = "logLoss",
+        numIterations: int = 100,
+        learningRate: float = 0.1,
+        maxDepth: int = 3,
+        maxBins: int = 32,
+    ) -> GradientBoostedTreesModel:
+        """
+        Train a gradient-boosted trees model for classification.
+
+        .. versionadded:: 1.3.0
+
+        Parameters
+        ----------
+        data : :py:class:`pyspark.RDD`
+            Training dataset: RDD of LabeledPoint. Labels should take values
+            {0, 1}.
+        categoricalFeaturesInfo : dict
+            Map storing arity of categorical features. An entry (n -> k)
+            indicates that feature n is categorical with k categories
+            indexed from 0: {0, 1, ..., k-1}.
+        loss : str, optional
+            Loss function used for minimization during gradient boosting.
+            Supported values: "logLoss", "leastSquaresError",
+            "leastAbsoluteError".
+            (default: "logLoss")
+        numIterations : int, optional
+            Number of iterations of boosting.
+            (default: 100)
+        learningRate : float, optional
+            Learning rate for shrinking the contribution of each estimator.
+            The learning rate should be between in the interval (0, 1].
+            (default: 0.1)
+        maxDepth : int, optional
+            Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1
+            means 1 internal node + 2 leaf nodes).
+            (default: 3)
+        maxBins : int, optional
+            Maximum number of bins used for splitting features. DecisionTree
+            requires maxBins >= max categories.
+            (default: 32)
+
+        Returns
+        -------
+        :py:class:`GradientBoostedTreesModel`
+            that can be used for prediction.
+
+        Examples
+        --------
+        >>> from pyspark.mllib.regression import LabeledPoint
+        >>> from pyspark.mllib.tree import GradientBoostedTrees
+        >>>
+        >>> data = [
+        ...     LabeledPoint(0.0, [0.0]),
+        ...     LabeledPoint(0.0, [1.0]),
+        ...     LabeledPoint(1.0, [2.0]),
+        ...     LabeledPoint(1.0, [3.0])
+        ... ]
+        >>>
+        >>> model = GradientBoostedTrees.trainClassifier(sc.parallelize(data), {}, numIterations=10)
+        >>> model.numTrees()
+        10
+        >>> model.totalNumNodes()
+        30
+        >>> print(model)  # it already has newline
+        TreeEnsembleModel classifier with 10 trees
+        >>> model.predict([2.0])
+        1.0
+        >>> model.predict([0.0])
+        0.0
+        >>> rdd = sc.parallelize([[2.0], [0.0]])
+        >>> model.predict(rdd).collect()
+        [1.0, 0.0]
+        """
+        return cls._train(
+            data,
+            "classification",
+            categoricalFeaturesInfo,
+            loss,
+            numIterations,
+            learningRate,
+            maxDepth,
+            maxBins,
+        )
+
+    @classmethod
+    def trainRegressor(
+        cls,
+        data: RDD[LabeledPoint],
+        categoricalFeaturesInfo: Dict[int, int],
+        loss: str = "leastSquaresError",
+        numIterations: int = 100,
+        learningRate: float = 0.1,
+        maxDepth: int = 3,
+        maxBins: int = 32,
+    ) -> GradientBoostedTreesModel:
+        """
+        Train a gradient-boosted trees model for regression.
+
+        .. versionadded:: 1.3.0
+
+        Parameters
+        ----------
+        data :
+            Training dataset: RDD of LabeledPoint. Labels are real numbers.
+        categoricalFeaturesInfo : dict
+            Map storing arity of categorical features. An entry (n -> k)
+            indicates that feature n is categorical with k categories
+            indexed from 0: {0, 1, ..., k-1}.
+        loss : str, optional
+            Loss function used for minimization during gradient boosting.
+            Supported values: "logLoss", "leastSquaresError",
+            "leastAbsoluteError".
+            (default: "leastSquaresError")
+        numIterations : int, optional
+            Number of iterations of boosting.
+            (default: 100)
+        learningRate : float, optional
+            Learning rate for shrinking the contribution of each estimator.
+            The learning rate should be between in the interval (0, 1].
+            (default: 0.1)
+        maxDepth : int, optional
+            Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1
+            means 1 internal node + 2 leaf nodes).
+            (default: 3)
+        maxBins : int, optional
+            Maximum number of bins used for splitting features. DecisionTree
+            requires maxBins >= max categories.
+            (default: 32)
+
+        Returns
+        -------
+        :py:class:`GradientBoostedTreesModel`
+            that can be used for prediction.
+
+        Examples
+        --------
+        >>> from pyspark.mllib.regression import LabeledPoint
+        >>> from pyspark.mllib.tree import GradientBoostedTrees
+        >>> from pyspark.mllib.linalg import SparseVector
+        >>>
+        >>> sparse_data = [
+        ...     LabeledPoint(0.0, SparseVector(2, {0: 1.0})),
+        ...     LabeledPoint(1.0, SparseVector(2, {1: 1.0})),
+        ...     LabeledPoint(0.0, SparseVector(2, {0: 1.0})),
+        ...     LabeledPoint(1.0, SparseVector(2, {1: 2.0}))
+        ... ]
+        >>>
+        >>> data = sc.parallelize(sparse_data)
+        >>> model = GradientBoostedTrees.trainRegressor(data, {}, numIterations=10)
+        >>> model.numTrees()
+        10
+        >>> model.totalNumNodes()
+        12
+        >>> model.predict(SparseVector(2, {1: 1.0}))
+        1.0
+        >>> model.predict(SparseVector(2, {0: 1.0}))
+        0.0
+        >>> rdd = sc.parallelize([[0.0, 1.0], [1.0, 0.0]])
+        >>> model.predict(rdd).collect()
+        [1.0, 0.0]
+        """
+        return cls._train(
+            data,
+            "regression",
+            categoricalFeaturesInfo,
+            loss,
+            numIterations,
+            learningRate,
+            maxDepth,
+            maxBins,
+        )
+
+
+def _test() -> None:
+    import doctest
+
+    globs = globals().copy()
+    from pyspark.sql import SparkSession
+
+    spark = SparkSession.builder.master("local[4]").appName("mllib.tree tests").getOrCreate()
+    globs["sc"] = spark.sparkContext
+    (failure_count, test_count) = doctest.testmod(
+        globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()