Issue-272: Add String Comparison operators (hablapps#355)

* Issue-272: Add String Comparison operators * Issue-272: Fix formatting
alfonsorr · May 16, 2023 · 619e228 · 619e228
1 parent 8eef33d
commit 619e228
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 0 deletions.
diff --git a/core/src/main/scala/doric/syntax/StringColumns.scala b/core/src/main/scala/doric/syntax/StringColumns.scala
@@ -2,6 +2,7 @@ package doric
 package syntax
 
 import cats.implicits._
+import doric.DoricColumn.sparkFunction
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.types.{DataType, StructType}
 import org.apache.spark.sql.{Column, functions => f}
@@ -689,5 +690,29 @@ protected trait StringColumns {
         .map(x => new Column(JsonTuple(x.map(_.expr))))
         .toDC
     }
+
+    /**
+      * @group Comparable Type
+      */
+    def <(other: StringColumn): BooleanColumn =
+      sparkFunction[String, Boolean](s, other, _ < _)
+
+    /**
+      * @group Comparable Type
+      */
+    def <=(other: StringColumn): BooleanColumn =
+      sparkFunction[String, Boolean](s, other, _ <= _)
+
+    /**
+      * @group Comparable Type
+      */
+    def >(other: StringColumn): BooleanColumn =
+      sparkFunction[String, Boolean](s, other, _ > _)
+
+    /**
+      * @group Comparable Type
+      */
+    def >=(other: StringColumn): BooleanColumn =
+      sparkFunction[String, Boolean](s, other, _ >= _)
   }
 }
diff --git a/core/src/test/scala/doric/syntax/StringColumnsSpec.scala b/core/src/test/scala/doric/syntax/StringColumnsSpec.scala
@@ -1048,4 +1048,44 @@ class StringColumnsSpec extends DoricTestElements {
     }
   }
 
+  describe("comparison operators") {
+    import spark.implicits._
+
+    val df =
+      List(("abc", "xyz"), ("abc", "abc"), ("xyz", "abc"), ("xyz", "xyz"))
+        .toDF("col1", "col2")
+
+    it("> should work as spark > function") {
+      df.testColumns2("col1", "col2")(
+        (col1, col2) => colString(col1) > colString(col2),
+        (col1, col2) => f.col(col1) > f.col(col2),
+        List(false, false, true, false).map(Option(_))
+      )
+    }
+
+    it(">= should work as spark >= function") {
+      df.testColumns2("col1", "col2")(
+        (col1, col2) => colString(col1) >= colString(col2),
+        (col1, col2) => f.col(col1) >= f.col(col2),
+        List(false, true, true, true).map(Option(_))
+      )
+    }
+
+    it("< should work as spark < function") {
+      df.testColumns2("col1", "col2")(
+        (col1, col2) => colString(col1) < colString(col2),
+        (col1, col2) => f.col(col1) < f.col(col2),
+        List(true, false, false, false).map(Option(_))
+      )
+    }
+
+    it("<= should work as spark <= function") {
+      df.testColumns2("col1", "col2")(
+        (col1, col2) => colString(col1) <= colString(col2),
+        (col1, col2) => f.col(col1) <= f.col(col2),
+        List(true, true, false, true).map(Option(_))
+      )
+    }
+  }
+
 }