Allow comparsion between Column and datetime objects (#274)

Resoves #273
zero323 · Dec 30, 2019 · 76db4a4 · 76db4a4
1 parent 84a927b
commit 76db4a4
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 5 deletions.
diff --git a/test-data/unit/sql-column.test b/test-data/unit/sql-column.test
@@ -0,0 +1,21 @@
+[case colDateTimeCompare]
+import datetime
+from pyspark.sql.functions import col
+
+today = datetime.date.today()
+now = datetime.datetime.now()
+a_col = col("")
+
+a_col < today
+a_col <= today
+a_col == today
+a_col >= today
+a_col > today
+
+
+a_col < now
+a_col <= now
+a_col == now
+a_col >= now
+a_col > now
+[out]
diff --git a/tests/unit/testpyspark.py b/tests/unit/testpyspark.py
@@ -10,6 +10,7 @@
         "ml-param.test",
         "ml-readable.test",
         "resultiterable.test",
+        "sql-column.test",
         "sql-readwriter.test",
         "sql-udf.test",
     ]
diff --git a/third_party/3/pyspark/sql/column.pyi b/third_party/3/pyspark/sql/column.pyi
@@ -27,12 +27,12 @@ class Column:
     def __rmod__(self, other: Union[bool, int, float, DecimalLiteral]) -> Column: ...
     def __pow__(self, other: Union[Column, LiteralType, DecimalLiteral]) -> Column: ...
     def __rpow__(self, other: Union[LiteralType, DecimalLiteral]) -> Column: ...
-    def __eq__(self, other: Union[Column, LiteralType, DecimalLiteral]) -> Column: ...  # type: ignore
+    def __eq__(self, other: Union[Column, LiteralType, DateTimeLiteral, DecimalLiteral]) -> Column: ...  # type: ignore
     def __ne__(self, other: Any) -> Column: ...  # type: ignore
-    def __lt__(self, other: Union[Column, LiteralType, DecimalLiteral]) -> Column: ...
-    def __le__(self, other: Union[Column, LiteralType, DecimalLiteral]) -> Column: ...
-    def __ge__(self, other: Union[Column, LiteralType, DecimalLiteral]) -> Column: ...
-    def __gt__(self, other: Union[Column, LiteralType, DecimalLiteral]) -> Column: ...
+    def __lt__(self, other: Union[Column, LiteralType, DateTimeLiteral, DecimalLiteral]) -> Column: ...
+    def __le__(self, other: Union[Column, LiteralType, DateTimeLiteral, DecimalLiteral]) -> Column: ...
+    def __ge__(self, other: Union[Column, LiteralType, DateTimeLiteral, DecimalLiteral]) -> Column: ...
+    def __gt__(self, other: Union[Column, LiteralType, DateTimeLiteral, DecimalLiteral]) -> Column: ...
     def eqNullSafe(self, other: Union[Column, LiteralType, DecimalLiteral]) -> Column: ...
     def __and__(self, other: Column) -> Column: ...
     def __or__(self, other: Column) -> Column: ...