TST: classifier doesn't have weightCol

facaiy · facaiy · commit 9ba0e2be0b6b · 2017-07-19T13:55:44.000+08:00
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala
@@ -158,9 +158,12 @@ class OneVsRestSuite extends SparkFunSuite with MLlibTestSparkContext with Defau
 
   test("SPARK-21306: OneVsRest should support setWeightCol") {
     val dataset2 = dataset.withColumn("weight", lit(1))
+    // classifier inherits hasWeightCol
     val ova = new OneVsRest().setWeightCol("weight").setClassifier(new LogisticRegression())
-    val ovaModel = ova.fit(dataset2)
-    assert(ovaModel !== null)
+    assert(ova.fit(dataset2) !== null)
+    // classifier doesn't inherit hasWeightCol
+    val ova2 = new OneVsRest().setWeightCol("weight").setClassifier(new DecisionTreeClassifier())
+    assert(ova2.fit(dataset2) !== null)
   }
 
   test("OneVsRest.copy and OneVsRestModel.copy") {
diff --git a/python/pyspark/ml/tests.py b/python/pyspark/ml/tests.py
@@ -1260,11 +1260,18 @@ def test_support_for_weightCol(self):
                                          (1.0, Vectors.sparse(2, [], []), 1.0),
                                          (2.0, Vectors.dense(0.5, 0.5), 1.0)],
                                         ["label", "features", "weight"])
+        # classifier inherits hasWeightCol
         lr = LogisticRegression(maxIter=5, regParam=0.01)
         ovr = OneVsRest(classifier=lr, weightCol="weight")
         self.assertIsNotNone(ovr.fit(df))
         ovr2 = OneVsRest(classifier=lr).setWeightCol("weight")
         self.assertIsNotNone(ovr2.fit(df))
+        # classifier doesn't inherit hasWeightCol
+        dt = DecisionTreeClassifier()
+        ovr3 = OneVsRest(classifier=dt, weightCol="weight")
+        self.assertIsNotNone(ovr3.fit(df))
+        ovr4 = OneVsRest(classifier=dt).setWeightCol("weight")
+        self.assertIsNotNone(ovr4.fit(df))
 
 
 class HashingTFTest(SparkSessionTestCase):
