From e55ecd52d5dafd152436f0750792812682ca94c9 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 14:34:40 +0800 Subject: [PATCH 1/5] [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 51692dab6bcf4..3652224e5a0fc 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 2caa7503eba13..4bc7118a837cf 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 8ad43bcbbe256..b4c1e32c90da8 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 9a78fa65ee322..126bc1e39685d 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 66c2020ee4db7..79fd08691a79b 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 f08ede67ab774..d064a02e8c2f2 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 98714e9fbb6a4..b918bf2b6c515 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 f0e4bedba174e..b80b98fe45bde 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 316bb44b55e23..7e59165f911cc 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 d8b9dcdf772b6..d85d1b082b910 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 01dfab186dddc..3681d0664ae1a 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 6a73b3cbd0771df7bdd096c4587d36f7ca3d5725 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 15:41:34 +0800 Subject: [PATCH 2/5] [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 126bc1e39685d..27580153251f8 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 d064a02e8c2f2..098161bb23b5f 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 b918bf2b6c515..ae341b66afcc5 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 b80b98fe45bde..3e885cc666a4f 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 9cc0f8ccdf677..0c70e2949605a 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 03c6eb81bcb7b95d5410dac32c7434d1d0907e14 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 16:27:59 +0800 Subject: [PATCH 3/5] [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 27580153251f8..1c076c982f84d 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 654b9ca14b3a33fcf826c8ecfd63de66c36b6e3b Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 17:40:04 +0800 Subject: [PATCH 4/5] [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 1c076c982f84d..6c47250b7ec29 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 098161bb23b5f..1d9739f74b857 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 66e5281bc3a848a17a4360149914e83f2e2fe208 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 20:46:26 +0800 Subject: [PATCH 5/5] [FLINK-40190][python] Refine DataFrame conversion docs and coverage Clarify that to_table does not execute a job, remove redundant result-page prose, and exercise from_pandas in the existing pandas round-trip integration smoke test. Generated-by: Codex (GPT-5) --- .../docs/reference/pyflink.dataframe/dataframe.rst | 4 ---- flink-python/pyflink/dataframe/dataframe.py | 2 ++ flink-python/pyflink/dataframe/tests/test_dataframe.py | 10 +++++----- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst index 4bc7118a837cf..4ab37c87b6b01 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 79fd08691a79b..3e58f5c8319e2 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 ae341b66afcc5..12eb8a0dd9382 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -685,17 +685,17 @@ 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": [1, 2], + "ts": pd.Series([timestamp, None], dtype="datetime64[ms]"), } ) result = ( - pf.from_arrow(arrow_table) + pf.from_pandas(pdf) .with_column("id_plus_one", pf.col("id") + 1) .to_pandas() )