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 5c422565ab1 [FLINK-40437][python] Add iteration and bounded retrieval
to DataFrame API
5c422565ab1 is described below
commit 5c422565ab1fa8a8b0d019aacec405ee44f4d116
Author: Federico Dolce <[email protected]>
AuthorDate: Wed Sep 30 16:27:33 2026 +0200
[FLINK-40437][python] Add iteration and bounded retrieval to DataFrame API
This closes #29381.
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 6 +
flink-python/pyflink/dataframe/__init__.py | 2 +
flink-python/pyflink/dataframe/dataframe.py | 217 ++++++++++++-
flink-python/pyflink/dataframe/iteration.py | 289 +++++++++++++++++
.../pyflink/dataframe/tests/test_dataframe.py | 2 +-
.../pyflink/dataframe/tests/test_iteration.py | 342 +++++++++++++++++++++
flink-python/pyflink/dataframe/validation.py | 56 ++++
7 files changed, 909 insertions(+), 5 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index dbeadf80fe9..46286497076 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -149,6 +149,12 @@ Results
DataFrame.collect
DataFrame.to_table
DataFrame.to_pandas
+ DataFrame.iter_rows
+ DataFrame.iter_batches
+ DataFrame.take
+ DataFrame.take_batch
+ CloseableIterator
+ CloseableIterator.close
Windowing
---------
diff --git a/flink-python/pyflink/dataframe/__init__.py
b/flink-python/pyflink/dataframe/__init__.py
index 0a30d752d19..0e67c1a9e54 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -66,11 +66,13 @@ from pyflink.dataframe.dataframe import DataFrame,
GroupedDataFrame, col, lit
from pyflink.dataframe.dataframe_config import config
from pyflink.dataframe.datatype import DataType
from pyflink.dataframe.io import read_catalog_table, read_generic, read_json,
read_parquet
+from pyflink.dataframe.iteration import CloseableIterator
from pyflink.dataframe.sql import sql
from pyflink.dataframe.udf import udf
from pyflink.dataframe.udtf import udtf
__all__ = [
+ "CloseableIterator",
"DataFrame",
"GroupedDataFrame",
"DataType",
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index 7a9679166c8..8698dda1367 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -25,6 +25,7 @@ from typing import (
Callable,
Dict,
List,
+ Literal,
Optional,
Set,
Tuple,
@@ -36,6 +37,7 @@ from typing import (
if TYPE_CHECKING:
import pandas
+ import pyarrow
from pyflink.dataframe.udf import _DataTypeLike
from pyflink.dataframe.udtf import _DataFrameUDTFWrapper
from pyflink.table.table_environment import TableEnvironment
@@ -43,6 +45,18 @@ if TYPE_CHECKING:
from pyflink.common import Row
from pyflink.dataframe.datatype import _INT_MAX, DataType
+from pyflink.dataframe.iteration import (
+ _BATCH_FORMATS,
+ CloseableIterator,
+ _build_arrow_schema,
+ _iterate_batches,
+ _iterate_rows,
+ _row_to_dict,
+ _rows_to_batch,
+ _take_rows,
+ _validate_row_kind_field,
+)
+from pyflink.dataframe.validation import _require_choice, _require_int,
_require_number
from pyflink.java_gateway import get_gateway
from pyflink.table.expression import Expression, _get_java_expression
from pyflink.table.expressions import (
@@ -64,10 +78,7 @@ T = TypeVar("T")
def _validate_row_count(n: int) -> None:
- if isinstance(n, bool) or not isinstance(n, int):
- raise TypeError("n must be an integer")
- if n < 0:
- raise ValueError("n must be non-negative")
+ _require_int(n, "n", 0)
if n > _INT_MAX:
raise ValueError(f"n must be less than or equal to {_INT_MAX}")
@@ -1720,6 +1731,204 @@ class DataFrame:
with self._table.execute().collect() as rows:
return list(rows)
+ @PublicEvolving()
+ def iter_rows(
+ self,
+ *,
+ include_row_kind: bool = False,
+ row_kind_field: str = "__row_kind__",
+ ) -> CloseableIterator[Dict[str, Any]]:
+ """
+ Execute this DataFrame and iterate over its rows as they arrive, one
dict per row.
+
+ Rows are fetched incrementally, so this also works on unbounded
sources. The job runs until
+ the iterator is exhausted or closed; use it in a ``with`` block.
+
+ On an updating DataFrame, such as a streaming aggregation, every
changelog entry is a
+ separate row. Set ``include_row_kind`` to tell insertions, updates and
deletions apart.
+
+ TIMESTAMP_LTZ columns are not supported yet: the row transfer shared
with :meth:`collect`
+ cannot serialize them.
+
+ :param include_row_kind: Whether to add each row's change kind, one of
``"+I"``,
+ ``"-U"``, ``"+U"`` and ``"-D"``.
+ :param row_kind_field: Key under which the change kind is added. It
must not clash with a
+ column name.
+ :return: An iterator of dicts mapping column names to values.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([{"id": 1}, {"id": 2}])
+ >>> with df.iter_rows() as rows:
+ ... for row in rows:
+ ... print(row["id"])
+
+ .. versionadded:: 2.4.0
+ """
+ columns = self.columns
+ _validate_row_kind_field(columns, include_row_kind, row_kind_field)
+ return _iterate_rows(self._table, columns, row_kind_field if
include_row_kind else None)
+
+ @PublicEvolving()
+ def iter_batches(
+ self,
+ *,
+ batch_size: int = 1000,
+ batch_format: Literal["pandas", "pyarrow"] = "pandas",
+ include_row_kind: bool = False,
+ row_kind_field: str = "__row_kind__",
+ ) -> CloseableIterator[Union["pandas.DataFrame", "pyarrow.Table"]]:
+ """
+ Execute this DataFrame and iterate over its rows in batches of
``batch_size``.
+
+ Every batch except possibly the last holds exactly ``batch_size``
rows, so on a slow
+ unbounded source a batch is only emitted once enough rows have
arrived. The job runs until
+ the iterator is exhausted or closed; use it in a ``with`` block.
+
+ TIMESTAMP_LTZ columns are not supported yet: the row transfer shared
with :meth:`collect`
+ cannot serialize them.
+
+ :param batch_size: Number of rows per batch.
+ :param batch_format: ``"pandas"`` for pandas DataFrames or
``"pyarrow"`` for PyArrow
+ Tables.
+ :param include_row_kind: Whether to add a column with each row's
change kind, one of
+ ``"+I"``, ``"-U"``, ``"+U"`` and ``"-D"``.
+ :param row_kind_field: Name of the change kind column. It must not
clash with a column
+ name.
+ :return: An iterator of batches in the requested format.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.range(10)
+ >>> with df.iter_batches(batch_size=4) as batches:
+ ... for batch in batches:
+ ... print(len(batch))
+
+ .. versionadded:: 2.4.0
+ """
+ _require_int(batch_size, "batch_size", 1)
+ _require_choice(batch_format, "batch_format", _BATCH_FORMATS)
+ schema = self._table.get_resolved_schema()
+ columns = schema.get_column_names()
+ data_types = schema.get_column_data_types()
+ _validate_row_kind_field(columns, include_row_kind, row_kind_field)
+ row_kind_field = row_kind_field if include_row_kind else None
+ return _iterate_batches(
+ self._table,
+ batch_size,
+ batch_format,
+ data_types,
+ _build_arrow_schema(columns, data_types, row_kind_field),
+ row_kind_field,
+ )
+
+ @PublicEvolving()
+ def take(
+ self,
+ n: int,
+ *,
+ timeout: Optional[float] = None,
+ include_row_kind: bool = False,
+ row_kind_field: str = "__row_kind__",
+ ) -> List[Dict[str, Any]]:
+ """
+ Execute this DataFrame and return its first ``n`` rows as dicts.
+
+ The job is cancelled once ``n`` rows have arrived if it is still
running, so this is a
+ safe way to peek at an unbounded source. Fewer rows are returned if
the result ends first
+ or ``timeout`` expires.
+ The query itself is not limited: on an updating DataFrame the rows are
the first ``n``
+ changelog entries.
+
+ TIMESTAMP_LTZ columns are not supported yet: the row transfer shared
with :meth:`collect`
+ cannot serialize them.
+
+ :param n: Maximum number of rows to return.
+ :param timeout: Maximum number of seconds to wait for rows once the
job is submitted.
+ ``None`` waits until ``n`` rows arrive or the result ends.
+ :param include_row_kind: Whether to add each row's change kind, one of
``"+I"``,
+ ``"-U"``, ``"+U"`` and ``"-D"``.
+ :param row_kind_field: Key under which the change kind is added. It
must not clash with a
+ column name.
+ :return: Up to ``n`` dicts mapping column names to values.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.range(100)
+ >>> df.take(2)
+ [{'id': 0}, {'id': 1}]
+
+ .. versionadded:: 2.4.0
+ """
+ _validate_row_count(n)
+ if timeout is not None:
+ _require_number(timeout, "timeout", 0)
+ columns = self.columns
+ _validate_row_kind_field(columns, include_row_kind, row_kind_field)
+ row_kind_field = row_kind_field if include_row_kind else None
+ return [
+ _row_to_dict(row, columns, row_kind_field)
+ for row in _take_rows(self._table, n, timeout)
+ ]
+
+ @PublicEvolving()
+ def take_batch(
+ self,
+ n: int,
+ *,
+ timeout: Optional[float] = None,
+ batch_format: Literal["pandas", "pyarrow"] = "pandas",
+ include_row_kind: bool = False,
+ row_kind_field: str = "__row_kind__",
+ ) -> Union["pandas.DataFrame", "pyarrow.Table"]:
+ """
+ Execute this DataFrame and return its first ``n`` rows as a single
batch.
+
+ Behaves like :meth:`take` but returns a pandas DataFrame or PyArrow
Table.
+
+ TIMESTAMP_LTZ columns are not supported yet: the row transfer shared
with :meth:`collect`
+ cannot serialize them.
+
+ :param n: Maximum number of rows to return.
+ :param timeout: Maximum number of seconds to wait for rows once the
job is submitted.
+ ``None`` waits until ``n`` rows arrive or the result ends.
+ :param batch_format: ``"pandas"`` for a pandas DataFrame or
``"pyarrow"`` for a PyArrow
+ Table.
+ :param include_row_kind: Whether to add a column with each row's
change kind, one of
+ ``"+I"``, ``"-U"``, ``"+U"`` and ``"-D"``.
+ :param row_kind_field: Name of the change kind column. It must not
clash with a column
+ name.
+ :return: A batch of up to ``n`` rows in the requested format.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.range(100)
+ >>> pdf = df.take_batch(5)
+
+ .. versionadded:: 2.4.0
+ """
+ _validate_row_count(n)
+ if timeout is not None:
+ _require_number(timeout, "timeout", 0)
+ _require_choice(batch_format, "batch_format", _BATCH_FORMATS)
+ schema = self._table.get_resolved_schema()
+ columns = schema.get_column_names()
+ data_types = schema.get_column_data_types()
+ _validate_row_kind_field(columns, include_row_kind, row_kind_field)
+ row_kind_field = row_kind_field if include_row_kind else None
+ arrow_schema = _build_arrow_schema(columns, data_types, row_kind_field)
+ return _rows_to_batch(
+ _take_rows(self._table, n, timeout),
+ data_types,
+ arrow_schema,
+ batch_format,
+ row_kind_field,
+ )
+
@PublicEvolving()
def to_table(self) -> Table:
"""
diff --git a/flink-python/pyflink/dataframe/iteration.py
b/flink-python/pyflink/dataframe/iteration.py
new file mode 100644
index 00000000000..53f6fd958c7
--- /dev/null
+++ b/flink-python/pyflink/dataframe/iteration.py
@@ -0,0 +1,289 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import itertools
+import queue
+import threading
+import time
+from typing import (
+ Any,
+ Callable,
+ Dict,
+ Iterator,
+ List,
+ Optional,
+ Protocol,
+ TypeVar,
+)
+
+from pyflink.common import Row
+from pyflink.dataframe.validation import _require_non_empty_str
+from pyflink.table.types import ArrayType, DataType, MapType, RowType,
create_arrow_schema
+from pyflink.util.api_stability_decorators import PublicEvolving
+
+__all__ = ["CloseableIterator"]
+
+T = TypeVar("T")
+T_co = TypeVar("T_co", covariant=True)
+
+_BATCH_FORMATS = ("pandas", "pyarrow")
+
+
+@PublicEvolving()
+class CloseableIterator(Iterator[T_co], Protocol):
+ """
+ An iterator over the results of a running Flink job.
+
+ Closing the iterator cancels the job if it is still running. The iterator
closes itself once
+ the results are exhausted or iteration raises an error, so only iterators
abandoned before
+ their end need an explicit :meth:`close`. Prefer a ``with`` block: on an
unbounded source the
+ job keeps running until the iterator is closed.
+
+ Instances are returned by :meth:`DataFrame.iter_rows` and
:meth:`DataFrame.iter_batches`.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([{"id": 1}, {"id": 2}])
+ >>> with df.iter_rows() as rows:
+ ... first = next(rows)
+
+ .. versionadded:: 2.4.0
+ """
+
+ def close(self) -> None:
+ """
+ Close the iterator and cancel the job if it is still running. Closing
twice is a no-op.
+ """
+ ...
+
+ def __enter__(self) -> "CloseableIterator[T_co]":
+ ...
+
+ def __exit__(self, exc_type, exc_value, traceback) -> None:
+ ...
+
+
+class _CloseableIterator(Iterator[T]):
+ def __init__(self, iterator: Iterator[T], close: Callable[[], None]):
+ self._iterator = iterator
+ self._close = close
+ self._closed = False
+
+ def __iter__(self) -> "_CloseableIterator[T]":
+ return self
+
+ def __next__(self) -> T:
+ if self._closed:
+ raise StopIteration
+ try:
+ return next(self._iterator)
+ except BaseException:
+ # Covers StopIteration and interrupts too: either way the job is
no longer needed.
+ self.close()
+ raise
+
+ def close(self) -> None:
+ if not self._closed:
+ self._closed = True
+ self._close()
+
+ def __enter__(self) -> "_CloseableIterator[T]":
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback) -> None:
+ self.close()
+
+
+def _validate_row_kind_field(
+ columns: List[str], include_row_kind: bool, row_kind_field: str
+) -> None:
+ if not include_row_kind:
+ return
+ _require_non_empty_str(row_kind_field, "row_kind_field")
+ if row_kind_field in columns:
+ raise ValueError(
+ f"row_kind_field '{row_kind_field}' conflicts with an existing
column; "
+ "pass a different row_kind_field"
+ )
+
+
+def _row_to_dict(
+ row: Row, columns: List[str], row_kind_field: Optional[str]
+) -> Dict[str, Any]:
+ result = dict(zip(columns, row))
+ if row_kind_field is not None:
+ result[row_kind_field] = str(row.get_row_kind())
+ return result
+
+
+def _iterate_rows(
+ table, columns: List[str], row_kind_field: Optional[str]
+) -> _CloseableIterator[Dict[str, Any]]:
+ rows = table.execute().collect()
+ return _CloseableIterator(
+ (_row_to_dict(row, columns, row_kind_field) for row in rows),
rows.close
+ )
+
+
+def _build_arrow_schema(
+ columns: List[str], data_types: List[DataType], row_kind_field:
Optional[str]
+) -> Any:
+ """
+ Build the Arrow schema of a batch, failing on unsupported column types
before any job runs.
+ """
+ import pyarrow as pa
+
+ schema = create_arrow_schema(columns, data_types, allow_nested=True)
+ if row_kind_field is not None:
+ schema = schema.append(pa.field(row_kind_field, pa.utf8(),
nullable=False))
+ return schema
+
+
+def _iterate_batches(
+ table,
+ batch_size: int,
+ batch_format: str,
+ data_types: List[DataType],
+ arrow_schema: Any,
+ row_kind_field: Optional[str],
+) -> _CloseableIterator[Any]:
+ rows = table.execute().collect()
+
+ def batches():
+ while True:
+ chunk = list(itertools.islice(rows, batch_size))
+ if not chunk:
+ return
+ yield _rows_to_batch(chunk, data_types, arrow_schema,
batch_format, row_kind_field)
+
+ return _CloseableIterator(batches(), rows.close)
+
+
+def _take_rows(table, n: int, timeout: Optional[float]) -> List[Row]:
+ if n == 0:
+ return []
+ rows = table.execute().collect()
+ if timeout is None:
+ with rows:
+ return list(itertools.islice(rows, n))
+ return _take_rows_until(rows, n, time.monotonic() + timeout)
+
+
+def _take_rows_until(rows, n: int, deadline: float) -> List[Row]:
+ """
+ Read up to ``n`` rows, returning early with what has arrived when
``deadline`` passes.
+
+ Fetching a row blocks inside the JVM and cannot be interrupted from
Python, so a worker
+ thread fetches while this thread waits on a queue with a timeout. Closing
the iterator
+ cancels the job; the JVM fetcher polls the job status between fetch
attempts, so the worker
+ unblocks within one retry interval, either with an end-of-results or with
an error. The
+ worker is a daemon thread and holds its own gateway connection, so it
cannot stall exit if
+ the cluster never answers.
+ """
+ results: "queue.Queue" = queue.Queue()
+ closed = threading.Event()
+
+ def fetch():
+ try:
+ for row in itertools.islice(rows, n):
+ if closed.is_set():
+ return
+ results.put(("row", row))
+ if not closed.is_set():
+ results.put(("end", None))
+ except BaseException as error:
+ # After close() the job is being cancelled, so a failing fetch is
the expected way
+ # for the worker to unblock and nobody is waiting for its result
any more. The
+ # caller sets `closed` before calling close(), so an error caused
by the
+ # cancellation is never mistaken for a genuine fetch failure.
+ if not closed.is_set():
+ results.put(("error", error))
+
+ taken: List[Row] = []
+ try:
+ threading.Thread(target=fetch, name="dataframe-take",
daemon=True).start()
+ while len(taken) < n:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ break
+ try:
+ kind, value = results.get(timeout=remaining)
+ except queue.Empty:
+ break
+ if kind == "row":
+ taken.append(value)
+ elif kind == "end":
+ break
+ else:
+ raise value
+ finally:
+ closed.set()
+ rows.close()
+ return taken
+
+
+def _rows_to_batch(
+ rows: List[Row],
+ data_types: List[DataType],
+ arrow_schema: Any,
+ batch_format: str,
+ row_kind_field: Optional[str],
+) -> Any:
+ import pyarrow as pa
+
+ arrays = [
+ pa.array(
+ [_to_arrow_value(row[index], data_type) for row in rows],
+ type=arrow_schema.field(index).type,
+ )
+ for index, data_type in enumerate(data_types)
+ ]
+ if row_kind_field is not None:
+ arrays.append(pa.array([str(row.get_row_kind()) for row in rows],
type=pa.utf8()))
+ batch = pa.Table.from_arrays(arrays, schema=arrow_schema)
+ return batch.to_pandas() if batch_format == "pandas" else batch
+
+
+def _to_arrow_value(value: Any, data_type: DataType) -> Any:
+ """
+ Reshape a collected value into what pyarrow expects for ``data_type``.
+
+ Collected nested rows are positional, while pyarrow builds structs from
dicts keyed by
+ field name. Maps become key-value pair lists, which every supported
pyarrow version accepts.
+ """
+ if value is None:
+ return None
+ if isinstance(data_type, RowType):
+ return {
+ name: _to_arrow_value(field_value, field_type)
+ for name, field_value, field_type in zip(
+ data_type.field_names(), value, data_type.field_types()
+ )
+ }
+ if isinstance(data_type, ArrayType):
+ return [_to_arrow_value(element, data_type.element_type) for element
in value]
+ if isinstance(data_type, MapType):
+ return [
+ (
+ _to_arrow_value(key, data_type.key_type),
+ _to_arrow_value(item, data_type.value_type),
+ )
+ for key, item in value.items()
+ ]
+ return value
diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py
b/flink-python/pyflink/dataframe/tests/test_dataframe.py
index 0716d2cc16c..b2833cfd8b5 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -232,7 +232,7 @@ class DataFrameSlicingTests(unittest.TestCase):
def test_rejects_negative_values(self):
for method_name in ("limit", "offset", "head"):
with self.subTest(method=method_name):
- with self.assertRaisesRegex(ValueError, "n must be
non-negative"):
+ with self.assertRaisesRegex(ValueError, "n must be at least
0"):
getattr(self.dataframe, method_name)(-1)
self.table.fetch.assert_not_called()
diff --git a/flink-python/pyflink/dataframe/tests/test_iteration.py
b/flink-python/pyflink/dataframe/tests/test_iteration.py
new file mode 100644
index 00000000000..9753b38cc59
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_iteration.py
@@ -0,0 +1,342 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import time
+import unittest
+
+import pandas as pd
+import pyarrow as pa
+from py4j.protocol import Py4JJavaError
+
+import pyflink.dataframe as pf
+from pyflink.common import Row, RowKind
+from pyflink.dataframe import DataType
+from pyflink.dataframe.iteration import _build_arrow_schema,
_CloseableIterator, _rows_to_batch
+from pyflink.table import DataTypes
+from pyflink.testing.test_case_utils import PyFlinkStreamDataFrameTestCase
+
+
+def _row(*values, kind=RowKind.INSERT):
+ row = Row(*values)
+ row.set_row_kind(kind)
+ return row
+
+
+class CloseableIteratorTests(unittest.TestCase):
+ def test_closes_once_exhausted(self):
+ closes = []
+ iterator = _CloseableIterator(iter([1, 2]), lambda:
closes.append(True))
+
+ self.assertEqual(list(iterator), [1, 2])
+ self.assertEqual(closes, [True])
+
+ def test_closes_and_reraises_when_iteration_fails(self):
+ def failing():
+ yield 1
+ raise RuntimeError("fetch failed")
+
+ closes = []
+ iterator = _CloseableIterator(failing(), lambda: closes.append(True))
+
+ self.assertEqual(next(iterator), 1)
+ with self.assertRaisesRegex(RuntimeError, "fetch failed"):
+ next(iterator)
+ self.assertEqual(closes, [True])
+
+ def test_context_manager_closes_early_and_close_is_idempotent(self):
+ closes = []
+ with _CloseableIterator(iter([1, 2, 3]), lambda: closes.append(True))
as iterator:
+ self.assertEqual(next(iterator), 1)
+ iterator.close()
+
+ self.assertEqual(closes, [True])
+ self.assertEqual(list(iterator), [])
+
+ def test_row_kind_strings_match_the_documented_short_forms(self):
+ # The docstrings promise these exact strings; they come from
RowKind.__str__.
+ self.assertEqual([str(kind) for kind in RowKind], ["+I", "-U", "+U",
"-D"])
+
+
+class RowsToBatchTests(unittest.TestCase):
+ _COLUMNS = ["id", "name"]
+ _TYPES = [DataTypes.BIGINT(), DataTypes.STRING()]
+
+ def _batch(self, rows, batch_format="pyarrow", row_kind_field=None,
+ columns=None, data_types=None):
+ columns = columns or self._COLUMNS
+ data_types = data_types or self._TYPES
+ schema = _build_arrow_schema(columns, data_types, row_kind_field)
+ return _rows_to_batch(rows, data_types, schema, batch_format,
row_kind_field)
+
+ def test_arrow_batch_uses_the_flink_schema(self):
+ batch = self._batch([_row(1, "Alice")])
+
+ self.assertIsInstance(batch, pa.Table)
+ self.assertEqual(batch.schema.types, [pa.int64(), pa.utf8()])
+ self.assertEqual(batch.to_pylist(), [{"id": 1, "name": "Alice"}])
+
+ def test_pandas_batch_with_row_kind_column(self):
+ batch = self._batch(
+ [_row(1, "Alice"), _row(1, "Alice", kind=RowKind.DELETE)],
+ batch_format="pandas",
+ row_kind_field="op",
+ )
+
+ self.assertIsInstance(batch, pd.DataFrame)
+ self.assertEqual(list(batch.columns), ["id", "name", "op"])
+ self.assertEqual(list(batch["op"]), ["+I", "-D"])
+
+ def test_empty_batch_keeps_the_schema(self):
+ batch = self._batch([], row_kind_field="op")
+
+ self.assertEqual(batch.schema.names, ["id", "name", "op"])
+ self.assertEqual(batch.num_rows, 0)
+
+ def test_nested_values_become_arrow_structs_lists_and_maps(self):
+ batch = self._batch(
+ [_row(Row(1, ["a"]), {"k": 2}), _row(None, None)],
+ columns=["point", "counts"],
+ data_types=[
+ DataTypes.ROW([
+ DataTypes.FIELD("x", DataTypes.INT()),
+ DataTypes.FIELD("tags",
DataTypes.ARRAY(DataTypes.STRING())),
+ ]),
+ DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT()),
+ ],
+ )
+
+ self.assertEqual(
+ batch.to_pylist(),
+ [
+ {"point": {"x": 1, "tags": ["a"]}, "counts": [("k", 2)]},
+ {"point": None, "counts": None},
+ ],
+ )
+
+
+class IterationTests(PyFlinkStreamDataFrameTestCase):
+ def _people(self):
+ return pf.from_records(
+ [(i, f"name-{i}") for i in range(5)], schema=["id", "name"]
+ )
+
+ def _unbounded_source(self, rows_per_second):
+ return pf.read_generic(
+ "datagen",
+ schema={"x": DataType.int64()},
+ options={"rows-per-second": str(rows_per_second)},
+ )
+
+ def _failing_source(self):
+ # Random datagen strings are not numbers, so the cast fails as soon as
a row flows.
+ self.t_env.execute_sql(
+ "CREATE TEMPORARY TABLE IF NOT EXISTS random_strings (s STRING) "
+ "WITH ('connector' = 'datagen', 'rows-per-second' = '10')"
+ )
+ return pf.from_table(
+ self.t_env.sql_query("SELECT CAST(s AS INT) AS x FROM
random_strings")
+ )
+
+ def _job_count(self):
+ return len(self.resource.getMiniCluster().listJobs().get())
+
+ def _assert_all_jobs_terminate(self, timeout_seconds=60):
+ mini_cluster = self.resource.getMiniCluster()
+ deadline = time.monotonic() + timeout_seconds
+ while True:
+ running = [
+ job for job in mini_cluster.listJobs().get()
+ if not job.getJobState().isGloballyTerminalState()
+ ]
+ if not running:
+ return
+ if time.monotonic() > deadline:
+ self.fail(f"jobs still running: {[str(job.getJobId()) for job
in running]}")
+ time.sleep(0.1)
+
+ # ------------------------------------------------------------------
iter_rows
+
+ def test_iter_rows_yields_dicts_keyed_by_column(self):
+ with self._people().iter_rows() as rows:
+ collected = sorted(rows, key=lambda row: row["id"])
+
+ self.assertEqual(collected, [{"id": i, "name": f"name-{i}"} for i in
range(5)])
+
+ def test_iter_rows_row_kind_field_is_only_checked_when_included(self):
+ with self._people().iter_rows(row_kind_field="name") as rows:
+ self.assertEqual(len(list(rows)), 5)
+
+ def test_iter_rows_reraises_when_the_job_fails(self):
+ with self._failing_source().iter_rows() as rows:
+ with self.assertRaises(Py4JJavaError):
+ list(rows)
+ self._assert_all_jobs_terminate()
+
+ def test_closing_iter_rows_early_cancels_an_unbounded_job(self):
+ with self._unbounded_source(100).iter_rows() as rows:
+ next(rows)
+
+ self._assert_all_jobs_terminate()
+
+ # ------------------------------------------------------------------
iter_batches
+
+ def test_iter_batches_match_to_pandas(self):
+ dataframe = pf.from_records(
+ [(i, f"name-{i}", i * 1.5) for i in range(7)], schema=["id",
"name", "score"]
+ )
+
+ with dataframe.iter_batches(batch_size=3) as batches:
+ batches = list(batches)
+
+ self.assertEqual([len(batch) for batch in batches], [3, 3, 1])
+ pd.testing.assert_frame_equal(
+ pd.concat(batches).sort_values("id").reset_index(drop=True),
+ dataframe.to_pandas().sort_values("id").reset_index(drop=True),
+ )
+
+ def test_iter_batches_in_pyarrow_format(self):
+ with self._people().iter_batches(batch_format="pyarrow") as batches:
+ (batch,) = batches
+
+ self.assertIsInstance(batch, pa.Table)
+ self.assertEqual(batch.num_rows, 5)
+
+ def test_iter_batches_reraises_when_the_job_fails(self):
+ with self._failing_source().iter_batches(batch_size=1) as batches:
+ with self.assertRaises(Py4JJavaError):
+ next(batches)
+ self._assert_all_jobs_terminate()
+
+ # ------------------------------------------------------------------ take
/ take_batch
+
+ def test_take_returns_first_n_rows(self):
+ taken = self._people().take(2)
+
+ self.assertEqual(len(taken), 2)
+ self.assertTrue(all(row.keys() == {"id", "name"} for row in taken))
+
+ def test_take_returns_all_rows_when_result_is_shorter(self):
+ self.assertEqual(len(self._people().take(10)), 5)
+ self.assertEqual(len(self._people().take(10, timeout=60)), 5)
+
+ def test_take_zero_rows_does_not_execute(self):
+ jobs_before = self._job_count()
+
+ self.assertEqual(self._people().take(0), [])
+ self.assertEqual(self._job_count(), jobs_before)
+
+ def test_row_kinds_of_an_updating_result(self):
+ dataframe = pf.from_records([(1, "a"), (2, "a")], schema=["id",
"name"])
+ counts = dataframe.group_by("name").agg(c=pf.col("id").count)
+
+ self.assertEqual(
+ counts.take(10, include_row_kind=True, row_kind_field="op"),
+ [
+ {"name": "a", "c": 1, "op": "+I"},
+ {"name": "a", "c": 1, "op": "-U"},
+ {"name": "a", "c": 2, "op": "+U"},
+ ],
+ )
+
+ def test_take_cancels_an_unbounded_job(self):
+ self.assertEqual(len(self._unbounded_source(100).take(3)), 3)
+
+ self._assert_all_jobs_terminate()
+
+ def
test_take_with_timeout_returns_partial_rows_from_an_unbounded_job(self):
+ rows = self._unbounded_source(1).take(1000, timeout=2)
+
+ self.assertLess(len(rows), 1000)
+ self._assert_all_jobs_terminate()
+
+ def test_take_reraises_when_the_job_fails(self):
+ with self.assertRaises(Py4JJavaError):
+ self._failing_source().take(5)
+ with self.assertRaises(Py4JJavaError):
+ self._failing_source().take(5, timeout=60)
+ self._assert_all_jobs_terminate()
+
+ def test_take_batch_returns_first_n_rows_as_one_batch(self):
+ batch = self._people().take_batch(3)
+
+ self.assertIsInstance(batch, pd.DataFrame)
+ self.assertEqual(len(batch), 3)
+
+ def
test_take_batch_with_timeout_returns_partial_rows_from_an_unbounded_job(self):
+ batch = self._unbounded_source(1).take_batch(
+ 1000, timeout=2, batch_format="pyarrow", include_row_kind=True
+ )
+
+ self.assertLess(batch.num_rows, 1000)
+ self.assertEqual(batch.schema.names, ["x", "__row_kind__"])
+ self._assert_all_jobs_terminate()
+
+ def test_take_batch_with_nested_types(self):
+ dataframe = pf.from_table(self.t_env.sql_query(
+ "SELECT ROW(1, 'a') AS point, ARRAY[1, 2] AS items, MAP['k', 3] AS
counts"
+ ))
+
+ batch = dataframe.take_batch(1, batch_format="pyarrow")
+
+ self.assertEqual(
+ batch.to_pylist(),
+ [{"point": {"EXPR$0": 1, "EXPR$1": "a"}, "items": [1, 2],
"counts": [("k", 3)]}],
+ )
+
+ # ------------------------------------------------------------------
argument checks
+
+ def test_invalid_arguments_are_rejected_before_executing(self):
+ dataframe = self._people()
+ multiset = pf.from_table(
+ self.t_env.sql_query("SELECT COLLECT(name) AS tags FROM (VALUES
('a')) AS t(name)")
+ )
+ cases = [
+ (ValueError, "conflicts", lambda: dataframe.iter_rows(
+ include_row_kind=True, row_kind_field="name")),
+ (ValueError, "row_kind_field", lambda: dataframe.iter_rows(
+ include_row_kind=True, row_kind_field="")),
+ (TypeError, "row_kind_field", lambda: dataframe.take(
+ 1, include_row_kind=True, row_kind_field=1)),
+ (ValueError, "batch_size", lambda:
dataframe.iter_batches(batch_size=0)),
+ (TypeError, "batch_size", lambda:
dataframe.iter_batches(batch_size=True)),
+ (ValueError, "batch_format", lambda:
dataframe.iter_batches(batch_format="arrow")),
+ (ValueError, "batch_format", lambda: dataframe.take_batch(1,
batch_format="polars")),
+ (ValueError, "n must be at least 0", lambda: dataframe.take(-1)),
+ (TypeError, "n must be an integer", lambda: dataframe.take(1.5)),
+ (TypeError, "n must be an integer", lambda: dataframe.take(True)),
+ (ValueError, "timeout", lambda: dataframe.take(1, timeout=-1)),
+ (ValueError, "timeout must be finite", lambda: dataframe.take(
+ 1, timeout=float("inf"))),
+ (ValueError, "timeout must be finite", lambda:
dataframe.take_batch(
+ 1, timeout=float("nan"))),
+ (TypeError, "timeout", lambda: dataframe.take(1, timeout="1")),
+ (TypeError, "timeout", lambda: dataframe.take_batch(1,
timeout=True)),
+ (ValueError, "not supported", lambda: multiset.iter_batches()),
+ (ValueError, "not supported", lambda: multiset.take_batch(1)),
+ ]
+ jobs_before = self._job_count()
+
+ for error, message, call in cases:
+ with self.subTest(message=message):
+ with self.assertRaisesRegex(error, message):
+ call()
+ self.assertEqual(self._job_count(), jobs_before)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/flink-python/pyflink/dataframe/validation.py
b/flink-python/pyflink/dataframe/validation.py
new file mode 100644
index 00000000000..49cc95d03aa
--- /dev/null
+++ b/flink-python/pyflink/dataframe/validation.py
@@ -0,0 +1,56 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+"""
+Argument checks shared by the DataFrame API.
+
+A wrong type raises :class:`TypeError` and a wrong value of the right type
raises
+:class:`ValueError`. ``bool`` is rejected wherever a number is expected, since
it is an ``int``
+subclass and ``take(True)`` is almost certainly a mistake.
+"""
+
+import math
+from typing import Any, Sequence
+
+
+def _require_int(value: Any, name: str, minimum: int) -> None:
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise TypeError(f"{name} must be an integer")
+ if value < minimum:
+ raise ValueError(f"{name} must be at least {minimum}, got {value}")
+
+
+def _require_number(value: Any, name: str, minimum: float) -> None:
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise TypeError(f"{name} must be a number")
+ if not math.isfinite(value):
+ raise ValueError(f"{name} must be finite, got {value}")
+ if value < minimum:
+ raise ValueError(f"{name} must be at least {minimum}, got {value}")
+
+
+def _require_non_empty_str(value: Any, name: str) -> None:
+ if not isinstance(value, str):
+ raise TypeError(f"{name} must be a string")
+ if not value:
+ raise ValueError(f"{name} must not be empty")
+
+
+def _require_choice(value: Any, name: str, choices: Sequence[str]) -> None:
+ if value not in choices:
+ raise ValueError(f"{name} must be one of {list(choices)}, got
{value!r}")
