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 008e440edc9 [FLINK-40422][python] Add slicing APIs to DataFrame API
008e440edc9 is described below
commit 008e440edc9203daa4cf7b43ee3f5a691cd92a0c
Author: Milesian111 <[email protected]>
AuthorDate: Fri Aug 28 14:39:06 2026 +0800
[FLINK-40422][python] Add slicing APIs to DataFrame API
This closes #29033.
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 7 +-
flink-python/pyflink/dataframe/convert.py | 3 +-
flink-python/pyflink/dataframe/dataframe.py | 92 ++++++++++++++-
flink-python/pyflink/dataframe/datatype.py | 4 +
.../pyflink/dataframe/tests/test_dataframe.py | 126 +++++++++++++++++++++
5 files changed, 228 insertions(+), 4 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index dc197a55e9a..4049d3dedd0 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -21,7 +21,9 @@ DataFrame
=========
A DataFrame provides a Pythonic interface for composing data transformations.
-Transformation methods return new DataFrames and support fluent chaining.
+Transformation methods return new DataFrames and support fluent chaining. They
build execution
+plans lazily without starting a Flink job; execution is triggered by an action
such as
+``DataFrame.collect`` or ``DataFrame.to_pandas``.
Example::
@@ -61,6 +63,9 @@ Transformations
DataFrame.drop_duplicates
DataFrame.distinct
DataFrame.unique
+ DataFrame.limit
+ DataFrame.offset
+ DataFrame.head
DataFrame.__getitem__
Aggregations
diff --git a/flink-python/pyflink/dataframe/convert.py
b/flink-python/pyflink/dataframe/convert.py
index a7265466dd1..84b31ba8f0d 100644
--- a/flink-python/pyflink/dataframe/convert.py
+++ b/flink-python/pyflink/dataframe/convert.py
@@ -38,6 +38,7 @@ if TYPE_CHECKING:
from pyflink.dataframe.context import get_or_create_table_environment
from pyflink.dataframe.dataframe import DataFrame
+from pyflink.dataframe.datatype import _BIGINT_MAX, _BIGINT_MIN
from pyflink.table import Schema, Table
from pyflink.table.types import (
_create_converter,
@@ -63,8 +64,6 @@ __all__ = [
]
_SCALAR_SEQUENCE_TYPES = (str, bytes, bytearray, memoryview)
-_BIGINT_MIN = -(1 << 63)
-_BIGINT_MAX = (1 << 63) - 1
class _WatermarkSpec(NamedTuple):
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index 0e542d0ab29..55b5e1610b6 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -35,7 +35,7 @@ if TYPE_CHECKING:
from pyflink.table.table_schema import TableSchema
from pyflink.common import Row
-from pyflink.dataframe.datatype import DataType
+from pyflink.dataframe.datatype import _INT_MAX, DataType
from pyflink.table.expression import Expression
from pyflink.table.expressions import (
and_,
@@ -51,6 +51,15 @@ __all__ = ["DataFrame", "GroupedDataFrame", "col", "lit"]
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")
+ if n > _INT_MAX:
+ raise ValueError(f"n must be less than or equal to {_INT_MAX}")
+
+
@PublicEvolving()
def col(name: str) -> Expression:
"""
@@ -542,6 +551,87 @@ class DataFrame:
distinct = drop_duplicates
unique = drop_duplicates
+ # ======================== Slicing ========================
+
+ @PublicEvolving()
+ def limit(self, n: int) -> "DataFrame":
+ """
+ Keep at most the first ``n`` rows.
+
+ This method builds a new DataFrame plan without executing a Flink job.
Execution is
+ triggered by an action such as :meth:`collect` or :meth:`to_pandas`.
Without an explicit
+ ordering on the underlying table, the selected rows and their order
are unspecified.
+ Changes to the underlying table content may also change the result.
+
+ :param n: Maximum number of rows to keep.
+ :return: A new DataFrame containing at most ``n`` rows.
+ :raises TypeError: If ``n`` is not an integer.
+ :raises ValueError: If ``n`` is negative.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([{"id": 1}, {"id": 2}, {"id": 3}])
+ >>> first_two = df.limit(2)
+
+ .. versionadded:: 2.4.0
+ """
+ _validate_row_count(n)
+ return DataFrame(self._table.fetch(n))
+
+ @PublicEvolving()
+ def offset(self, n: int) -> "DataFrame":
+ """
+ Skip the first ``n`` rows.
+
+ This method builds a new DataFrame plan without executing a Flink job.
Execution is
+ triggered by an action such as :meth:`collect` or :meth:`to_pandas`.
Without an explicit
+ ordering on the underlying table, the skipped rows and their order are
unspecified.
+ Changes to the underlying table content may also change the result.
Combine this method
+ with :meth:`limit` for pagination.
+
+ :param n: Number of rows to skip.
+ :return: A new DataFrame without the first ``n`` rows.
+ :raises TypeError: If ``n`` is not an integer.
+ :raises ValueError: If ``n`` is negative.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([{"id": 1}, {"id": 2}, {"id": 3}])
+ >>> page = df.offset(1).limit(2)
+
+ .. versionadded:: 2.4.0
+ """
+ _validate_row_count(n)
+ return DataFrame(self._table.offset(n))
+
+ @PublicEvolving()
+ def head(self, n: int) -> "DataFrame":
+ """
+ Keep at most the first ``n`` rows.
+
+ This method builds a new DataFrame plan without executing a Flink job
and delegates to
+ :meth:`limit`. Execution is triggered by an action such as
:meth:`collect` or
+ :meth:`to_pandas`. Without an explicit ordering on the underlying
table, the selected rows
+ and their order are unspecified. Changes to the underlying table
content may also change
+ the result.
+
+ :param n: Maximum number of rows to keep.
+ :return: A new DataFrame containing at most ``n`` rows.
+ :raises TypeError: If ``n`` is not an integer.
+ :raises ValueError: If ``n`` is negative.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([{"id": 1}, {"id": 2}, {"id": 3}])
+ >>> first_two = df.head(2)
+
+ .. versionadded:: 2.4.0
+ """
+ return self.limit(n)
+
# ======================== Aggregation ========================
@PublicEvolving()
diff --git a/flink-python/pyflink/dataframe/datatype.py
b/flink-python/pyflink/dataframe/datatype.py
index 84676885a6f..d604f2c7fda 100644
--- a/flink-python/pyflink/dataframe/datatype.py
+++ b/flink-python/pyflink/dataframe/datatype.py
@@ -27,6 +27,10 @@ from pyflink.util.api_stability_decorators import
PublicEvolving
__all__ = ["DataType"]
+_INT_MIN = -(1 << 31)
+_INT_MAX = (1 << 31) - 1
+_BIGINT_MIN = -(1 << 63)
+_BIGINT_MAX = (1 << 63) - 1
_PEP_604_UNION_TYPE = getattr(types, "UnionType", None)
_BASIC_TYPE_HINT_FACTORIES: Dict[Any, Callable[[], TableDataType]] = {
diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py
b/flink-python/pyflink/dataframe/tests/test_dataframe.py
index dc34f28229e..36d5c8efeda 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -23,6 +23,7 @@ import unittest
from py4j.protocol import Py4JJavaError
from datetime import date, datetime, time, timedelta, timezone
from typing import NamedTuple
+from unittest.mock import Mock, patch
import pandas as pd
import pyarrow as pa
@@ -151,6 +152,97 @@ class DataFrameCompositionTests(unittest.TestCase):
self.assertIs(pf.DataFrame.rename, pf.DataFrame.rename_columns)
+class DataFrameSlicingTests(unittest.TestCase):
+ def setUp(self):
+ self.table = Mock()
+ self.dataframe = pf.DataFrame(self.table)
+
+ def test_limit_is_lazy_and_returns_new_dataframe(self):
+ limited_table = Mock()
+ self.table.fetch.return_value = limited_table
+
+ result = self.dataframe.limit(3)
+
+ self.assertIsInstance(result, pf.DataFrame)
+ self.assertIs(result.to_table(), limited_table)
+ self.assertIs(self.dataframe.to_table(), self.table)
+ self.table.fetch.assert_called_once_with(3)
+ self.table.execute.assert_not_called()
+
+ def test_offset_is_lazy_and_returns_new_dataframe(self):
+ offset_table = Mock()
+ self.table.offset.return_value = offset_table
+
+ result = self.dataframe.offset(2)
+
+ self.assertIsInstance(result, pf.DataFrame)
+ self.assertIs(result.to_table(), offset_table)
+ self.assertIs(self.dataframe.to_table(), self.table)
+ self.table.offset.assert_called_once_with(2)
+ self.table.execute.assert_not_called()
+
+ def test_offset_and_limit_compose(self):
+ offset_table = Mock()
+ limited_table = Mock()
+ self.table.offset.return_value = offset_table
+ offset_table.fetch.return_value = limited_table
+
+ result = self.dataframe.offset(2).limit(3)
+
+ self.assertIs(result.to_table(), limited_table)
+ self.table.offset.assert_called_once_with(2)
+ offset_table.fetch.assert_called_once_with(3)
+ self.table.execute.assert_not_called()
+
+ def test_head_delegates_to_limit(self):
+ expected = pf.DataFrame(Mock())
+
+ with patch.object(pf.DataFrame, "limit", autospec=True) as limit:
+ limit.return_value = expected
+
+ result = self.dataframe.head(3)
+
+ self.assertIs(result, expected)
+ limit.assert_called_once_with(self.dataframe, 3)
+ self.table.execute.assert_not_called()
+
+ def test_zero_is_supported(self):
+ limited_table = Mock()
+ offset_table = Mock()
+ self.table.fetch.return_value = limited_table
+ self.table.offset.return_value = offset_table
+
+ self.assertIs(self.dataframe.limit(0).to_table(), limited_table)
+ self.assertIs(self.dataframe.head(0).to_table(), limited_table)
+ self.assertIs(self.dataframe.offset(0).to_table(), offset_table)
+
+ self.assertEqual(self.table.fetch.call_count, 2)
+ self.table.fetch.assert_called_with(0)
+ self.table.offset.assert_called_once_with(0)
+ self.table.execute.assert_not_called()
+
+ 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"):
+ getattr(self.dataframe, method_name)(-1)
+
+ self.table.fetch.assert_not_called()
+ self.table.offset.assert_not_called()
+ self.table.execute.assert_not_called()
+
+ def test_rejects_unsupported_types(self):
+ for method_name in ("limit", "offset", "head"):
+ for value in (True, 1.5, "1", None):
+ with self.subTest(method=method_name, value=value):
+ with self.assertRaisesRegex(TypeError, "n must be an
integer"):
+ getattr(self.dataframe, method_name)(value)
+
+ self.table.fetch.assert_not_called()
+ self.table.offset.assert_not_called()
+ self.table.execute.assert_not_called()
+
+
class DataFrameCreationTests(PyFlinkDataFrameUTTestCase):
def test_from_dict_uses_insertion_order_without_schema(self):
dataframe = pf.from_dict({"name": ["Alice"], "id": [1]})
@@ -1804,6 +1896,40 @@ class DataFrameBatchITTests(PyFlinkITTestCase):
self.addCleanup(pf.set_table_environment, previous_environment)
self.t_env =
TableEnvironment.create(EnvironmentSettings.in_batch_mode())
+ def _ordered_dataframe(self):
+ table = self.t_env.sql_query(
+ "SELECT * FROM (VALUES (3, 'C'), (1, 'A'), (4, 'D'), (2, 'B')) "
+ "AS T(id, name)"
+ )
+ return pf.from_table(table.order_by(table.id))
+
+ def test_limit_returns_first_rows(self):
+ self.assertEqual(
+ self._ordered_dataframe().limit(2).collect(),
+ [Row(1, "A"), Row(2, "B")],
+ )
+
+ def test_offset_and_limit_compose_for_pagination(self):
+ self.assertEqual(
+ self._ordered_dataframe().offset(1).limit(2).collect(),
+ [Row(2, "B"), Row(3, "C")],
+ )
+
+ def test_head_and_limit_are_equivalent(self):
+ dataframe = self._ordered_dataframe()
+
+ self.assertEqual(dataframe.head(3).collect(),
dataframe.limit(3).collect())
+
+ def test_zero_slicing(self):
+ dataframe = self._ordered_dataframe()
+
+ self.assertEqual(dataframe.limit(0).collect(), [])
+ self.assertEqual(dataframe.head(0).collect(), [])
+ self.assertEqual(
+ dataframe.offset(0).collect(),
+ [Row(1, "A"), Row(2, "B"), Row(3, "C"), Row(4, "D")],
+ )
+
def test_from_records_with_batch_table_environment(self):
pf.set_table_environment(self.t_env)
