Skip to content

Commit ae8b993

Browse files
committed
updated some commented tests to use the new ModelExporter object
reordered the imports
1 parent df8a89e commit ae8b993

File tree

1 file changed

+7
-8
lines changed

1 file changed

+7
-8
lines changed

mllib/src/test/scala/org/apache/spark/mllib/export/pmml/KMeansPMMLModelExportSuite.scala

+7-8
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,13 @@
1717

1818
package org.apache.spark.mllib.export.pmml
1919

20-
import org.scalatest.FunSuite
21-
import org.apache.spark.mllib.linalg.Vectors
22-
import org.apache.spark.mllib.export.ModelExportFactory
20+
import scala.annotation.varargs
2321
import org.apache.spark.mllib.clustering.KMeansModel
22+
import org.apache.spark.mllib.export.ModelExportFactory
2423
import org.apache.spark.mllib.export.ModelExportType
24+
import org.apache.spark.mllib.linalg.Vectors
25+
import org.scalatest.FunSuite
2526
import org.dmg.pmml.ClusteringModel
26-
import javax.xml.parsers.DocumentBuilderFactory
27-
import java.io.ByteArrayOutputStream
2827

2928
class KMeansPMMLModelExportSuite extends FunSuite{
3029

@@ -43,7 +42,7 @@ class KMeansPMMLModelExportSuite extends FunSuite{
4342

4443
//assert that the PMML format is as expected
4544
assert(modelExport.isInstanceOf[PMMLModelExport])
46-
var pmml = modelExport.asInstanceOf[PMMLModelExport].getPmml()
45+
val pmml = modelExport.asInstanceOf[PMMLModelExport].getPmml()
4746
assert(pmml.getHeader().getDescription() === "k-means clustering")
4847
//check that the number of fields match the single vector size
4948
assert(pmml.getDataDictionary().getNumberOfFields() === clusterCenters(0).size)
@@ -52,8 +51,8 @@ class KMeansPMMLModelExportSuite extends FunSuite{
5251
assert(pmml.getModels().get(0).asInstanceOf[ClusteringModel].getNumberOfClusters() === clusterCenters.size)
5352

5453
//manual checking
55-
//modelExport.asInstanceOf[PMMLModelExport].save(System.out)
56-
//modelExport.asInstanceOf[PMMLModelExport].saveLocalFile("/tmp/kmeans.xml")
54+
//ModelExporter.toPMML(kmeansModel,new StreamResult(System.out))
55+
//ModelExporter.toPMML(kmeansModel,"/tmp/kmeans.xml")
5756

5857
}
5958

0 commit comments

Comments
 (0)