fixed cluster export

vanzin · Oct 30, 2014 · aba5ee1 · aba5ee1
1 parent cd6c07c
commit aba5ee1
Showing 1 changed file with 13 additions and 15 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/export/pmml/KMeansPMMLModelExport.scala b/mllib/src/main/scala/org/apache/spark/mllib/export/pmml/KMeansPMMLModelExport.scala
@@ -61,32 +61,30 @@ class KMeansPMMLModelExport(model : KMeansModel) extends PMMLModelExport{
 
        var miningSchema = new MiningSchema()
 
-       for ( i <- 0 to (clusterCenter.size - 1)) {
-         fields(i) = FieldName.create("field_" + i)
-         dataDictionary
-            .withDataFields(new DataField(fields(i), OpType.CONTINUOUS, DataType.DOUBLE))
-         miningSchema
-            .withMiningFields(new MiningField(fields(i))
-            .withUsageType(FieldUsageType.ACTIVE))
-         }
-
        var comparisonMeasure = new ComparisonMeasure()
             .withKind(Kind.DISTANCE)
             .withMeasure(new SquaredEuclidean()
        );
 
-       dataDictionary.withNumberOfFields((dataDictionary.getDataFields()).size());
-
-       pmml.setDataDictionary(dataDictionary);
-
        var clusteringModel = new ClusteringModel(miningSchema, comparisonMeasure, 
         MiningFunctionType.CLUSTERING, ModelClass.CENTER_BASED, model.clusterCenters.length)
         .withModelName("k-means");
 
        for ( i <- 0 to (clusterCenter.size - 1)) {
+         fields(i) = FieldName.create("field_" + i)
+         dataDictionary
+            .withDataFields(new DataField(fields(i), OpType.CONTINUOUS, DataType.DOUBLE))
+         miningSchema
+            .withMiningFields(new MiningField(fields(i))
+            .withUsageType(FieldUsageType.ACTIVE))
          clusteringModel.withClusteringFields(
              new ClusteringField(fields(i)).withCompareFunction(CompareFunctionType.ABS_DIFF)
-         )
+         )     
+       }
+
+       dataDictionary.withNumberOfFields((dataDictionary.getDataFields()).size());
+
+       for ( i <- 0 to (model.clusterCenters.size - 1)) {
          var cluster = new Cluster()
             .withName("cluster_" + i)
             .withArray(new org.dmg.pmml.Array()
@@ -95,10 +93,10 @@ class KMeansPMMLModelExport(model : KMeansModel) extends PMMLModelExport{
             .withValue(model.clusterCenters(i).toArray.mkString(" ")))
             // we don't have the size of the single cluster but only the centroids (withValue)
             // .withSize(value)
-
          clusteringModel.withClusters(cluster)
        }
 
+       pmml.setDataDictionary(dataDictionary);
        pmml.withModels(clusteringModel);
 
      }