[python] Normalize nested naive timestamps when writing Avro to object stores and HDFS (#10356)

Author: jackylee-ch

paimon-python/pypaimon/common/file_io.py

@@ -19,6 +19,7 @@
 import os
 import uuid
 from abc import ABC, abstractmethod
+from datetime import datetime, timezone
 from pathlib import Path
 from typing import Dict, List, Optional
 
@@ -30,6 +31,27 @@
 _LOG = logging.getLogger(__name__)
 
 
+def normalize_naive_datetimes(value):
+    """Recursively attach UTC to naive datetimes, including ones nested inside
+    ROW (dict) and ARRAY/MAP (list/tuple) values.
+
+    fastavro converts a naive datetime using the host timezone, which corrupts
+    the stored instant on a non-UTC host; Paimon timestamps are UTC-based. The
+    Avro ``write_avro`` paths must normalize nested timestamps, not just
+    top-level column values, so a naive datetime inside a struct or list is not
+    written shifted by the host offset.
+    """
+    if isinstance(value, datetime):
+        if value.tzinfo is None:
+            return value.replace(tzinfo=timezone.utc)
+        return value
+    if isinstance(value, dict):
+        return {key: normalize_naive_datetimes(val) for key, val in value.items()}
+    if isinstance(value, (list, tuple)):
+        return type(value)(normalize_naive_datetimes(val) for val in value)
+    return value
+
+
 def supports_pread(stream) -> bool:
     """Check if the stream supports position-based reads (thread-safe I/O)."""
     if hasattr(stream, 'read_at'):

paimon-python/pypaimon/filesystem/hdfs_native_file_io.py

@@ -28,7 +28,7 @@
 import pyarrow
 import pyarrow.fs as pafs
 
-from pypaimon.common.file_io import FileIO
+from pypaimon.common.file_io import FileIO, normalize_naive_datetimes
 from pypaimon.common.options import Options
 from pypaimon.common.options.config import HdfsOptions, SecurityOptions
 from pypaimon.common.uri_reader import UriReaderFactory
@@ -594,9 +594,7 @@ def record_generator():
                 record = {}
                 for col in records_dict.keys():
                     value = records_dict[col][i]
-                    if isinstance(value, datetime) and value.tzinfo is None:
-                        value = value.replace(tzinfo=timezone.utc)
-                    record[col] = value
+                    record[col] = normalize_naive_datetimes(value)
                 yield record
 
         codec_map = {

paimon-python/pypaimon/filesystem/local_file_io.py

@@ -19,42 +19,20 @@
 import os
 import shutil
 import threading
-from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, Dict, Optional
 from urllib.parse import unquote, urlparse
 
 import pyarrow
 import pyarrow.fs as pafs
 
-from pypaimon.common.file_io import FileIO, create_temp_path
+from pypaimon.common.file_io import FileIO, create_temp_path, normalize_naive_datetimes
 from pypaimon.common.options import Options
 from pypaimon.common.uri_reader import UriReaderFactory
 from pypaimon.filesystem.local import PaimonLocalFileSystem
 from pypaimon.schema.data_types import PyarrowFieldParser
 
 
-def _normalize_naive_datetimes(value):
-    """Recursively attach UTC to naive datetimes, including ones nested inside
-    ROW (dict) and ARRAY/MAP (list/tuple) values.
-
-    fastavro converts a naive datetime using the host timezone, which corrupts
-    the stored instant on a non-UTC host; Paimon timestamps are UTC-based.
-    ``write_avro`` previously normalized only top-level column values, so a
-    naive datetime inside a struct or list was written shifted by the host
-    offset. Normalizing recursively keeps nested timestamps correct.
-    """
-    if isinstance(value, datetime):
-        if value.tzinfo is None:
-            return value.replace(tzinfo=timezone.utc)
-        return value
-    if isinstance(value, dict):
-        return {key: _normalize_naive_datetimes(val) for key, val in value.items()}
-    if isinstance(value, (list, tuple)):
-        return type(value)(_normalize_naive_datetimes(val) for val in value)
-    return value
-
-
 def _file_uri_path(parsed, windows=None) -> str:
     """Decode a file URI path, preserving Windows drives and UNC hosts."""
     if windows is None:
@@ -378,7 +356,7 @@ def record_generator():
                 record = {}
                 for col in records_dict.keys():
                     value = records_dict[col][i]
-                    record[col] = _normalize_naive_datetimes(value)
+                    record[col] = normalize_naive_datetimes(value)
                 yield record
 
         records = record_generator()

paimon-python/pypaimon/filesystem/pyarrow_file_io.py

@@ -25,7 +25,6 @@
 import tempfile
 import threading
 import time
-from datetime import datetime, timezone
 from itertools import islice
 from pathlib import PurePosixPath
 from typing import Any, Dict, Iterable, List, Optional
@@ -36,7 +35,7 @@
 from packaging.version import parse
 from pyarrow._fs import FileSystem
 
-from pypaimon.common.file_io import FileIO, create_temp_path
+from pypaimon.common.file_io import FileIO, create_temp_path, normalize_naive_datetimes
 from pypaimon.common.options import Options
 from pypaimon.common.options.config import OssOptions, S3Options, SecurityOptions
 from pypaimon.common.options.options_utils import OptionsUtils
@@ -957,9 +956,7 @@ def record_generator():
                 record = {}
                 for col in records_dict.keys():
                     value = records_dict[col][i]
-                    if isinstance(value, datetime) and value.tzinfo is None:
-                        value = value.replace(tzinfo=timezone.utc)
-                    record[col] = value
+                    record[col] = normalize_naive_datetimes(value)
                 yield record
 
         records = record_generator()

paimon-python/pypaimon/tests/data_types_test.py

@@ -432,3 +432,53 @@ def test_avro_nested_timestamp_roundtrips_on_non_utc_host(self):
         self.assertEqual(rows[0]["ts"].replace(tzinfo=None), ts)
         self.assertEqual(rows[0]["r"]["ts"].replace(tzinfo=None), ts)
         self.assertEqual(rows[0]["arr"][0].replace(tzinfo=None), ts)
+
+    def test_pyarrow_file_io_avro_normalizes_nested_naive_timestamps(self):
+        # Same contract as the LocalFileIO case, but for the object-store base
+        # PyArrowFileIO (OSS/S3/GCS): it previously normalized only top-level
+        # column values, so a naive TIMESTAMP nested in a ROW/ARRAY written as
+        # Avro to object storage on a non-UTC host was shifted by the host
+        # offset. The output stream is mocked to a local file so the real
+        # write_avro record path runs without cloud credentials.
+        import datetime
+        import os
+        import tempfile
+        import time
+        from unittest import mock
+
+        import fastavro
+
+        from pypaimon.filesystem.pyarrow_file_io import PyArrowFileIO
+
+        if not hasattr(time, "tzset"):
+            self.skipTest("time.tzset is unavailable on this platform")
+
+        ts = datetime.datetime(2024, 1, 2, 3, 4, 5)
+        table = pa.table({
+            "ts": pa.array([ts], pa.timestamp('s')),
+            "r": pa.array([{"ts": ts}], pa.struct([("ts", pa.timestamp('s'))])),
+            "arr": pa.array([[ts]], pa.list_(pa.timestamp('s'))),
+        })
+        file_io = PyArrowFileIO.__new__(PyArrowFileIO)
+        previous_tz = os.environ.get("TZ")
+        os.environ["TZ"] = "Asia/Shanghai"
+        time.tzset()
+        try:
+            with tempfile.TemporaryDirectory() as d:
+                path = os.path.join(d, "data.avro")
+                with mock.patch.object(
+                        PyArrowFileIO, "new_output_stream",
+                        side_effect=lambda p: open(p, "wb")):
+                    file_io.write_avro(path, table)
+                with open(path, 'rb') as f:
+                    rows = list(fastavro.reader(f))
+        finally:
+            if previous_tz is None:
+                os.environ.pop("TZ", None)
+            else:
+                os.environ["TZ"] = previous_tz
+            time.tzset()
+
+        self.assertEqual(rows[0]["ts"].replace(tzinfo=None), ts)
+        self.assertEqual(rows[0]["r"]["ts"].replace(tzinfo=None), ts)
+        self.assertEqual(rows[0]["arr"][0].replace(tzinfo=None), ts)