improve static mode parsing for certain types in postgres input (#9969)

GitOrigin-RevId: b923ecece38ba91c60c67826b532a0f722587977

Author: zxqfd555

integration_tests/db_connectors/test_postgres.py

@@ -436,11 +436,11 @@ def on_change(
                 2025, 4, 23, 10, 13, tzinfo=datetime.timezone.utc
             )
             assert row["i"].value == SimpleObject("test")
-            assert row["j"] == pd.Timedelta("4 days 2 seconds 123 us").value // 1_000
+            assert row["j"] == pd.Timedelta("4 days 2 seconds 123 us")
             assert row["k"] == ("abc", "def", "ghi")
             assert row["l"] == (
-                pd.Timedelta("4 days 2 seconds 123 us").value // 1_000,
-                pd.Timedelta("1 days 2 seconds 3 us").value // 1_000,
+                pd.Timedelta("4 days 2 seconds 123 us"),
+                pd.Timedelta("1 days 2 seconds 3 us"),
             )
             assert row["m"] == (("a", "b"), ("c", "d"))
 
@@ -2117,6 +2117,95 @@ def stream_target():
         assert oor_id not in ids_out, f"Out-of-range row {oor_id} must be skipped"
 
 
+@pytest.mark.parametrize("pg_type", ["BIGINT", "INTEGER"])
+def test_static_int_read_as_duration(tmp_path, postgres, pg_type):
+    """Integral columns declared as pw.Duration must be returned as Duration values
+    (nanoseconds in jsonlines output) in static read mode."""
+
+    class InputSchema(pw.Schema):
+        id: int = pw.column_definition(primary_key=True)
+        duration_col: pw.Duration
+
+    output_path = tmp_path / "output.jsonl"
+    table_name = postgres.random_table_name()
+
+    postgres.execute_sql(
+        f"""
+        CREATE TABLE {table_name} (
+            id {pg_type} PRIMARY KEY,
+            duration_col {pg_type} NOT NULL
+        );
+        """
+    )
+    # 10 seconds expressed in microseconds, as the output connector writes pw.Duration
+    microseconds = 10_000_000
+    postgres.execute_sql(
+        f"INSERT INTO {table_name} (id, duration_col) VALUES (1, {microseconds});"
+    )
+
+    table = pw.io.postgres.read(
+        postgres_settings=POSTGRES_SETTINGS,
+        table_name=table_name,
+        schema=InputSchema,
+        mode="static",
+    )
+    pw.io.jsonlines.write(table, output_path)
+    run()
+
+    rows = []
+    with open(output_path) as f:
+        for line in f:
+            rows.append(json.loads(line))
+
+    assert len(rows) == 1
+    # pw.Duration is serialized to jsonlines as nanoseconds
+    assert rows[0]["duration_col"] == microseconds * 1_000
+
+
+def test_static_bigint_array_read_as_duration_list(tmp_path, postgres):
+    """BIGINT[] columns declared as list[pw.Duration] must be returned as lists of
+    Duration values (each element in nanoseconds in jsonlines) in static read mode."""
+
+    class InputSchema(pw.Schema):
+        id: int = pw.column_definition(primary_key=True)
+        durations: list[pw.Duration]
+
+    output_path = tmp_path / "output.jsonl"
+    table_name = postgres.random_table_name()
+
+    postgres.execute_sql(
+        f"""
+        CREATE TABLE {table_name} (
+            id BIGINT PRIMARY KEY,
+            durations BIGINT[] NOT NULL
+        );
+        """
+    )
+    # Three durations in microseconds: 1s, 2s, 3s
+    postgres.execute_sql(
+        f"INSERT INTO {table_name} (id, durations)"
+        f" VALUES (1, ARRAY[1000000, 2000000, 3000000]::BIGINT[]);"
+    )
+
+    table = pw.io.postgres.read(
+        postgres_settings=POSTGRES_SETTINGS,
+        table_name=table_name,
+        schema=InputSchema,
+        mode="static",
+    )
+    pw.io.jsonlines.write(table, output_path)
+    run()
+
+    rows = []
+    with open(output_path) as f:
+        for line in f:
+            rows.append(json.loads(line))
+
+    assert len(rows) == 1
+    # Each Duration element serialized as nanoseconds
+    assert rows[0]["durations"] == [1_000_000_000, 2_000_000_000, 3_000_000_000]
+
+
 def test_no_publication(tmp_path, postgres):
     class InputSchema(pw.Schema):
         value: str

integration_tests/db_connectors/test_postgres_parsing.py

@@ -1,4 +1,4 @@
-from typing import Any, get_args
+from typing import Any, Union, get_args, get_origin
 
 import numpy as np
 import pandas as pd
@@ -155,11 +155,15 @@ def streaming_target():
             output_table_name,
             init_mode="create_if_not_exists",
         )
+        observer, type_errors = _make_type_check_observer(ItemType)
+        pw.io.python.write(table, observer)
         wait_result_with_checker(
             PostgresRowCountChecker(postgres, output_table_name, 2 * len(items)), 60
         )
         streaming_thread.join()
 
+        assert not type_errors, "\n".join(type_errors)
+
         output_rows = postgres.get_table_contents(
             output_table_name, InputSchemaWithPkey.column_names()
         )
@@ -176,6 +180,60 @@ def streaming_target():
         _compare_input_and_output(ItemType, input_rows, output_rows)
 
 
+def _get_expected_python_type(ItemType: type) -> type | tuple:
+    import types as builtin_types
+
+    origin = get_origin(ItemType)
+
+    if origin is Union or (
+        hasattr(builtin_types, "UnionType") and origin is builtin_types.UnionType
+    ):
+        args = get_args(ItemType)
+        non_none_args = [a for a in args if a is not type(None)]
+        inner = _get_expected_python_type(non_none_args[0])
+        if isinstance(inner, tuple):
+            return inner + (type(None),)
+        return (inner, type(None))
+
+    if origin is not None:
+        return origin
+
+    if ItemType is pw.Duration:
+        return pd.Timedelta
+    if ItemType in (pw.DateTimeNaive, pw.DateTimeUtc):
+        return pd.Timestamp
+
+    return ItemType
+
+
+def _make_type_check_observer(
+    ItemType: type,
+) -> tuple[pw.io.python.ConnectorObserver, list[str]]:
+    type_errors: list[str] = []
+    expected_type = _get_expected_python_type(ItemType)
+
+    class TypeCheckObserver(pw.io.python.ConnectorObserver):
+        def on_change(self, key, row, time, is_addition):
+            if is_addition:
+                value = row["item"]
+                if not isinstance(value, expected_type):
+                    # tuple is acceptable when the schema type is list
+                    if isinstance(value, tuple) and (
+                        expected_type is list
+                        or (isinstance(expected_type, tuple) and list in expected_type)
+                    ):
+                        return
+                    type_errors.append(
+                        f"item value {value!r} has type {type(value)}, "
+                        f"expected {expected_type}"
+                    )
+
+        def on_end(self):
+            pass
+
+    return TypeCheckObserver(), type_errors
+
+
 def _test_postgres_static(
     postgres, ItemType: type, items: list[Any], create_table: Any = _create_table
 ) -> str:
@@ -205,8 +263,13 @@ def _test_postgres_static(
     pw.io.postgres.write(
         table, POSTGRES_SETTINGS, output_table_name, init_mode="create_if_not_exists"
     )
+
+    observer, type_errors = _make_type_check_observer(ItemType)
+    pw.io.python.write(table, observer)
     run()
 
+    assert not type_errors, "\n".join(type_errors)
+
     output_rows = postgres.get_table_contents(
         output_table_name, InputSchemaWithPkey.column_names()
     )