From 4cd599a5efa1679ce44117e75063abea4347683b Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 14:34:40 +0800 Subject: [PATCH 01/11] [FLINK-40190][python] Add DataFrame creation and conversion APIs Add pandas, Arrow, Table, and range creators with strict schema and watermark validation. Add DataFrame conversion wrappers and schema-aware in-memory and Arrow source paths. Generated-by: Codex (GPT-5) --- .../reference/pyflink.dataframe/creation.rst | 41 +- .../reference/pyflink.dataframe/dataframe.rst | 6 + flink-python/pyflink/dataframe/__init__.py | 13 +- flink-python/pyflink/dataframe/convert.py | 351 +++++++++++++++++- flink-python/pyflink/dataframe/dataframe.py | 43 ++- .../pyflink/dataframe/tests/test_convert.py | 86 +++++ .../pyflink/dataframe/tests/test_dataframe.py | 175 +++++++++ .../pyflink/table/table_environment.py | 50 ++- flink-python/pyflink/table/types.py | 2 +- .../flink/table/runtime/arrow/ArrowUtils.java | 6 +- .../table/utils/python/PythonTableUtils.java | 18 +- 11 files changed, 773 insertions(+), 18 deletions(-) diff --git a/flink-python/docs/reference/pyflink.dataframe/creation.rst b/flink-python/docs/reference/pyflink.dataframe/creation.rst index 51692dab6bcf4d..3652224e5a0fcb 100644 --- a/flink-python/docs/reference/pyflink.dataframe/creation.rst +++ b/flink-python/docs/reference/pyflink.dataframe/creation.rst @@ -20,7 +20,20 @@ DataFrame Creation ================== -Functions for creating DataFrames from row-oriented or column-oriented Python data. +Functions for creating DataFrames from row-oriented and column-oriented Python data, pandas +DataFrames, PyArrow tables, PyFlink Tables, and integer ranges. + +``schema`` is an optional list of column names. For dictionaries and mapping records it selects +and reorders named fields. For pandas and Arrow inputs it renames columns positionally and must +contain exactly one name per input column. Names must be non-empty strings and must be unique. + +Dictionary and record inputs must contain at least one row. Empty pandas and Arrow inputs are +supported when their column types can be inferred from pandas dtypes or the Arrow schema. An empty +:func:`range` still has one ``id BIGINT`` column. + +The native data creators accept an optional ``watermark=(column, expression)`` declaration. The +column must exist and have a timestamp-compatible type. Watermark columns are normalized to +``TIMESTAMP(3)`` or ``TIMESTAMP_LTZ(3)``; sub-millisecond precision is truncated. Example:: @@ -30,6 +43,28 @@ Example:: ... {"id": 2, "name": "Bob"}, ... ]) >>> users = pf.from_dict({"id": [1, 2], "name": ["Alice", "Bob"]}) + >>> identifiers = pf.range(1, 5) + +Pandas and Arrow inputs can be renamed positionally:: + + >>> import pandas as pd + >>> import pyarrow as pa + >>> pandas_users = pf.from_pandas( + ... pd.DataFrame({"identifier": [1], "display_name": ["Alice"]}), + ... schema=["id", "name"], + ... ) + >>> arrow_users = pf.from_arrow( + ... pa.table({"identifier": [1], "display_name": ["Alice"]}), + ... schema=["id", "name"], + ... ) + +A watermark can be attached while creating event data:: + + >>> from datetime import datetime + >>> events = pf.from_records( + ... [{"id": 1, "ts": datetime(2026, 1, 1)}], + ... watermark=("ts", "ts - INTERVAL '5' SECOND"), + ... ) .. currentmodule:: pyflink.dataframe @@ -38,3 +73,7 @@ Example:: from_records from_dict + from_pandas + from_arrow + from_table + range diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst index 2caa7503eba132..4bc7118a837cf1 100644 --- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst +++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst @@ -57,12 +57,18 @@ Transformations Results ------- +``to_pandas()`` executes the DataFrame and transfers every result row to the client. Use it only +when the complete result fits in client memory. ``to_table()`` returns the exact underlying +PyFlink Table without executing or copying it. + .. currentmodule:: pyflink.dataframe .. autosummary:: :toctree: api/ DataFrame.collect + DataFrame.to_table + DataFrame.to_pandas Expressions ----------- diff --git a/flink-python/pyflink/dataframe/__init__.py b/flink-python/pyflink/dataframe/__init__.py index 8ad43bcbbe2560..b4c1e32c90da8a 100644 --- a/flink-python/pyflink/dataframe/__init__.py +++ b/flink-python/pyflink/dataframe/__init__.py @@ -38,7 +38,14 @@ """ -from pyflink.dataframe.convert import from_dict, from_records +from pyflink.dataframe.convert import ( + from_arrow, + from_dict, + from_pandas, + from_records, + from_table, + range, +) from pyflink.dataframe.context import ( get_or_create_table_environment, get_table_environment, @@ -52,8 +59,12 @@ "DataType", "col", "lit", + "from_arrow", "from_dict", + "from_pandas", "from_records", + "from_table", + "range", "set_table_environment", "get_table_environment", "get_or_create_table_environment", diff --git a/flink-python/pyflink/dataframe/convert.py b/flink-python/pyflink/dataframe/convert.py index 9a78fa65ee322e..126bc1e39685d6 100644 --- a/flink-python/pyflink/dataframe/convert.py +++ b/flink-python/pyflink/dataframe/convert.py @@ -16,6 +16,7 @@ # limitations under the License. ################################################################################ +import builtins from enum import Enum from typing import ( Any, @@ -31,9 +32,28 @@ from pyflink.dataframe.context import get_or_create_table_environment from pyflink.dataframe.dataframe import DataFrame +from pyflink.table import Schema, Table +from pyflink.table.types import ( + _create_converter, + _create_type_verifier, + _infer_schema_from_data, + DataTypes, + LocalZonedTimestampType, + RowField, + RowType, + TimestampType, + from_arrow_type, +) from pyflink.util.api_stability_decorators import PublicEvolving -__all__ = ["from_dict", "from_records"] +__all__ = [ + "from_arrow", + "from_dict", + "from_pandas", + "from_records", + "from_table", + "range", +] _SCALAR_SEQUENCE_TYPES = (str, bytes, bytearray, memoryview) @@ -109,7 +129,7 @@ def normalize_record( return tuple(getattr(record, name) for name in schema) -def _validate_schema(schema: List[str]) -> None: +def _validate_schema(schema: Any) -> None: if not isinstance(schema, list) or any(not isinstance(name, str) for name in schema): raise TypeError("schema must be a list of strings") if not schema: @@ -120,10 +140,256 @@ def _validate_schema(schema: List[str]) -> None: raise ValueError("schema field names must be unique") +def _resolve_column_names( + input_names: Sequence[str], schema: Optional[List[str]] +) -> List[str]: + if schema is None: + column_names = list(input_names) + else: + _validate_schema(schema) + if len(schema) != len(input_names): + raise ValueError( + f"schema has {len(schema)} fields but data has " + f"{len(input_names)} columns" + ) + column_names = schema + _validate_schema(column_names) + return column_names + + +def _validate_watermark( + watermark: Optional[Tuple[str, str]], +) -> Optional[Tuple[str, str]]: + if watermark is None: + return None + if not isinstance(watermark, tuple) or len(watermark) != 2: + raise TypeError("watermark must be a tuple of (column, expression)") + if any(not isinstance(value, str) or not value.strip() for value in watermark): + raise TypeError("watermark column and expression must be non-empty strings") + return watermark + + +def _normalize_watermark_row_type( + row_type: RowType, watermark: Tuple[str, str] +) -> RowType: + column_name = watermark[0] + matching_fields = [field for field in row_type.fields if field.name == column_name] + if not matching_fields: + raise ValueError(f"watermark column {column_name!r} is not present in data") + + watermark_type = matching_fields[0].data_type + if not isinstance(watermark_type, (TimestampType, LocalZonedTimestampType)): + raise ValueError( + f"watermark column {column_name!r} must have a timestamp type" + ) + + fields = [] + for field in row_type.fields: + data_type = field.data_type + if field.name == column_name and data_type.precision != 3: + data_type = type(data_type)(3, data_type._nullable) + fields.append(RowField(field.name, data_type, field.description)) + return RowType(fields, row_type._nullable) + + +def _resolve_watermark_schema( + row_type: RowType, watermark: Optional[Tuple[str, str]] +) -> Tuple[RowType, Optional[Schema]]: + watermark = _validate_watermark(watermark) + if watermark is None: + return row_type, None + + row_type = _normalize_watermark_row_type(row_type, watermark) + table_schema = ( + Schema.new_builder() + .from_row_data_type(row_type) + .watermark(*watermark) + .build() + ) + return row_type, table_schema + + +def _from_rows( + rows: Sequence[Sequence[Any]], + row_type: RowType, + watermark: Optional[Tuple[str, str]] = None, +) -> DataFrame: + verify_row = _create_type_verifier(row_type) + verified_rows = [] + for row in rows: + verify_row(row) + verified_rows.append(row_type.to_sql_type(row)) + + _, table_schema = _resolve_watermark_schema(row_type, watermark) + table = get_or_create_table_environment()._from_elements( + verified_rows, row_type, table_schema + ) + return DataFrame(table) + + +def _infer_row_type( + rows: Sequence[Sequence[Any]], schema: List[str] +) -> Tuple[List[Sequence[Any]], RowType]: + row_type = _infer_schema_from_data(rows, names=schema) + converter = _create_converter(row_type) + return [converter(row) for row in rows], row_type + + +def _timestamp_precision(unit: str) -> int: + return {"s": 0, "ms": 3, "us": 6, "ns": 9}[unit] + + +def _row_type_from_arrow_schema(arrow_schema: Any, names: List[str]) -> RowType: + import pyarrow as pa + + fields = [] + for name, arrow_field in zip(names, arrow_schema): + if pa.types.is_timestamp(arrow_field.type) and arrow_field.type.tz is not None: + data_type = LocalZonedTimestampType( + _timestamp_precision(arrow_field.type.unit), arrow_field.nullable + ) + else: + data_type = from_arrow_type(arrow_field.type, arrow_field.nullable) + fields.append(RowField(name, data_type)) + return RowType(fields) + + +@PublicEvolving() +def from_table(table: Table) -> DataFrame: + """ + Create a DataFrame that wraps a PyFlink Table. + + :param table: Table to wrap without copying or converting it. + :return: A DataFrame backed by the exact supplied Table. + :raises TypeError: If ``table`` is not a :class:`~pyflink.table.Table`. + + Example:: + + >>> import pyflink.dataframe as pf + >>> table = table_env.from_elements([(1, "Alice")], ["id", "name"]) + >>> dataframe = pf.from_table(table) + >>> dataframe.to_table() is table + True + + .. versionadded:: 2.4.0 + """ + if not isinstance(table, Table): + raise TypeError("table must be a pyflink.table.Table") + return DataFrame(table) + + +@PublicEvolving() +def from_pandas( + pdf: Any, + schema: Optional[List[str]] = None, + watermark: Optional[Tuple[str, str]] = None, +) -> DataFrame: + """ + Create a DataFrame from a pandas DataFrame. + + Types are inferred from the Arrow representation of the pandas columns. An explicit ``schema`` + renames columns positionally and must contain exactly one unique, non-empty name per input + column. Empty inputs are supported when their pandas dtypes can be converted to Flink types. + + ``watermark`` declares an event-time column and its SQL watermark expression. The selected + column must have a timestamp-compatible type. Its precision is normalized to milliseconds; + values with finer precision are truncated to ``TIMESTAMP(3)`` or ``TIMESTAMP_LTZ(3)``. + + :param pdf: pandas DataFrame to convert. + :param schema: Optional list of positional result column names. + :param watermark: Optional ``(column, expression)`` watermark declaration. + :return: A DataFrame containing the pandas rows. + :raises TypeError: If the input, schema, watermark, or inferred types are invalid. + :raises ValueError: If schema width or watermark column requirements are not met. + + Example:: + + >>> import pandas as pd + >>> import pyflink.dataframe as pf + >>> pdf = pd.DataFrame({"identifier": [1, 2], "name": ["Alice", "Bob"]}) + >>> dataframe = pf.from_pandas(pdf, schema=["id", "name"]) + + .. versionadded:: 2.4.0 + """ + import pandas as pd + + if not isinstance(pdf, pd.DataFrame): + raise TypeError( + f"data must be a pandas.DataFrame, but was {type(pdf).__name__}" + ) + watermark = _validate_watermark(watermark) + + import pyarrow as pa + + arrow_table = pa.Table.from_pandas(pdf, preserve_index=False) + names = _resolve_column_names(arrow_table.column_names, schema) + row_type = _row_type_from_arrow_schema(arrow_table.schema, names) + resolved_row_type, table_schema = _resolve_watermark_schema(row_type, watermark) + table_environment = get_or_create_table_environment() + + if len(pdf) > 0 and watermark is None: + return DataFrame(table_environment.from_pandas(pdf, schema)) + return DataFrame( + table_environment._from_arrow( + arrow_table, resolved_row_type, table_schema + ) + ) + + +@PublicEvolving() +def from_arrow( + table: Any, + schema: Optional[List[str]] = None, + watermark: Optional[Tuple[str, str]] = None, +) -> DataFrame: + """ + Create a DataFrame from a PyArrow Table without converting through pandas. + + An explicit ``schema`` renames columns positionally and must contain exactly one unique, + non-empty name per input column. Empty tables are supported when their Arrow field types can be + converted to Flink types. + + ``watermark`` declares an event-time column and its SQL watermark expression. The selected + column must have a timestamp-compatible type. Its precision is normalized to milliseconds; + values with finer precision are truncated to ``TIMESTAMP(3)`` or ``TIMESTAMP_LTZ(3)``. + + :param table: PyArrow Table to convert. + :param schema: Optional list of positional result column names. + :param watermark: Optional ``(column, expression)`` watermark declaration. + :return: A DataFrame containing the Arrow rows. + :raises TypeError: If the input, schema, watermark, or inferred types are invalid. + :raises ValueError: If schema width or watermark column requirements are not met. + + Example:: + + >>> import pyarrow as pa + >>> import pyflink.dataframe as pf + >>> table = pa.table({"id": [1, 2], "name": ["Alice", "Bob"]}) + >>> dataframe = pf.from_arrow(table) + + .. versionadded:: 2.4.0 + """ + import pyarrow as pa + + if not isinstance(table, pa.Table): + raise TypeError( + f"data must be a pyarrow.Table, but was {type(table).__name__}" + ) + watermark = _validate_watermark(watermark) + names = _resolve_column_names(table.column_names, schema) + row_type = _row_type_from_arrow_schema(table.schema, names) + resolved_row_type, table_schema = _resolve_watermark_schema(row_type, watermark) + result = get_or_create_table_environment()._from_arrow( + table, resolved_row_type, table_schema + ) + return DataFrame(result) + + @PublicEvolving() def from_records( data: Sequence[Union[Sequence[Any], Mapping[str, Any]]], schema: Optional[List[str]] = None, + watermark: Optional[Tuple[str, str]] = None, ) -> DataFrame: """ Create a DataFrame from row-oriented records. @@ -137,8 +403,13 @@ def from_records( Field types are inferred from the record values. + ``watermark`` declares an event-time column and its SQL watermark expression. The selected + column must have a timestamp-compatible type. Its precision is normalized to milliseconds; + values with finer precision are truncated to ``TIMESTAMP(3)`` or ``TIMESTAMP_LTZ(3)``. + :param data: Non-empty sequence of mapping or sequence records. :param schema: Optional non-empty list of field names. + :param watermark: Optional ``(column, expression)`` watermark declaration. :return: A DataFrame containing the records. :raises TypeError: If a record or schema has an invalid type. :raises ValueError: If data or schema is empty, schema field names are invalid, a required @@ -163,6 +434,11 @@ def from_records( >>> selected_users = pf.from_records( ... [User(1, "Alice")], schema=["name", "id"] ... ) + >>> from datetime import datetime + >>> events = pf.from_records( + ... [{"id": 1, "ts": datetime(2026, 1, 1)}], + ... watermark=("ts", "ts - INTERVAL '5' SECOND"), + ... ) .. versionadded:: 2.4.0 """ @@ -172,6 +448,7 @@ def from_records( ) if not data: raise ValueError("data must not be empty") + watermark = _validate_watermark(watermark) first_record = data[0] try: @@ -204,14 +481,17 @@ def from_records( raise ValueError(f"invalid record at index {index}") from error rows.append(row) - return DataFrame( - get_or_create_table_environment().from_elements(rows, schema) - ) + if watermark is not None: + converted_rows, row_type = _infer_row_type(rows, schema) + return _from_rows(converted_rows, row_type, watermark) + return DataFrame(get_or_create_table_environment().from_elements(rows, schema)) @PublicEvolving() def from_dict( - data: Mapping[str, Sequence[Any]], schema: Optional[List[str]] = None + data: Mapping[str, Sequence[Any]], + schema: Optional[List[str]] = None, + watermark: Optional[Tuple[str, str]] = None, ) -> DataFrame: """ Create a DataFrame from a column-oriented dictionary. @@ -219,8 +499,13 @@ def from_dict( All selected columns must contain the same non-zero number of values. ``schema`` can select a subset of columns and controls their order. If omitted, dictionary insertion order is used. + ``watermark`` declares an event-time column and its SQL watermark expression. The selected + column must have a timestamp-compatible type. Its precision is normalized to milliseconds; + values with finer precision are truncated to ``TIMESTAMP(3)`` or ``TIMESTAMP_LTZ(3)``. + :param data: Non-empty mapping of column names to value sequences. :param schema: Optional non-empty list of selected column names. + :param watermark: Optional ``(column, expression)`` watermark declaration. :return: A DataFrame containing the selected columns. :raises TypeError: If ``data`` is not a mapping, or the selected schema or a selected column value has an invalid type. @@ -241,6 +526,7 @@ def from_dict( raise TypeError("data must be a mapping") if not data: raise ValueError("data must not be empty") + watermark = _validate_watermark(watermark) if schema is None: schema = list(data.keys()) _validate_schema(schema) @@ -263,8 +549,53 @@ def from_dict( raise ValueError("data must contain at least one row") rows = [ tuple(data[name][row_index] for name in schema) - for row_index in range(row_count) + for row_index in builtins.range(row_count) ] - return DataFrame( - get_or_create_table_environment().from_elements(rows, schema) - ) + if watermark is not None: + converted_rows, row_type = _infer_row_type(rows, schema) + return _from_rows(converted_rows, row_type, watermark) + return DataFrame(get_or_create_table_environment().from_elements(rows, schema)) + + +@PublicEvolving() +def range(start_or_end: int, end: Optional[int] = None, step: int = 1) -> DataFrame: + """ + Create a DataFrame containing an integer range in one ``id`` column. + + The arguments follow Python's built-in :func:`range` semantics. The result always has an + ``id BIGINT`` column, including when the requested range is empty. + + :param start_or_end: End value when ``end`` is omitted, otherwise the start value. + :param end: Optional exclusive end value. + :param step: Distance between adjacent values; must not be zero. + :return: A DataFrame with one ``id`` column. + :raises TypeError: If an argument is not an integer. + :raises ValueError: If ``step`` is zero. + + Example:: + + >>> import pyflink.dataframe as pf + >>> identifiers = pf.range(1, 6, 2) + >>> identifiers.collect() + [, , ] + + .. versionadded:: 2.4.0 + """ + if not isinstance(start_or_end, int): + raise TypeError("start_or_end must be an integer") + if end is not None and not isinstance(end, int): + raise TypeError("end must be an integer") + if not isinstance(step, int): + raise TypeError("step must be an integer") + if step == 0: + raise ValueError("step must not be zero") + + if end is None: + start = 0 + stop = start_or_end + else: + start = start_or_end + stop = end + rows = [(value,) for value in builtins.range(start, stop, step)] + row_type = DataTypes.ROW([DataTypes.FIELD("id", DataTypes.BIGINT())]) + return _from_rows(rows, row_type) diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 66c2020ee4db74..79fd08691a79b4 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -16,7 +16,10 @@ # limitations under the License. ################################################################################ -from typing import Any, Callable, List, Optional, Tuple, Union, overload +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union, overload + +if TYPE_CHECKING: + import pandas from pyflink.common import Row from pyflink.dataframe.datatype import DataType @@ -348,3 +351,41 @@ def collect(self) -> List[Row]: """ with self._table.execute().collect() as rows: return list(rows) + + @PublicEvolving() + def to_table(self) -> Table: + """ + Return the underlying PyFlink Table without copying or converting it. + + :return: The exact Table wrapped by this DataFrame. + + Example:: + + >>> import pyflink.dataframe as pf + >>> table = table_env.from_elements([(1,)], ["id"]) + >>> dataframe = pf.from_table(table) + >>> dataframe.to_table() is table + True + + .. versionadded:: 2.4.0 + """ + return self._table + + @PublicEvolving() + def to_pandas(self) -> "pandas.DataFrame": + """ + Execute this DataFrame and collect its rows into a pandas DataFrame. + + All results are transferred to the client and must fit in client memory. + + :return: A pandas DataFrame containing all result rows. + + Example:: + + >>> import pyflink.dataframe as pf + >>> dataframe = pf.from_records([{"id": 1}, {"id": 2}]) + >>> pdf = dataframe.to_pandas() + + .. versionadded:: 2.4.0 + """ + return self._table.to_pandas() diff --git a/flink-python/pyflink/dataframe/tests/test_convert.py b/flink-python/pyflink/dataframe/tests/test_convert.py index f08ede67ab7749..d064a02e8c2f20 100644 --- a/flink-python/pyflink/dataframe/tests/test_convert.py +++ b/flink-python/pyflink/dataframe/tests/test_convert.py @@ -17,9 +17,14 @@ ################################################################################ import unittest +from datetime import datetime +from unittest.mock import Mock, patch from typing import NamedTuple +import pandas as pd +import pyarrow as pa import pyflink.dataframe as pf +from pyflink.table.types import BigIntType, RowType class _Point(NamedTuple): @@ -224,5 +229,86 @@ def test_rejects_duplicate_schema_field_names(self): with self.assertRaisesRegex(ValueError, "schema field names must be unique"): pf.from_dict({"id": [1]}, schema=["id", "id"]) + +class CreationValidationTests(unittest.TestCase): + def test_rejects_invalid_watermarks(self): + invalid_watermarks = [ + ("ts", "watermark must be a tuple"), + (("ts",), "watermark must be a tuple"), + (("ts", "ts", "extra"), "watermark must be a tuple"), + (("", "ts"), "must be non-empty strings"), + (("ts", ""), "must be non-empty strings"), + ((1, "ts"), "must be non-empty strings"), + ] + for watermark, message in invalid_watermarks: + with self.subTest(watermark=watermark): + with self.assertRaisesRegex(TypeError, message): + pf.from_dict( + {"ts": [datetime(2026, 1, 1)]}, watermark=watermark + ) + + def test_pandas_and_arrow_reject_invalid_positional_schemas(self): + inputs = [ + (pf.from_pandas, pd.DataFrame({"left": [1], "right": [2]})), + (pf.from_arrow, pa.table({"left": [1], "right": [2]})), + ] + invalid_schemas = [ + ("names", TypeError, "schema must be a list of strings"), + (["left", 2], TypeError, "schema must be a list of strings"), + (["left"], ValueError, "schema has 1 fields but data has 2 columns"), + (["left", "left"], ValueError, "schema field names must be unique"), + ] + for creator, data in inputs: + for schema, error_type, message in invalid_schemas: + with self.subTest(creator=creator.__name__, schema=schema): + with self.assertRaisesRegex(error_type, message): + creator(data, schema=schema) + + def test_rejects_invalid_table_and_columnar_inputs(self): + invalid_inputs = [ + (pf.from_table, object(), "pyflink.table.Table"), + (pf.from_pandas, object(), "pandas.DataFrame"), + (pf.from_arrow, object(), "pyarrow.Table"), + ] + for creator, data, message in invalid_inputs: + with self.subTest(creator=creator.__name__): + with self.assertRaisesRegex(TypeError, message): + creator(data) + + +class RangeTests(unittest.TestCase): + def test_matches_python_range_and_preserves_bigint_schema_when_empty(self): + cases = [ + ((4,), [(0,), (1,), (2,), (3,)]), + ((4, -1, -2), [(4,), (2,), (0,)]), + ((2, 2), []), + ] + for arguments, expected_rows in cases: + table_environment = Mock() + table_environment._from_elements.return_value = object() + with self.subTest(arguments=arguments), patch( + "pyflink.dataframe.convert.get_or_create_table_environment", + return_value=table_environment, + ): + pf.range(*arguments) + + rows, row_type = table_environment._from_elements.call_args.args[:2] + self.assertEqual([row[1:] for row in rows], expected_rows) + self.assertIsInstance(row_type, RowType) + self.assertEqual(row_type.field_names(), ["id"]) + self.assertIsInstance(row_type.field_types()[0], BigIntType) + + def test_rejects_invalid_arguments(self): + invalid_arguments = [ + ((1.5,), TypeError, "start_or_end must be an integer"), + ((0, 1.5), TypeError, "end must be an integer"), + ((0, 1, 1.5), TypeError, "step must be an integer"), + ((0, 1, 0), ValueError, "step must not be zero"), + ] + for arguments, error_type, message in invalid_arguments: + with self.subTest(arguments=arguments): + with self.assertRaisesRegex(error_type, message): + pf.range(*arguments) + if __name__ == "__main__": unittest.main() diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index 98714e9fbb6a47..b918bf2b6c5153 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -17,8 +17,12 @@ ################################################################################ import unittest +from datetime import datetime, timezone +from unittest.mock import patch from typing import NamedTuple +import pandas as pd +import pyarrow as pa import pyflink.dataframe as pf from py4j.protocol import Py4JJavaError from pyflink.common import Row @@ -28,6 +32,7 @@ TableEnvironment, ) from pyflink.table.expression import Expression +from pyflink.table.types import LocalZonedTimestampType, TimestampType from pyflink.testing.test_case_utils import ( PyFlinkDataFrameUTTestCase, PyFlinkITTestCase, @@ -77,6 +82,17 @@ def execute(self): return _TableResult(self._iterator) +class _PandasTable: + def __init__(self, result=None, error=None): + self._result = result + self._error = error + + def to_pandas(self): + if self._error is not None: + raise self._error + return self._result + + class DataFrameCollectTests(unittest.TestCase): def test_collect_returns_all_rows_and_closes_iterator(self): iterator = _CloseableIterator([Row(1, "Alice")]) @@ -95,6 +111,24 @@ def test_collect_closes_iterator_when_iteration_fails(self): self.assertTrue(iterator.closed) +class DataFrameConversionTests(unittest.TestCase): + def test_to_table_returns_underlying_table(self): + table = _PandasTable() + + self.assertIs(pf.DataFrame(table).to_table(), table) + + def test_to_pandas_delegates_to_underlying_table(self): + expected = pd.DataFrame({"id": [1]}) + + self.assertIs(pf.DataFrame(_PandasTable(expected)).to_pandas(), expected) + + def test_to_pandas_propagates_errors(self): + with self.assertRaisesRegex(RuntimeError, "conversion failed"): + pf.DataFrame( + _PandasTable(error=RuntimeError("conversion failed")) + ).to_pandas() + + class DataFrameCreationTests(PyFlinkDataFrameUTTestCase): def test_from_dict_uses_insertion_order_without_schema(self): dataframe = pf.from_dict({"name": ["Alice"], "id": [1]}) @@ -192,6 +226,126 @@ def test_from_records_selects_named_tuple_fields_with_explicit_schema(self): [TableDataTypes.STRING(), TableDataTypes.BIGINT()], ) + def test_from_pandas_and_arrow_rename_columns_positionally(self): + inputs = [ + pd.DataFrame( + {"original_id": [1], "original_ts": [datetime(2026, 1, 1)]} + ), + pa.table( + { + "original_id": pa.array([1], type=pa.int64()), + "original_ts": pa.array( + [datetime(2026, 1, 1)], type=pa.timestamp("us") + ), + } + ), + ] + for creator, data in zip((pf.from_pandas, pf.from_arrow), inputs): + with self.subTest(creator=creator.__name__): + dataframe = creator(data, schema=["id", "ts"]) + self.assert_dataframe_schema(dataframe, ["id", "ts"]) + + def test_empty_pandas_and_arrow_inputs_preserve_inferred_types(self): + inputs = [ + ( + pf.from_pandas, + pd.DataFrame({"id": pd.Series([], dtype="int64")}), + ), + ( + pf.from_arrow, + pa.table({"id": pa.array([], type=pa.int64())}), + ), + ] + for creator, data in inputs: + with self.subTest(creator=creator.__name__): + dataframe = creator(data) + self.assert_dataframe_schema( + dataframe, + ["id"], + [TableDataTypes.BIGINT()], + ) + + def test_from_arrow_does_not_use_pandas_conversion(self): + with patch.object( + self.t_env, + "from_pandas", + side_effect=AssertionError("from_pandas must not be called"), + ): + dataframe = pf.from_arrow(pa.table({"id": [1]})) + + self.assert_dataframe_schema( + dataframe, + ["id"], + [TableDataTypes.BIGINT()], + ) + + def test_creators_attach_and_normalize_watermarks(self): + timestamp = datetime(2026, 1, 1, 0, 0, 0, 123456) + creators = [ + ( + lambda: pf.from_dict( + {"ts": [timestamp]}, + watermark=("ts", "ts - INTERVAL '1' SECOND"), + ), + LocalZonedTimestampType, + ), + ( + lambda: pf.from_records( + [{"ts": timestamp}], + watermark=("ts", "ts - INTERVAL '1' SECOND"), + ), + LocalZonedTimestampType, + ), + ( + lambda: pf.from_pandas( + pd.DataFrame({"ts": [timestamp]}), + watermark=("ts", "ts - INTERVAL '1' SECOND"), + ), + TimestampType, + ), + ( + lambda: pf.from_arrow( + pa.table( + { + "ts": pa.array( + [timestamp.replace(tzinfo=timezone.utc)], + type=pa.timestamp("us", tz="UTC"), + ) + } + ), + watermark=("ts", "ts - INTERVAL '1' SECOND"), + ), + LocalZonedTimestampType, + ), + ] + for creator, expected_type in creators: + with self.subTest(creator=creator): + resolved_schema = creator().to_table().get_resolved_schema() + timestamp_type = resolved_schema.get_column_data_types()[0] + self.assertIsInstance(timestamp_type, expected_type) + self.assertEqual(timestamp_type.precision, 3) + watermark_specs = resolved_schema.get_watermark_specs() + self.assertEqual(len(watermark_specs), 1) + self.assertEqual(watermark_specs[0].get_rowtime_attribute(), "ts") + + def test_watermark_requires_existing_timestamp_column(self): + invalid_watermarks = [ + (("missing", "ts"), "watermark column 'missing' is not present"), + (("id", "id"), "watermark column 'id' must have a timestamp type"), + ] + for watermark, message in invalid_watermarks: + with self.subTest(watermark=watermark): + with self.assertRaisesRegex(ValueError, message): + pf.from_records( + [{"id": 1, "ts": datetime(2026, 1, 1)}], + watermark=watermark, + ) + + def test_from_table_and_to_table_preserve_identity(self): + table = self.t_env.from_elements([(1,)], ["id"]) + + self.assertIs(pf.from_table(table).to_table(), table) + class DataFrameSelectTests(PyFlinkDataFrameUTTestCase): def setUp(self): @@ -480,6 +634,27 @@ def test_from_records(self): [Row(1, "Alice"), Row(2, "Bob")], ) + def test_arrow_to_pandas_round_trip(self): + timestamp = datetime(2026, 1, 1, 0, 0, 0, 123000) + arrow_table = pa.table( + { + "id": pa.array([1, 2], type=pa.int64()), + "ts": pa.array([timestamp, None], type=pa.timestamp("ms")), + } + ) + + result = ( + pf.from_arrow(arrow_table) + .with_column("id_plus_one", pf.col("id") + 1) + .to_pandas() + ) + + self.assertEqual(list(result.columns), ["id", "ts", "id_plus_one"]) + self.assertEqual(result["id"].tolist(), [1, 2]) + self.assertEqual(result["id_plus_one"].tolist(), [2, 3]) + self.assertEqual(result["ts"].isna().tolist(), [False, True]) + self.assertEqual(result.loc[0, "ts"].to_pydatetime(), timestamp) + def test_basic_functionality(self): df = pf.from_dict( { diff --git a/flink-python/pyflink/table/table_environment.py b/flink-python/pyflink/table/table_environment.py index f0e4bedba174e8..b80b98fe45bdea 100644 --- a/flink-python/pyflink/table/table_environment.py +++ b/flink-python/pyflink/table/table_environment.py @@ -1469,11 +1469,17 @@ def verify_obj(obj): elements = [schema.to_sql_type(element) for element in elements] return self._from_elements(elements, schema) - def _from_elements(self, elements: List, schema: DataType) -> Table: + def _from_elements( + self, + elements: List, + schema: DataType, + table_schema: Schema = None) -> Table: """ Creates a table from a collection of elements. :param elements: The elements to create a table from. + :param schema: Data type used to serialize the elements. + :param table_schema: Optional declarative schema for the resulting source table. :return: The result :class:`~pyflink.table.Table`. """ # serializes to a file, and we read the file in java @@ -1482,7 +1488,8 @@ def _from_elements(self, elements: List, schema: DataType) -> Table: try: with temp_file: serializer.serialize(elements, temp_file) - j_schema = _to_java_data_type(schema) + j_schema = (table_schema._j_schema if table_schema is not None + else _to_java_data_type(schema)) gateway = get_gateway() PythonTableUtils = gateway.jvm \ .org.apache.flink.table.utils.python.PythonTableUtils @@ -1492,6 +1499,45 @@ def _from_elements(self, elements: List, schema: DataType) -> Table: finally: atexit.register(lambda: os.unlink(temp_file.name)) + def _from_arrow( + self, + table, + row_type: RowType, + table_schema: Schema = None) -> Table: + """Creates a table from a PyArrow Table through the Arrow table source.""" + import pyarrow as pa + + if not isinstance(table, pa.Table): + raise TypeError(f"table must be a pyarrow.Table, but was {type(table).__name__}") + + arrow_schema = create_arrow_schema(row_type.field_names(), row_type.field_types()) + try: + compatible_table = table.rename_columns(row_type.field_names()).cast( + arrow_schema, safe=False) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError, ValueError) as e: + raise TypeError( + f"Could not convert pyarrow.Table to the inferred Flink schema: {row_type}" + ) from e + + temp_file = tempfile.NamedTemporaryFile(delete=False, dir=tempfile.mkdtemp()) + try: + with temp_file: + with pa.ipc.new_stream(temp_file, arrow_schema) as writer: + writer.write_table(compatible_table) + + jvm = get_gateway().jvm + if table_schema is None: + source_schema = _to_java_data_type(row_type).notNull() + source_schema = source_schema.bridgedTo( + load_java_class('org.apache.flink.table.data.RowData')) + else: + source_schema = table_schema._j_schema + descriptor = jvm.org.apache.flink.table.runtime.arrow.ArrowUtils \ + .createArrowTableSourceDesc(source_schema, temp_file.name) + return Table(getattr(self._j_tenv, "from")(descriptor), self) + finally: + os.unlink(temp_file.name) + def from_pandas(self, pdf: 'pandas.DataFrame', schema: Union[RowType, List[str], Tuple[str], List[DataType], Tuple[DataType]] = None, diff --git a/flink-python/pyflink/table/types.py b/flink-python/pyflink/table/types.py index 316bb44b55e23e..7e59165f911cc2 100644 --- a/flink-python/pyflink/table/types.py +++ b/flink-python/pyflink/table/types.py @@ -2286,7 +2286,7 @@ def from_arrow_type(arrow_type, nullable: bool = True) -> DataType: elif types.is_null(arrow_type): return NullType() else: - raise TypeError("Unsupported data type to convert to Arrow type: " + str(dt)) + raise TypeError("Unsupported data type to convert from Arrow type: " + str(arrow_type)) def to_arrow_type(data_type: DataType): diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java index d8b9dcdf772b66..d85d1b082b910b 100644 --- a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java +++ b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java @@ -488,14 +488,18 @@ public static TableDescriptor createArrowTableSourceDesc(DataType dataType, Stri for (int i = 0; i < fieldNames.size(); i++) { schemaBuilder.column(fieldNames.get(i), fieldTypes.get(i)); } + return createArrowTableSourceDesc(schemaBuilder.build(), fileName); + } + public static TableDescriptor createArrowTableSourceDesc( + org.apache.flink.table.api.Schema schema, String fileName) { try { byte[][] data = readArrowBatches(fileName); return TableDescriptor.forConnector(ArrowTableSourceFactory.IDENTIFIER) .option( ArrowTableSourceOptions.DATA, ByteArrayUtils.twoDimByteArrayToString(data)) - .schema(schemaBuilder.build()) + .schema(schema) .build(); } catch (Throwable e) { throw new TableException("Failed to read the arrow data from " + fileName, e); diff --git a/flink-python/src/main/java/org/apache/flink/table/utils/python/PythonTableUtils.java b/flink-python/src/main/java/org/apache/flink/table/utils/python/PythonTableUtils.java index 01dfab186dddcf..3681d0664ae1a3 100644 --- a/flink-python/src/main/java/org/apache/flink/table/utils/python/PythonTableUtils.java +++ b/flink-python/src/main/java/org/apache/flink/table/utils/python/PythonTableUtils.java @@ -100,11 +100,27 @@ private PythonTableUtils() {} */ public static Table createTableFromElement( TableEnvironment tEnv, String filePath, DataType schema, boolean batched) { + return createTableFromElement( + tEnv, filePath, Schema.newBuilder().fromRowDataType(schema).build(), batched); + } + + /** + * Create a table from {@link PythonDynamicTableSource} that reads data from an input file with + * the given declarative {@link Schema}. + * + * @param tEnv The TableEnvironment to create the table. + * @param filePath the file path of the input data. + * @param schema the schema of the table, including time attributes when present. + * @param batched Whether to read data in a batch. + * @return Table backed by the input file. + */ + public static Table createTableFromElement( + TableEnvironment tEnv, String filePath, Schema schema, boolean batched) { TableDescriptor.Builder builder = TableDescriptor.forConnector(PythonDynamicTableFactory.IDENTIFIER) .option(PythonDynamicTableOptions.INPUT_FILE_PATH, filePath) .option(PythonDynamicTableOptions.BATCH_MODE, batched) - .schema(Schema.newBuilder().fromRowDataType(schema).build()); + .schema(schema); return tEnv.from(builder.build()); } From 6625267b92423f25617ea5600c905d547f5178e8 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 15:41:34 +0800 Subject: [PATCH 02/11] [FLINK-40190][python] Refine DataFrame conversion internals Align timezone-aware Arrow timestamps with existing Table API semantics, delegate pandas creation to the Arrow path, and add split-aware Arrow IPC serialization. Refine watermark and row helpers with focused tests. Generated-by: Codex (GPT-5) --- flink-python/pyflink/dataframe/convert.py | 120 ++++++++---------- .../pyflink/dataframe/tests/test_convert.py | 9 ++ .../pyflink/dataframe/tests/test_dataframe.py | 65 +++++++++- .../pyflink/table/table_environment.py | 22 +++- .../table/tests/test_pandas_conversion.py | 41 ++++++ 5 files changed, 181 insertions(+), 76 deletions(-) diff --git a/flink-python/pyflink/dataframe/convert.py b/flink-python/pyflink/dataframe/convert.py index 126bc1e39685d6..27580153251f82 100644 --- a/flink-python/pyflink/dataframe/convert.py +++ b/flink-python/pyflink/dataframe/convert.py @@ -23,6 +23,7 @@ Collection, List, Mapping, + NamedTuple, Optional, Sequence, Tuple, @@ -58,6 +59,11 @@ _SCALAR_SEQUENCE_TYPES = (str, bytes, bytearray, memoryview) +class _WatermarkSpec(NamedTuple): + column: str + expression: str + + class _RecordType(Enum): NAMED_TUPLE = "named_tuple" MAPPING = "mapping" @@ -143,36 +149,34 @@ def _validate_schema(schema: Any) -> None: def _resolve_column_names( input_names: Sequence[str], schema: Optional[List[str]] ) -> List[str]: - if schema is None: - column_names = list(input_names) - else: - _validate_schema(schema) - if len(schema) != len(input_names): - raise ValueError( - f"schema has {len(schema)} fields but data has " - f"{len(input_names)} columns" - ) - column_names = schema + column_names = list(input_names) if schema is None else schema + if ( + schema is not None + and isinstance(schema, list) + and len(schema) != len(input_names) + ): + raise ValueError( + f"schema has {len(schema)} fields but data has " + f"{len(input_names)} columns" + ) _validate_schema(column_names) return column_names -def _validate_watermark( +def _parse_watermark( watermark: Optional[Tuple[str, str]], -) -> Optional[Tuple[str, str]]: +) -> Optional[_WatermarkSpec]: if watermark is None: return None if not isinstance(watermark, tuple) or len(watermark) != 2: raise TypeError("watermark must be a tuple of (column, expression)") if any(not isinstance(value, str) or not value.strip() for value in watermark): raise TypeError("watermark column and expression must be non-empty strings") - return watermark + return _WatermarkSpec(*watermark) -def _normalize_watermark_row_type( - row_type: RowType, watermark: Tuple[str, str] -) -> RowType: - column_name = watermark[0] +def _normalize_watermark_row_type(row_type: RowType, watermark: _WatermarkSpec) -> RowType: + column_name = watermark.column matching_fields = [field for field in row_type.fields if field.name == column_name] if not matching_fields: raise ValueError(f"watermark column {column_name!r} is not present in data") @@ -193,9 +197,8 @@ def _normalize_watermark_row_type( def _resolve_watermark_schema( - row_type: RowType, watermark: Optional[Tuple[str, str]] + row_type: RowType, watermark: Optional[_WatermarkSpec] ) -> Tuple[RowType, Optional[Schema]]: - watermark = _validate_watermark(watermark) if watermark is None: return row_type, None @@ -203,16 +206,16 @@ def _resolve_watermark_schema( table_schema = ( Schema.new_builder() .from_row_data_type(row_type) - .watermark(*watermark) + .watermark(watermark.column, watermark.expression) .build() ) return row_type, table_schema -def _from_rows( +def _create_dataframe_from_rows( rows: Sequence[Sequence[Any]], row_type: RowType, - watermark: Optional[Tuple[str, str]] = None, + watermark: Optional[_WatermarkSpec] = None, ) -> DataFrame: verify_row = _create_type_verifier(row_type) verified_rows = [] @@ -227,7 +230,7 @@ def _from_rows( return DataFrame(table) -def _infer_row_type( +def _infer_row_type_and_convert_rows( rows: Sequence[Sequence[Any]], schema: List[str] ) -> Tuple[List[Sequence[Any]], RowType]: row_type = _infer_schema_from_data(rows, names=schema) @@ -235,23 +238,13 @@ def _infer_row_type( return [converter(row) for row in rows], row_type -def _timestamp_precision(unit: str) -> int: - return {"s": 0, "ms": 3, "us": 6, "ns": 9}[unit] - - def _row_type_from_arrow_schema(arrow_schema: Any, names: List[str]) -> RowType: - import pyarrow as pa - - fields = [] - for name, arrow_field in zip(names, arrow_schema): - if pa.types.is_timestamp(arrow_field.type) and arrow_field.type.tz is not None: - data_type = LocalZonedTimestampType( - _timestamp_precision(arrow_field.type.unit), arrow_field.nullable - ) - else: - data_type = from_arrow_type(arrow_field.type, arrow_field.nullable) - fields.append(RowField(name, data_type)) - return RowType(fields) + return RowType( + [ + RowField(name, from_arrow_type(arrow_field.type, arrow_field.nullable)) + for name, arrow_field in zip(names, arrow_schema) + ] + ) @PublicEvolving() @@ -308,6 +301,10 @@ def from_pandas( >>> import pyflink.dataframe as pf >>> pdf = pd.DataFrame({"identifier": [1, 2], "name": ["Alice", "Bob"]}) >>> dataframe = pf.from_pandas(pdf, schema=["id", "name"]) + >>> events = pf.from_pandas( + ... pd.DataFrame({"ts": pd.to_datetime(["2026-01-01T00:00:00Z"])}), + ... watermark=("ts", "ts - INTERVAL '5' SECOND"), + ... ) .. versionadded:: 2.4.0 """ @@ -317,22 +314,13 @@ def from_pandas( raise TypeError( f"data must be a pandas.DataFrame, but was {type(pdf).__name__}" ) - watermark = _validate_watermark(watermark) import pyarrow as pa - arrow_table = pa.Table.from_pandas(pdf, preserve_index=False) - names = _resolve_column_names(arrow_table.column_names, schema) - row_type = _row_type_from_arrow_schema(arrow_table.schema, names) - resolved_row_type, table_schema = _resolve_watermark_schema(row_type, watermark) - table_environment = get_or_create_table_environment() - - if len(pdf) > 0 and watermark is None: - return DataFrame(table_environment.from_pandas(pdf, schema)) - return DataFrame( - table_environment._from_arrow( - arrow_table, resolved_row_type, table_schema - ) + return from_arrow( + pa.Table.from_pandas(pdf, preserve_index=False), + schema=schema, + watermark=watermark, ) @@ -366,6 +354,10 @@ def from_arrow( >>> import pyflink.dataframe as pf >>> table = pa.table({"id": [1, 2], "name": ["Alice", "Bob"]}) >>> dataframe = pf.from_arrow(table) + >>> events = pf.from_arrow( + ... pa.table({"ts": pa.array([0], type=pa.timestamp("ms"))}), + ... watermark=("ts", "ts - INTERVAL '5' SECOND"), + ... ) .. versionadded:: 2.4.0 """ @@ -375,10 +367,12 @@ def from_arrow( raise TypeError( f"data must be a pyarrow.Table, but was {type(table).__name__}" ) - watermark = _validate_watermark(watermark) + watermark_spec = _parse_watermark(watermark) names = _resolve_column_names(table.column_names, schema) row_type = _row_type_from_arrow_schema(table.schema, names) - resolved_row_type, table_schema = _resolve_watermark_schema(row_type, watermark) + resolved_row_type, table_schema = _resolve_watermark_schema( + row_type, watermark_spec + ) result = get_or_create_table_environment()._from_arrow( table, resolved_row_type, table_schema ) @@ -448,7 +442,7 @@ def from_records( ) if not data: raise ValueError("data must not be empty") - watermark = _validate_watermark(watermark) + watermark_spec = _parse_watermark(watermark) first_record = data[0] try: @@ -481,10 +475,8 @@ def from_records( raise ValueError(f"invalid record at index {index}") from error rows.append(row) - if watermark is not None: - converted_rows, row_type = _infer_row_type(rows, schema) - return _from_rows(converted_rows, row_type, watermark) - return DataFrame(get_or_create_table_environment().from_elements(rows, schema)) + converted_rows, row_type = _infer_row_type_and_convert_rows(rows, schema) + return _create_dataframe_from_rows(converted_rows, row_type, watermark_spec) @PublicEvolving() @@ -526,7 +518,7 @@ def from_dict( raise TypeError("data must be a mapping") if not data: raise ValueError("data must not be empty") - watermark = _validate_watermark(watermark) + watermark_spec = _parse_watermark(watermark) if schema is None: schema = list(data.keys()) _validate_schema(schema) @@ -551,10 +543,8 @@ def from_dict( tuple(data[name][row_index] for name in schema) for row_index in builtins.range(row_count) ] - if watermark is not None: - converted_rows, row_type = _infer_row_type(rows, schema) - return _from_rows(converted_rows, row_type, watermark) - return DataFrame(get_or_create_table_environment().from_elements(rows, schema)) + converted_rows, row_type = _infer_row_type_and_convert_rows(rows, schema) + return _create_dataframe_from_rows(converted_rows, row_type, watermark_spec) @PublicEvolving() @@ -598,4 +588,4 @@ def range(start_or_end: int, end: Optional[int] = None, step: int = 1) -> DataFr stop = end rows = [(value,) for value in builtins.range(start, stop, step)] row_type = DataTypes.ROW([DataTypes.FIELD("id", DataTypes.BIGINT())]) - return _from_rows(rows, row_type) + return _create_dataframe_from_rows(rows, row_type) diff --git a/flink-python/pyflink/dataframe/tests/test_convert.py b/flink-python/pyflink/dataframe/tests/test_convert.py index d064a02e8c2f20..098161bb23b5f4 100644 --- a/flink-python/pyflink/dataframe/tests/test_convert.py +++ b/flink-python/pyflink/dataframe/tests/test_convert.py @@ -24,6 +24,7 @@ import pandas as pd import pyarrow as pa import pyflink.dataframe as pf +import pyflink.dataframe.convert as dataframe_convert from pyflink.table.types import BigIntType, RowType @@ -231,6 +232,14 @@ def test_rejects_duplicate_schema_field_names(self): class CreationValidationTests(unittest.TestCase): + def test_parses_watermark_into_semantic_specification(self): + watermark = dataframe_convert._parse_watermark( + ("ts", "ts - INTERVAL '5' SECOND") + ) + + self.assertEqual(watermark.column, "ts") + self.assertEqual(watermark.expression, "ts - INTERVAL '5' SECOND") + def test_rejects_invalid_watermarks(self): invalid_watermarks = [ ("ts", "watermark must be a tuple"), diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index b918bf2b6c5153..ae341b66afcc50 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -265,18 +265,69 @@ def test_empty_pandas_and_arrow_inputs_preserve_inferred_types(self): [TableDataTypes.BIGINT()], ) - def test_from_arrow_does_not_use_pandas_conversion(self): + def test_columnar_creators_do_not_use_table_environment_from_pandas(self): with patch.object( self.t_env, "from_pandas", side_effect=AssertionError("from_pandas must not be called"), ): - dataframe = pf.from_arrow(pa.table({"id": [1]})) + for creator, data in [ + (pf.from_pandas, pd.DataFrame({"id": [1]})), + (pf.from_arrow, pa.table({"id": [1]})), + ]: + with self.subTest(creator=creator.__name__): + dataframe = creator(data) + self.assert_dataframe_schema( + dataframe, + ["id"], + [TableDataTypes.BIGINT()], + ) - self.assert_dataframe_schema( - dataframe, - ["id"], - [TableDataTypes.BIGINT()], + def test_from_pandas_matches_table_environment_schema(self): + pdf = pd.DataFrame( + { + "original_id": [1.0, None], + "original_name": ["Alice", None], + "original_ts": pd.Series( + pd.to_datetime( + ["2026-01-01T00:00:00Z", "2026-01-02T00:00:00Z"] + ) + ), + } + ) + names = ["id", "name", "ts"] + + dataframe_schema = ( + pf.from_pandas(pdf, schema=names).to_table().get_resolved_schema() + ) + table_schema = self.t_env.from_pandas( + pdf, schema=names + ).get_resolved_schema() + + self.assertEqual( + table_schema.get_column_names(), dataframe_schema.get_column_names() + ) + self.assertEqual( + table_schema.get_column_data_types(), + dataframe_schema.get_column_data_types(), + ) + + empty_pdf = pd.DataFrame( + { + "original_id": pd.Series([], dtype="float64"), + "original_name": pd.Series([], dtype="string"), + "original_ts": pd.Series([], dtype="datetime64[ns, UTC]"), + } + ) + empty_schema = pf.from_pandas( + empty_pdf, schema=names + ).to_table().get_resolved_schema() + self.assertEqual( + table_schema.get_column_names(), empty_schema.get_column_names() + ) + self.assertEqual( + table_schema.get_column_data_types(), + empty_schema.get_column_data_types(), ) def test_creators_attach_and_normalize_watermarks(self): @@ -315,7 +366,7 @@ def test_creators_attach_and_normalize_watermarks(self): ), watermark=("ts", "ts - INTERVAL '1' SECOND"), ), - LocalZonedTimestampType, + TimestampType, ), ] for creator, expected_type in creators: diff --git a/flink-python/pyflink/table/table_environment.py b/flink-python/pyflink/table/table_environment.py index b80b98fe45bdea..3e885cc666a4f8 100644 --- a/flink-python/pyflink/table/table_environment.py +++ b/flink-python/pyflink/table/table_environment.py @@ -19,7 +19,7 @@ import os import sys import tempfile -from typing import Union, List, Tuple, Iterable, Optional, TYPE_CHECKING +from typing import BinaryIO, Union, List, Tuple, Iterable, Optional, TYPE_CHECKING if TYPE_CHECKING: import pandas @@ -61,6 +61,20 @@ ] +def _serialize_arrow_table(table, stream: BinaryIO, splits_num: int) -> None: + if isinstance(splits_num, bool) or not isinstance(splits_num, int): + raise TypeError("splits_num must be an integer") + if splits_num <= 0: + raise ValueError("splits_num must be greater than 0") + + import pyarrow as pa + + with pa.ipc.new_stream(stream, table.schema) as writer: + if table.num_rows > 0: + max_chunksize = -(-table.num_rows // splits_num) + writer.write_table(table, max_chunksize=max_chunksize) + + @PublicEvolving() class TableEnvironment(object): """ @@ -1503,7 +1517,8 @@ def _from_arrow( self, table, row_type: RowType, - table_schema: Schema = None) -> Table: + table_schema: Schema = None, + splits_num: int = 1) -> Table: """Creates a table from a PyArrow Table through the Arrow table source.""" import pyarrow as pa @@ -1522,8 +1537,7 @@ def _from_arrow( temp_file = tempfile.NamedTemporaryFile(delete=False, dir=tempfile.mkdtemp()) try: with temp_file: - with pa.ipc.new_stream(temp_file, arrow_schema) as writer: - writer.write_table(compatible_table) + _serialize_arrow_table(compatible_table, temp_file, splits_num) jvm = get_gateway().jvm if table_schema is None: diff --git a/flink-python/pyflink/table/tests/test_pandas_conversion.py b/flink-python/pyflink/table/tests/test_pandas_conversion.py index 9cc0f8ccdf6772..0c70e2949605a4 100644 --- a/flink-python/pyflink/table/tests/test_pandas_conversion.py +++ b/flink-python/pyflink/table/tests/test_pandas_conversion.py @@ -17,16 +17,57 @@ ################################################################################ import datetime import decimal +import io +import unittest from pandas.testing import assert_frame_equal +import pyarrow as pa from pyflink.common import Row +from pyflink.table import table_environment from pyflink.table.types import DataTypes from pyflink.testing import source_sink_utils from pyflink.testing.test_case_utils import PyFlinkBatchTableTestCase, \ PyFlinkStreamTableTestCase +class ArrowTableSerializationTests(unittest.TestCase): + + def test_serializes_expected_batch_sizes(self): + table = pa.table({"id": [1, 2, 3, 4, 5]}) + stream = io.BytesIO() + + table_environment._serialize_arrow_table(table, stream, splits_num=2) + + reader = pa.ipc.open_stream(stream.getvalue()) + self.assertEqual([3, 2], [batch.num_rows for batch in reader]) + + def test_serializes_empty_table_with_schema(self): + table = pa.table({"id": pa.array([], type=pa.int64())}) + stream = io.BytesIO() + + table_environment._serialize_arrow_table(table, stream, splits_num=1) + + reader = pa.ipc.open_stream(stream.getvalue()) + self.assertEqual(table.schema, reader.schema) + self.assertEqual([], list(reader)) + + def test_rejects_invalid_split_counts(self): + table = pa.table({"id": [1]}) + invalid_splits = [ + (True, TypeError, "splits_num must be an integer"), + (1.5, TypeError, "splits_num must be an integer"), + (0, ValueError, "splits_num must be greater than 0"), + (-1, ValueError, "splits_num must be greater than 0"), + ] + for splits_num, error_type, message in invalid_splits: + with self.subTest(splits_num=splits_num): + with self.assertRaisesRegex(error_type, message): + table_environment._serialize_arrow_table( + table, io.BytesIO(), splits_num + ) + + class PandasConversionTestBase(object): @classmethod From eee94a05022ad1e36c458a1d95e18e9875c794ce Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 16:27:59 +0800 Subject: [PATCH 03/11] [FLINK-40190][python] Simplify DataFrame row creation Combine inferred-schema row conversion with DataFrame creation and let range use its known BIGINT schema directly. Generated-by: Codex (GPT-5) --- flink-python/pyflink/dataframe/convert.py | 37 +++++++++++------------ 1 file changed, 17 insertions(+), 20 deletions(-) diff --git a/flink-python/pyflink/dataframe/convert.py b/flink-python/pyflink/dataframe/convert.py index 27580153251f82..1c076c982f84da 100644 --- a/flink-python/pyflink/dataframe/convert.py +++ b/flink-python/pyflink/dataframe/convert.py @@ -212,32 +212,27 @@ def _resolve_watermark_schema( return row_type, table_schema -def _create_dataframe_from_rows( +def _infer_schema_and_create_dataframe( rows: Sequence[Sequence[Any]], - row_type: RowType, + column_names: List[str], watermark: Optional[_WatermarkSpec] = None, ) -> DataFrame: + row_type = _infer_schema_from_data(rows, names=column_names) + row_type, table_schema = _resolve_watermark_schema(row_type, watermark) + converter = _create_converter(row_type) verify_row = _create_type_verifier(row_type) - verified_rows = [] + sql_rows = [] for row in rows: + row = converter(row) verify_row(row) - verified_rows.append(row_type.to_sql_type(row)) + sql_rows.append(row_type.to_sql_type(row)) - _, table_schema = _resolve_watermark_schema(row_type, watermark) table = get_or_create_table_environment()._from_elements( - verified_rows, row_type, table_schema + sql_rows, row_type, table_schema ) return DataFrame(table) -def _infer_row_type_and_convert_rows( - rows: Sequence[Sequence[Any]], schema: List[str] -) -> Tuple[List[Sequence[Any]], RowType]: - row_type = _infer_schema_from_data(rows, names=schema) - converter = _create_converter(row_type) - return [converter(row) for row in rows], row_type - - def _row_type_from_arrow_schema(arrow_schema: Any, names: List[str]) -> RowType: return RowType( [ @@ -475,8 +470,7 @@ def from_records( raise ValueError(f"invalid record at index {index}") from error rows.append(row) - converted_rows, row_type = _infer_row_type_and_convert_rows(rows, schema) - return _create_dataframe_from_rows(converted_rows, row_type, watermark_spec) + return _infer_schema_and_create_dataframe(rows, schema, watermark_spec) @PublicEvolving() @@ -543,8 +537,7 @@ def from_dict( tuple(data[name][row_index] for name in schema) for row_index in builtins.range(row_count) ] - converted_rows, row_type = _infer_row_type_and_convert_rows(rows, schema) - return _create_dataframe_from_rows(converted_rows, row_type, watermark_spec) + return _infer_schema_and_create_dataframe(rows, schema, watermark_spec) @PublicEvolving() @@ -586,6 +579,10 @@ def range(start_or_end: int, end: Optional[int] = None, step: int = 1) -> DataFr else: start = start_or_end stop = end - rows = [(value,) for value in builtins.range(start, stop, step)] row_type = DataTypes.ROW([DataTypes.FIELD("id", DataTypes.BIGINT())]) - return _create_dataframe_from_rows(rows, row_type) + sql_rows = [ + row_type.to_sql_type((value,)) + for value in builtins.range(start, stop, step) + ] + table = get_or_create_table_environment()._from_elements(sql_rows, row_type) + return DataFrame(table) From 729d161463b4f4284bb53d30d96fc5f75d2f0047 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 17:40:04 +0800 Subject: [PATCH 04/11] [FLINK-40190][python] Validate DataFrame range BIGINT bounds Reject ranges whose emitted values exceed signed BIGINT bounds before creating the underlying table. Cover valid boundary values and ascending and descending overflow cases. Generated-by: Codex (GPT-5) --- flink-python/pyflink/dataframe/convert.py | 18 +++++++++++----- .../pyflink/dataframe/tests/test_convert.py | 21 +++++++++++++++++++ 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/flink-python/pyflink/dataframe/convert.py b/flink-python/pyflink/dataframe/convert.py index 1c076c982f84da..6c47250b7ec293 100644 --- a/flink-python/pyflink/dataframe/convert.py +++ b/flink-python/pyflink/dataframe/convert.py @@ -57,6 +57,8 @@ ] _SCALAR_SEQUENCE_TYPES = (str, bytes, bytearray, memoryview) +_BIGINT_MIN = -(1 << 63) +_BIGINT_MAX = (1 << 63) - 1 class _WatermarkSpec(NamedTuple): @@ -553,7 +555,8 @@ def range(start_or_end: int, end: Optional[int] = None, step: int = 1) -> DataFr :param step: Distance between adjacent values; must not be zero. :return: A DataFrame with one ``id`` column. :raises TypeError: If an argument is not an integer. - :raises ValueError: If ``step`` is zero. + :raises ValueError: If ``step`` is zero or the range contains values outside the signed + ``BIGINT`` bounds. Example:: @@ -579,10 +582,15 @@ def range(start_or_end: int, end: Optional[int] = None, step: int = 1) -> DataFr else: start = start_or_end stop = end + values = builtins.range(start, stop, step) + has_values = start < stop if step > 0 else start > stop + if has_values and not ( + _BIGINT_MIN <= values[0] <= _BIGINT_MAX + and _BIGINT_MIN <= values[-1] <= _BIGINT_MAX + ): + raise ValueError("range values must fit in signed BIGINT") + row_type = DataTypes.ROW([DataTypes.FIELD("id", DataTypes.BIGINT())]) - sql_rows = [ - row_type.to_sql_type((value,)) - for value in builtins.range(start, stop, step) - ] + sql_rows = [row_type.to_sql_type((value,)) for value in values] table = get_or_create_table_environment()._from_elements(sql_rows, row_type) return DataFrame(table) diff --git a/flink-python/pyflink/dataframe/tests/test_convert.py b/flink-python/pyflink/dataframe/tests/test_convert.py index 098161bb23b5f4..1d9739f74b8578 100644 --- a/flink-python/pyflink/dataframe/tests/test_convert.py +++ b/flink-python/pyflink/dataframe/tests/test_convert.py @@ -291,6 +291,8 @@ def test_matches_python_range_and_preserves_bigint_schema_when_empty(self): ((4,), [(0,), (1,), (2,), (3,)]), ((4, -1, -2), [(4,), (2,), (0,)]), ((2, 2), []), + ((2**63 - 1, 2**63), [(2**63 - 1,)]), + ((-(2**63), -(2**63) + 1), [(-(2**63),)]), ] for arguments, expected_rows in cases: table_environment = Mock() @@ -307,6 +309,25 @@ def test_matches_python_range_and_preserves_bigint_schema_when_empty(self): self.assertEqual(row_type.field_names(), ["id"]) self.assertIsInstance(row_type.field_types()[0], BigIntType) + def test_rejects_values_outside_bigint_bounds(self): + invalid_ranges = [ + (2**63, 2**63 + 1), + (2**63 - 1, 2**63 + 2), + (-(2**63) - 1, -(2**63) - 2, -1), + (-(2**63), -(2**63) - 3, -1), + ] + table_environment = Mock() + for arguments in invalid_ranges: + with self.subTest(arguments=arguments), patch( + "pyflink.dataframe.convert.get_or_create_table_environment", + return_value=table_environment, + ) as get_table_environment: + with self.assertRaisesRegex( + ValueError, "range values must fit in signed BIGINT" + ): + pf.range(*arguments) + get_table_environment.assert_not_called() + def test_rejects_invalid_arguments(self): invalid_arguments = [ ((1.5,), TypeError, "start_or_end must be an integer"), From 950ef610114125d5084fb6e3ef48770265e2a298 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 20:46:26 +0800 Subject: [PATCH 05/11] [FLINK-40190][python] Refine DataFrame conversion docs and coverage Clarify that to_table does not execute a job, simplify the DataFrame creation and results reference pages, and exercise from_pandas through filtering, projection, and conversion in the existing pandas round-trip integration smoke test. Generated-by: Codex (GPT-5) --- .../reference/pyflink.dataframe/creation.rst | 35 +++---------------- .../reference/pyflink.dataframe/dataframe.rst | 4 --- flink-python/pyflink/dataframe/dataframe.py | 2 ++ .../pyflink/dataframe/tests/test_dataframe.py | 14 ++++---- 4 files changed, 15 insertions(+), 40 deletions(-) diff --git a/flink-python/docs/reference/pyflink.dataframe/creation.rst b/flink-python/docs/reference/pyflink.dataframe/creation.rst index 3652224e5a0fcb..9a4b577a12e8cd 100644 --- a/flink-python/docs/reference/pyflink.dataframe/creation.rst +++ b/flink-python/docs/reference/pyflink.dataframe/creation.rst @@ -20,20 +20,7 @@ DataFrame Creation ================== -Functions for creating DataFrames from row-oriented and column-oriented Python data, pandas -DataFrames, PyArrow tables, PyFlink Tables, and integer ranges. - -``schema`` is an optional list of column names. For dictionaries and mapping records it selects -and reorders named fields. For pandas and Arrow inputs it renames columns positionally and must -contain exactly one name per input column. Names must be non-empty strings and must be unique. - -Dictionary and record inputs must contain at least one row. Empty pandas and Arrow inputs are -supported when their column types can be inferred from pandas dtypes or the Arrow schema. An empty -:func:`range` still has one ``id BIGINT`` column. - -The native data creators accept an optional ``watermark=(column, expression)`` declaration. The -column must exist and have a timestamp-compatible type. Watermark columns are normalized to -``TIMESTAMP(3)`` or ``TIMESTAMP_LTZ(3)``; sub-millisecond precision is truncated. +Functions for creating DataFrames from row-oriented or column-oriented Python data. Example:: @@ -43,28 +30,16 @@ Example:: ... {"id": 2, "name": "Bob"}, ... ]) >>> users = pf.from_dict({"id": [1, 2], "name": ["Alice", "Bob"]}) - >>> identifiers = pf.range(1, 5) - -Pandas and Arrow inputs can be renamed positionally:: - >>> import pandas as pd >>> import pyarrow as pa >>> pandas_users = pf.from_pandas( - ... pd.DataFrame({"identifier": [1], "display_name": ["Alice"]}), - ... schema=["id", "name"], + ... pd.DataFrame({"id": [1, 2], "name": ["Alice", "Bob"]}) ... ) >>> arrow_users = pf.from_arrow( - ... pa.table({"identifier": [1], "display_name": ["Alice"]}), - ... schema=["id", "name"], - ... ) - -A watermark can be attached while creating event data:: - - >>> from datetime import datetime - >>> events = pf.from_records( - ... [{"id": 1, "ts": datetime(2026, 1, 1)}], - ... watermark=("ts", "ts - INTERVAL '5' SECOND"), + ... pa.table({"id": [1, 2], "name": ["Alice", "Bob"]}) ... ) + >>> table_users = pf.from_table(users.to_table()) + >>> identifiers = pf.range(5) .. currentmodule:: pyflink.dataframe diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst index 4bc7118a837cf1..4ab37c87b6b010 100644 --- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst +++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst @@ -57,10 +57,6 @@ Transformations Results ------- -``to_pandas()`` executes the DataFrame and transfers every result row to the client. Use it only -when the complete result fits in client memory. ``to_table()`` returns the exact underlying -PyFlink Table without executing or copying it. - .. currentmodule:: pyflink.dataframe .. autosummary:: diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 79fd08691a79b4..3e58f5c8319e29 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -357,6 +357,8 @@ def to_table(self) -> Table: """ Return the underlying PyFlink Table without copying or converting it. + This method does not trigger job execution. + :return: The exact Table wrapped by this DataFrame. Example:: diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index ae341b66afcc50..3bdb81d2b75e06 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -685,22 +685,24 @@ def test_from_records(self): [Row(1, "Alice"), Row(2, "Bob")], ) - def test_arrow_to_pandas_round_trip(self): + def test_pandas_to_pandas_round_trip(self): timestamp = datetime(2026, 1, 1, 0, 0, 0, 123000) - arrow_table = pa.table( + pdf = pd.DataFrame( { - "id": pa.array([1, 2], type=pa.int64()), - "ts": pa.array([timestamp, None], type=pa.timestamp("ms")), + "id": [0, 1, 2], + "ts": pd.Series([None, timestamp, None], dtype="datetime64[ms]"), } ) result = ( - pf.from_arrow(arrow_table) + pf.from_pandas(pdf) + .filter(pf.col("id") > 0) .with_column("id_plus_one", pf.col("id") + 1) + .select("id", "id_plus_one", "ts") .to_pandas() ) - self.assertEqual(list(result.columns), ["id", "ts", "id_plus_one"]) + self.assertEqual(list(result.columns), ["id", "id_plus_one", "ts"]) self.assertEqual(result["id"].tolist(), [1, 2]) self.assertEqual(result["id_plus_one"].tolist(), [2, 3]) self.assertEqual(result["ts"].isna().tolist(), [False, True]) From 6151e0f27e21567e1b6f880d20aa92ea1c6ffb98 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 7 Aug 2026 11:26:01 +0800 Subject: [PATCH 06/11] [FLINK-40190][python] Inline Arrow table serialization Inline the single-use Arrow IPC writer into TableEnvironment._from_arrow and remove its helper-specific tests. This also avoids the NamedTemporaryFile wrapper type mismatch reported by mypy. Generated-by: Codex (GPT-5) --- .../pyflink/table/table_environment.py | 25 ++++------- .../table/tests/test_pandas_conversion.py | 41 ------------------- 2 files changed, 9 insertions(+), 57 deletions(-) diff --git a/flink-python/pyflink/table/table_environment.py b/flink-python/pyflink/table/table_environment.py index 3e885cc666a4f8..917c768d88f839 100644 --- a/flink-python/pyflink/table/table_environment.py +++ b/flink-python/pyflink/table/table_environment.py @@ -19,7 +19,7 @@ import os import sys import tempfile -from typing import BinaryIO, Union, List, Tuple, Iterable, Optional, TYPE_CHECKING +from typing import Union, List, Tuple, Iterable, Optional, TYPE_CHECKING if TYPE_CHECKING: import pandas @@ -61,20 +61,6 @@ ] -def _serialize_arrow_table(table, stream: BinaryIO, splits_num: int) -> None: - if isinstance(splits_num, bool) or not isinstance(splits_num, int): - raise TypeError("splits_num must be an integer") - if splits_num <= 0: - raise ValueError("splits_num must be greater than 0") - - import pyarrow as pa - - with pa.ipc.new_stream(stream, table.schema) as writer: - if table.num_rows > 0: - max_chunksize = -(-table.num_rows // splits_num) - writer.write_table(table, max_chunksize=max_chunksize) - - @PublicEvolving() class TableEnvironment(object): """ @@ -1524,6 +1510,10 @@ def _from_arrow( if not isinstance(table, pa.Table): raise TypeError(f"table must be a pyarrow.Table, but was {type(table).__name__}") + if isinstance(splits_num, bool) or not isinstance(splits_num, int): + raise TypeError("splits_num must be an integer") + if splits_num <= 0: + raise ValueError("splits_num must be greater than 0") arrow_schema = create_arrow_schema(row_type.field_names(), row_type.field_types()) try: @@ -1537,7 +1527,10 @@ def _from_arrow( temp_file = tempfile.NamedTemporaryFile(delete=False, dir=tempfile.mkdtemp()) try: with temp_file: - _serialize_arrow_table(compatible_table, temp_file, splits_num) + with pa.ipc.new_stream(temp_file, compatible_table.schema) as writer: + if compatible_table.num_rows > 0: + max_chunksize = -(-compatible_table.num_rows // splits_num) + writer.write_table(compatible_table, max_chunksize=max_chunksize) jvm = get_gateway().jvm if table_schema is None: diff --git a/flink-python/pyflink/table/tests/test_pandas_conversion.py b/flink-python/pyflink/table/tests/test_pandas_conversion.py index 0c70e2949605a4..9cc0f8ccdf6772 100644 --- a/flink-python/pyflink/table/tests/test_pandas_conversion.py +++ b/flink-python/pyflink/table/tests/test_pandas_conversion.py @@ -17,57 +17,16 @@ ################################################################################ import datetime import decimal -import io -import unittest from pandas.testing import assert_frame_equal -import pyarrow as pa from pyflink.common import Row -from pyflink.table import table_environment from pyflink.table.types import DataTypes from pyflink.testing import source_sink_utils from pyflink.testing.test_case_utils import PyFlinkBatchTableTestCase, \ PyFlinkStreamTableTestCase -class ArrowTableSerializationTests(unittest.TestCase): - - def test_serializes_expected_batch_sizes(self): - table = pa.table({"id": [1, 2, 3, 4, 5]}) - stream = io.BytesIO() - - table_environment._serialize_arrow_table(table, stream, splits_num=2) - - reader = pa.ipc.open_stream(stream.getvalue()) - self.assertEqual([3, 2], [batch.num_rows for batch in reader]) - - def test_serializes_empty_table_with_schema(self): - table = pa.table({"id": pa.array([], type=pa.int64())}) - stream = io.BytesIO() - - table_environment._serialize_arrow_table(table, stream, splits_num=1) - - reader = pa.ipc.open_stream(stream.getvalue()) - self.assertEqual(table.schema, reader.schema) - self.assertEqual([], list(reader)) - - def test_rejects_invalid_split_counts(self): - table = pa.table({"id": [1]}) - invalid_splits = [ - (True, TypeError, "splits_num must be an integer"), - (1.5, TypeError, "splits_num must be an integer"), - (0, ValueError, "splits_num must be greater than 0"), - (-1, ValueError, "splits_num must be greater than 0"), - ] - for splits_num, error_type, message in invalid_splits: - with self.subTest(splits_num=splits_num): - with self.assertRaisesRegex(error_type, message): - table_environment._serialize_arrow_table( - table, io.BytesIO(), splits_num - ) - - class PandasConversionTestBase(object): @classmethod From d73b437403494891a0bff2b219f3be2f8c36e6cb Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 11 Aug 2026 17:07:31 +0800 Subject: [PATCH 07/11] [FLINK-40190][python] Refine watermark handling Encapsulate watermark parsing and row-type normalization, inline schema construction, and verify watermark unpacking order. Generated-by: Codex (GPT-5) --- flink-python/pyflink/dataframe/convert.py | 130 +++++++++--------- .../pyflink/dataframe/tests/test_convert.py | 9 +- 2 files changed, 72 insertions(+), 67 deletions(-) diff --git a/flink-python/pyflink/dataframe/convert.py b/flink-python/pyflink/dataframe/convert.py index 6c47250b7ec293..738ac2bd842eeb 100644 --- a/flink-python/pyflink/dataframe/convert.py +++ b/flink-python/pyflink/dataframe/convert.py @@ -65,6 +65,41 @@ class _WatermarkSpec(NamedTuple): column: str expression: str + @classmethod + def parse(cls, watermark: Optional[Tuple[str, str]]) -> Optional["_WatermarkSpec"]: + if watermark is None: + return None + if not isinstance(watermark, tuple) or len(watermark) != 2: + raise TypeError("watermark must be a tuple of (column, expression)") + if any(not isinstance(value, str) or not value.strip() for value in watermark): + raise TypeError( + "watermark column and expression must be non-empty strings" + ) + return cls(*watermark) + + def normalize_row_type(self, row_type: RowType) -> RowType: + matching_fields = [ + field for field in row_type.fields if field.name == self.column + ] + if not matching_fields: + raise ValueError( + f"watermark column {self.column!r} is not present in data" + ) + + watermark_type = matching_fields[0].data_type + if not isinstance(watermark_type, (TimestampType, LocalZonedTimestampType)): + raise ValueError( + f"watermark column {self.column!r} must have a timestamp type" + ) + + fields = [] + for field in row_type.fields: + data_type = field.data_type + if field.name == self.column and data_type.precision != 3: + data_type = type(data_type)(3, data_type._nullable) + fields.append(RowField(field.name, data_type, field.description)) + return RowType(fields, row_type._nullable) + class _RecordType(Enum): NAMED_TUPLE = "named_tuple" @@ -165,62 +200,22 @@ def _resolve_column_names( return column_names -def _parse_watermark( - watermark: Optional[Tuple[str, str]], -) -> Optional[_WatermarkSpec]: - if watermark is None: - return None - if not isinstance(watermark, tuple) or len(watermark) != 2: - raise TypeError("watermark must be a tuple of (column, expression)") - if any(not isinstance(value, str) or not value.strip() for value in watermark): - raise TypeError("watermark column and expression must be non-empty strings") - return _WatermarkSpec(*watermark) - - -def _normalize_watermark_row_type(row_type: RowType, watermark: _WatermarkSpec) -> RowType: - column_name = watermark.column - matching_fields = [field for field in row_type.fields if field.name == column_name] - if not matching_fields: - raise ValueError(f"watermark column {column_name!r} is not present in data") - - watermark_type = matching_fields[0].data_type - if not isinstance(watermark_type, (TimestampType, LocalZonedTimestampType)): - raise ValueError( - f"watermark column {column_name!r} must have a timestamp type" - ) - - fields = [] - for field in row_type.fields: - data_type = field.data_type - if field.name == column_name and data_type.precision != 3: - data_type = type(data_type)(3, data_type._nullable) - fields.append(RowField(field.name, data_type, field.description)) - return RowType(fields, row_type._nullable) - - -def _resolve_watermark_schema( - row_type: RowType, watermark: Optional[_WatermarkSpec] -) -> Tuple[RowType, Optional[Schema]]: - if watermark is None: - return row_type, None - - row_type = _normalize_watermark_row_type(row_type, watermark) - table_schema = ( - Schema.new_builder() - .from_row_data_type(row_type) - .watermark(watermark.column, watermark.expression) - .build() - ) - return row_type, table_schema - - def _infer_schema_and_create_dataframe( rows: Sequence[Sequence[Any]], column_names: List[str], watermark: Optional[_WatermarkSpec] = None, ) -> DataFrame: row_type = _infer_schema_from_data(rows, names=column_names) - row_type, table_schema = _resolve_watermark_schema(row_type, watermark) + if watermark is None: + table_schema = None + else: + row_type = watermark.normalize_row_type(row_type) + table_schema = ( + Schema.new_builder() + .from_row_data_type(row_type) + .watermark(*watermark) + .build() + ) converter = _create_converter(row_type) verify_row = _create_type_verifier(row_type) sql_rows = [] @@ -235,15 +230,6 @@ def _infer_schema_and_create_dataframe( return DataFrame(table) -def _row_type_from_arrow_schema(arrow_schema: Any, names: List[str]) -> RowType: - return RowType( - [ - RowField(name, from_arrow_type(arrow_field.type, arrow_field.nullable)) - for name, arrow_field in zip(names, arrow_schema) - ] - ) - - @PublicEvolving() def from_table(table: Table) -> DataFrame: """ @@ -364,14 +350,26 @@ def from_arrow( raise TypeError( f"data must be a pyarrow.Table, but was {type(table).__name__}" ) - watermark_spec = _parse_watermark(watermark) + watermark_spec = _WatermarkSpec.parse(watermark) names = _resolve_column_names(table.column_names, schema) - row_type = _row_type_from_arrow_schema(table.schema, names) - resolved_row_type, table_schema = _resolve_watermark_schema( - row_type, watermark_spec + row_type = RowType( + [ + RowField(name, from_arrow_type(field.type, field.nullable)) + for name, field in zip(names, table.schema) + ] ) + if watermark_spec is None: + table_schema = None + else: + row_type = watermark_spec.normalize_row_type(row_type) + table_schema = ( + Schema.new_builder() + .from_row_data_type(row_type) + .watermark(*watermark_spec) + .build() + ) result = get_or_create_table_environment()._from_arrow( - table, resolved_row_type, table_schema + table, row_type, table_schema ) return DataFrame(result) @@ -439,7 +437,7 @@ def from_records( ) if not data: raise ValueError("data must not be empty") - watermark_spec = _parse_watermark(watermark) + watermark_spec = _WatermarkSpec.parse(watermark) first_record = data[0] try: @@ -514,7 +512,7 @@ def from_dict( raise TypeError("data must be a mapping") if not data: raise ValueError("data must not be empty") - watermark_spec = _parse_watermark(watermark) + watermark_spec = _WatermarkSpec.parse(watermark) if schema is None: schema = list(data.keys()) _validate_schema(schema) diff --git a/flink-python/pyflink/dataframe/tests/test_convert.py b/flink-python/pyflink/dataframe/tests/test_convert.py index 1d9739f74b8578..6a0ff6dc39a34c 100644 --- a/flink-python/pyflink/dataframe/tests/test_convert.py +++ b/flink-python/pyflink/dataframe/tests/test_convert.py @@ -233,13 +233,20 @@ def test_rejects_duplicate_schema_field_names(self): class CreationValidationTests(unittest.TestCase): def test_parses_watermark_into_semantic_specification(self): - watermark = dataframe_convert._parse_watermark( + watermark = dataframe_convert._WatermarkSpec.parse( ("ts", "ts - INTERVAL '5' SECOND") ) self.assertEqual(watermark.column, "ts") self.assertEqual(watermark.expression, "ts - INTERVAL '5' SECOND") + def test_watermark_spec_unpacks_column_before_expression(self): + watermark = dataframe_convert._WatermarkSpec( + "ts", "ts - INTERVAL '5' SECOND" + ) + + self.assertEqual(tuple(watermark), ("ts", "ts - INTERVAL '5' SECOND")) + def test_rejects_invalid_watermarks(self): invalid_watermarks = [ ("ts", "watermark must be a tuple"), From 48020b96448b40655152ef6d08df50afb1d1f358 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 13 Aug 2026 00:11:01 +0800 Subject: [PATCH 08/11] [FLINK-40190][python] Address DataFrame conversion review feedback Handle duplicate pandas column names, reject unresolved Arrow null types, preserve timezone-aware instants across conversion, and refine focused coverage. Generated-by: Codex (GPT-5) --- flink-python/pyflink/dataframe/convert.py | 27 +++++- .../pyflink/dataframe/tests/test_convert.py | 36 ++++++++ .../pyflink/dataframe/tests/test_dataframe.py | 83 +++++++++++++------ .../pyflink/table/tests/test_types.py | 24 +++++- flink-python/pyflink/table/types.py | 13 ++- flink-python/pyflink/table/utils.py | 8 +- 6 files changed, 157 insertions(+), 34 deletions(-) diff --git a/flink-python/pyflink/dataframe/convert.py b/flink-python/pyflink/dataframe/convert.py index 738ac2bd842eeb..ddee331a580758 100644 --- a/flink-python/pyflink/dataframe/convert.py +++ b/flink-python/pyflink/dataframe/convert.py @@ -37,6 +37,7 @@ from pyflink.table.types import ( _create_converter, _create_type_verifier, + _has_nulltype, _infer_schema_from_data, DataTypes, LocalZonedTimestampType, @@ -266,6 +267,8 @@ def from_pandas( Types are inferred from the Arrow representation of the pandas columns. An explicit ``schema`` renames columns positionally and must contain exactly one unique, non-empty name per input column. Empty inputs are supported when their pandas dtypes can be converted to Flink types. + Timezone-aware timestamps are represented as ``TIMESTAMP_LTZ`` in the TableEnvironment's + configured local timezone; timezone-naive timestamps are represented as ``TIMESTAMP``. ``watermark`` declares an event-time column and its SQL watermark expression. The selected column must have a timestamp-compatible type. Its precision is normalized to milliseconds; @@ -300,9 +303,15 @@ def from_pandas( import pyarrow as pa + input_names = list(pdf.columns) + arrow_pdf = pdf.copy(deep=False) + arrow_pdf.columns = [ + f"__pyflink_dataframe_column_{index}" + for index in builtins.range(len(input_names)) + ] return from_arrow( - pa.Table.from_pandas(pdf, preserve_index=False), - schema=schema, + pa.Table.from_pandas(arrow_pdf, preserve_index=False), + schema=input_names if schema is None else schema, watermark=watermark, ) @@ -358,6 +367,15 @@ def from_arrow( for name, field in zip(names, table.schema) ] ) + null_field_names = [ + field.name for field in row_type.fields if _has_nulltype(field.data_type) + ] + if null_field_names: + columns = ", ".join(repr(name) for name in null_field_names) + raise TypeError( + f"Cannot infer Flink data types for columns with Arrow null types: {columns}. " + "Use explicit pandas or Arrow dtypes for these columns." + ) if watermark_spec is None: table_schema = None else: @@ -500,11 +518,16 @@ def from_dict( Example:: + >>> from datetime import datetime >>> import pyflink.dataframe as pf >>> users = pf.from_dict( ... {"name": ["Alice", "Bob"], "id": [1, 2]}, ... schema=["id", "name"], ... ) + >>> events = pf.from_dict( + ... {"id": [1], "ts": [datetime(2026, 1, 1)]}, + ... watermark=("ts", "ts - INTERVAL '5' SECOND"), + ... ) .. versionadded:: 2.4.0 """ diff --git a/flink-python/pyflink/dataframe/tests/test_convert.py b/flink-python/pyflink/dataframe/tests/test_convert.py index 6a0ff6dc39a34c..a204f7a6611ba2 100644 --- a/flink-python/pyflink/dataframe/tests/test_convert.py +++ b/flink-python/pyflink/dataframe/tests/test_convert.py @@ -280,6 +280,42 @@ def test_pandas_and_arrow_reject_invalid_positional_schemas(self): with self.assertRaisesRegex(error_type, message): creator(data, schema=schema) + def test_pandas_rejects_duplicate_columns_without_schema(self): + pdf = pd.DataFrame([[1, 2]], columns=["value", "value"]) + + with self.assertRaisesRegex(ValueError, "schema field names must be unique"): + pf.from_pandas(pdf) + + def test_rejects_columnar_fields_containing_null_type(self): + inputs = [ + ( + lambda: pf.from_pandas(pd.DataFrame({"value": [None]})), + "columns with Arrow null types: 'value'", + ), + ( + lambda: pf.from_arrow( + pa.table({"left": pa.nulls(1), "right": pa.nulls(1)}) + ), + "columns with Arrow null types: 'left', 'right'", + ), + ( + lambda: pf.from_arrow( + pa.table( + { + "value": pa.array( + [[None]], type=pa.list_(pa.null()) + ) + } + ) + ), + "columns with Arrow null types: 'value'", + ), + ] + for creator, message in inputs: + with self.subTest(creator=creator): + with self.assertRaisesRegex(TypeError, message): + creator() + def test_rejects_invalid_table_and_columnar_inputs(self): invalid_inputs = [ (pf.from_table, object(), "pyflink.table.Table"), diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index 3bdb81d2b75e06..2d6b62714df605 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -32,7 +32,7 @@ TableEnvironment, ) from pyflink.table.expression import Expression -from pyflink.table.types import LocalZonedTimestampType, TimestampType +from pyflink.table.types import LocalZonedTimestampType from pyflink.testing.test_case_utils import ( PyFlinkDataFrameUTTestCase, PyFlinkITTestCase, @@ -245,6 +245,16 @@ def test_from_pandas_and_arrow_rename_columns_positionally(self): dataframe = creator(data, schema=["id", "ts"]) self.assert_dataframe_schema(dataframe, ["id", "ts"]) + duplicate_pdf = pd.DataFrame( + [[1, "Alice"], [2, "Bob"]], columns=["value", "value"] + ) + dataframe = pf.from_pandas(duplicate_pdf, schema=["id", "name"]) + self.assert_dataframe_schema( + dataframe, + ["id", "name"], + [TableDataTypes.BIGINT(), TableDataTypes.STRING()], + ) + def test_empty_pandas_and_arrow_inputs_preserve_inferred_types(self): inputs = [ ( @@ -311,6 +321,9 @@ def test_from_pandas_matches_table_environment_schema(self): table_schema.get_column_data_types(), dataframe_schema.get_column_data_types(), ) + self.assertIsInstance( + dataframe_schema.get_column_data_types()[2], LocalZonedTimestampType + ) empty_pdf = pd.DataFrame( { @@ -349,10 +362,17 @@ def test_creators_attach_and_normalize_watermarks(self): ), ( lambda: pf.from_pandas( - pd.DataFrame({"ts": [timestamp]}), + pd.DataFrame( + { + "ts": pd.Series( + [timestamp.replace(tzinfo=timezone.utc)], + dtype="datetime64[us, UTC]", + ) + } + ), watermark=("ts", "ts - INTERVAL '1' SECOND"), ), - TimestampType, + LocalZonedTimestampType, ), ( lambda: pf.from_arrow( @@ -366,7 +386,7 @@ def test_creators_attach_and_normalize_watermarks(self): ), watermark=("ts", "ts - INTERVAL '1' SECOND"), ), - TimestampType, + LocalZonedTimestampType, ), ] for creator, expected_type in creators: @@ -686,27 +706,42 @@ def test_from_records(self): ) def test_pandas_to_pandas_round_trip(self): - timestamp = datetime(2026, 1, 1, 0, 0, 0, 123000) - pdf = pd.DataFrame( - { - "id": [0, 1, 2], - "ts": pd.Series([None, timestamp, None], dtype="datetime64[ms]"), - } - ) + original_timezone = self.t_env.get_config().get_local_timezone() + self.t_env.get_config().set_local_timezone("America/New_York") + try: + first_fold = pd.Timestamp("2026-11-01T05:30:00.123Z") + second_fold = pd.Timestamp("2026-11-01T06:30:00.123Z") + pdf = pd.DataFrame( + { + "id": [0, 1, 2, 3], + "ts": pd.Series( + [None, first_fold, second_fold, None], + dtype="datetime64[ms, UTC]", + ), + } + ) - result = ( - pf.from_pandas(pdf) - .filter(pf.col("id") > 0) - .with_column("id_plus_one", pf.col("id") + 1) - .select("id", "id_plus_one", "ts") - .to_pandas() - ) - - self.assertEqual(list(result.columns), ["id", "id_plus_one", "ts"]) - self.assertEqual(result["id"].tolist(), [1, 2]) - self.assertEqual(result["id_plus_one"].tolist(), [2, 3]) - self.assertEqual(result["ts"].isna().tolist(), [False, True]) - self.assertEqual(result.loc[0, "ts"].to_pydatetime(), timestamp) + result = ( + pf.from_pandas(pdf) + .filter(pf.col("id") > 0) + .with_column("id_plus_one", pf.col("id") + 1) + .select("id", "id_plus_one", "ts") + .to_pandas() + ) + + self.assertEqual(list(result.columns), ["id", "id_plus_one", "ts"]) + self.assertEqual(result["id"].tolist(), [1, 2, 3]) + self.assertEqual(result["id_plus_one"].tolist(), [2, 3, 4]) + self.assertEqual(result["ts"].isna().tolist(), [False, False, True]) + self.assertEqual( + result["ts"].tolist()[:2], + [ + first_fold.tz_convert("America/New_York"), + second_fold.tz_convert("America/New_York"), + ], + ) + finally: + self.t_env.get_config().set_local_timezone(original_timezone) def test_basic_functionality(self): df = pf.from_dict( diff --git a/flink-python/pyflink/table/tests/test_types.py b/flink-python/pyflink/table/tests/test_types.py index 400a3729ee9212..681884beeeb871 100644 --- a/flink-python/pyflink/table/tests/test_types.py +++ b/flink-python/pyflink/table/tests/test_types.py @@ -24,6 +24,8 @@ import tempfile import unittest +import pyarrow as pa + from pyflink.pyflink_gateway_server import on_windows from pyflink.serializers import BatchedSerializer, PickleSerializer @@ -35,7 +37,8 @@ _create_type_verifier, UserDefinedType, DataTypes, Row, RowField, RowType, ArrayType, BigIntType, VarCharType, MapType, DataType, _from_java_data_type, ZonedTimestampType, - LocalZonedTimestampType, _to_java_data_type) + LocalZonedTimestampType, TimestampType, _to_java_data_type, + from_arrow_type) from pyflink.testing.test_case_utils import PyFlinkTestCase @@ -126,6 +129,25 @@ def dst(self, dt): return self.OFFSET +class ArrowTypeConversionTests(unittest.TestCase): + + def test_timestamp_type_mapping(self): + units_and_precisions = [("s", 0), ("ms", 3), ("us", 6), ("ns", 9)] + for timezone_id, expected_type in [ + (None, TimestampType), + ("Asia/Shanghai", LocalZonedTimestampType), + ]: + for unit, precision in units_and_precisions: + with self.subTest(timezone_id=timezone_id, unit=unit): + data_type = from_arrow_type( + pa.timestamp(unit, tz=timezone_id), nullable=False + ) + + self.assertIsInstance(data_type, expected_type) + self.assertEqual(data_type.precision, precision) + self.assertFalse(data_type._nullable) + + class TypesTests(PyFlinkTestCase): def test_row_type_repr_includes_nullability(self): diff --git a/flink-python/pyflink/table/types.py b/flink-python/pyflink/table/types.py index 7e59165f911cc2..6627be47bd9d14 100644 --- a/flink-python/pyflink/table/types.py +++ b/flink-python/pyflink/table/types.py @@ -2263,14 +2263,19 @@ def from_arrow_type(arrow_type, nullable: bool = True) -> DataType: else: return TimeType(9, nullable) elif types.is_timestamp(arrow_type): + timestamp_type = ( + LocalZonedTimestampType + if arrow_type.tz is not None + else TimestampType + ) if arrow_type.unit == 's': - return TimestampType(0, nullable) + return timestamp_type(0, nullable) elif arrow_type.unit == 'ms': - return TimestampType(3, nullable) + return timestamp_type(3, nullable) elif arrow_type.unit == 'us': - return TimestampType(6, nullable) + return timestamp_type(6, nullable) else: - return TimestampType(9, nullable) + return timestamp_type(9, nullable) elif types.is_map(arrow_type): return MapType(from_arrow_type(arrow_type.key_type), from_arrow_type(arrow_type.item_type), diff --git a/flink-python/pyflink/table/utils.py b/flink-python/pyflink/table/utils.py index bc4444e66ebb96..89f906ca687eb6 100644 --- a/flink-python/pyflink/table/utils.py +++ b/flink-python/pyflink/table/utils.py @@ -82,7 +82,7 @@ def tz_convert_from_internal(s, t: DataType, local_tz): returns a converted series. """ if type(t) == LocalZonedTimestampType: - return s.dt.tz_localize(local_tz) + return s.dt.tz_localize(datetime.timezone.utc).dt.tz_convert(local_tz) else: return s @@ -95,9 +95,11 @@ def tz_convert_to_internal(s, t: DataType, local_tz): 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) + return s.dt.tz_localize(local_tz).dt.tz_convert( + datetime.timezone.utc + ).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(datetime.timezone.utc).dt.tz_localize(None) return s From 9a5e938156801a410bf7ce645b46004286ffd29f Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 13 Aug 2026 10:17:41 +0800 Subject: [PATCH 09/11] [FLINK-40190][python] Support Arrow local-zoned timestamps Serialize TIMESTAMP_LTZ values as instants when collecting Python table results and enable their use in nested Arrow types. Generated-by: Codex (GPT-5) --- .../pyflink/dataframe/tests/test_dataframe.py | 20 ++++++ .../pyflink/table/tests/test_types.py | 16 ++++- flink-python/pyflink/table/types.py | 4 +- flink-python/pyflink/table/utils.py | 5 ++ .../api/common/python/PythonBridgeUtils.java | 24 +++++++ .../common/python/PythonBridgeUtilsTest.java | 64 +++++++++++++++++++ 6 files changed, 130 insertions(+), 3 deletions(-) create mode 100644 flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index 2d6b62714df605..f7bda3da2f4674 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -743,6 +743,26 @@ def test_pandas_to_pandas_round_trip(self): finally: self.t_env.get_config().set_local_timezone(original_timezone) + def test_collect_timezone_aware_arrow_timestamps(self): + timestamp = pd.Timestamp("2026-01-01T00:00:00.123Z") + arrow_type = pa.timestamp("ms", tz="UTC") + dataframe = pf.from_arrow( + pa.table( + { + "timestamp": pa.array([timestamp], type=arrow_type), + "timestamps": pa.array( + [[timestamp]], type=pa.list_(arrow_type) + ), + } + ) + ) + expected = timestamp.to_pydatetime() + + self.assertEqual( + dataframe.collect(), + [Row(expected, [expected])], + ) + def test_basic_functionality(self): df = pf.from_dict( { diff --git a/flink-python/pyflink/table/tests/test_types.py b/flink-python/pyflink/table/tests/test_types.py index 681884beeeb871..c7b6fae9a19f62 100644 --- a/flink-python/pyflink/table/tests/test_types.py +++ b/flink-python/pyflink/table/tests/test_types.py @@ -38,7 +38,7 @@ RowType, ArrayType, BigIntType, VarCharType, MapType, DataType, _from_java_data_type, ZonedTimestampType, LocalZonedTimestampType, TimestampType, _to_java_data_type, - from_arrow_type) + from_arrow_type, to_arrow_type) from pyflink.testing.test_case_utils import PyFlinkTestCase @@ -147,6 +147,20 @@ def test_timestamp_type_mapping(self): self.assertEqual(data_type.precision, precision) self.assertFalse(data_type._nullable) + def test_nested_local_zoned_timestamp_type_mapping(self): + timestamp_type = pa.timestamp("ms", tz="UTC") + + self.assertEqual( + to_arrow_type(from_arrow_type(pa.list_(timestamp_type))), + pa.list_(pa.timestamp("ms")), + ) + self.assertEqual( + to_arrow_type( + from_arrow_type(pa.struct([pa.field("ts", timestamp_type)])) + ), + pa.struct([pa.field("ts", pa.timestamp("ms"))]), + ) + class TypesTests(PyFlinkTestCase): diff --git a/flink-python/pyflink/table/types.py b/flink-python/pyflink/table/types.py index 6627be47bd9d14..de5abba8a1cebc 100644 --- a/flink-python/pyflink/table/types.py +++ b/flink-python/pyflink/table/types.py @@ -2344,13 +2344,13 @@ def to_arrow_type(data_type: DataType): elif isinstance(data_type, MapType): return pa.map_(to_arrow_type(data_type.key_type), to_arrow_type(data_type.value_type)) elif isinstance(data_type, ArrayType): - if type(data_type.element_type) in [LocalZonedTimestampType, RowType]: + if isinstance(data_type.element_type, RowType): raise ValueError("%s is not supported to be used as the element type of ArrayType." % data_type.element_type) return pa.list_(to_arrow_type(data_type.element_type)) elif isinstance(data_type, RowType): for field in data_type: - if type(field.data_type) in [LocalZonedTimestampType, RowType]: + if isinstance(field.data_type, RowType): raise TypeError("%s is not supported to be used as the field type of RowType" % field.data_type) fields = [pa.field(field.name, to_arrow_type(field.data_type), field.data_type._nullable) diff --git a/flink-python/pyflink/table/utils.py b/flink-python/pyflink/table/utils.py index 89f906ca687eb6..8533c03a895d94 100644 --- a/flink-python/pyflink/table/utils.py +++ b/flink-python/pyflink/table/utils.py @@ -134,6 +134,11 @@ def pickled_bytes_to_python_converter(data, field_type: DataType): return field_type.from_sql_type(data) elif isinstance(field_type, TimestampType): return field_type.from_sql_type(int(data.timestamp() * 10**6)) + elif isinstance(field_type, LocalZonedTimestampType): + seconds, nanoseconds = data + return datetime.datetime.fromtimestamp( + seconds, datetime.timezone.utc + ).replace(microsecond=nanoseconds // 1000) elif isinstance(field_type, MapType): key_type = field_type.key_type value_type = field_type.value_type diff --git a/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java b/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java index 1205e79d2a8a83..cbb6b2810c86fd 100644 --- a/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java +++ b/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java @@ -36,10 +36,12 @@ import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos; import org.apache.flink.core.memory.DataOutputViewStreamWrapper; import org.apache.flink.streaming.api.typeinfo.python.PickledByteArrayTypeInfo; +import org.apache.flink.table.data.TimestampData; import org.apache.flink.table.types.DataType; import org.apache.flink.table.types.logical.ArrayType; import org.apache.flink.table.types.logical.DateType; import org.apache.flink.table.types.logical.FloatType; +import org.apache.flink.table.types.logical.LocalZonedTimestampType; import org.apache.flink.table.types.logical.LogicalType; import org.apache.flink.table.types.logical.MapType; import org.apache.flink.table.types.logical.RowType; @@ -58,6 +60,7 @@ import java.sql.Date; import java.sql.Time; import java.sql.Timestamp; +import java.time.Instant; import java.time.LocalDate; import java.time.LocalDateTime; import java.time.LocalTime; @@ -193,6 +196,9 @@ private static Object getPickledBytesFromJavaObject(Object obj, LogicalType data } else { return pickler.dumps(obj); } + } else if (dataType instanceof LocalZonedTimestampType) { + Instant instant = toInstant(obj); + return pickler.dumps(Arrays.asList(instant.getEpochSecond(), instant.getNano())); } else if (dataType instanceof RowType) { Row tmpRow = (Row) obj; LogicalType[] tmpRowFieldTypes = @@ -238,6 +244,24 @@ private static Object getPickledBytesFromJavaObject(Object obj, LogicalType data } } + private static Instant toInstant(Object value) { + if (value instanceof Instant) { + return (Instant) value; + } else if (value instanceof Integer) { + return Instant.ofEpochSecond(((Integer) value).longValue()); + } else if (value instanceof Long) { + return Instant.ofEpochMilli((Long) value); + } else if (value instanceof TimestampData) { + return ((TimestampData) value).toInstant(); + } else if (value instanceof Timestamp) { + return ((Timestamp) value).toInstant(); + } + throw new IllegalArgumentException( + String.format( + "Unsupported value class for TIMESTAMP_LTZ: %s", + value.getClass().getName())); + } + public static Object getPickledBytesFromJavaObject(Object obj, TypeInformation dataType) throws IOException { Pickler pickler = new Pickler(); diff --git a/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java b/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java new file mode 100644 index 00000000000000..c5227bf5170f03 --- /dev/null +++ b/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.api.common.python; + +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.types.DataType; +import org.apache.flink.types.Row; + +import net.razorvine.pickle.Unpickler; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.time.Instant; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +class PythonBridgeUtilsTest { + + @Test + void testGetPickledBytesFromRowWithLocalZonedTimestamp() throws IOException { + Instant instant = Instant.parse("2026-01-01T00:00:00.123456789Z"); + Row row = Row.of(instant, new Object[] {instant}); + DataType timestampType = DataTypes.TIMESTAMP_LTZ(9); + + Object serialized = + PythonBridgeUtils.getPickledBytesFromRow( + row, new DataType[] {timestampType, DataTypes.ARRAY(timestampType)}); + + assertThat(serialized).isInstanceOf(List.class); + List fields = (List) serialized; + assertSerializedInstant(fields.get(1), instant); + + Object serializedArray = new Unpickler().loads((byte[]) fields.get(2)); + assertThat(serializedArray).isInstanceOf(List.class); + List array = (List) serializedArray; + assertSerializedInstant(array.get(0), instant); + } + + private static void assertSerializedInstant(Object serialized, Instant expected) + throws IOException { + Object deserialized = new Unpickler().loads((byte[]) serialized); + assertThat(deserialized).isInstanceOf(List.class); + List instant = (List) deserialized; + assertThat(((Number) instant.get(0)).longValue()).isEqualTo(expected.getEpochSecond()); + assertThat(((Number) instant.get(1)).intValue()).isEqualTo(expected.getNano()); + } +} From 7ccce59604c98d9c22241e7085a66c2ed4e43677 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 13 Aug 2026 22:28:33 +0800 Subject: [PATCH 10/11] [FLINK-40190][python] Revert Arrow local-zoned timestamp support Revert the initial Python-side Arrow TIMESTAMP_LTZ handling before replacing it with reader-side support. Generated-by: Codex (GPT-5) --- .../pyflink/dataframe/tests/test_dataframe.py | 20 ------ .../pyflink/table/tests/test_types.py | 16 +---- flink-python/pyflink/table/types.py | 4 +- flink-python/pyflink/table/utils.py | 5 -- .../api/common/python/PythonBridgeUtils.java | 24 ------- .../common/python/PythonBridgeUtilsTest.java | 64 ------------------- 6 files changed, 3 insertions(+), 130 deletions(-) delete mode 100644 flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index f7bda3da2f4674..2d6b62714df605 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -743,26 +743,6 @@ def test_pandas_to_pandas_round_trip(self): finally: self.t_env.get_config().set_local_timezone(original_timezone) - def test_collect_timezone_aware_arrow_timestamps(self): - timestamp = pd.Timestamp("2026-01-01T00:00:00.123Z") - arrow_type = pa.timestamp("ms", tz="UTC") - dataframe = pf.from_arrow( - pa.table( - { - "timestamp": pa.array([timestamp], type=arrow_type), - "timestamps": pa.array( - [[timestamp]], type=pa.list_(arrow_type) - ), - } - ) - ) - expected = timestamp.to_pydatetime() - - self.assertEqual( - dataframe.collect(), - [Row(expected, [expected])], - ) - def test_basic_functionality(self): df = pf.from_dict( { diff --git a/flink-python/pyflink/table/tests/test_types.py b/flink-python/pyflink/table/tests/test_types.py index c7b6fae9a19f62..681884beeeb871 100644 --- a/flink-python/pyflink/table/tests/test_types.py +++ b/flink-python/pyflink/table/tests/test_types.py @@ -38,7 +38,7 @@ RowType, ArrayType, BigIntType, VarCharType, MapType, DataType, _from_java_data_type, ZonedTimestampType, LocalZonedTimestampType, TimestampType, _to_java_data_type, - from_arrow_type, to_arrow_type) + from_arrow_type) from pyflink.testing.test_case_utils import PyFlinkTestCase @@ -147,20 +147,6 @@ def test_timestamp_type_mapping(self): self.assertEqual(data_type.precision, precision) self.assertFalse(data_type._nullable) - def test_nested_local_zoned_timestamp_type_mapping(self): - timestamp_type = pa.timestamp("ms", tz="UTC") - - self.assertEqual( - to_arrow_type(from_arrow_type(pa.list_(timestamp_type))), - pa.list_(pa.timestamp("ms")), - ) - self.assertEqual( - to_arrow_type( - from_arrow_type(pa.struct([pa.field("ts", timestamp_type)])) - ), - pa.struct([pa.field("ts", pa.timestamp("ms"))]), - ) - class TypesTests(PyFlinkTestCase): diff --git a/flink-python/pyflink/table/types.py b/flink-python/pyflink/table/types.py index de5abba8a1cebc..6627be47bd9d14 100644 --- a/flink-python/pyflink/table/types.py +++ b/flink-python/pyflink/table/types.py @@ -2344,13 +2344,13 @@ def to_arrow_type(data_type: DataType): elif isinstance(data_type, MapType): return pa.map_(to_arrow_type(data_type.key_type), to_arrow_type(data_type.value_type)) elif isinstance(data_type, ArrayType): - if isinstance(data_type.element_type, RowType): + if type(data_type.element_type) in [LocalZonedTimestampType, RowType]: raise ValueError("%s is not supported to be used as the element type of ArrayType." % data_type.element_type) return pa.list_(to_arrow_type(data_type.element_type)) elif isinstance(data_type, RowType): for field in data_type: - if isinstance(field.data_type, RowType): + if type(field.data_type) in [LocalZonedTimestampType, RowType]: raise TypeError("%s is not supported to be used as the field type of RowType" % field.data_type) fields = [pa.field(field.name, to_arrow_type(field.data_type), field.data_type._nullable) diff --git a/flink-python/pyflink/table/utils.py b/flink-python/pyflink/table/utils.py index 8533c03a895d94..89f906ca687eb6 100644 --- a/flink-python/pyflink/table/utils.py +++ b/flink-python/pyflink/table/utils.py @@ -134,11 +134,6 @@ def pickled_bytes_to_python_converter(data, field_type: DataType): return field_type.from_sql_type(data) elif isinstance(field_type, TimestampType): return field_type.from_sql_type(int(data.timestamp() * 10**6)) - elif isinstance(field_type, LocalZonedTimestampType): - seconds, nanoseconds = data - return datetime.datetime.fromtimestamp( - seconds, datetime.timezone.utc - ).replace(microsecond=nanoseconds // 1000) elif isinstance(field_type, MapType): key_type = field_type.key_type value_type = field_type.value_type diff --git a/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java b/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java index cbb6b2810c86fd..1205e79d2a8a83 100644 --- a/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java +++ b/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java @@ -36,12 +36,10 @@ import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos; import org.apache.flink.core.memory.DataOutputViewStreamWrapper; import org.apache.flink.streaming.api.typeinfo.python.PickledByteArrayTypeInfo; -import org.apache.flink.table.data.TimestampData; import org.apache.flink.table.types.DataType; import org.apache.flink.table.types.logical.ArrayType; import org.apache.flink.table.types.logical.DateType; import org.apache.flink.table.types.logical.FloatType; -import org.apache.flink.table.types.logical.LocalZonedTimestampType; import org.apache.flink.table.types.logical.LogicalType; import org.apache.flink.table.types.logical.MapType; import org.apache.flink.table.types.logical.RowType; @@ -60,7 +58,6 @@ import java.sql.Date; import java.sql.Time; import java.sql.Timestamp; -import java.time.Instant; import java.time.LocalDate; import java.time.LocalDateTime; import java.time.LocalTime; @@ -196,9 +193,6 @@ private static Object getPickledBytesFromJavaObject(Object obj, LogicalType data } else { return pickler.dumps(obj); } - } else if (dataType instanceof LocalZonedTimestampType) { - Instant instant = toInstant(obj); - return pickler.dumps(Arrays.asList(instant.getEpochSecond(), instant.getNano())); } else if (dataType instanceof RowType) { Row tmpRow = (Row) obj; LogicalType[] tmpRowFieldTypes = @@ -244,24 +238,6 @@ private static Object getPickledBytesFromJavaObject(Object obj, LogicalType data } } - private static Instant toInstant(Object value) { - if (value instanceof Instant) { - return (Instant) value; - } else if (value instanceof Integer) { - return Instant.ofEpochSecond(((Integer) value).longValue()); - } else if (value instanceof Long) { - return Instant.ofEpochMilli((Long) value); - } else if (value instanceof TimestampData) { - return ((TimestampData) value).toInstant(); - } else if (value instanceof Timestamp) { - return ((Timestamp) value).toInstant(); - } - throw new IllegalArgumentException( - String.format( - "Unsupported value class for TIMESTAMP_LTZ: %s", - value.getClass().getName())); - } - public static Object getPickledBytesFromJavaObject(Object obj, TypeInformation dataType) throws IOException { Pickler pickler = new Pickler(); diff --git a/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java b/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java deleted file mode 100644 index c5227bf5170f03..00000000000000 --- a/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java +++ /dev/null @@ -1,64 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.flink.api.common.python; - -import org.apache.flink.table.api.DataTypes; -import org.apache.flink.table.types.DataType; -import org.apache.flink.types.Row; - -import net.razorvine.pickle.Unpickler; -import org.junit.jupiter.api.Test; - -import java.io.IOException; -import java.time.Instant; -import java.util.List; - -import static org.assertj.core.api.Assertions.assertThat; - -class PythonBridgeUtilsTest { - - @Test - void testGetPickledBytesFromRowWithLocalZonedTimestamp() throws IOException { - Instant instant = Instant.parse("2026-01-01T00:00:00.123456789Z"); - Row row = Row.of(instant, new Object[] {instant}); - DataType timestampType = DataTypes.TIMESTAMP_LTZ(9); - - Object serialized = - PythonBridgeUtils.getPickledBytesFromRow( - row, new DataType[] {timestampType, DataTypes.ARRAY(timestampType)}); - - assertThat(serialized).isInstanceOf(List.class); - List fields = (List) serialized; - assertSerializedInstant(fields.get(1), instant); - - Object serializedArray = new Unpickler().loads((byte[]) fields.get(2)); - assertThat(serializedArray).isInstanceOf(List.class); - List array = (List) serializedArray; - assertSerializedInstant(array.get(0), instant); - } - - private static void assertSerializedInstant(Object serialized, Instant expected) - throws IOException { - Object deserialized = new Unpickler().loads((byte[]) serialized); - assertThat(deserialized).isInstanceOf(List.class); - List instant = (List) deserialized; - assertThat(((Number) instant.get(0)).longValue()).isEqualTo(expected.getEpochSecond()); - assertThat(((Number) instant.get(1)).intValue()).isEqualTo(expected.getNano()); - } -} From 790290f4ab0c02b0dbda9664e8426db7f927b772 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 13 Aug 2026 23:22:42 +0800 Subject: [PATCH 11/11] [FLINK-40190][python] Support timezone-aware Arrow timestamp reading Preserve timezone-aware Arrow timestamp vectors through the Java reader and serialize TIMESTAMP_LTZ values as instants when collecting results. Generated-by: Codex (GPT-5) --- .../pyflink/dataframe/tests/test_dataframe.py | 25 +++++ .../pyflink/table/table_environment.py | 34 +++++-- flink-python/pyflink/table/utils.py | 5 + .../api/common/python/PythonBridgeUtils.java | 24 +++++ .../flink/table/runtime/arrow/ArrowUtils.java | 12 ++- .../vectors/ArrowTimestampColumnVector.java | 36 +++---- .../common/python/PythonBridgeUtilsTest.java | 64 ++++++++++++ .../table/runtime/arrow/ArrowUtilsTest.java | 99 +++++++++++++++++++ 8 files changed, 270 insertions(+), 29 deletions(-) create mode 100644 flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index 2d6b62714df605..342522119870f3 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -743,6 +743,31 @@ def test_pandas_to_pandas_round_trip(self): finally: self.t_env.get_config().set_local_timezone(original_timezone) + def test_collect_timezone_aware_arrow_timestamps(self): + first_fold = pd.Timestamp("2026-11-01T05:30:00.123Z") + second_fold = pd.Timestamp("2026-11-01T06:30:00.123Z") + arrow_type = pa.timestamp("ms", tz="America/New_York") + dataframe = pf.from_arrow( + pa.table( + { + "timestamp": pa.array( + [first_fold, second_fold], type=arrow_type + ), + "timestamps": pa.array( + [[first_fold], [second_fold]], type=pa.list_(arrow_type) + ), + } + ) + ) + + self.assertEqual( + dataframe.collect(), + [ + Row(first_fold.to_pydatetime(), [first_fold.to_pydatetime()]), + Row(second_fold.to_pydatetime(), [second_fold.to_pydatetime()]), + ], + ) + def test_basic_functionality(self): df = pf.from_dict( { diff --git a/flink-python/pyflink/table/table_environment.py b/flink-python/pyflink/table/table_environment.py index 917c768d88f839..0d1fd20e0cb934 100644 --- a/flink-python/pyflink/table/table_environment.py +++ b/flink-python/pyflink/table/table_environment.py @@ -45,9 +45,9 @@ from pyflink.table.table_config import TableConfig from pyflink.table.table_descriptor import TableDescriptor from pyflink.table.table_result import TableResult -from pyflink.table.types import _create_type_verifier, RowType, DataType, \ - _infer_schema_from_data, _create_converter, from_arrow_type, RowField, create_arrow_schema, \ - _to_java_data_type +from pyflink.table.types import _create_type_verifier, RowType, DataType, TimestampType, \ + LocalZonedTimestampType, _infer_schema_from_data, _create_converter, from_arrow_type, \ + RowField, create_arrow_schema, to_arrow_type, _to_java_data_type from pyflink.table.udf import UserDefinedFunctionWrapper, AggregateFunction, udaf, \ udtaf, TableAggregateFunction from pyflink.table.utils import to_expression_jarray @@ -1515,10 +1515,32 @@ def _from_arrow( if splits_num <= 0: raise ValueError("splits_num must be greater than 0") - arrow_schema = create_arrow_schema(row_type.field_names(), row_type.field_types()) + field_names = row_type.field_names() + field_types = row_type.field_types() try: - compatible_table = table.rename_columns(row_type.field_names()).cast( - arrow_schema, safe=False) + compatible_table = table.rename_columns(field_names) + for index, (field, data_type) in enumerate(zip(table.schema, field_types)): + if not isinstance( + data_type, (TimestampType, LocalZonedTimestampType) + ): + continue + target_type = to_arrow_type(data_type) + if isinstance(data_type, LocalZonedTimestampType): + timezone = field.type.tz if pa.types.is_timestamp(field.type) else None + target_type = pa.timestamp(target_type.unit, tz=timezone) + if field.type == target_type: + continue + target_field = pa.field( + field_names[index], + target_type, + field.nullable, + field.metadata, + ) + compatible_table = compatible_table.set_column( + index, + target_field, + compatible_table.column(index).cast(target_type, safe=False), + ) except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError, ValueError) as e: raise TypeError( f"Could not convert pyarrow.Table to the inferred Flink schema: {row_type}" diff --git a/flink-python/pyflink/table/utils.py b/flink-python/pyflink/table/utils.py index 89f906ca687eb6..8533c03a895d94 100644 --- a/flink-python/pyflink/table/utils.py +++ b/flink-python/pyflink/table/utils.py @@ -134,6 +134,11 @@ def pickled_bytes_to_python_converter(data, field_type: DataType): return field_type.from_sql_type(data) elif isinstance(field_type, TimestampType): return field_type.from_sql_type(int(data.timestamp() * 10**6)) + elif isinstance(field_type, LocalZonedTimestampType): + seconds, nanoseconds = data + return datetime.datetime.fromtimestamp( + seconds, datetime.timezone.utc + ).replace(microsecond=nanoseconds // 1000) elif isinstance(field_type, MapType): key_type = field_type.key_type value_type = field_type.value_type diff --git a/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java b/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java index 1205e79d2a8a83..cbb6b2810c86fd 100644 --- a/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java +++ b/flink-python/src/main/java/org/apache/flink/api/common/python/PythonBridgeUtils.java @@ -36,10 +36,12 @@ import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos; import org.apache.flink.core.memory.DataOutputViewStreamWrapper; import org.apache.flink.streaming.api.typeinfo.python.PickledByteArrayTypeInfo; +import org.apache.flink.table.data.TimestampData; import org.apache.flink.table.types.DataType; import org.apache.flink.table.types.logical.ArrayType; import org.apache.flink.table.types.logical.DateType; import org.apache.flink.table.types.logical.FloatType; +import org.apache.flink.table.types.logical.LocalZonedTimestampType; import org.apache.flink.table.types.logical.LogicalType; import org.apache.flink.table.types.logical.MapType; import org.apache.flink.table.types.logical.RowType; @@ -58,6 +60,7 @@ import java.sql.Date; import java.sql.Time; import java.sql.Timestamp; +import java.time.Instant; import java.time.LocalDate; import java.time.LocalDateTime; import java.time.LocalTime; @@ -193,6 +196,9 @@ private static Object getPickledBytesFromJavaObject(Object obj, LogicalType data } else { return pickler.dumps(obj); } + } else if (dataType instanceof LocalZonedTimestampType) { + Instant instant = toInstant(obj); + return pickler.dumps(Arrays.asList(instant.getEpochSecond(), instant.getNano())); } else if (dataType instanceof RowType) { Row tmpRow = (Row) obj; LogicalType[] tmpRowFieldTypes = @@ -238,6 +244,24 @@ private static Object getPickledBytesFromJavaObject(Object obj, LogicalType data } } + private static Instant toInstant(Object value) { + if (value instanceof Instant) { + return (Instant) value; + } else if (value instanceof Integer) { + return Instant.ofEpochSecond(((Integer) value).longValue()); + } else if (value instanceof Long) { + return Instant.ofEpochMilli((Long) value); + } else if (value instanceof TimestampData) { + return ((TimestampData) value).toInstant(); + } else if (value instanceof Timestamp) { + return ((Timestamp) value).toInstant(); + } + throw new IllegalArgumentException( + String.format( + "Unsupported value class for TIMESTAMP_LTZ: %s", + value.getClass().getName())); + } + public static Object getPickledBytesFromJavaObject(Object obj, TypeInformation dataType) throws IOException { Pickler pickler = new Pickler(); diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java index d85d1b082b910b..31460a1fc59412 100644 --- a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java +++ b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java @@ -445,9 +445,15 @@ public static ColumnVector createColumnVector(ValueVector vector, LogicalType fi || vector instanceof TimeMicroVector || vector instanceof TimeNanoVector) { return new ArrowTimeColumnVector(vector); - } else if (vector instanceof TimeStampVector - && ((ArrowType.Timestamp) vector.getField().getType()).getTimezone() == null) { - return new ArrowTimestampColumnVector(vector); + } else if (vector instanceof TimeStampVector) { + String timezone = ((ArrowType.Timestamp) vector.getField().getType()).getTimezone(); + if (timezone != null && !(fieldType instanceof LocalZonedTimestampType)) { + throw new UnsupportedOperationException( + String.format( + "Arrow timestamp with timezone '%s' cannot be read as %s.", + timezone, fieldType)); + } + return new ArrowTimestampColumnVector((TimeStampVector) vector); } else if (vector instanceof MapVector) { MapVector mapVector = (MapVector) vector; LogicalType keyType = ((MapType) fieldType).getKeyType(); diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowTimestampColumnVector.java b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowTimestampColumnVector.java index 297f21d3978192..7da2ca485bebc9 100644 --- a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowTimestampColumnVector.java +++ b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowTimestampColumnVector.java @@ -23,12 +23,8 @@ import org.apache.flink.table.data.columnar.vector.TimestampColumnVector; import org.apache.flink.util.Preconditions; -import org.apache.arrow.vector.TimeStampMicroVector; -import org.apache.arrow.vector.TimeStampMilliVector; -import org.apache.arrow.vector.TimeStampNanoVector; -import org.apache.arrow.vector.TimeStampSecVector; import org.apache.arrow.vector.TimeStampVector; -import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; /** Arrow column vector for Timestamp. */ @@ -36,28 +32,28 @@ public final class ArrowTimestampColumnVector implements TimestampColumnVector { /** Container which is used to store the sequence of timestamp values of a column to read. */ - private final ValueVector valueVector; + private final TimeStampVector valueVector; - public ArrowTimestampColumnVector(ValueVector valueVector) { + private final TimeUnit timeUnit; + + public ArrowTimestampColumnVector(TimeStampVector valueVector) { this.valueVector = Preconditions.checkNotNull(valueVector); - Preconditions.checkState( - valueVector instanceof TimeStampVector - && ((ArrowType.Timestamp) valueVector.getField().getType()).getTimezone() - == null); + this.timeUnit = ((ArrowType.Timestamp) valueVector.getField().getType()).getUnit(); } @Override public TimestampData getTimestamp(int i, int precision) { - if (valueVector instanceof TimeStampSecVector) { - return TimestampData.fromEpochMillis(((TimeStampSecVector) valueVector).get(i) * 1000); - } else if (valueVector instanceof TimeStampMilliVector) { - return TimestampData.fromEpochMillis(((TimeStampMilliVector) valueVector).get(i)); - } else if (valueVector instanceof TimeStampMicroVector) { - long micros = ((TimeStampMicroVector) valueVector).get(i); - return TimestampData.fromEpochMillis(micros / 1000, (int) (micros % 1000) * 1000); + long value = valueVector.get(i); + if (timeUnit == TimeUnit.SECOND) { + return TimestampData.fromEpochMillis(value * 1000); + } else if (timeUnit == TimeUnit.MILLISECOND) { + return TimestampData.fromEpochMillis(value); + } else if (timeUnit == TimeUnit.MICROSECOND) { + return TimestampData.fromEpochMillis( + Math.floorDiv(value, 1000), (int) Math.floorMod(value, 1000) * 1000); } else { - long nanos = ((TimeStampNanoVector) valueVector).get(i); - return TimestampData.fromEpochMillis(nanos / 1_000_000, (int) (nanos % 1_000_000)); + return TimestampData.fromEpochMillis( + Math.floorDiv(value, 1_000_000), (int) Math.floorMod(value, 1_000_000)); } } diff --git a/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java b/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java new file mode 100644 index 00000000000000..b410563dd316e9 --- /dev/null +++ b/flink-python/src/test/java/org/apache/flink/api/common/python/PythonBridgeUtilsTest.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.api.common.python; + +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.types.DataType; +import org.apache.flink.types.Row; + +import net.razorvine.pickle.Unpickler; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.time.Instant; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +class PythonBridgeUtilsTest { + + @Test + void testGetPickledBytesFromRowWithLocalZonedTimestamp() throws IOException { + Instant instant = Instant.parse("2026-11-01T05:30:00.123456789Z"); + Row row = Row.of(instant, new Object[] {instant}); + DataType timestampType = DataTypes.TIMESTAMP_LTZ(9); + + Object serialized = + PythonBridgeUtils.getPickledBytesFromRow( + row, new DataType[] {timestampType, DataTypes.ARRAY(timestampType)}); + + assertThat(serialized).isInstanceOf(List.class); + List fields = (List) serialized; + assertSerializedInstant(fields.get(1), instant); + + Object serializedArray = new Unpickler().loads((byte[]) fields.get(2)); + assertThat(serializedArray).isInstanceOf(List.class); + List array = (List) serializedArray; + assertSerializedInstant(array.get(0), instant); + } + + private static void assertSerializedInstant(Object serialized, Instant expected) + throws IOException { + Object deserialized = new Unpickler().loads((byte[]) serialized); + assertThat(deserialized).isInstanceOf(List.class); + List instant = (List) deserialized; + assertThat(((Number) instant.get(0)).longValue()).isEqualTo(expected.getEpochSecond()); + assertThat(((Number) instant.get(1)).intValue()).isEqualTo(expected.getNano()); + } +} diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowUtilsTest.java b/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowUtilsTest.java index e7b3a81d1f4aeb..025b28675016a4 100644 --- a/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowUtilsTest.java +++ b/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowUtilsTest.java @@ -21,6 +21,7 @@ import org.apache.flink.api.java.tuple.Tuple5; import org.apache.flink.table.data.GenericRowData; import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.TimestampData; import org.apache.flink.table.data.columnar.vector.ColumnVector; import org.apache.flink.table.runtime.arrow.vectors.ArrowArrayColumnVector; import org.apache.flink.table.runtime.arrow.vectors.ArrowBigIntColumnVector; @@ -74,6 +75,8 @@ import org.apache.flink.shaded.guava33.com.google.common.collect.Lists; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.TimeStampMilliTZVector; +import org.apache.arrow.vector.TimeStampVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowStreamWriter; import org.apache.arrow.vector.types.DateUnit; @@ -81,6 +84,7 @@ import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -89,11 +93,14 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; +import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link ArrowUtils}. */ class ArrowUtilsTest { @@ -352,6 +359,98 @@ void testCreateArrowReader() { } } + @Test + void testCreateArrowReaderForTimezoneAwareLocalZonedTimestamps() { + ArrowType.Timestamp timestampType = + new ArrowType.Timestamp(TimeUnit.MILLISECOND, "America/New_York"); + Field timestampField = new Field("timestamp", FieldType.nullable(timestampType), null); + Field timestampsField = + new Field( + "timestamps", + FieldType.nullable(ArrowType.List.INSTANCE), + Collections.singletonList( + new Field("element", FieldType.nullable(timestampType), null))); + RowType timestampRowType = + RowType.of( + new LocalZonedTimestampType(3), + new ArrayType(new LocalZonedTimestampType(3))); + + try (VectorSchemaRoot root = + VectorSchemaRoot.create( + new Schema(Arrays.asList(timestampField, timestampsField)), allocator)) { + TimeStampMilliTZVector timestampVector = + (TimeStampMilliTZVector) root.getVector("timestamp"); + long firstFold = Instant.parse("2026-11-01T05:30:00.123Z").toEpochMilli(); + long secondFold = Instant.parse("2026-11-01T06:30:00.123Z").toEpochMilli(); + timestampVector.setSafe(0, firstFold); + timestampVector.setSafe(1, secondFold); + timestampVector.setValueCount(2); + root.setRowCount(2); + + ArrowReader reader = ArrowUtils.createArrowReader(root, timestampRowType); + + assertThat(reader.getColumnVectors()[0]).isInstanceOf(ArrowTimestampColumnVector.class); + assertThat(reader.getColumnVectors()[1]).isInstanceOf(ArrowArrayColumnVector.class); + assertThat(reader.read(0).getTimestamp(0, 3).getMillisecond()).isEqualTo(firstFold); + assertThat(reader.read(1).getTimestamp(0, 3).getMillisecond()).isEqualTo(secondFold); + } + } + + @Test + void testCreateArrowReaderForPreEpochSubMillisecondTimestamps() { + List fields = + Arrays.asList( + new Field( + "micros", + FieldType.nullable( + new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)), + null), + new Field( + "nanos", + FieldType.nullable( + new ArrowType.Timestamp(TimeUnit.NANOSECOND, null)), + null)); + RowType timestampRowType = RowType.of(new TimestampType(6), new TimestampType(9)); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(fields), allocator)) { + TimeStampVector micros = (TimeStampVector) root.getVector("micros"); + TimeStampVector nanos = (TimeStampVector) root.getVector("nanos"); + micros.setSafe(0, -1); + nanos.setSafe(0, -1); + micros.setValueCount(1); + nanos.setValueCount(1); + root.setRowCount(1); + + RowData row = ArrowUtils.createArrowReader(root, timestampRowType).read(0); + + assertThat(row.getTimestamp(0, 6)) + .isEqualTo(TimestampData.fromEpochMillis(-1, 999_000)); + assertThat(row.getTimestamp(1, 9)) + .isEqualTo(TimestampData.fromEpochMillis(-1, 999_999)); + } + } + + @Test + void testCreateArrowReaderRejectsTimezoneAwareTimestamp() { + Field timestampField = + new Field( + "timestamp", + FieldType.nullable( + new ArrowType.Timestamp(TimeUnit.MILLISECOND, "America/New_York")), + null); + + try (VectorSchemaRoot root = + VectorSchemaRoot.create( + new Schema(Collections.singletonList(timestampField)), allocator)) { + assertThatThrownBy( + () -> + ArrowUtils.createArrowReader( + root, RowType.of(new TimestampType(3)))) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("cannot be read as TIMESTAMP(3)"); + } + } + @Test void testCreateArrowWriter() { VectorSchemaRoot root =