This is an automated email from the ASF dual-hosted git repository. dianfu pushed a commit to branch release-2.2 in repository https://gitbox.apache.org/repos/asf/flink.git
commit 8cdc4cc145c6c8b0d7909cef35a9b4f4636c7661 Author: Biao Geng <[email protected]> AuthorDate: Tue Oct 27 09:49:59 2026 +0800 [FLINK-41196][python] Preserve TIMESTAMP_LTZ instants in Python conversions (#29603) --- .../pyflink/fn_execution/coder_impl_fast.pyx | 13 +- .../pyflink/fn_execution/coder_impl_slow.py | 14 ++- .../pyflink/fn_execution/tests/test_coders.py | 131 +++++++++++++++++++++ .../pyflink/table/tests/test_pandas_conversion.py | 56 +++++++++ .../pyflink/table/tests/test_pandas_udf.py | 33 +++++- flink-python/pyflink/table/tests/test_udf.py | 48 ++++++-- flink-python/pyflink/table/utils.py | 16 +-- 7 files changed, 283 insertions(+), 28 deletions(-) diff --git a/flink-python/pyflink/fn_execution/coder_impl_fast.pyx b/flink-python/pyflink/fn_execution/coder_impl_fast.pyx index 92dff893fe9..9e36ff321ef 100644 --- a/flink-python/pyflink/fn_execution/coder_impl_fast.pyx +++ b/flink-python/pyflink/fn_execution/coder_impl_fast.pyx @@ -21,6 +21,7 @@ # cython: boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True from libc.stdint cimport int32_t, int64_t from libc.stdlib cimport free, malloc +cimport cython import datetime import decimal @@ -685,7 +686,8 @@ 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()) + timestamp_seconds = <int64_t> ( + value.replace(microsecond=0, tzinfo=datetime.timezone.utc).timestamp()) microseconds_of_second = value.microsecond timestamp_milliseconds = timestamp_seconds * 1000 + microseconds_of_second // 1000 nanoseconds = microseconds_of_second % 1000 * 1000 @@ -698,6 +700,7 @@ cdef class TimestampCoderImpl(FieldCoderImpl): cpdef decode_from_stream(self, InputStream in_stream, size_t size): return self._decode_timestamp_data_from_stream(in_stream) + @cython.cdivision(False) cdef _decode_timestamp_data_from_stream(self, InputStream in_stream): cdef int64_t milliseconds cdef int32_t nanoseconds, seconds, microseconds @@ -721,7 +724,13 @@ cdef class LocalZonedTimestampCoderImpl(TimestampCoderImpl): self._timezone = timezone cpdef decode_from_stream(self, InputStream in_stream, size_t size): - return self._timezone.localize(self._decode_timestamp_data_from_stream(in_stream)) + return self._decode_timestamp_data_from_stream(in_stream) \ + .replace(tzinfo=datetime.timezone.utc).astimezone(self._timezone) + + cpdef encode_to_stream(self, value, OutputStream out_stream): + if value.tzinfo is not None: + value = value.astimezone(datetime.timezone.utc) + TimestampCoderImpl.encode_to_stream(self, value, out_stream) cdef class InstantCoderImpl(FieldCoderImpl): """ diff --git a/flink-python/pyflink/fn_execution/coder_impl_slow.py b/flink-python/pyflink/fn_execution/coder_impl_slow.py index 769720dc277..5987348354d 100644 --- a/flink-python/pyflink/fn_execution/coder_impl_slow.py +++ b/flink-python/pyflink/fn_execution/coder_impl_slow.py @@ -572,7 +572,7 @@ class TimestampCoderImpl(FieldCoderImpl): @staticmethod def timestamp_to_internal(timestamp): - seconds = int(timestamp.replace(tzinfo=datetime.timezone.utc).timestamp()) + seconds = int(timestamp.replace(microsecond=0, tzinfo=datetime.timezone.utc).timestamp()) microseconds_of_second = timestamp.microsecond milliseconds = seconds * 1000 + microseconds_of_second // 1000 nanoseconds = microseconds_of_second % 1000 * 1000 @@ -594,9 +594,15 @@ class LocalZonedTimestampCoderImpl(TimestampCoderImpl): self.timezone = timezone 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) + + @staticmethod + def timestamp_to_internal(timestamp): + if timestamp.tzinfo is not None: + timestamp = timestamp.astimezone(datetime.timezone.utc) + return TimestampCoderImpl.timestamp_to_internal(timestamp) class InstantCoderImpl(FieldCoderImpl): diff --git a/flink-python/pyflink/fn_execution/tests/test_coders.py b/flink-python/pyflink/fn_execution/tests/test_coders.py index f04a33f1649..bf8be497a80 100644 --- a/flink-python/pyflink/fn_execution/tests/test_coders.py +++ b/flink-python/pyflink/fn_execution/tests/test_coders.py @@ -21,6 +21,9 @@ import decimal import logging import unittest +import pyarrow as pa +import pytz + from pyflink.fn_execution.coders import BigIntCoder, TinyIntCoder, BooleanCoder, \ SmallIntCoder, IntCoder, FloatCoder, DoubleCoder, BinaryCoder, CharCoder, DateCoder, \ TimeCoder, TimestampCoder, GenericArrayCoder, MapCoder, DecimalCoder, FlattenRowCoder, \ @@ -30,6 +33,134 @@ from pyflink.datastream.window import TimeWindow, CountWindow from pyflink.testing.test_case_utils import PyFlinkTestCase +class ArrowCodersTests(unittest.TestCase): + from pyflink.fn_execution import coder_impl_slow as implementation + + def test_pandas_ltz_decode_preserves_instants(self): + import pandas as pd + from pandas.testing import assert_series_equal + from pyflink.table import DataTypes + from pyflink.table.types import create_arrow_schema + + epochs = [-1, 1704067200123, 1730611800123, 1730615400123, None] + row_type = DataTypes.ROW([DataTypes.FIELD('ltz', DataTypes.TIMESTAMP_LTZ(3))]) + schema = create_arrow_schema(row_type.field_names(), row_type.field_types()) + for zone in ['UTC', 'Asia/Shanghai', 'America/New_York']: + timezone = pytz.timezone(zone) + for size in [0, len(epochs)]: + with self.subTest(timezone=zone, size=size): + batch = pa.record_batch( + [pa.array(epochs[:size], type=pa.timestamp('ms'))], schema=schema) + with pa.BufferOutputStream() as output: + with pa.ipc.new_stream(output, schema) as writer: + writer.write_batch(batch) + data = output.getvalue().to_pybytes() + coder = self.implementation.ArrowCoderImpl(schema, row_type, timezone) + result = coder.decode(data)[0] + expected = pd.Series(pd.to_datetime(epochs[:size], unit='ms', utc=True), + name='ltz') \ + .dt.tz_convert(timezone) + self.assertEqual(zone, str(result.dt.tz)) + assert_series_equal(result, expected.astype(result.dtype)) + + def test_pandas_ltz_encode_preserves_epoch_values(self): + import pandas as pd + from pyflink.table import DataTypes + from pyflink.table.types import create_arrow_schema + + epochs = [-1, 1704067200123, 1730611800123, 1730615400123, None] + row_type = DataTypes.ROW([DataTypes.FIELD('ltz', DataTypes.TIMESTAMP_LTZ(3))]) + schema = create_arrow_schema(row_type.field_names(), row_type.field_types()) + aware = pd.Series(pd.to_datetime(epochs, unit='ms', utc=True)) \ + .dt.tz_convert('Asia/Shanghai') + naive = pd.Series(pd.to_datetime(epochs, unit='ms')) + for zone in ['UTC', 'Asia/Shanghai', 'America/New_York']: + for values in [aware, naive]: + for size in [0, len(epochs)]: + with self.subTest(timezone=zone, dtype=values.dtype, size=size): + coder = self.implementation.ArrowCoderImpl( + schema, row_type, pytz.timezone(zone)) + encoded = coder.encode([values.iloc[:size]]) + result = pa.ipc.open_stream(encoded).read_all().column(0) + self.assertEqual(epochs[:size], result.cast(pa.int64()).to_pylist()) + + +try: + from pyflink.fn_execution import coder_impl_fast +except ImportError: + coder_impl_fast = None + + [email protected](coder_impl_fast is None, "Compiled coders are not installed") +class FastArrowCodersTests(ArrowCodersTests): + implementation = coder_impl_fast + + +class LocalZonedTimestampCodersTests(unittest.TestCase): + from pyflink.fn_execution import coder_impl_slow as implementation + + def test_timestamp_before_epoch(self): + import datetime + import struct + + for precision, micros in [(3, 999000), (6, 999456)]: + with self.subTest(precision=precision): + coder = self.implementation.TimestampCoderImpl(precision) + value = datetime.datetime(1969, 12, 31, 23, 59, 59, micros) + data = struct.pack('>q', -1) + if precision == 6: + data += struct.pack('>i', 456000) + self.assertEqual(data, coder.encode(value)) + self.assertEqual(value, coder.decode(data)) + + def test_decode_preserves_instants(self): + import datetime + import struct + + for zone in ['UTC', 'Asia/Shanghai', 'America/New_York']: + for precision in [3, 6]: + for millis in [-1001, -1000, -1, 0, 1, 1704067200123, 1710052200123, 1710055800123, + 1730611800123, 1730615400123]: + with self.subTest(zone=zone, precision=precision, millis=millis): + nanos = 456000 if precision == 6 else 0 + data = struct.pack('>q', millis) + if precision == 6: + data += struct.pack('>i', nanos) + coder = self.implementation.LocalZonedTimestampCoderImpl( + precision, pytz.timezone(zone)) + result = coder.decode(data) + expected = datetime.datetime(1970, 1, 1, tzinfo=pytz.UTC) + \ + datetime.timedelta(milliseconds=millis, microseconds=nanos // 1000) + self.assertEqual(expected, result) + self.assertEqual(zone, str(result.tzinfo)) + + def test_encode_preserves_epoch_values(self): + import datetime + import struct + + for zone in ['UTC', 'Asia/Shanghai', 'America/New_York']: + for precision in [3, 6]: + for millis in [-1001, -1000, -1, 0, 1, 1704067200123, 1730611800123, 1730615400123]: + nanos = 456000 if precision == 6 else 0 + instant = datetime.datetime(1970, 1, 1, tzinfo=pytz.UTC) + \ + datetime.timedelta(milliseconds=millis, microseconds=nanos // 1000) + expected = struct.pack('>q', millis) + if precision == 6: + expected += struct.pack('>i', nanos) + coder = self.implementation.LocalZonedTimestampCoderImpl( + precision, pytz.timezone(zone)) + for value in [instant.astimezone(pytz.timezone('Asia/Shanghai')), + instant.astimezone(pytz.timezone('America/New_York')), + instant.replace(tzinfo=None)]: + with self.subTest(zone=zone, precision=precision, value=value): + self.assertEqual(expected, coder.encode(value)) + + [email protected](coder_impl_fast is None, "Compiled coders are not installed") +class FastLocalZonedTimestampCodersTests(LocalZonedTimestampCodersTests): + implementation = coder_impl_fast + + class CodersTest(PyFlinkTestCase): def check_coder(self, coder, *values): diff --git a/flink-python/pyflink/table/tests/test_pandas_conversion.py b/flink-python/pyflink/table/tests/test_pandas_conversion.py index 9cc0f8ccdf6..f856f4cd249 100644 --- a/flink-python/pyflink/table/tests/test_pandas_conversion.py +++ b/flink-python/pyflink/table/tests/test_pandas_conversion.py @@ -166,6 +166,62 @@ class PandasConversionITTests(PandasConversionTestBase): pdf = table.filter(table.f1 < 0).to_pandas() self.assertTrue(pdf.empty) + def test_to_pandas_preserves_local_zoned_timestamp_instants(self): + import pandas as pd + + timestamps = pd.to_datetime([ + '1969-12-31 23:59:59.999', + '2024-01-01 00:00:00.123', + '2024-03-10 06:30:00.000', + '2024-03-10 07:30:00.000', + '2024-11-03 05:30:00.000', + '2024-11-03 06:30:00.000', + ], utc=True) + values = ', '.join('(%d, %d)' % (i, timestamp.value // 1000000) + for i, timestamp in enumerate(timestamps)) + values += ', (%d, CAST(NULL AS BIGINT))' % len(timestamps) + original_timezone = self.t_env.get_config().get_local_timezone() + self.addCleanup(self.t_env.get_config().set_local_timezone, original_timezone) + for timezone in ['UTC', 'Asia/Shanghai', 'America/New_York']: + with self.subTest(timezone=timezone): + self.t_env.get_config().set_local_timezone(timezone) + table = self.t_env.sql_query(""" + SELECT id, TO_TIMESTAMP_LTZ(epoch_millis, 3) AS ltz, + TIMESTAMP '2024-01-01 00:00:00.123' AS ts + FROM (VALUES %s) AS v(id, epoch_millis) + """ % values) + result = table.to_pandas().sort_values('id') + + self.assertEqual(len(timestamps) + 1, len(result)) + self.assertEqual(timezone, str(result['ltz'].dt.tz)) + self.assertEqual(list(timestamps), result['ltz'].iloc[:-1].tolist()) + self.assertTrue(pd.isna(result['ltz'].iloc[-1])) + self.assertIsNone(result['ts'].dt.tz) + self.assertTrue((result['ts'] == pd.Timestamp('2024-01-01 00:00:00.123')).all()) + + def test_from_pandas_preserves_local_zoned_timestamp_instants(self): + import pandas as pd + + epochs = [-1, 1704067200123, 1730611800123, 1730615400123, None] + pdf = pd.DataFrame({'ltz': pd.to_datetime(epochs, unit='ms', utc=True) + .tz_convert('Asia/Shanghai')}) + schema = DataTypes.ROW([DataTypes.FIELD('ltz', DataTypes.TIMESTAMP_LTZ(3))]) + original_timezone = self.t_env.get_config().get_local_timezone() + self.addCleanup(self.t_env.get_config().set_local_timezone, original_timezone) + for timezone in ['UTC', 'Asia/Shanghai', 'America/New_York']: + with self.subTest(timezone=timezone): + self.t_env.get_config().set_local_timezone(timezone) + table = self.t_env.from_pandas(pdf, schema) + results = table._j_table.execute().collect() + actual = [] + try: + while results.hasNext(): + timestamp = results.next().getField(0) + actual.append(None if timestamp is None else timestamp.toEpochMilli()) + finally: + results.close() + self.assertCountEqual(epochs, actual) + def test_to_pandas_for_retract_table(self): table = self.t_env.from_pandas(self.pdf, self.data_type) result_pdf = table.group_by(table.f1).select(table.f2.max.alias('f2')).to_pandas() diff --git a/flink-python/pyflink/table/tests/test_pandas_udf.py b/flink-python/pyflink/table/tests/test_pandas_udf.py index 7a6b36f5b17..079e623c1f2 100644 --- a/flink-python/pyflink/table/tests/test_pandas_udf.py +++ b/flink-python/pyflink/table/tests/test_pandas_udf.py @@ -38,6 +38,31 @@ class PandasUDFTests(PyFlinkTestCase): class PandasUDFITTests(object): + def test_local_zoned_timestamp_instants(self): + original_timezone = self.t_env.get_config().get_local_timezone() + self.addCleanup(self.t_env.get_config().set_local_timezone, original_timezone) + for timezone in ['Asia/Shanghai', 'America/New_York']: + with self.subTest(timezone=timezone): + self.t_env.get_config().set_local_timezone(timezone) + + @udf(result_type=DataTypes.TIMESTAMP_LTZ(3), func_type='pandas') + def add_hour(values): + import pandas as pd + + assert str(values.dt.tz) == timezone + assert (values == pd.Timestamp(1730611800123, unit='ms', tz='UTC')).all() + return values + pd.Timedelta(hours=1) + + table = self.t_env.sql_query( + 'SELECT TO_TIMESTAMP_LTZ(1730611800123, 3) AS ltz') + results = table.select(add_hour(table.ltz))._j_table.execute().collect() + try: + self.assertTrue(results.hasNext()) + self.assertEqual(1730615400123, results.next().getField(0).toEpochMilli()) + self.assertFalse(results.hasNext()) + finally: + results.close() + def test_basic_functionality(self): # pandas UDF add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT(), func_type="pandas") @@ -353,8 +378,8 @@ class PandasUDFITTests(object): import pandas as pd timezone = self.t_env.get_config().get_local_timezone() - local_datetime = pytz.timezone(timezone).localize( - datetime.datetime(1970, 1, 2, 0, 0, 0, 123000)) + local_datetime = datetime.datetime( + 1970, 1, 2, 0, 0, 0, 123000, tzinfo=pytz.UTC).astimezone(pytz.timezone(timezone)) @udf(result_type=DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3), func_type="pandas") def local_zoned_timestamp_func(local_zoned_timestamp_param): @@ -372,9 +397,7 @@ class PandasUDFITTests(object): """ self.t_env.execute_sql(sink_table_ddl) - t = self.t_env.from_elements( - [(local_datetime,)], - DataTypes.ROW([DataTypes.FIELD("a", DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3))])) + t = self.t_env.sql_query('SELECT TO_TIMESTAMP_LTZ(86400123, 3) AS a') t.select(local_zoned_timestamp_func(local_zoned_timestamp_func(t.a))) \ .execute_insert("Results_test_data_types").wait() diff --git a/flink-python/pyflink/table/tests/test_udf.py b/flink-python/pyflink/table/tests/test_udf.py index 4667702bbbd..da92776f41c 100644 --- a/flink-python/pyflink/table/tests/test_udf.py +++ b/flink-python/pyflink/table/tests/test_udf.py @@ -35,6 +35,38 @@ def generate_random_table_name(): return "Table{0}".format(str(uuid.uuid1()).replace("-", "_")) +class ProcessModeLocalZonedTimestampTests(object): + + def test_local_zoned_timestamp_instants(self): + original_timezone = self.t_env.get_config().get_local_timezone() + self.addCleanup(self.t_env.get_config().set_local_timezone, original_timezone) + for timezone in ['Asia/Shanghai', 'America/New_York']: + with self.subTest(timezone=timezone): + self.t_env.get_config().set_local_timezone(timezone) + + @udf(result_type=DataTypes.TIMESTAMP_LTZ(3)) + def add_hour(value): + if value is None: + return None + assert str(value.tzinfo) == timezone + assert value == datetime.datetime(2024, 11, 3, 5, 30, 0, 123000, + tzinfo=pytz.UTC) + return value + datetime.timedelta(hours=1) + + table = self.t_env.sql_query( + 'SELECT TO_TIMESTAMP_LTZ(epoch_millis, 3) AS ltz ' + 'FROM (VALUES (1730611800123), (CAST(NULL AS BIGINT))) AS v(epoch_millis)') + results = table.select(add_hour(table.ltz))._j_table.execute().collect() + actual = [] + try: + while results.hasNext(): + value = results.next().getField(0) + actual.append(None if value is None else value.toEpochMilli()) + finally: + results.close() + self.assertCountEqual([1730615400123, None], actual) + + class UserDefinedFunctionTests(object): def test_scalar_function(self): @@ -904,8 +936,8 @@ def float_equal(a, b, rel_tol=1e-09, abs_tol=0.0): return abs(a - b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol) -class PyFlinkStreamUserDefinedFunctionTests(UserDefinedFunctionTests, - PyFlinkStreamTableTestCase): +class PyFlinkStreamUserDefinedFunctionTests( + ProcessModeLocalZonedTimestampTests, UserDefinedFunctionTests, PyFlinkStreamTableTestCase): def test_deterministic(self): add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT()) self.assertTrue(add_one._deterministic) @@ -972,8 +1004,8 @@ class PyFlinkStreamUserDefinedFunctionTests(UserDefinedFunctionTests, def test_data_types(self): timezone = self.t_env.get_config().get_local_timezone() - local_datetime = pytz.timezone(timezone).localize( - datetime.datetime(1970, 1, 1, 0, 0, 0, 123000)) + local_datetime = datetime.datetime( + 1970, 1, 1, 0, 0, 0, 123000, tzinfo=pytz.UTC).astimezone(pytz.timezone(timezone)) @udf(result_type=DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3)) def local_zoned_timestamp_func(local_zoned_timestamp_param): @@ -987,9 +1019,7 @@ class PyFlinkStreamUserDefinedFunctionTests(UserDefinedFunctionTests, """ self.t_env.execute_sql(sink_table_ddl) - t = self.t_env.from_elements( - [(local_datetime,)], - DataTypes.ROW([DataTypes.FIELD("a", DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3))])) + t = self.t_env.sql_query('SELECT TO_TIMESTAMP_LTZ(123, 3) AS a') t.select(local_zoned_timestamp_func(local_zoned_timestamp_func(t.a))) \ .execute_insert(sink_table) \ @@ -1088,8 +1118,8 @@ class PyFlinkStreamUserDefinedFunctionTests(UserDefinedFunctionTests, self.assert_equals(actual, ['+I[42]', '+I[5]', '+I[1000]', '+I[1000]']) -class PyFlinkBatchUserDefinedFunctionTests(UserDefinedFunctionTests, - PyFlinkBatchTableTestCase): +class PyFlinkBatchUserDefinedFunctionTests( + ProcessModeLocalZonedTimestampTests, UserDefinedFunctionTests, PyFlinkBatchTableTestCase): pass diff --git a/flink-python/pyflink/table/utils.py b/flink-python/pyflink/table/utils.py index bc4444e66eb..ed92ad230af 100644 --- a/flink-python/pyflink/table/utils.py +++ b/flink-python/pyflink/table/utils.py @@ -53,7 +53,7 @@ def pandas_to_arrow(schema, timezone, field_types, series): arrays.append(pa.StructArray.from_arrays(struct_arrays, struct_names)) else: arrays.append(create_array( - tz_convert_to_internal(s, field_type, timezone), schema_type)) + tz_convert_to_internal(s, field_type), schema_type)) return pa.RecordBatch.from_arrays(arrays, schema=schema) @@ -76,28 +76,28 @@ def arrow_to_pandas(timezone, field_types, batches): @Internal() def tz_convert_from_internal(s, t: DataType, local_tz): """ - Converts the timestamp series from internal according to the specified local timezone. + Converts UTC-backed TIMESTAMP_LTZ series to the specified local timezone. - Returns the same series if the series is not a timestamp series. Otherwise, - returns a converted series. + Returns the same series for other data types. """ if type(t) == LocalZonedTimestampType: - return s.dt.tz_localize(local_tz) + return s.dt.tz_localize('UTC').dt.tz_convert(local_tz) else: return s @Internal() -def tz_convert_to_internal(s, t: DataType, local_tz): +def tz_convert_to_internal(s, t: DataType): """ - Converts the timestamp series to internal according to the specified local timezone. + Converts timezone-aware TIMESTAMP_LTZ series to UTC and removes the timezone. + Timezone-naive series and other data types are left unchanged. """ if type(t) == LocalZonedTimestampType: from pandas.api.types import is_datetime64_dtype, is_datetime64tz_dtype if is_datetime64_dtype(s.dtype): return s.dt.tz_localize(None) elif is_datetime64tz_dtype(s.dtype): - return s.dt.tz_convert(local_tz).dt.tz_localize(None) + return s.dt.tz_convert('UTC').dt.tz_localize(None) return s
