Merge pull request #841 from praj-0/u/praj/math_binary_input_validation

Allow only one of input or default in MathBinary

Author: austinzh

mleap-spark-extension/src/main/scala/org/apache/spark/ml/mleap/feature/MathBinary.scala

@@ -33,6 +33,12 @@ class MathBinary(override val uid: String = Identifiable.randomUID("math_binary"
 
   @org.apache.spark.annotation.Since("2.0.0")
   override def transform(dataset: Dataset[_]): DataFrame = {
+    // Check this condition at runtime else the input schema inferred by MathBinaryModel might be wrong
+    // and this will cause the transform to fail at inference time
+    if((isSet(inputA) && model.da.isDefined) || (isSet(inputB) && model.db.isDefined)) {
+      throw new RuntimeException("Only one of input column or default value can be present.")
+    }
+
     val binaryUdfA = udf {
       a: Double => model(Some(a), None)
     }

mleap-spark-extension/src/test/scala/org/apache/spark/ml/mleap/parity/feature/MathBinaryParitySpec.scala

@@ -22,5 +22,25 @@ class MathBinaryParitySpec extends SparkParityBase {
       setOutputCol("bin_out")
   )).fit(dataset)
 
+  describe("has valid inputs") {
+    val model = MathBinaryModel(Multiply, da=Some(4.0), db=Some(5.0))
+    it("Only one of inputA or defaultA") {
+      val invalidSparkTransformer: Transformer = new MathBinary(uid = "math_bin", model = model).
+        setInputA("dti").
+        setOutputCol("bin_out")
+      assertThrows[RuntimeException] {
+        invalidSparkTransformer.transform(dataset)
+      }
+    }
+    it("Only one of inputB or defaultB") {
+      val invalidSparkTransformer: Transformer = new MathBinary(uid = "math_bin", model = model).
+        setInputB("dti").
+        setOutputCol("bin_out")
+      assertThrows[RuntimeException] {
+        invalidSparkTransformer.transform(dataset)
+      }
+    }
+  }
+
   override val unserializedParams = Set("stringOrderType")
 }