Adding sync between Python and Java default timezones

Author: BryanCutler

python/pyspark/sql/tests.py

@@ -3000,6 +3000,14 @@ class ArrowTests(ReusedPySparkTestCase):
     def setUpClass(cls):
         from datetime import datetime
         ReusedPySparkTestCase.setUpClass()
+
+        # Synchronize default timezone between Python and Java
+        tz = "America/Los_Angeles"
+        os.environ["TZ"] = tz
+        time.tzset()
+        cls.old_tz = cls.sc._jvm.org.apache.spark.sql.catalyst.util.DateTimeTestUtils\
+            .setDefaultTimeZone(tz)
+
         cls.spark = SparkSession(cls.sc)
         cls.spark.conf.set("spark.sql.execution.arrow.enable", "true")
         cls.schema = StructType([
@@ -3014,6 +3022,13 @@ def setUpClass(cls):
                     ("b", 2, 20, 0.4, 4.0, datetime(2012, 2, 2), datetime(2012, 2, 2, 2, 2, 2)),
                     ("c", 3, 30, 0.8, 6.0, datetime(2100, 3, 3), datetime(2100, 3, 3, 3, 3, 3))]
 
+    @classmethod
+    def tearDownClass(cls):
+        del os.environ["TZ"]
+        time.tzset()
+        cls.sc._jvm.org.apache.spark.sql.catalyst.util.DateTimeTestUtils\
+            .setDefaultTimeZone(cls.old_tz)
+
     def assertFramesEqual(self, df_with_arrow, df_without):
         msg = ("DataFrame from Arrow is not equal" +
                ("\n\nWith Arrow:\n%s\n%s" % (df_with_arrow, df_with_arrow.dtypes)) +

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeTestUtils.scala

@@ -37,4 +37,11 @@ object DateTimeTestUtils {
       DateTimeUtils.resetThreadLocals()
     }
   }
+
+  def setDefaultTimeZone(id: String): String = {
+    val originalDefaultTimeZone = DateTimeUtils.defaultTimeZone().getID
+    DateTimeUtils.resetThreadLocals()
+    TimeZone.setDefault(TimeZone.getTimeZone(id))
+    originalDefaultTimeZone
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala

@@ -20,7 +20,7 @@ import java.io.File
 import java.nio.charset.StandardCharsets
 import java.sql.{Date, Timestamp}
 import java.text.SimpleDateFormat
-import java.util.Locale
+import java.util.{Locale, TimeZone}
 
 import com.google.common.io.Files
 import org.apache.arrow.memory.RootAllocator
@@ -31,7 +31,7 @@ import org.scalatest.BeforeAndAfterAll
 
 import org.apache.spark.SparkException
 import org.apache.spark.sql.{DataFrame, Row}
-import org.apache.spark.sql.catalyst.util.DateTimeUtils
+import org.apache.spark.sql.catalyst.util.{DateTimeTestUtils, DateTimeUtils}
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types.{BinaryType, StructField, StructType}
 import org.apache.spark.util.Utils
@@ -841,52 +841,54 @@ class ArrowConvertersSuite extends SharedSQLContext with BeforeAndAfterAll {
   }
 
   test("timestamp type conversion") {
-    val json =
-      s"""
-         |{
-         |  "schema" : {
-         |    "fields" : [ {
-         |      "name" : "timestamp",
-         |      "type" : {
-         |        "name" : "timestamp",
-         |        "unit" : "MICROSECOND",
-         |        "timezone" : "${DateTimeUtils.defaultTimeZone().getID}"
-         |      },
-         |      "nullable" : true,
-         |      "children" : [ ],
-         |      "typeLayout" : {
-         |        "vectors" : [ {
-         |          "type" : "VALIDITY",
-         |          "typeBitWidth" : 1
-         |        }, {
-         |          "type" : "DATA",
-         |          "typeBitWidth" : 64
-         |        } ]
-         |      }
-         |    } ]
-         |  },
-         |  "batches" : [ {
-         |    "count" : 4,
-         |    "columns" : [ {
-         |      "name" : "timestamp",
-         |      "count" : 4,
-         |      "VALIDITY" : [ 1, 1, 1, 1 ],
-         |      "DATA" : [ -1234, 0, 1365383415567000, 33057298500000000 ]
-         |    } ]
-         |  } ]
-         |}
-       """.stripMargin
-
-    val sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS z", Locale.US)
-    val ts1 = DateTimeUtils.toJavaTimestamp(-1234L)
-    val ts2 = DateTimeUtils.toJavaTimestamp(0L)
-    val ts3 = new Timestamp(sdf.parse("2013-04-08 01:10:15.567 UTC").getTime)
-    val ts4 = new Timestamp(sdf.parse("3017-07-18 14:55:00.000 UTC").getTime)
-    val data = Seq(ts1, ts2, ts3, ts4)
-
-    val df = data.toDF("timestamp")
-
-    collectAndValidate(df, json, "timestampData.json")
+    DateTimeTestUtils.withDefaultTimeZone(TimeZone.getTimeZone("America/Los_Angeles")) {
+      val json =
+        s"""
+           |{
+           |  "schema" : {
+           |    "fields" : [ {
+           |      "name" : "timestamp",
+           |      "type" : {
+           |        "name" : "timestamp",
+           |        "unit" : "MICROSECOND",
+           |        "timezone" : "${DateTimeUtils.defaultTimeZone().getID}"
+           |      },
+           |      "nullable" : true,
+           |      "children" : [ ],
+           |      "typeLayout" : {
+           |        "vectors" : [ {
+           |          "type" : "VALIDITY",
+           |          "typeBitWidth" : 1
+           |        }, {
+           |          "type" : "DATA",
+           |          "typeBitWidth" : 64
+           |        } ]
+           |      }
+           |    } ]
+           |  },
+           |  "batches" : [ {
+           |    "count" : 4,
+           |    "columns" : [ {
+           |      "name" : "timestamp",
+           |      "count" : 4,
+           |      "VALIDITY" : [ 1, 1, 1, 1 ],
+           |      "DATA" : [ -1234, 0, 1365383415567000, 33057298500000000 ]
+           |    } ]
+           |  } ]
+           |}
+         """.stripMargin
+
+      val sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS z", Locale.US)
+      val ts1 = DateTimeUtils.toJavaTimestamp(-1234L)
+      val ts2 = DateTimeUtils.toJavaTimestamp(0L)
+      val ts3 = new Timestamp(sdf.parse("2013-04-08 01:10:15.567 UTC").getTime)
+      val ts4 = new Timestamp(sdf.parse("3017-07-18 14:55:00.000 UTC").getTime)
+      val data = Seq(ts1, ts2, ts3, ts4)
+
+      val df = data.toDF("timestamp")
+
+      collectAndValidate(df, json, "timestampData.json")
+    }
   }
 
   test("floating-point NaN") {