From 080e983bca0bbc311a4cf0b8e88f90caba3243b9 Mon Sep 17 00:00:00 2001
From: Kousuke Saruta <sarutak@oss.nttdata.co.jp>
Date: Sat, 19 Dec 2015 04:59:34 +0900
Subject: [PATCH] Modified test

---
 .../org/apache/spark/ml/param/params.scala    |  5 +----
 .../apache/spark/ml/param/ParamsSuite.scala   | 21 +++++++++++++++++--
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/param/params.scala b/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
index 4c900f46a327a..c0546695e487b 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
@@ -863,10 +863,7 @@ final class ParamMap private[ml] (private val map: mutable.Map[Param[Any], Any])
     // returns the instance of collections.Map, not mutable.Map.
     // Otherwise, we get ClassCastException.
     // Not using filterKeys also avoid SI-6654
-    val filtered = map.filter {
-      case (k, _) =>
-        k.parent == parent.uid
-    }
+    val filtered = map.filter { case (k, _) => k.parent == parent.uid }
     new ParamMap(filtered)
   }
 
diff --git a/mllib/src/test/scala/org/apache/spark/ml/param/ParamsSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/param/ParamsSuite.scala
index ddcedc4121cc9..f00e096bf33a6 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/param/ParamsSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/param/ParamsSuite.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.ml.param
 
+import java.io.{ByteArrayOutputStream, NotSerializableException, ObjectOutputStream}
+
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.util.MyParams
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
@@ -367,8 +369,23 @@ class ParamsSuite extends SparkFunSuite {
         assert(p.parent === params1.uid)
     }
 
-    // Following assertion is to avoid SI-6654
-    assert(filteredParamMap.isInstanceOf[Serializable])
+    // At the previous implementation of ParamMap#filter,
+    // mutable.Map#filterKeys was used internally but
+    // the return type of the method is not serializable (see SI-6654).
+    // Now mutable.Map#filter is used instead of filterKeys and the return type is serializable.
+    // So let's ensure serializability.
+    val objOut = new ObjectOutputStream(new ByteArrayOutputStream())
+    try {
+      objOut.writeObject(filteredParamMap)
+    } catch {
+      case _: NotSerializableException =>
+        fail("The field of ParamMap 'map' may not be serializable. " +
+          "See SI-6654 and the implementation of ParamMap#filter")
+      case e: Exception =>
+        fail(s"Exception was thrown unexpectedly during the serializability test: ${e.getMessage}")
+    } finally {
+      objOut.close()
+    }
   }
 }