@@ -131,10 +131,10 @@ object DecisionTreeModel extends Loader[DecisionTreeModel] {
131131 }
132132
133133 case class SplitData (
134- feature : Int ,
135- threshold : Double ,
136- featureType : Int ,
137- categories : Seq [Double ]) { // TODO: Change to List once SPARK-3365 is fixed
134+ feature : Int ,
135+ threshold : Double ,
136+ featureType : Int ,
137+ categories : Seq [Double ]) { // TODO: Change to List once SPARK-3365 is fixed
138138 def toSplit : Split = {
139139 new Split (feature, threshold, FeatureType (featureType), categories.toList)
140140 }
@@ -152,15 +152,15 @@ object DecisionTreeModel extends Loader[DecisionTreeModel] {
152152
153153 /** Model data for model import/export */
154154 case class NodeData (
155- treeId : Int ,
156- nodeId : Int ,
157- predict : PredictData ,
158- impurity : Double ,
159- isLeaf : Boolean ,
160- split : Option [SplitData ],
161- leftNodeId : Option [Int ],
162- rightNodeId : Option [Int ],
163- infoGain : Option [Double ])
155+ treeId : Int ,
156+ nodeId : Int ,
157+ predict : PredictData ,
158+ impurity : Double ,
159+ isLeaf : Boolean ,
160+ split : Option [SplitData ],
161+ leftNodeId : Option [Int ],
162+ rightNodeId : Option [Int ],
163+ infoGain : Option [Double ])
164164
165165 object NodeData {
166166 def apply (treeId : Int , n : Node ): NodeData = {
@@ -174,8 +174,8 @@ object DecisionTreeModel extends Loader[DecisionTreeModel] {
174174 val leftNodeId = if (r.isNullAt(6 )) None else Some (r.getInt(6 ))
175175 val rightNodeId = if (r.isNullAt(7 )) None else Some (r.getInt(7 ))
176176 val infoGain = if (r.isNullAt(8 )) None else Some (r.getDouble(8 ))
177- NodeData (r.getInt(0 ), r.getInt(1 ), PredictData (r.getStruct(2 )), r.getDouble(3 ), r.getBoolean( 4 ),
178- split, leftNodeId, rightNodeId, infoGain)
177+ NodeData (r.getInt(0 ), r.getInt(1 ), PredictData (r.getStruct(2 )), r.getDouble(3 ),
178+ r.getBoolean( 4 ), split, leftNodeId, rightNodeId, infoGain)
179179 }
180180 }
181181
@@ -210,9 +210,8 @@ object DecisionTreeModel extends Loader[DecisionTreeModel] {
210210 assert(trees.size == 1 ,
211211 " Decision tree should contain exactly one tree but got ${trees.size} trees." )
212212 val model = new DecisionTreeModel (trees(0 ), Algo .fromString(algo))
213- assert(model.numNodes == numNodes,
214- s " Unable to load DecisionTreeModel data from: $datapath. " +
215- s " Expected $numNodes nodes but found ${model.numNodes}" )
213+ assert(model.numNodes == numNodes, s " Unable to load DecisionTreeModel data from: $datapath. " +
214+ s " Expected $numNodes nodes but found ${model.numNodes}" )
216215 model
217216 }
218217
0 commit comments