support datetime with timezone

Author: Davies Liu

python/pyspark/sql/_types.py

@@ -19,6 +19,7 @@
 import decimal
 import time
 import datetime
+import calendar
 import keyword
 import warnings
 import json
@@ -638,6 +639,8 @@ def _need_python_to_sql_conversion(dataType):
     elif isinstance(dataType, MapType):
         return _need_python_to_sql_conversion(dataType.keyType) or \
             _need_python_to_sql_conversion(dataType.valueType)
+    elif isinstance(dataType, TimestampType):
+        return True
     elif isinstance(dataType, UserDefinedType):
         return True
     else:
@@ -691,6 +694,16 @@ def converter(obj):
         key_converter = _python_to_sql_converter(dataType.keyType)
         value_converter = _python_to_sql_converter(dataType.valueType)
         return lambda m: dict([(key_converter(k), value_converter(v)) for k, v in m.items()])
+
+    elif isinstance(dataType, TimestampType):
+
+        def to_posix_timstamp(dt):
+            if dt.tzinfo is None:
+                return time.mktime(dt.timetuple()) + dt.microsecond / 1e6
+            else:
+                return calendar.timegm(dt.utctimetuple()) + dt.microsecond / 1e6
+        return to_posix_timstamp
+
     elif isinstance(dataType, UserDefinedType):
         return lambda obj: dataType.serialize(obj)
     else:

python/pyspark/sql/tests.py

@@ -26,6 +26,7 @@
 import tempfile
 import pickle
 import functools
+import time
 import datetime
 
 import py4j
@@ -46,6 +47,20 @@
 from pyspark.sql.functions import UserDefinedFunction
 
 
+class UTC(datetime.tzinfo):
+    """UTC"""
+    ZERO = datetime.timedelta(0)
+
+    def utcoffset(self, dt):
+        return self.ZERO
+
+    def tzname(self, dt):
+        return "UTC"
+
+    def dst(self, dt):
+        return self.ZERO
+
+
 class ExamplePointUDT(UserDefinedType):
     """
     User-defined type (UDT) for ExamplePoint.
@@ -571,6 +586,20 @@ def test_filter_with_datetime(self):
         self.assertEqual(0, df.filter(df.date > date).count())
         self.assertEqual(0, df.filter(df.time > time).count())
 
+    def test_time_with_timezone(self):
+        now = datetime.datetime.now()
+        ts = time.mktime(now.timetuple()) + now.microsecond / 1e6
+        # class in __main__ is not serializable
+        from pyspark.sql.tests import UTC
+        utc = UTC()
+        utcnow = datetime.datetime.fromtimestamp(ts, utc)
+        df = self.sqlCtx.createDataFrame([(now, utcnow)])
+        now1, utcnow1 = df.first()
+        # Spark SQL does not support microsecond, the error should be
+        # less than 1 millisecond
+        self.assertTrue(now1 - now < datetime.timedelta(0.001))
+        self.assertTrue(utcnow1 - now < datetime.timedelta(0.001))
+
     def test_dropna(self):
         schema = StructType([
             StructField("name", StringType(), True),

sql/core/src/main/scala/org/apache/spark/sql/execution/pythonUdfs.scala

@@ -28,8 +28,7 @@ import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.api.python.{PythonBroadcast, PythonRDD}
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.catalyst.expressions.Row
-import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.{Row, _}
 import org.apache.spark.sql.catalyst.plans.logical
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -184,6 +183,8 @@ object EvaluatePython {
 
     case (c: java.util.Calendar, TimestampType) =>
       new java.sql.Timestamp(c.getTime().getTime())
+    case (c: Double, TimestampType) =>
+      new java.sql.Timestamp((c * 1000).toLong)
 
     case (_, udt: UserDefinedType[_]) =>
       fromJava(obj, udt.sqlType)