diff options
author | Jayant Shekar <jayant@user-MBPMBA-3.local> | 2015-10-23 08:45:13 -0700 |
---|---|---|
committer | Xiangrui Meng <meng@databricks.com> | 2015-10-23 08:45:13 -0700 |
commit | 4e38defae13b2b13e196b4d172722ef5e6266c66 (patch) | |
tree | 727d246ccd43d9860249a947aee76984c2ab930f /mllib/src/test/scala/org/apache | |
parent | 282a15f78e08f0dc9e696945be4fc973011a96d9 (diff) | |
download | spark-4e38defae13b2b13e196b4d172722ef5e6266c66.tar.gz spark-4e38defae13b2b13e196b4d172722ef5e6266c66.tar.bz2 spark-4e38defae13b2b13e196b4d172722ef5e6266c66.zip |
[SPARK-6723] [MLLIB] Model import/export for ChiSqSelector
This is a PR for Parquet-based model import/export.
* Added save/load for ChiSqSelectorModel
* Updated the test suite ChiSqSelectorSuite
Author: Jayant Shekar <jayant@user-MBPMBA-3.local>
Closes #6785 from jayantshekhar/SPARK-6723.
Diffstat (limited to 'mllib/src/test/scala/org/apache')
-rw-r--r-- | mllib/src/test/scala/org/apache/spark/mllib/feature/ChiSqSelectorSuite.scala | 26 |
1 files changed, 26 insertions, 0 deletions
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/feature/ChiSqSelectorSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/feature/ChiSqSelectorSuite.scala index 889727fb55..734800a9af 100644 --- a/mllib/src/test/scala/org/apache/spark/mllib/feature/ChiSqSelectorSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/mllib/feature/ChiSqSelectorSuite.scala @@ -21,6 +21,7 @@ import org.apache.spark.SparkFunSuite import org.apache.spark.mllib.linalg.Vectors import org.apache.spark.mllib.regression.LabeledPoint import org.apache.spark.mllib.util.MLlibTestSparkContext +import org.apache.spark.util.Utils class ChiSqSelectorSuite extends SparkFunSuite with MLlibTestSparkContext { @@ -63,4 +64,29 @@ class ChiSqSelectorSuite extends SparkFunSuite with MLlibTestSparkContext { }.collect().toSet assert(filteredData == preFilteredData) } + + test("model load / save") { + val model = ChiSqSelectorSuite.createModel() + val tempDir = Utils.createTempDir() + val path = tempDir.toURI.toString + try { + model.save(sc, path) + val sameModel = ChiSqSelectorModel.load(sc, path) + ChiSqSelectorSuite.checkEqual(model, sameModel) + } finally { + Utils.deleteRecursively(tempDir) + } + } +} + +object ChiSqSelectorSuite extends SparkFunSuite { + + def createModel(): ChiSqSelectorModel = { + val arr = Array(1, 2, 3, 4) + new ChiSqSelectorModel(arr) + } + + def checkEqual(a: ChiSqSelectorModel, b: ChiSqSelectorModel): Unit = { + assert(a.selectedFeatures.deep == b.selectedFeatures.deep) + } } |