17
17
18
18
package org .apache .spark .mllib .export .pmml
19
19
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
23
21
import org .apache .spark .mllib .clustering .KMeansModel
22
+ import org .apache .spark .mllib .export .ModelExportFactory
24
23
import org .apache .spark .mllib .export .ModelExportType
24
+ import org .apache .spark .mllib .linalg .Vectors
25
+ import org .scalatest .FunSuite
25
26
import org .dmg .pmml .ClusteringModel
26
- import javax .xml .parsers .DocumentBuilderFactory
27
- import java .io .ByteArrayOutputStream
28
27
29
28
class KMeansPMMLModelExportSuite extends FunSuite {
30
29
@@ -43,7 +42,7 @@ class KMeansPMMLModelExportSuite extends FunSuite{
43
42
44
43
// assert that the PMML format is as expected
45
44
assert(modelExport.isInstanceOf [PMMLModelExport ])
46
- var pmml = modelExport.asInstanceOf [PMMLModelExport ].getPmml()
45
+ val pmml = modelExport.asInstanceOf [PMMLModelExport ].getPmml()
47
46
assert(pmml.getHeader().getDescription() === " k-means clustering" )
48
47
// check that the number of fields match the single vector size
49
48
assert(pmml.getDataDictionary().getNumberOfFields() === clusterCenters(0 ).size)
@@ -52,8 +51,8 @@ class KMeansPMMLModelExportSuite extends FunSuite{
52
51
assert(pmml.getModels().get(0 ).asInstanceOf [ClusteringModel ].getNumberOfClusters() === clusterCenters.size)
53
52
54
53
// 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")
57
56
58
57
}
59
58
0 commit comments