feat: [+] #271 unary functions (-) & (!) (#276)

hablapps · Sep 7, 2022 · c0b4ee0 · c0b4ee0
1 parent 86620fa
commit c0b4ee0
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 21 deletions.
diff --git a/core/src/main/scala/doric/syntax/BooleanColumns.scala b/core/src/main/scala/doric/syntax/BooleanColumns.scala
@@ -15,21 +15,18 @@ private[syntax] trait BooleanColumns {
     */
   def not(col: BooleanColumn): BooleanColumn = col.elem.map(f.not).toDC
 
-  /**
-    * Inversion of boolean expression, i.e. NOT.
-    *
-    * @group Boolean Type
-    * @see [[not]]
-    */
-  @inline def !(col: BooleanColumn): BooleanColumn = not(col)
-
-  /**
-    * @group Boolean Type
-    */
   implicit class BooleanOperationsSyntax(
       column: DoricColumn[Boolean]
   ) {
 
+    /**
+      * Inversion of boolean expression, i.e. NOT.
+      *
+      * @group Boolean Type
+      * @see [[not]]
+      */
+    def unary_! : DoricColumn[Boolean] = not(column)
+
     /**
       * Boolean AND
       *

diff --git a/core/src/main/scala/doric/syntax/NumericColumns.scala b/core/src/main/scala/doric/syntax/NumericColumns.scala
@@ -4,8 +4,8 @@ package syntax
 import cats.implicits._
 import doric.DoricColumn.sparkFunction
 import doric.types.{CollectionType, NumericType}
+import org.apache.spark.sql.catalyst.expressions.{BRound, FormatNumber, FromUnixTime, Rand, Randn, Round, UnaryMinus}
 import org.apache.spark.sql.{Column, functions => f}
-import org.apache.spark.sql.catalyst.expressions.{BRound, FormatNumber, FromUnixTime, Rand, Randn, Round, ShiftLeft, ShiftRight, ShiftRightUnsigned}
 
 private[syntax] trait NumericColumns {
 
@@ -101,6 +101,14 @@ private[syntax] trait NumericColumns {
       column: DoricColumn[T]
   ) {
 
+    /**
+      * Unary minus, i.e. negate the expression.
+      *
+      * @group Numeric Type
+      */
+    def unary_- : DoricColumn[T] =
+      column.elem.map(x => new Column(UnaryMinus(x.expr))).toDC
+
     /**
       * @group Numeric Type
       */

diff --git a/core/src/test/scala/doric/syntax/BooleanColumnsSpec.scala b/core/src/test/scala/doric/syntax/BooleanColumnsSpec.scala
@@ -1,15 +1,9 @@
 package doric
 package syntax
 
-import org.scalatest.EitherValues
-import org.scalatest.matchers.should.Matchers
-
 import org.apache.spark.sql.{functions => f}
 
-class BooleanColumnsSpec
-    extends DoricTestElements
-    with EitherValues
-    with Matchers {
+class BooleanColumnsSpec extends DoricTestElements {
 
   describe("Boolean columns") {
     import spark.implicits._
@@ -28,8 +22,8 @@ class BooleanColumnsSpec
 
     it("should be inverted by !") {
       df.testColumns("col1")(
-        c => BoolF.!(colBoolean(c)),
-        c => f.not(f.col(c)),
+        c => !colBoolean(c),
+        c => !f.col(c),
         List(Some(false), Some(true), None)
       )
     }

diff --git a/core/src/test/scala/doric/syntax/NumericOperationsSpec.scala b/core/src/test/scala/doric/syntax/NumericOperationsSpec.scala
@@ -344,6 +344,15 @@ trait NumericOperationsSpec
           f.negate
         )
       }
+
+      it(s"unary (-) function $numTypeStr") {
+        testDoricSpark[T, T](
+          List(Some(-1), Some(1), Some(2), None),
+          List(Some(1), Some(-1), Some(-2), None),
+          x => -x,
+          x => -x
+        )
+      }
     }
   }