IGNITE-11654: [ML] Memory leak in KNNClassificationModel (#6392)

apache · Apr 3, 2019 · db38352 · db38352
1 parent 4daa681
commit db38352
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 4 deletions.
diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/clustering/kmeans/KMeansModel.java b/modules/ml/src/main/java/org/apache/ignite/ml/clustering/kmeans/KMeansModel.java
@@ -68,7 +68,7 @@ public DistanceMeasure distanceMeasure() {
      *
      * @param vec Vector.
      */
-    public Integer predict(Vector vec) {
+    @Override public Integer predict(Vector vec) {
         int res = -1;
         double minDist = Double.POSITIVE_INFINITY;
 

diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/inference/Model.java b/modules/ml/src/main/java/org/apache/ignite/ml/inference/Model.java
@@ -33,5 +33,5 @@ public interface Model<I, O> extends AutoCloseable {
     public O predict(I input);
 
     /** {@inheritDoc} */
-    public void close();
-}
+    @Override public void close();
+}
diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/knn/NNClassificationModel.java b/modules/ml/src/main/java/org/apache/ignite/ml/knn/NNClassificationModel.java
@@ -229,5 +229,5 @@ protected void copyParametersFrom(NNClassificationModel mdl) {
     }
 
     /** */
-    public abstract <P> void saveModel(Exporter<KNNModelFormat, P> exporter, P path);
+    @Override public abstract <P> void saveModel(Exporter<KNNModelFormat, P> exporter, P path);
 }
diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/knn/classification/KNNClassificationModel.java b/modules/ml/src/main/java/org/apache/ignite/ml/knn/classification/KNNClassificationModel.java
@@ -38,6 +38,9 @@
 
 /**
  * kNN algorithm model to solve multi-class classification task.
+ *
+ * NOTE: This model is especial, because it takes dataset as an input.
+ * The dataset (or datasets in case of model updating) are closing automatically.
  */
 public class KNNClassificationModel extends NNClassificationModel implements Exportable<KNNModelFormat> {
     /** */
@@ -144,4 +147,16 @@ public void copyStateFrom(KNNClassificationModel mdl) {
         this.copyParametersFrom(mdl);
         datasets.addAll(mdl.datasets);
     }
+
+    /** {@inheritDoc} */
+    @Override public void close() {
+        for (int i = 0; i < datasets.size(); i++) {
+            try {
+                datasets.get(i).close();
+            }
+            catch (Exception e) {
+                e.printStackTrace();
+            }
+        }
+    }
 }