This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new ffd8d8836df [FLINK-40190][python] Introduce PyFlink DataFrame creation
and conversion APIs (#28934)
ffd8d8836df is described below
commit ffd8d8836df6da3c6ddd92a4ba2d291017588a11
Author: Liu Liu <[email protected]>
AuthorDate: Tue Aug 18 14:31:18 2026 +0800
[FLINK-40190][python] Introduce PyFlink DataFrame creation and conversion
APIs (#28934)
---
.../docs/reference/pyflink.dataframe/creation.rst | 14 +
.../docs/reference/pyflink.dataframe/dataframe.rst | 2 +
flink-python/pyflink/dataframe/__init__.py | 13 +-
flink-python/pyflink/dataframe/convert.py | 384 ++++++++++++++++++++-
flink-python/pyflink/dataframe/dataframe.py | 45 ++-
.../pyflink/dataframe/tests/test_convert.py | 159 +++++++++
.../pyflink/dataframe/tests/test_dataframe.py | 267 ++++++++++++++
flink-python/pyflink/table/table_environment.py | 161 ++++++++-
.../pyflink/table/tests/test_pandas_conversion.py | 21 ++
.../table/tests/test_table_environment_api.py | 117 +++++++
flink-python/pyflink/table/tests/test_types.py | 34 +-
flink-python/pyflink/table/types.py | 19 +-
.../flink/table/runtime/arrow/ArrowUtils.java | 118 ++++++-
.../arrow/vectors/ArrowTimestampColumnVector.java | 9 +-
.../flink/table/utils/python/PythonTableUtils.java | 18 +-
.../flink/table/runtime/arrow/ArrowUtilsTest.java | 159 +++++++++
16 files changed, 1507 insertions(+), 33 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/creation.rst
b/flink-python/docs/reference/pyflink.dataframe/creation.rst
index 51692dab6bc..9a4b577a12e 100644
--- a/flink-python/docs/reference/pyflink.dataframe/creation.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/creation.rst
@@ -30,6 +30,16 @@ Example::
... {"id": 2, "name": "Bob"},
... ])
>>> users = pf.from_dict({"id": [1, 2], "name": ["Alice", "Bob"]})
+ >>> import pandas as pd
+ >>> import pyarrow as pa
+ >>> pandas_users = pf.from_pandas(
+ ... pd.DataFrame({"id": [1, 2], "name": ["Alice", "Bob"]})
+ ... )
+ >>> arrow_users = pf.from_arrow(
+ ... pa.table({"id": [1, 2], "name": ["Alice", "Bob"]})
+ ... )
+ >>> table_users = pf.from_table(users.to_table())
+ >>> identifiers = pf.range(5)
.. currentmodule:: pyflink.dataframe
@@ -38,3 +48,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 19ebb0ad9c0..29c09bee435 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -76,6 +76,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 4efe9fa1c8c..88e7ca2aad5 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -38,7 +38,14 @@ Example::
<Row(1, 'Alice', 31)>
"""
-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,
@@ -54,8 +61,12 @@ __all__ = [
"DataType",
"col",
"lit",
+ "from_arrow",
"from_dict",
+ "from_pandas",
"from_records",
+ "from_table",
+ "range",
"read_generic",
"set_table_environment",
"get_table_environment",
diff --git a/flink-python/pyflink/dataframe/convert.py
b/flink-python/pyflink/dataframe/convert.py
index 9a78fa65ee3..a7265466dd1 100644
--- a/flink-python/pyflink/dataframe/convert.py
+++ b/flink-python/pyflink/dataframe/convert.py
@@ -16,26 +16,95 @@
# limitations under the License.
################################################################################
+import builtins
from enum import Enum
from typing import (
Any,
Collection,
List,
Mapping,
+ NamedTuple,
Optional,
Sequence,
+ TYPE_CHECKING,
Tuple,
Union,
cast,
)
+if TYPE_CHECKING:
+ import pandas
+ import pyarrow
+
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,
+ _has_nulltype,
+ _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
+
+ @classmethod
+ def parse(cls, watermark: Optional[Tuple[str, str]]) ->
Optional["_WatermarkSpec"]:
+ if watermark is None:
+ return None
+ if not isinstance(watermark, tuple) or len(watermark) != 2:
+ raise TypeError("watermark must be a tuple of (column,
expression)")
+ if any(not isinstance(value, str) or not value.strip() for value in
watermark):
+ raise TypeError(
+ "watermark column and expression must be non-empty strings"
+ )
+ return cls(*watermark)
+
+ def normalize_row_type(self, row_type: RowType) -> RowType:
+ matching_fields = [
+ field for field in row_type.fields if field.name == self.column
+ ]
+ if not matching_fields:
+ raise ValueError(
+ f"watermark column {self.column!r} is not present in data"
+ )
+
+ watermark_type = matching_fields[0].data_type
+ if not isinstance(watermark_type, (TimestampType,
LocalZonedTimestampType)):
+ raise ValueError(
+ f"watermark column {self.column!r} must have a timestamp type"
+ )
+
+ fields = []
+ for field in row_type.fields:
+ data_type = field.data_type
+ if field.name == self.column and data_type.precision != 3:
+ data_type = type(data_type)(3, data_type._nullable)
+ fields.append(RowField(field.name, data_type, field.description))
+ return RowType(fields, row_type._nullable)
class _RecordType(Enum):
@@ -120,10 +189,233 @@ 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 _normalize_pandas_column_name(name: Any) -> str:
+ """Normalize a pandas column label to the corresponding Arrow field
name."""
+ if isinstance(name, str):
+ return name
+ if isinstance(name, bytes):
+ return name.decode("utf-8")
+ if isinstance(name, tuple):
+ return str(tuple(_normalize_pandas_column_name(value) for value in
name))
+ return str(name)
+
+
+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)
+ if watermark is None:
+ table_schema = None
+ else:
+ row_type = watermark.normalize_row_type(row_type)
+ table_schema = (
+ Schema.new_builder()
+ .from_row_data_type(row_type)
+ .watermark(*watermark)
+ .build()
+ )
+ converter = _create_converter(row_type)
+ verify_row = _create_type_verifier(row_type)
+ sql_rows = []
+ 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)
+
+
+@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: "pandas.DataFrame",
+ 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.
+ Timezone-aware timestamps are represented as ``TIMESTAMP`` in the
TableEnvironment's
+ configured local timezone; timezone-naive timestamps are unchanged.
+
+ ``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)``.
+
+ :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
+
+ input_names = [_normalize_pandas_column_name(name) for name in pdf.columns]
+ arrow_pdf = pdf.copy(deep=False)
+ arrow_pdf.columns = [
+ f"__pyflink_dataframe_column_{index}"
+ for index in builtins.range(len(input_names))
+ ]
+ return from_arrow(
+ pa.Table.from_pandas(arrow_pdf, preserve_index=False),
+ schema=input_names if schema is None else schema,
+ watermark=watermark,
+ )
+
+
+@PublicEvolving()
+def from_arrow(
+ table: "pyarrow.Table",
+ 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)``.
+
+ :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 = _WatermarkSpec.parse(watermark)
+ names = _resolve_column_names(table.column_names, schema)
+ row_type = RowType(
+ [
+ RowField(
+ name,
+ from_arrow_type(field.type, field.nullable),
+ )
+ for name, field in zip(names, table.schema)
+ ]
+ )
+ null_field_names = [
+ field.name for field in row_type.fields if
_has_nulltype(field.data_type)
+ ]
+ if null_field_names:
+ columns = ", ".join(repr(name) for name in null_field_names)
+ raise TypeError(
+ f"Cannot infer Flink data types for columns with Arrow null types:
{columns}. "
+ "Use explicit pandas or Arrow dtypes for these columns."
+ )
+ if watermark_spec is None:
+ table_schema = None
+ else:
+ row_type = watermark_spec.normalize_row_type(row_type)
+ table_schema = (
+ Schema.new_builder()
+ .from_row_data_type(row_type)
+ .watermark(*watermark_spec)
+ .build()
+ )
+ result = get_or_create_table_environment()._from_arrow(
+ table, 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 +429,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 +460,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 +474,7 @@ def from_records(
)
if not data:
raise ValueError("data must not be empty")
+ watermark_spec = _WatermarkSpec.parse(watermark)
first_record = data[0]
try:
@@ -204,14 +507,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 +522,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.
@@ -229,11 +537,16 @@ def from_dict(
Example::
+ >>> from datetime import datetime
>>> import pyflink.dataframe as pf
>>> users = pf.from_dict(
... {"name": ["Alice", "Bob"], "id": [1, 2]},
... schema=["id", "name"],
... )
+ >>> events = pf.from_dict(
+ ... {"id": [1], "ts": [datetime(2026, 1, 1)]},
+ ... watermark=("ts", "ts - INTERVAL '5' SECOND"),
+ ... )
.. versionadded:: 2.4.0
"""
@@ -241,6 +554,7 @@ def from_dict(
raise TypeError("data must be a mapping")
if not data:
raise ValueError("data must not be empty")
+ watermark_spec = _WatermarkSpec.parse(watermark)
if schema is None:
schema = list(data.keys())
_validate_schema(schema)
@@ -263,8 +577,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()
+ [<Row(1)>, <Row(3)>, <Row(5)>]
+
+ .. 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 ba12d06399f..35f03af35f8 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, Dict, List, Optional, Tuple, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple,
Union, overload
+
+if TYPE_CHECKING:
+ import pandas
from pyflink.common import Row
from pyflink.dataframe.datatype import DataType
@@ -434,6 +437,46 @@ class DataFrame:
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()
+
# ======================== I/O ========================
@PublicEvolving()
diff --git a/flink-python/pyflink/dataframe/tests/test_convert.py
b/flink-python/pyflink/dataframe/tests/test_convert.py
index f08ede67ab7..a204f7a6611 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,158 @@ class FromDictTests(unittest.TestCase):
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._WatermarkSpec.parse(
+ ("ts", "ts - INTERVAL '5' SECOND")
+ )
+
+ self.assertEqual(watermark.column, "ts")
+ self.assertEqual(watermark.expression, "ts - INTERVAL '5' SECOND")
+
+ def test_watermark_spec_unpacks_column_before_expression(self):
+ watermark = dataframe_convert._WatermarkSpec(
+ "ts", "ts - INTERVAL '5' SECOND"
+ )
+
+ self.assertEqual(tuple(watermark), ("ts", "ts - INTERVAL '5' SECOND"))
+
+ def test_rejects_invalid_watermarks(self):
+ invalid_watermarks = [
+ ("ts", "watermark must be a tuple"),
+ (("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_pandas_rejects_duplicate_columns_without_schema(self):
+ pdf = pd.DataFrame([[1, 2]], columns=["value", "value"])
+
+ with self.assertRaisesRegex(ValueError, "schema field names must be
unique"):
+ pf.from_pandas(pdf)
+
+ def test_rejects_columnar_fields_containing_null_type(self):
+ inputs = [
+ (
+ lambda: pf.from_pandas(pd.DataFrame({"value": [None]})),
+ "columns with Arrow null types: 'value'",
+ ),
+ (
+ lambda: pf.from_arrow(
+ pa.table({"left": pa.nulls(1), "right": pa.nulls(1)})
+ ),
+ "columns with Arrow null types: 'left', 'right'",
+ ),
+ (
+ lambda: pf.from_arrow(
+ pa.table(
+ {
+ "value": pa.array(
+ [[None]], type=pa.list_(pa.null())
+ )
+ }
+ )
+ ),
+ "columns with Arrow null types: 'value'",
+ ),
+ ]
+ for creator, message in inputs:
+ with self.subTest(creator=creator):
+ with self.assertRaisesRegex(TypeError, message):
+ creator()
+
+ def test_rejects_invalid_table_and_columnar_inputs(self):
+ invalid_inputs = [
+ (pf.from_table, object(), "pyflink.table.Table"),
+ (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 8f5b7620cab..a3d095f2ab8 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -17,8 +17,11 @@
################################################################################
import unittest
+from datetime import datetime, timezone
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 +31,7 @@ from pyflink.table import (
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 +81,17 @@ class _Table:
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 +110,24 @@ class DataFrameCollectTests(unittest.TestCase):
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 +225,202 @@ class DataFrameCreationTests(PyFlinkDataFrameUTTestCase):
[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"])
+
+ duplicate_pdf = pd.DataFrame(
+ [[1, "Alice"], [2, "Bob"]], columns=["value", "value"]
+ )
+ dataframe = pf.from_pandas(duplicate_pdf, schema=["id", "name"])
+ self.assert_dataframe_schema(
+ dataframe,
+ ["id", "name"],
+ [TableDataTypes.BIGINT(), TableDataTypes.STRING()],
+ )
+
+ def test_from_pandas_normalizes_inferred_column_names(self):
+ dataframe = pf.from_pandas(pd.DataFrame([[1, 2]]))
+
+ self.assert_dataframe_schema(
+ dataframe,
+ ["0", "1"],
+ [TableDataTypes.BIGINT(), TableDataTypes.BIGINT()],
+ )
+
+ 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_pandas_schema_inference(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(
+ dataframe_schema.get_column_names(),
empty_schema.get_column_names()
+ )
+ self.assertEqual(
+ dataframe_schema.get_column_data_types(),
+ empty_schema.get_column_data_types(),
+ )
+
+ def test_timezone_aware_creation_supports_java_timezone_ids(self):
+ original_timezone = self.t_env.get_config().get_local_timezone()
+ self.t_env.get_config().set_local_timezone("SystemV/PST8PDT")
+ try:
+ dataframe = pf.from_arrow(
+ pa.table({
+ "ts": pa.array([0], type=pa.timestamp("ms", tz="UTC")),
+ })
+ )
+ self.assert_dataframe_schema(
+ dataframe,
+ ["ts"],
+ [TableDataTypes.TIMESTAMP(3)],
+ )
+ finally:
+ self.t_env.get_config().set_local_timezone(original_timezone)
+
+ 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": pd.Series(
+ [timestamp.replace(tzinfo=timezone.utc)],
+ dtype="datetime64[us, UTC]",
+ )
+ }
+ ),
+ 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):
@@ -572,6 +801,44 @@ class DataFrameITTests(PyFlinkStreamDataFrameTestCase):
[Row(1, "Alice"), Row(2, "Bob")],
)
+ def test_pandas_to_pandas_round_trip(self):
+ original_timezone = self.t_env.get_config().get_local_timezone()
+ self.t_env.get_config().set_local_timezone("America/New_York")
+ try:
+ first_fold = pd.Timestamp("2026-11-01T05:30:00.123Z")
+ second_fold = pd.Timestamp("2026-11-01T06:30:00.123Z")
+ pdf = pd.DataFrame(
+ {
+ "id": [0, 1, 2, 3],
+ "ts": pd.Series(
+ [None, first_fold, second_fold, None],
+ dtype="datetime64[ms, UTC]",
+ ),
+ }
+ )
+
+ result = (
+ pf.from_pandas(pdf)
+ .filter(pf.col("id") > 0)
+ .with_column("id_plus_one", pf.col("id") + 1)
+ .select("id", "id_plus_one", "ts")
+ .to_pandas()
+ .sort_values("id")
+ .reset_index(drop=True)
+ )
+
+ self.assertEqual(list(result.columns), ["id", "id_plus_one", "ts"])
+ self.assertEqual(result["id"].tolist(), [1, 2, 3])
+ self.assertEqual(result["id_plus_one"].tolist(), [2, 3, 4])
+ self.assertEqual(result["ts"].isna().tolist(), [False, False,
True])
+ local_fold = pd.Timestamp("2026-11-01T01:30:00.123")
+ self.assertEqual(
+ result["ts"].tolist()[:2],
+ [local_fold, local_fold],
+ )
+ finally:
+ self.t_env.get_config().set_local_timezone(original_timezone)
+
def test_basic_functionality(self):
df = pf.from_dict(
{
diff --git a/flink-python/pyflink/table/table_environment.py
b/flink-python/pyflink/table/table_environment.py
index f0e4bedba17..a1d90c41db1 100644
--- a/flink-python/pyflink/table/table_environment.py
+++ b/flink-python/pyflink/table/table_environment.py
@@ -45,9 +45,9 @@ from pyflink.table.statement_set import StatementSet
from pyflink.table.table_config import TableConfig
from pyflink.table.table_descriptor import TableDescriptor
from pyflink.table.table_result import TableResult
-from pyflink.table.types import _create_type_verifier, RowType, DataType, \
+from pyflink.table.types import _create_type_verifier, RowType, DataType,
TimestampType, \
_infer_schema_from_data, _create_converter, from_arrow_type, RowField,
create_arrow_schema, \
- _to_java_data_type
+ to_arrow_type, _to_java_data_type
from pyflink.table.udf import UserDefinedFunctionWrapper, AggregateFunction,
udaf, \
udtaf, TableAggregateFunction
from pyflink.table.utils import to_expression_jarray
@@ -61,6 +61,64 @@ __all__ = [
]
+_TIMESTAMP_UNITS_PER_SECOND = {
+ 's': 1,
+ 'ms': 1_000,
+ 'us': 1_000_000,
+ 'ns': 1_000_000_000,
+}
+
+
+def _cast_arrow_timestamp(array, target_type):
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ if isinstance(array, pa.ChunkedArray):
+ chunks = [
+ _cast_arrow_timestamp(chunk, target_type) for chunk in array.chunks
+ ]
+ return pa.chunked_array(chunks, type=target_type)
+ if array.type == target_type:
+ return array
+
+ source_units = _TIMESTAMP_UNITS_PER_SECOND[array.type.unit]
+ target_units = _TIMESTAMP_UNITS_PER_SECOND[target_type.unit]
+ values = array.view(pa.int64())
+ if source_units > target_units:
+ divisor = source_units // target_units
+ quotient = pc.divide_checked(values, divisor)
+ remainder = pc.subtract_checked(
+ values, pc.multiply_checked(quotient, divisor))
+ # Arrow integer division truncates toward zero. Timestamp precision
reduction
+ # needs floor division to preserve the local date and time before the
epoch.
+ values = pc.subtract_checked(
+ quotient,
+ pc.cast(pc.less(remainder, 0), pa.int64()),
+ )
+ elif source_units < target_units:
+ values = pc.multiply_checked(values, target_units // source_units)
+ return values.view(target_type)
+
+
+def _contains_timezone_aware_timestamp(arrow_type):
+ import pyarrow as pa
+
+ if pa.types.is_timestamp(arrow_type):
+ return arrow_type.tz is not None
+ if pa.types.is_struct(arrow_type):
+ return any(
+ _contains_timezone_aware_timestamp(field.type) for field in
arrow_type
+ )
+ if pa.types.is_list(arrow_type):
+ return _contains_timezone_aware_timestamp(arrow_type.value_type)
+ if pa.types.is_map(arrow_type):
+ return (
+ _contains_timezone_aware_timestamp(arrow_type.key_type)
+ or _contains_timezone_aware_timestamp(arrow_type.item_type)
+ )
+ return False
+
+
@PublicEvolving()
class TableEnvironment(object):
"""
@@ -1469,11 +1527,17 @@ class TableEnvironment(object):
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 +1546,8 @@ class TableEnvironment(object):
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 +1557,94 @@ class TableEnvironment(object):
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__}")
+ 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")
+
+ field_names = row_type.field_names()
+ field_types = row_type.field_types()
+ try:
+ compatible_table = table.rename_columns(field_names)
+ for index, (field, data_type) in enumerate(zip(table.schema,
field_types)):
+ source_column = compatible_table.column(index)
+ column = source_column
+ target_type = column.type
+ if isinstance(data_type, TimestampType):
+ target_type = to_arrow_type(data_type)
+ if column.type.tz is not None:
+ target_type = pa.timestamp(
+ target_type.unit, tz=column.type.tz)
+ if column.type != target_type:
+ if pa.types.is_timestamp(column.type):
+ column = _cast_arrow_timestamp(column, target_type)
+ else:
+ column = column.cast(target_type, safe=False)
+ if column is source_column:
+ continue
+ target_field = pa.field(
+ field_names[index],
+ target_type,
+ field.nullable,
+ field.metadata,
+ )
+ compatible_table = compatible_table.set_column(
+ index,
+ target_field,
+ column,
+ )
+ except (
+ pa.ArrowInvalid,
+ pa.ArrowNotImplementedError,
+ pa.ArrowTypeError,
+ TypeError,
+ ValueError,
+ ) as e:
+ raise TypeError(
+ f"Could not convert pyarrow.Table to the inferred Flink
schema: {row_type}"
+ ) from e
+
+ with tempfile.TemporaryDirectory() as temp_dir:
+ with tempfile.NamedTemporaryFile(delete=False, dir=temp_dir) as
temp_file:
+ with pa.ipc.new_stream(temp_file, compatible_table.schema) as
writer:
+ if compatible_table.num_rows > 0:
+ max_chunksize = -(-compatible_table.num_rows //
splits_num)
+ writer.write_table(
+ compatible_table, max_chunksize=max_chunksize)
+
+ jvm = get_gateway().jvm
+ if table_schema is None:
+ 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
+ create_descriptor =
jvm.org.apache.flink.table.runtime.arrow.ArrowUtils \
+ .createArrowTableSourceDesc
+ if any(
+ _contains_timezone_aware_timestamp(field.type)
+ for field in compatible_table.schema
+ ):
+ descriptor = create_descriptor(
+ source_schema,
+ temp_file.name,
+ self.get_config().get_local_timezone(),
+ )
+ else:
+ descriptor = create_descriptor(source_schema, temp_file.name)
+ return Table(getattr(self._j_tenv, "from")(descriptor), self)
+
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 9cc0f8ccdf6..6ebce2f1005 100644
--- a/flink-python/pyflink/table/tests/test_pandas_conversion.py
+++ b/flink-python/pyflink/table/tests/test_pandas_conversion.py
@@ -113,6 +113,27 @@ class PandasConversionTests(PandasConversionTestBase):
table = self.t_env.from_pandas(self.pdf, schema=tuple(new_types))
self.assertEqual(new_types, table.get_schema().get_field_data_types())
+ def test_from_pandas_with_nested_timezone_aware_timestamps(self):
+ import pandas as pd
+
+ timestamp = pd.Timestamp("2026-01-01T00:00:00Z")
+ pdf = pd.DataFrame(
+ {
+ "payload": [{"ts": timestamp}],
+ "timestamps": [[timestamp]],
+ }
+ )
+
+ table = self.t_env.from_pandas(pdf)
+
+ self.assertEqual(
+ [
+ DataTypes.ROW([DataTypes.FIELD("ts", DataTypes.TIMESTAMP(6))]),
+ DataTypes.ARRAY(DataTypes.TIMESTAMP(6)),
+ ],
+ table.get_schema().get_field_data_types(),
+ )
+
class PandasConversionITTests(PandasConversionTestBase):
diff --git a/flink-python/pyflink/table/tests/test_table_environment_api.py
b/flink-python/pyflink/table/tests/test_table_environment_api.py
index 505d1dbd26f..492c20bb8ff 100644
--- a/flink-python/pyflink/table/tests/test_table_environment_api.py
+++ b/flink-python/pyflink/table/tests/test_table_environment_api.py
@@ -17,9 +17,13 @@
################################################################################
import datetime
import decimal
+import os
import sys
+import tempfile
import unittest
+from unittest.mock import MagicMock, patch
+import pyarrow as pa
from py4j.protocol import Py4JJavaError
from typing import Iterable
@@ -38,6 +42,7 @@ from pyflink.table.catalog import ObjectPath,
CatalogBaseTable, CatalogDescripto
from pyflink.table.explain_detail import ExplainDetail
from pyflink.table.expressions import col, source_watermark
from pyflink.table.table_descriptor import TableDescriptor
+from pyflink.table.table_environment import TableEnvironment
from pyflink.table.types import RowType, Row, UserDefinedType
from pyflink.table.udf import udf
from pyflink.testing import source_sink_utils
@@ -46,6 +51,118 @@ from pyflink.testing.test_case_utils import
(PyFlinkStreamTableTestCase, PyFlink
from pyflink.util.java_utils import get_j_env_configuration
+class ArrowTimestampConversionTests(unittest.TestCase):
+
+ def _write_from_arrow(
+ self,
+ table,
+ row_type,
+ splits_num=1):
+ table_environment = object.__new__(TableEnvironment)
+ table_environment._j_tenv = MagicMock()
+ setattr(table_environment._j_tenv, 'from',
MagicMock(return_value=object()))
+
+ batches = []
+ arrow_schemas = []
+ gateway = MagicMock()
+ descriptor_factory = gateway.jvm.org.apache.flink.table.runtime.arrow \
+ .ArrowUtils.createArrowTableSourceDesc
+
+ def capture_batches(source_schema, file_name, *args):
+ with pa.ipc.open_stream(file_name) as reader:
+ arrow_schemas.append(reader.schema)
+ batches.extend(reader)
+ return object()
+
+ descriptor_factory.side_effect = capture_batches
+ with patch(
+ 'pyflink.table.table_environment.get_gateway', return_value=gateway
+ ):
+ table_environment._from_arrow(
+ table,
+ row_type,
+ table_schema=MagicMock(),
+ splits_num=splits_num,
+ )
+ return batches, arrow_schemas[0]
+
+ def test_from_arrow_preserves_existing_batches(self):
+ table = pa.table({
+ 'id': pa.chunked_array(
+ [pa.array([value], type=pa.int64()) for value in range(4)]
+ )
+ })
+ row_type = DataTypes.ROW([DataTypes.FIELD('id', DataTypes.BIGINT())])
+ batches, _ = self._write_from_arrow(
+ table, row_type, splits_num=2)
+
+ self.assertEqual([batch.num_rows for batch in batches], [1, 1, 1, 1])
+
+ def test_from_arrow_handles_split_boundaries(self):
+ row_type = DataTypes.ROW([DataTypes.FIELD('id', DataTypes.BIGINT())])
+ cases = [
+ (pa.table({'id': range(5)}), 2, [3, 2]),
+ (pa.table({'id': range(2)}), 5, [1, 1]),
+ ]
+ for table, splits_num, expected_batch_sizes in cases:
+ with self.subTest(splits_num=splits_num, rows=table.num_rows):
+ batches, _ = self._write_from_arrow(
+ table, row_type, splits_num=splits_num)
+ self.assertEqual(
+ [batch.num_rows for batch in batches],
expected_batch_sizes)
+
+ def test_from_arrow_rejects_invalid_split_counts(self):
+ table = pa.table({'id': [1]})
+ row_type = DataTypes.ROW([DataTypes.FIELD('id', DataTypes.BIGINT())])
+ invalid_splits = [
+ (True, TypeError, 'must be an integer'),
+ (1.5, TypeError, 'must be an integer'),
+ (0, ValueError, 'must be greater than 0'),
+ (-1, ValueError, '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):
+ self._write_from_arrow(table, row_type,
splits_num=splits_num)
+
+ def test_from_arrow_writes_schema_only_empty_stream(self):
+ table = pa.table({
+ 'id': pa.array([], type=pa.int64()),
+ })
+ row_type = DataTypes.ROW([DataTypes.FIELD('id', DataTypes.BIGINT())])
+
+ batches, arrow_schema = self._write_from_arrow(
+ table, row_type, splits_num=5)
+
+ self.assertEqual(batches, [])
+ self.assertEqual(arrow_schema, table.schema)
+
+ def test_from_arrow_floors_pre_epoch_timestamp_precision(self):
+ table = pa.table({
+ 'ts': pa.array([-1], type=pa.timestamp('us')),
+ })
+ row_type = DataTypes.ROW([
+ DataTypes.FIELD('ts', DataTypes.TIMESTAMP(3)),
+ ])
+
+ batches, _ = self._write_from_arrow(table, row_type)
+
+ self.assertEqual(
+ batches[0].column(0).to_pylist(),
+ [datetime.datetime.fromisoformat('1969-12-31T23:59:59.999')],
+ )
+
+ def test_from_arrow_cleans_temporary_directory(self):
+ table = pa.table({'id': [1]})
+ row_type = DataTypes.ROW([DataTypes.FIELD('id', DataTypes.BIGINT())])
+
+ with tempfile.TemporaryDirectory() as temp_root:
+ with patch('pyflink.table.table_environment.tempfile.tempdir',
temp_root):
+ self._write_from_arrow(table, row_type)
+
+ self.assertEqual(os.listdir(temp_root), [])
+
+
class TableEnvironmentTest(PyFlinkUTTestCase):
def test_set_sys_executable_for_local_mode(self):
diff --git a/flink-python/pyflink/table/tests/test_types.py
b/flink-python/pyflink/table/tests/test_types.py
index 400a3729ee9..80d08d5b1f0 100644
--- a/flink-python/pyflink/table/tests/test_types.py
+++ b/flink-python/pyflink/table/tests/test_types.py
@@ -24,6 +24,8 @@ import sys
import tempfile
import unittest
+import pyarrow as pa
+
from pyflink.pyflink_gateway_server import on_windows
from pyflink.serializers import BatchedSerializer, PickleSerializer
@@ -35,7 +37,7 @@ from pyflink.table.types import (_infer_schema_from_data,
_infer_type,
_create_type_verifier, UserDefinedType,
DataTypes, Row, RowField,
RowType, ArrayType, BigIntType, VarCharType,
MapType, DataType,
_from_java_data_type, ZonedTimestampType,
- LocalZonedTimestampType, _to_java_data_type)
+ LocalZonedTimestampType, _to_java_data_type,
from_arrow_type)
from pyflink.testing.test_case_utils import PyFlinkTestCase
@@ -126,6 +128,36 @@ class UTCOffsetTimezone(datetime.tzinfo):
return self.OFFSET
+class ArrowTypeConversionTests(unittest.TestCase):
+
+ def test_nested_field_nullability(self):
+ try:
+ arrow_map_type = pa.map_(
+ pa.string(),
+ pa.field("value", pa.int64(), nullable=False),
+ )
+ except TypeError:
+ arrow_map_type = pa.map_(pa.string(), pa.int64())
+
+ map_type = from_arrow_type(arrow_map_type)
+ self.assertFalse(map_type.key_type._nullable)
+ item_field = getattr(arrow_map_type, "item_field", None)
+ expected_item_nullable = (
+ item_field.nullable if item_field is not None else True
+ )
+ self.assertEqual(expected_item_nullable, map_type.value_type._nullable)
+
+ array_type = from_arrow_type(
+ pa.list_(pa.field("item", pa.int64(), nullable=False))
+ )
+ self.assertFalse(array_type.element_type._nullable)
+
+ row_type = from_arrow_type(
+ pa.struct([pa.field("value", pa.int64())]), nullable=False
+ )
+ self.assertFalse(row_type._nullable)
+
+
class TypesTests(PyFlinkTestCase):
def test_row_type_repr_includes_nullability(self):
diff --git a/flink-python/pyflink/table/types.py
b/flink-python/pyflink/table/types.py
index 316bb44b55e..16d92c3c1ff 100644
--- a/flink-python/pyflink/table/types.py
+++ b/flink-python/pyflink/table/types.py
@@ -2272,21 +2272,28 @@ def from_arrow_type(arrow_type, nullable: bool = True)
-> DataType:
else:
return TimestampType(9, nullable)
elif types.is_map(arrow_type):
- return MapType(from_arrow_type(arrow_type.key_type),
- from_arrow_type(arrow_type.item_type),
- nullable)
+ item_field = getattr(arrow_type, 'item_field', None)
+ item_nullable = item_field.nullable if item_field is not None else True
+ key_type = from_arrow_type(arrow_type.key_type, nullable=False)
+ value_type = from_arrow_type(arrow_type.item_type,
nullable=item_nullable)
+ return MapType(key_type, value_type, nullable)
elif types.is_list(arrow_type):
- return ArrayType(from_arrow_type(arrow_type.value_type), nullable)
+ return ArrayType(
+ from_arrow_type(
+ arrow_type.value_type, nullable=arrow_type.value_field.nullable
+ ),
+ nullable,
+ )
elif types.is_struct(arrow_type):
if any(types.is_struct(field.type) for field in arrow_type):
raise TypeError("Nested RowType is not supported in conversion
from Arrow: " +
str(arrow_type))
return RowType([RowField(field.name, from_arrow_type(field.type,
field.nullable))
- for field in arrow_type])
+ for field in arrow_type], nullable)
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 d8b9dcdf772..87f70e8a0c2 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
@@ -128,12 +128,15 @@ import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.VectorUnloader;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.MapVector;
import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.ipc.ArrowStreamReader;
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
import org.apache.arrow.vector.ipc.ReadChannel;
import org.apache.arrow.vector.ipc.WriteChannel;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.ipc.message.MessageMetadataResult;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.DateUnit;
@@ -154,6 +157,9 @@ import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.ReadableByteChannel;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.zone.ZoneRules;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
@@ -481,6 +487,15 @@ public final class ArrowUtils {
}
public static TableDescriptor createArrowTableSourceDesc(DataType
dataType, String fileName) {
+ return createArrowTableSourceDesc(createTableSchema(dataType),
fileName);
+ }
+
+ public static TableDescriptor createArrowTableSourceDesc(
+ DataType dataType, String fileName, String localTimeZoneId) {
+ return createArrowTableSourceDesc(createTableSchema(dataType),
fileName, localTimeZoneId);
+ }
+
+ private static org.apache.flink.table.api.Schema
createTableSchema(DataType dataType) {
List<String> fieldNames = getFieldNames(dataType);
List<DataType> fieldTypes = dataType.getChildren();
org.apache.flink.table.api.Schema.Builder schemaBuilder =
@@ -488,20 +503,109 @@ public final class ArrowUtils {
for (int i = 0; i < fieldNames.size(); i++) {
schemaBuilder.column(fieldNames.get(i), fieldTypes.get(i));
}
+ return schemaBuilder.build();
+ }
+ 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())
- .build();
+ return createArrowTableSourceDesc(schema,
readArrowBatches(fileName));
} catch (Throwable e) {
throw new TableException("Failed to read the arrow data from " +
fileName, e);
}
}
+ public static TableDescriptor createArrowTableSourceDesc(
+ org.apache.flink.table.api.Schema schema, String fileName, String
localTimeZoneId) {
+ try {
+ return createArrowTableSourceDesc(
+ schema, readArrowBatchesInLocalTimeZone(fileName,
ZoneId.of(localTimeZoneId)));
+ } catch (Throwable e) {
+ throw new TableException("Failed to read the arrow data from " +
fileName, e);
+ }
+ }
+
+ private static TableDescriptor createArrowTableSourceDesc(
+ org.apache.flink.table.api.Schema schema, byte[][] data) throws
IOException {
+ return TableDescriptor.forConnector(ArrowTableSourceFactory.IDENTIFIER)
+ .option(ArrowTableSourceOptions.DATA,
ByteArrayUtils.twoDimByteArrayToString(data))
+ .schema(schema)
+ .build();
+ }
+
+ private static byte[][] readArrowBatchesInLocalTimeZone(String fileName,
ZoneId localTimeZone)
+ throws IOException {
+ checkArrowUsable();
+ List<byte[]> results = new ArrayList<>();
+ try (BufferAllocator allocator =
+ getRootAllocator()
+ .newChildAllocator(
+
"ArrowUtils#readArrowBatchesInLocalTimeZone",
+ 0,
+ Long.MAX_VALUE);
+ FileInputStream inputStream = new FileInputStream(fileName);
+ ArrowStreamReader reader = new ArrowStreamReader(inputStream,
allocator)) {
+ while (reader.loadNextBatch()) {
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ convertTimezoneAwareTimestampsToLocalTime(root, localTimeZone);
+ try (ArrowRecordBatch batch = new
VectorUnloader(root).getRecordBatch();
+ ByteArrayOutputStream outputStream = new
ByteArrayOutputStream()) {
+ MessageSerializer.serialize(
+ new
WriteChannel(Channels.newChannel(outputStream)), batch);
+ results.add(outputStream.toByteArray());
+ }
+ }
+ }
+ return results.toArray(new byte[0][]);
+ }
+
+ static void convertTimezoneAwareTimestampsToLocalTime(
+ VectorSchemaRoot root, ZoneId localTimeZone) {
+ ZoneRules rules = localTimeZone.getRules();
+ for (FieldVector vector : root.getFieldVectors()) {
+ convertTimezoneAwareTimestampsToLocalTime(vector, rules);
+ }
+ }
+
+ private static void convertTimezoneAwareTimestampsToLocalTime(
+ FieldVector vector, ZoneRules rules) {
+ if (vector instanceof TimeStampVector
+ && ((ArrowType.Timestamp)
vector.getField().getType()).getTimezone() != null) {
+ TimeStampVector timestampVector = (TimeStampVector) vector;
+ TimeUnit unit = ((ArrowType.Timestamp)
vector.getField().getType()).getUnit();
+ long unitsPerSecond = getTimestampUnitsPerSecond(unit);
+ for (int i = 0; i < timestampVector.getValueCount(); i++) {
+ if (timestampVector.isNull(i)) {
+ continue;
+ }
+ long value = timestampVector.get(i);
+ Instant instant = Instant.ofEpochSecond(Math.floorDiv(value,
unitsPerSecond));
+ long offset = rules.getOffset(instant).getTotalSeconds();
+ timestampVector.set(
+ i, Math.addExact(value, Math.multiplyExact(offset,
unitsPerSecond)));
+ }
+ return;
+ }
+ for (FieldVector child : vector.getChildrenFromFields()) {
+ convertTimezoneAwareTimestampsToLocalTime(child, rules);
+ }
+ }
+
+ private static long getTimestampUnitsPerSecond(TimeUnit unit) {
+ switch (unit) {
+ case SECOND:
+ return 1L;
+ case MILLISECOND:
+ return 1_000L;
+ case MICROSECOND:
+ return 1_000_000L;
+ case NANOSECOND:
+ return 1_000_000_000L;
+ default:
+ throw new IllegalArgumentException("Unsupported timestamp
unit: " + unit);
+ }
+ }
+
public static byte[][] readArrowBatches(String fileName) throws
IOException {
try (FileInputStream fis = new FileInputStream(fileName)) {
return readArrowBatches(fis.getChannel());
diff --git
a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowTimestampColumnVector.java
b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowTimestampColumnVector.java
index 297f21d3978..612056db14f 100644
---
a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowTimestampColumnVector.java
+++
b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowTimestampColumnVector.java
@@ -49,15 +49,18 @@ public final class ArrowTimestampColumnVector implements
TimestampColumnVector {
@Override
public TimestampData getTimestamp(int i, int precision) {
if (valueVector instanceof TimeStampSecVector) {
- return TimestampData.fromEpochMillis(((TimeStampSecVector)
valueVector).get(i) * 1000);
+ return TimestampData.fromEpochMillis(
+ Math.multiplyExact(((TimeStampSecVector)
valueVector).get(i), 1000));
} else if (valueVector instanceof TimeStampMilliVector) {
return TimestampData.fromEpochMillis(((TimeStampMilliVector)
valueVector).get(i));
} else if (valueVector instanceof TimeStampMicroVector) {
long micros = ((TimeStampMicroVector) valueVector).get(i);
- return TimestampData.fromEpochMillis(micros / 1000, (int) (micros
% 1000) * 1000);
+ return TimestampData.fromEpochMillis(
+ Math.floorDiv(micros, 1000), (int) Math.floorMod(micros,
1000) * 1000);
} else {
long nanos = ((TimeStampNanoVector) valueVector).get(i);
- return TimestampData.fromEpochMillis(nanos / 1_000_000, (int)
(nanos % 1_000_000));
+ return TimestampData.fromEpochMillis(
+ Math.floorDiv(nanos, 1_000_000), (int)
Math.floorMod(nanos, 1_000_000));
}
}
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 01dfab186dd..3681d0664ae 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 @@ public final class 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());
}
diff --git
a/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowUtilsTest.java
b/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowUtilsTest.java
index e7b3a81d1f4..971e0f560fb 100644
---
a/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowUtilsTest.java
+++
b/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowUtilsTest.java
@@ -21,6 +21,7 @@ package org.apache.flink.table.runtime.arrow;
import org.apache.flink.api.java.tuple.Tuple5;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.TimestampData;
import org.apache.flink.table.data.columnar.vector.ColumnVector;
import org.apache.flink.table.runtime.arrow.vectors.ArrowArrayColumnVector;
import org.apache.flink.table.runtime.arrow.vectors.ArrowBigIntColumnVector;
@@ -74,13 +75,16 @@ import org.apache.flink.table.types.logical.VarCharType;
import org.apache.flink.shaded.guava33.com.google.common.collect.Lists;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.TimeStampVector;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -89,11 +93,16 @@ import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.channels.Channels;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.Arrays;
+import java.util.Collections;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for {@link ArrowUtils}. */
class ArrowUtilsTest {
@@ -352,6 +361,156 @@ class ArrowUtilsTest {
}
}
+ @Test
+ void testCreateArrowReaderForPreEpochSubMillisecondTimestamps() {
+ List<Field> fields =
+ Arrays.asList(
+ new Field(
+ "micros",
+ FieldType.nullable(
+ new
ArrowType.Timestamp(TimeUnit.MICROSECOND, null)),
+ null),
+ new Field(
+ "nanos",
+ FieldType.nullable(
+ new
ArrowType.Timestamp(TimeUnit.NANOSECOND, null)),
+ null));
+ RowType timestampRowType = RowType.of(new TimestampType(6), new
TimestampType(9));
+
+ try (VectorSchemaRoot root = VectorSchemaRoot.create(new
Schema(fields), allocator)) {
+ TimeStampVector micros = (TimeStampVector)
root.getVector("micros");
+ TimeStampVector nanos = (TimeStampVector) root.getVector("nanos");
+ micros.setSafe(0, -1);
+ nanos.setSafe(0, -1);
+ micros.setValueCount(1);
+ nanos.setValueCount(1);
+ root.setRowCount(1);
+
+ RowData row = ArrowUtils.createArrowReader(root,
timestampRowType).read(0);
+
+ assertThat(row.getTimestamp(0, 6))
+ .isEqualTo(TimestampData.fromEpochMillis(-1, 999_000));
+ assertThat(row.getTimestamp(1, 9))
+ .isEqualTo(TimestampData.fromEpochMillis(-1, 999_999));
+ }
+ }
+
+ @Test
+ void testCreateArrowReaderRejectsSecondTimestampOverflow() {
+ Field timestampField =
+ new Field(
+ "ts",
+ FieldType.nullable(new
ArrowType.Timestamp(TimeUnit.SECOND, null)),
+ null);
+ RowType timestampRowType = RowType.of(new TimestampType(0));
+
+ try (VectorSchemaRoot root =
+ VectorSchemaRoot.create(
+ new Schema(Collections.singletonList(timestampField)),
allocator)) {
+ TimeStampVector timestamp = (TimeStampVector) root.getVector("ts");
+ timestamp.setValueCount(1);
+ root.setRowCount(1);
+ ArrowReader reader = ArrowUtils.createArrowReader(root,
timestampRowType);
+
+ for (long value : new long[] {Long.MIN_VALUE, Long.MAX_VALUE}) {
+ timestamp.setSafe(0, value);
+ assertThatThrownBy(() -> reader.read(0).getTimestamp(0, 0))
+ .isInstanceOf(ArithmeticException.class);
+ }
+ }
+ }
+
+ @Test
+ void testConvertTimezoneAwareTimestampsUsingJavaZoneRules() {
+ Field timestampField =
+ new Field(
+ "ts",
+ FieldType.nullable(new
ArrowType.Timestamp(TimeUnit.SECOND, "UTC")),
+ null);
+ Field nestedField =
+ new Field(
+ "nested",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Collections.singletonList(timestampField));
+ long instant = Instant.parse("2050-01-01T00:00:00Z").getEpochSecond();
+ ZoneId localTimeZone = ZoneId.of("America/Vancouver");
+
+ try (VectorSchemaRoot root =
+ VectorSchemaRoot.create(
+ new Schema(Arrays.asList(timestampField,
nestedField)), allocator)) {
+ root.allocateNew();
+ TimeStampVector timestamp = (TimeStampVector) root.getVector("ts");
+ timestamp.setSafe(0, instant);
+ timestamp.setValueCount(1);
+ StructVector nested = (StructVector) root.getVector("nested");
+ TimeStampVector nestedTimestamp = (TimeStampVector)
nested.getChild("ts");
+ nestedTimestamp.setSafe(0, instant);
+ nestedTimestamp.setValueCount(1);
+ nested.setIndexDefined(0);
+ nested.setValueCount(1);
+ root.setRowCount(1);
+
+ ArrowUtils.convertTimezoneAwareTimestampsToLocalTime(root,
localTimeZone);
+
+ long expected =
+ Instant.ofEpochSecond(instant)
+ .atZone(localTimeZone)
+ .toLocalDateTime()
+ .toEpochSecond(ZoneOffset.UTC);
+ assertThat(timestamp.get(0)).isEqualTo(expected);
+ assertThat(nestedTimestamp.get(0)).isEqualTo(expected);
+ }
+ }
+
+ @Test
+ void testConvertTimezoneAwareTimestampsWithJavaOnlyTimeZoneId() {
+ Field timestampField =
+ new Field(
+ "ts",
+ FieldType.nullable(new
ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")),
+ null);
+
+ try (VectorSchemaRoot root =
+ VectorSchemaRoot.create(
+ new Schema(Collections.singletonList(timestampField)),
allocator)) {
+ root.allocateNew();
+ TimeStampVector timestamp = (TimeStampVector) root.getVector("ts");
+ timestamp.setSafe(0, 0);
+ timestamp.setValueCount(1);
+ root.setRowCount(1);
+
+ ArrowUtils.convertTimezoneAwareTimestampsToLocalTime(
+ root, ZoneId.of("SystemV/PST8PDT"));
+
+ assertThat(timestamp.get(0)).isEqualTo(-8 * 60 * 60 * 1000L);
+ }
+ }
+
+ @Test
+ void testRejectTimezoneConversionOverflow() {
+ Field timestampField =
+ new Field(
+ "ts",
+ FieldType.nullable(new
ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")),
+ null);
+
+ try (VectorSchemaRoot root =
+ VectorSchemaRoot.create(
+ new Schema(Collections.singletonList(timestampField)),
allocator)) {
+ root.allocateNew();
+ TimeStampVector timestamp = (TimeStampVector) root.getVector("ts");
+ timestamp.setSafe(0, Long.MAX_VALUE);
+ timestamp.setValueCount(1);
+ root.setRowCount(1);
+
+ assertThatThrownBy(
+ () ->
+
ArrowUtils.convertTimezoneAwareTimestampsToLocalTime(
+ root, ZoneId.of("Asia/Shanghai")))
+ .isInstanceOf(ArithmeticException.class);
+ }
+ }
+
@Test
void testCreateArrowWriter() {
VectorSchemaRoot root =
