Skip to content

Commit aba5ee1

Browse files
committed
fixed cluster export
1 parent cd6c07c commit aba5ee1

File tree

1 file changed

+13
-15
lines changed

1 file changed

+13
-15
lines changed

mllib/src/main/scala/org/apache/spark/mllib/export/pmml/KMeansPMMLModelExport.scala

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -61,32 +61,30 @@ class KMeansPMMLModelExport(model : KMeansModel) extends PMMLModelExport{
6161

6262
var miningSchema = new MiningSchema()
6363

64-
for ( i <- 0 to (clusterCenter.size - 1)) {
65-
fields(i) = FieldName.create("field_" + i)
66-
dataDictionary
67-
.withDataFields(new DataField(fields(i), OpType.CONTINUOUS, DataType.DOUBLE))
68-
miningSchema
69-
.withMiningFields(new MiningField(fields(i))
70-
.withUsageType(FieldUsageType.ACTIVE))
71-
}
72-
7364
var comparisonMeasure = new ComparisonMeasure()
7465
.withKind(Kind.DISTANCE)
7566
.withMeasure(new SquaredEuclidean()
7667
);
7768

78-
dataDictionary.withNumberOfFields((dataDictionary.getDataFields()).size());
79-
80-
pmml.setDataDictionary(dataDictionary);
81-
8269
var clusteringModel = new ClusteringModel(miningSchema, comparisonMeasure,
8370
MiningFunctionType.CLUSTERING, ModelClass.CENTER_BASED, model.clusterCenters.length)
8471
.withModelName("k-means");
8572

8673
for ( i <- 0 to (clusterCenter.size - 1)) {
74+
fields(i) = FieldName.create("field_" + i)
75+
dataDictionary
76+
.withDataFields(new DataField(fields(i), OpType.CONTINUOUS, DataType.DOUBLE))
77+
miningSchema
78+
.withMiningFields(new MiningField(fields(i))
79+
.withUsageType(FieldUsageType.ACTIVE))
8780
clusteringModel.withClusteringFields(
8881
new ClusteringField(fields(i)).withCompareFunction(CompareFunctionType.ABS_DIFF)
89-
)
82+
)
83+
}
84+
85+
dataDictionary.withNumberOfFields((dataDictionary.getDataFields()).size());
86+
87+
for ( i <- 0 to (model.clusterCenters.size - 1)) {
9088
var cluster = new Cluster()
9189
.withName("cluster_" + i)
9290
.withArray(new org.dmg.pmml.Array()
@@ -95,10 +93,10 @@ class KMeansPMMLModelExport(model : KMeansModel) extends PMMLModelExport{
9593
.withValue(model.clusterCenters(i).toArray.mkString(" ")))
9694
// we don't have the size of the single cluster but only the centroids (withValue)
9795
// .withSize(value)
98-
9996
clusteringModel.withClusters(cluster)
10097
}
10198

99+
pmml.setDataDictionary(dataDictionary);
102100
pmml.withModels(clusteringModel);
103101

104102
}

0 commit comments

Comments
 (0)