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..4ab37c87b6b010 100644 --- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst +++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst @@ -63,6 +63,8 @@ Results :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..6c47250b7ec293 100644 --- a/flink-python/pyflink/dataframe/convert.py +++ b/flink-python/pyflink/dataframe/convert.py @@ -16,12 +16,14 @@ # limitations under the License. ################################################################################ +import builtins from enum import Enum from typing import ( Any, Collection, List, Mapping, + NamedTuple, Optional, Sequence, Tuple, @@ -31,11 +33,37 @@ 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) +_BIGINT_MIN = -(1 << 63) +_BIGINT_MAX = (1 << 63) - 1 + + +class _WatermarkSpec(NamedTuple): + column: str + expression: str class _RecordType(Enum): @@ -109,7 +137,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 +148,239 @@ 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]: + 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 _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) + converter = _create_converter(row_type) + verify_row = _create_type_verifier(row_type) + sql_rows = [] + for row in rows: + row = converter(row) + verify_row(row) + sql_rows.append(row_type.to_sql_type(row)) + + table = get_or_create_table_environment()._from_elements( + sql_rows, row_type, table_schema + ) + 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: + """ + 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"]) + >>> 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 + """ + import pandas as pd + + if not isinstance(pdf, pd.DataFrame): + raise TypeError( + f"data must be a pandas.DataFrame, but was {type(pdf).__name__}" + ) + + import pyarrow as pa + + return from_arrow( + pa.Table.from_pandas(pdf, preserve_index=False), + schema=schema, + watermark=watermark, + ) + + +@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) + >>> events = pf.from_arrow( + ... pa.table({"ts": pa.array([0], type=pa.timestamp("ms"))}), + ... watermark=("ts", "ts - INTERVAL '5' SECOND"), + ... ) + + .. 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_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_spec + ) + 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 +394,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 +425,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 +439,7 @@ def from_records( ) if not data: raise ValueError("data must not be empty") + watermark_spec = _parse_watermark(watermark) first_record = data[0] try: @@ -204,14 +472,14 @@ 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) - ) + return _infer_schema_and_create_dataframe(rows, schema, watermark_spec) @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 +487,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 +514,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) if schema is None: schema = list(data.keys()) _validate_schema(schema) @@ -263,8 +537,60 @@ 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) - ) + return _infer_schema_and_create_dataframe(rows, schema, watermark_spec) + + +@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 or the range contains values outside the signed + ``BIGINT`` bounds. + + 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 + 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 values] + table = get_or_create_table_environment()._from_elements(sql_rows, row_type) + return DataFrame(table) diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 66c2020ee4db74..3e58f5c8319e29 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,43 @@ 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. + + This method does not trigger job execution. + + :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..1d9739f74b8578 100644 --- a/flink-python/pyflink/dataframe/tests/test_convert.py +++ b/flink-python/pyflink/dataframe/tests/test_convert.py @@ -17,9 +17,15 @@ ################################################################################ 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 +import pyflink.dataframe.convert as dataframe_convert +from pyflink.table.types import BigIntType, RowType class _Point(NamedTuple): @@ -224,5 +230,115 @@ 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_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"), + (("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), []), + ((2**63 - 1, 2**63), [(2**63 - 1,)]), + ((-(2**63), -(2**63) + 1), [(-(2**63),)]), + ] + 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_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"), + ((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..12eb8a0dd9382e 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,177 @@ 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_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"), + ): + 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()], + ) + + 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): + 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"), + ), + TimestampType, + ), + ] + 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 +685,27 @@ def test_from_records(self): [Row(1, "Alice"), Row(2, "Bob")], ) + def test_pandas_to_pandas_round_trip(self): + timestamp = datetime(2026, 1, 1, 0, 0, 0, 123000) + pdf = pd.DataFrame( + { + "id": [1, 2], + "ts": pd.Series([timestamp, None], dtype="datetime64[ms]"), + } + ) + + result = ( + pf.from_pandas(pdf) + .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..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): """ @@ -1469,11 +1483,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 +1502,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 +1513,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, + splits_num: int = 1) -> 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: + _serialize_arrow_table(compatible_table, temp_file, splits_num) + + 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/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 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()); }