Github user mengxr commented on a diff in the pull request:
https://github.com/apache/spark/pull/1798#discussion_r15854074
--- Diff:
mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala ---
@@ -300,6 +293,198 @@ object DecisionTree extends Serializable with Logging
{
new DecisionTree(strategy).train(input)
}
+ /**
+ * Method to train a decision tree model.
+ * The method supports binary and multiclass classification and
regression.
+ * This version takes basic types, for consistency with Python API.
+ *
+ * @param input Training dataset: RDD of
[[org.apache.spark.mllib.regression.LabeledPoint]].
+ * For classification, labels should take values {0, 1,
..., numClasses-1}.
+ * For regression, labels are real numbers.
+ * @param algo "classification" or "regression"
+ * @param numClassesForClassification number of classes for
classification. Default value of 2.
+ * @param categoricalFeaturesInfo Map storing arity of categorical
features.
+ * E.g., an entry (n -> k) indicates that
feature n is categorical
+ * with k categories indexed from 0: {0,
1, ..., k-1}.
+ * @param impurity criterion used for information gain calculation
+ * @param maxDepth Maximum depth of the tree.
+ * E.g., depth 0 means 1 leaf node; depth 1 means 1
internal node + 2 leaf nodes.
+ * @param maxBins maximum number of bins used for splitting features
+ * (default Python value = 100)
+ * @return DecisionTreeModel that can be used for prediction
+ */
+ def train(
+ input: RDD[LabeledPoint],
+ algo: String,
+ numClassesForClassification: Int,
+ categoricalFeaturesInfo: Map[Int, Int],
+ impurity: String,
+ maxDepth: Int,
+ maxBins: Int): DecisionTreeModel = {
+ val algoType = Algo.stringToAlgo(algo)
+ val impurityType = Impurities.stringToImpurity(impurity)
+ train(input, algoType, impurityType, maxDepth,
numClassesForClassification, maxBins, Sort,
+ categoricalFeaturesInfo)
+ }
+
+ /**
+ * Method to train a decision tree model.
+ * The method supports binary and multiclass classification and
regression.
+ * This version takes basic types, for consistency with Python API.
+ * This version is Java-friendly, taking a Java map for
categoricalFeaturesInfo.
+ *
+ * @param input Training dataset: RDD of
[[org.apache.spark.mllib.regression.LabeledPoint]].
+ * For classification, labels should take values {0, 1,
..., numClasses-1}.
+ * For regression, labels are real numbers.
+ * @param algo "classification" or "regression"
+ * @param numClassesForClassification number of classes for
classification. Default value of 2.
+ * @param categoricalFeaturesInfo Map storing arity of categorical
features.
+ * E.g., an entry (n -> k) indicates that
feature n is categorical
+ * with k categories indexed from 0: {0,
1, ..., k-1}.
+ * @param impurity criterion used for information gain calculation
+ * @param maxDepth Maximum depth of the tree.
+ * E.g., depth 0 means 1 leaf node; depth 1 means 1
internal node + 2 leaf nodes.
+ * @param maxBins maximum number of bins used for splitting features
+ * (default Python value = 100)
+ * @return DecisionTreeModel that can be used for prediction
+ */
+ def train(
+ input: RDD[LabeledPoint],
+ algo: String,
+ numClassesForClassification: Int,
+ categoricalFeaturesInfo: java.util.Map[java.lang.Integer,
java.lang.Integer],
+ impurity: String,
+ maxDepth: Int,
+ maxBins: Int): DecisionTreeModel = {
+ train(input, algo, numClassesForClassification,
+ categoricalFeaturesInfo.asInstanceOf[java.util.Map[Int,
Int]].asScala.toMap,
+ impurity, maxDepth, maxBins)
+ }
+
+ /**
+ * Method to train a decision tree model for binary or multiclass
classification.
+ * This version takes basic types, for consistency with Python API.
+ *
+ * @param input Training dataset: RDD of
[[org.apache.spark.mllib.regression.LabeledPoint]].
+ * Labels should take values {0, 1, ..., numClasses-1}.
+ * @param numClassesForClassification number of classes for
classification.
+ * @param categoricalFeaturesInfo Map storing arity of categorical
features.
+ * E.g., an entry (n -> k) indicates that
feature n is categorical
+ * with k categories indexed from 0: {0,
1, ..., k-1}.
+ * (default Python value = {}, i.e., no
categorical features)
--- End diff --
We don't need to mention Python's default value here. Btw, having the empty
map as the default value may be error-prone. What do you think?
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
