[FLINK-38525][python] Fix coder doesn't process TIMESTAMP_LTZ correctly

Author: fsk119

flink-python/pyflink/fn_execution/coder_impl_fast.pxd

@@ -148,6 +148,10 @@ cdef class TimestampCoderImpl(FieldCoderImpl):
 
     cdef _decode_timestamp_data_from_stream(self, InputStream in_stream)
 
+    cdef _to_utc_timestamp(self, value)
+
+    cdef _to_datetime(self, int64_t seconds, int32_t microseconds)
+
 cdef class LocalZonedTimestampCoderImpl(TimestampCoderImpl):
     cdef object _timezone
 

flink-python/pyflink/fn_execution/coder_impl_fast.pyx

@@ -685,8 +685,9 @@ cdef class TimestampCoderImpl(FieldCoderImpl):
     cpdef encode_to_stream(self, value, OutputStream out_stream):
         cdef int32_t microseconds_of_second, nanoseconds
         cdef int64_t timestamp_seconds, timestamp_milliseconds
-        timestamp_seconds = <int64_t> (value.replace(tzinfo=datetime.timezone.utc).timestamp())
-        microseconds_of_second = value.microsecond
+        utc_ts = self._to_utc_timestamp(value)
+        timestamp_seconds = <int64_t> (utc_ts.timestamp())
+        microseconds_of_second = utc_ts.microsecond
         timestamp_milliseconds = timestamp_seconds * 1000 + microseconds_of_second // 1000
         nanoseconds = microseconds_of_second % 1000 * 1000
         if self._is_compact:
@@ -709,7 +710,15 @@ cdef class TimestampCoderImpl(FieldCoderImpl):
             nanoseconds = in_stream.read_int32()
         seconds = milliseconds // 1000
         microseconds = milliseconds % 1000 * 1000 + nanoseconds // 1000
-        return datetime.datetime.utcfromtimestamp(seconds).replace(microsecond=microseconds)
+        return self._to_datetime(seconds, microseconds)
+
+    cdef _to_utc_timestamp(self, value):
+        return value.replace(tzinfo=datetime.timezone.utc)
+
+    cdef _to_datetime(self, int64_t seconds, int32_t microseconds):
+        datetime.datetime.utcfromtimestamp(seconds).replace(microsecond=microseconds)
+
+    cdef _to_datetime(self, int64_t seconds, int32_t microseconds)
 
 cdef class LocalZonedTimestampCoderImpl(TimestampCoderImpl):
     """
@@ -723,6 +732,13 @@ cdef class LocalZonedTimestampCoderImpl(TimestampCoderImpl):
     cpdef decode_from_stream(self, InputStream in_stream, size_t size):
         return self._timezone.localize(self._decode_timestamp_data_from_stream(in_stream))
 
+    cpdef _to_utc_timestamp(self, value):
+        return value.astimezone(tzinfo=datetime.timezone.utc)
+
+    cdef _to_datetime(self, int64_t seconds, int32_t microseconds):
+        (datetime.datetime.fromtimestamp(seconds, tz=datetime.timezone.utc)
+         .replace(microsecond=microseconds).astimezone(self._timezone))
+
 cdef class InstantCoderImpl(FieldCoderImpl):
     """
     A coder for Instant.

flink-python/pyflink/fn_execution/coder_impl_slow.py

@@ -570,8 +570,7 @@ def decode_from_stream(self, in_stream: InputStream, length=0):
             nanoseconds = in_stream.read_int32()
         return self.internal_to_timestamp(milliseconds, nanoseconds)
 
-    @staticmethod
-    def timestamp_to_internal(timestamp):
+    def timestamp_to_internal(self, timestamp):
         seconds = int(timestamp.replace(tzinfo=datetime.timezone.utc).timestamp())
         microseconds_of_second = timestamp.microsecond
         milliseconds = seconds * 1000 + microseconds_of_second // 1000
@@ -593,10 +592,18 @@ def __init__(self, precision, timezone):
         super(LocalZonedTimestampCoderImpl, self).__init__(precision)
         self.timezone = timezone
 
+    def timestamp_to_internal(self, timestamp):
+        seconds = int(timestamp.astimezone(datetime.timezone.utc).timestamp())
+        microseconds_of_second = timestamp.microsecond
+        milliseconds = seconds * 1000 + microseconds_of_second // 1000
+        nanoseconds = microseconds_of_second % 1000 * 1000
+        return milliseconds, nanoseconds
+
     def internal_to_timestamp(self, milliseconds, nanoseconds):
-        return self.timezone.localize(
-            super(LocalZonedTimestampCoderImpl, self).internal_to_timestamp(
-                milliseconds, nanoseconds))
+        return (super(LocalZonedTimestampCoderImpl, self).internal_to_timestamp(
+                milliseconds, nanoseconds)
+                .replace(tzinfo=datetime.timezone.utc)
+                .astimezone(self.timezone))
 
 
 class InstantCoderImpl(FieldCoderImpl):

flink-python/pyflink/table/tests/test_udf.py

@@ -364,6 +364,13 @@ def timestamp_func(timestamp_param):
                 'timestamp_param is wrong value %s !' % timestamp_param
             return timestamp_param
 
+        @udf(result_type=DataTypes.TIMESTAMP_LTZ(3))
+        def timestamp_ltz_func(timestamp_ltz_param):
+            from datetime import datetime, timezone
+            assert timestamp_ltz_param == datetime(2018, 3, 11, 3, 0, 0, 123000, timezone.utc), \
+                'timestamp_ltz_param is wrong value %s !' % timestamp_ltz_param
+            return timestamp_ltz_param
+
         @udf(result_type=DataTypes.ARRAY(DataTypes.BIGINT()))
         def array_func(array_param):
             assert array_param == [[1, 2, 3]] or array_param == ((1, 2, 3),), \
@@ -427,7 +434,8 @@ def varchar_func(varchar_param):
                 q DECIMAL(38, 18),
                 r BINARY(5),
                 s CHAR(7),
-                t VARCHAR(10)
+                t VARCHAR(10),
+                u TIMESTAMP_LTZ(3)
             ) WITH ('connector'='test-sink')
         """
         self.t_env.execute_sql(sink_table_ddl)
@@ -441,7 +449,8 @@ def varchar_func(varchar_param):
               datetime.datetime(2018, 3, 11, 3, 0, 0, 123000), [[1, 2, 3]],
               {1: 'flink', 2: 'pyflink'}, decimal.Decimal('1000000000000000000.05'),
               decimal.Decimal('1000000000000000000.05999999999999999899999999999'),
-              bytearray(b'flink'), 'pyflink', 'pyflink')],
+              bytearray(b'flink'), 'pyflink', 'pyflink',
+             datetime.datetime(2018, 3, 11, 3, 0, 0, 123000, datetime.timezone.utc))],
             DataTypes.ROW(
                 [DataTypes.FIELD("a", DataTypes.BIGINT()),
                  DataTypes.FIELD("b", DataTypes.BIGINT()),
@@ -462,7 +471,8 @@ def varchar_func(varchar_param):
                  DataTypes.FIELD("q", DataTypes.DECIMAL(38, 18)),
                  DataTypes.FIELD("r", DataTypes.BINARY(5)),
                  DataTypes.FIELD("s", DataTypes.CHAR(7)),
-                 DataTypes.FIELD("t", DataTypes.VARCHAR(10))]))
+                 DataTypes.FIELD("t", DataTypes.VARCHAR(10)),
+                 DataTypes.FIELD("u", DataTypes.TIMESTAMP_LTZ(3))]))
 
         t.select(
             bigint_func(t.a),
@@ -484,7 +494,8 @@ def varchar_func(varchar_param):
             decimal_cut_func(t.q),
             binary_func(t.r),
             char_func(t.s),
-            varchar_func(t.t)) \
+            varchar_func(t.t),
+            timestamp_ltz_func(t.u)) \
             .execute_insert(sink_table).wait()
         actual = source_sink_utils.results()
         # Currently the sink result precision of DataTypes.TIME(precision) only supports 0.
@@ -494,7 +505,7 @@ def varchar_func(varchar_param):
                             "2018-03-11T03:00:00.123, [1, 2, 3], "
                             "{1=flink, 2=pyflink}, 1000000000000000000.050000000000000000, "
                             "1000000000000000000.059999999999999999, [102, 108, 105, 110, 107], "
-                            "pyflink, pyflink]"])
+                            "pyflink, pyflink, 2018-03-11T03:00:00.123Z]"])
 
     def test_all_data_types(self):
         def boolean_func(bool_param):
@@ -995,7 +1006,7 @@ def local_zoned_timestamp_func(local_zoned_timestamp_param):
             .execute_insert(sink_table) \
             .wait()
         actual = source_sink_utils.results()
-        self.assert_equals(actual, ["+I[1970-01-01T00:00:00.123Z]"])
+        self.assert_equals(actual, ["+I[1969-12-31T16:00:01.123Z]"])
 
     def test_execute_from_json_plan(self):
         # create source file path
@@ -1161,6 +1172,8 @@ def echo(i: str):
 if __name__ == '__main__':
     import unittest
 
+    os.environ['_python_worker_execution_mode'] = "loopback"
+
     try:
         import xmlrunner
 

flink-python/pyflink/table/types.py

@@ -33,6 +33,7 @@
 
 from pyflink.common.types import _create_row
 from pyflink.util.api_stability_decorators import PublicEvolving
+from pyflink.util.exceptions import TableException
 from pyflink.util.java_utils import to_jarray, is_instance_of
 from pyflink.java_gateway import get_gateway
 from pyflink.common import Row, RowKind
@@ -498,14 +499,19 @@ def need_conversion(self):
 
     def to_sql_type(self, dt):
         if dt is not None:
-            seconds = (calendar.timegm(dt.utctimetuple()) if dt.tzinfo
-                       else time.mktime(dt.timetuple()))
-            return int(seconds) * 10 ** 6 + dt.microsecond + self.EPOCH_ORDINAL
+            if dt.tzinfo is None:
+                raise TableException(
+                    f"""The input field {dt} does not specify time zone but its SQL type \
+                TIMESTAMP_LTZ requires TIME ZONE. Please use TIMESTAMP instead or use CAST \
+                function to cast TIMESTAMP as TIMESTAMP_LTZ."""
+                )
+            seconds = calendar.timegm(dt.utctimetuple())
+            return int(seconds) * 10 ** 6 + dt.microsecond
 
     def from_sql_type(self, ts):
         if ts is not None:
-            ts = ts - self.EPOCH_ORDINAL
-            return datetime.datetime.fromtimestamp(ts // 10 ** 6).replace(microsecond=ts % 10 ** 6)
+            return (datetime.datetime.fromtimestamp(ts // 10 ** 6, datetime.timezone.utc)
+                    .replace(microsecond=ts % 10 ** 6))
 
 
 class ZonedTimestampType(AtomicType):

flink-python/pyflink/testing/test_case_utils.py

@@ -164,6 +164,7 @@ def setUpClass(cls):
         super(PyFlinkStreamTableTestCase, cls).setUpClass()
         cls.env.set_runtime_mode(RuntimeExecutionMode.STREAMING)
         cls.env.set_parallelism(2)
+        os.environ['_python_worker_execution_mode'] = "loopback"
         cls.t_env = StreamTableEnvironment.create(cls.env)
         cls.t_env.get_config().set("python.fn-execution.bundle.size", "1")