Add column method 'isNaN' (#313)

* Add column method 'isNaN'

* Add a type class to restrict types that can be NaN

Author: prestonph

core/src/main/scala/frameless/CatalystNaN.scala

@@ -0,0 +1,16 @@
+package frameless
+
+import scala.annotation.implicitNotFound
+
+/** Spark does NaN check only for these types */
+@implicitNotFound("Columns of type ${A} cannot be NaN.")
+trait CatalystNaN[A]
+
+object CatalystNaN {
+  private[this] val theInstance = new CatalystNaN[Any] {}
+  private[this] def of[A]: CatalystNaN[A] = theInstance.asInstanceOf[CatalystNaN[A]]
+
+  implicit val framelessFloatNaN     : CatalystNaN[Float]      = of[Float]
+  implicit val framelessDoubleNaN    : CatalystNaN[Double]     = of[Double]
+}
+

dataset/src/main/scala/frameless/TypedColumn.scala

@@ -139,6 +139,13 @@ abstract class AbstractTypedColumn[T, U]
   def isNotNone(implicit i0: U <:< Option[_]): ThisType[T, Boolean] =
     typed(Not(equalsTo(lit(None.asInstanceOf[U])).expr))
 
+  /** True if the current expression is a fractional number and is not NaN.
+    *
+    * apache/spark
+    */
+  def isNaN(implicit n: CatalystNaN[U]): ThisType[T, Boolean] =
+    typed(self.untyped.isNaN)
+
   /** Convert an Optional column by providing a default value
     * {{{
     *   df( df('opt).getOrElse(df('defaultValue)) )

dataset/src/test/scala/frameless/NumericTests.scala

@@ -1,7 +1,10 @@
 package frameless
 
-import org.scalacheck.Prop
+import org.apache.spark.sql.Encoder
+import org.scalacheck.{Arbitrary, Gen, Prop}
 import org.scalacheck.Prop._
+import org.scalatest.Matchers._
+
 import scala.reflect.ClassTag
 
 class NumericTests extends TypedDatasetSuite {
@@ -169,4 +172,35 @@ class NumericTests extends TypedDatasetSuite {
     check(prop[Short ] _)
     check(prop[BigDecimal] _)
   }
+
+  test("isNaN") {
+    val spark = session
+    import spark.implicits._
+
+    implicit val doubleWithNaN = Arbitrary {
+      implicitly[Arbitrary[Double]].arbitrary.flatMap(Gen.oneOf(_, Double.NaN))
+    }
+    implicit val x1 = Arbitrary{ doubleWithNaN.arbitrary.map(X1(_)) }
+
+    def prop[A : TypedEncoder : Encoder : CatalystNaN](data: List[X1[A]]): Prop = {
+      val ds = TypedDataset.create(data)
+
+      val expected = ds.toDF().filter(!$"a".isNaN).map(_.getAs[A](0)).collect().toSeq
+      val rs = ds.filter(!ds('a).isNaN).collect().run().map(_.a)
+
+      rs ?= expected
+    }
+
+    check(forAll(prop[Float] _))
+    check(forAll(prop[Double] _))
+  }
+
+  test("isNaN with non-nan types should not compile") {
+    val ds = TypedDataset.create((1, false, 'a, "b") :: Nil)
+
+    "ds.filter(ds('_1).isNaN)" shouldNot typeCheck
+    "ds.filter(ds('_2).isNaN)" shouldNot typeCheck
+    "ds.filter(ds('_3).isNaN)" shouldNot typeCheck
+    "ds.filter(ds('_4).isNaN)" shouldNot typeCheck
+  }
 }