@@ -70,7 +70,7 @@ object DecisionTreeExample {
70
70
val parser = new OptionParser [Params ](" DecisionTreeExample" ) {
71
71
head(" DecisionTreeExample: an example decision tree app." )
72
72
opt[String ](" algo" )
73
- .text(s " algorithm (Classification, Regression ), default: ${defaultParams.algo}" )
73
+ .text(s " algorithm (classification, regression ), default: ${defaultParams.algo}" )
74
74
.action((x, c) => c.copy(algo = x))
75
75
opt[Int ](" maxDepth" )
76
76
.text(s " max depth of the tree, default: ${defaultParams.maxDepth}" )
@@ -221,18 +221,23 @@ object DecisionTreeExample {
221
221
// (1) For classification, re-index classes.
222
222
val labelColName = if (algo == " classification" ) " indexedLabel" else " label"
223
223
if (algo == " classification" ) {
224
- val labelIndexer = new StringIndexer ().setInputCol(" labelString" ).setOutputCol(labelColName)
224
+ val labelIndexer = new StringIndexer ()
225
+ .setInputCol(" labelString" )
226
+ .setOutputCol(labelColName)
225
227
stages += labelIndexer
226
228
}
227
229
// (2) Identify categorical features using VectorIndexer.
228
230
// Features with more than maxCategories values will be treated as continuous.
229
- val featuresIndexer = new VectorIndexer ().setInputCol(" features" )
230
- .setOutputCol(" indexedFeatures" ).setMaxCategories(10 )
231
+ val featuresIndexer = new VectorIndexer ()
232
+ .setInputCol(" features" )
233
+ .setOutputCol(" indexedFeatures" )
234
+ .setMaxCategories(10 )
231
235
stages += featuresIndexer
232
236
// (3) Learn DecisionTree
233
237
val dt = algo match {
234
238
case " classification" =>
235
- new DecisionTreeClassifier ().setFeaturesCol(" indexedFeatures" )
239
+ new DecisionTreeClassifier ()
240
+ .setFeaturesCol(" indexedFeatures" )
236
241
.setLabelCol(labelColName)
237
242
.setMaxDepth(params.maxDepth)
238
243
.setMaxBins(params.maxBins)
@@ -241,7 +246,8 @@ object DecisionTreeExample {
241
246
.setCacheNodeIds(params.cacheNodeIds)
242
247
.setCheckpointInterval(params.checkpointInterval)
243
248
case " regression" =>
244
- new DecisionTreeRegressor ().setFeaturesCol(" indexedFeatures" )
249
+ new DecisionTreeRegressor ()
250
+ .setFeaturesCol(" indexedFeatures" )
245
251
.setLabelCol(labelColName)
246
252
.setMaxDepth(params.maxDepth)
247
253
.setMaxBins(params.maxBins)
0 commit comments