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 f83a73c19a4 [FLINK-40421][python] Add sorting APIs to DataFrame API
f83a73c19a4 is described below
commit f83a73c19a41a066360cb353877b1ca145a18929
Author: Milesian111 <[email protected]>
AuthorDate: Thu Sep 3 17:09:23 2026 +0800
[FLINK-40421][python] Add sorting APIs to DataFrame API
This closes #29078.
Generated-by: Codex GPT-5
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 1 +
flink-python/pyflink/dataframe/dataframe.py | 130 ++++++++++++++++++-
.../pyflink/dataframe/tests/test_dataframe.py | 144 ++++++++++++++++++++-
3 files changed, 270 insertions(+), 5 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index d43dd2a2be4..ca7cecc499f 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -63,6 +63,7 @@ Transformations
DataFrame.drop_duplicates
DataFrame.distinct
DataFrame.unique
+ DataFrame.sort
DataFrame.top_n
DataFrame.limit
DataFrame.offset
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index 38f1971abcf..ed2cf3dab1c 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -36,6 +36,7 @@ if TYPE_CHECKING:
from pyflink.common import Row
from pyflink.dataframe.datatype import _INT_MAX, DataType
+from pyflink.java_gateway import get_gateway
from pyflink.table.expression import Expression
from pyflink.table.expressions import (
and_,
@@ -620,6 +621,80 @@ class DataFrame:
distinct = drop_duplicates
unique = drop_duplicates
+ # ======================== Filtering & Ordering ========================
+
+ @PublicEvolving()
+ def sort(
+ self,
+ by: Union[str, Expression, List[Union[str, Expression]]],
+ *,
+ descending: Union[bool, List[bool]] = False,
+ nulls_first: Union[bool, List[bool]] = None,
+ ) -> "DataFrame":
+ """
+ Sort rows globally by one or more columns or expressions.
+
+ This method builds a new DataFrame plan without executing a Flink job.
The ``by``
+ expressions must not already specify ``asc`` or ``desc``; use
``descending`` to control
+ their direction. When ``nulls_first`` is omitted, the Table API
default is used: NULLs
+ are ordered last for ascending keys and first for descending keys.
+
+ The result is globally sorted across all parallel partitions. For
unbounded tables, the
+ first sort key must be an ascending time attribute unless the sort is
followed by
+ :meth:`limit`.
+
+ :param by: Column name or expression, or a list of them, used as sort
keys.
+ :param descending: Whether to sort in descending order, either for all
keys or once per
+ key.
+ :param nulls_first: Whether to place NULLs first, either for all keys
or once per key. When
+ omitted, the Table API default applies.
+ :return: A new sorted DataFrame.
+ :raises TypeError: If ``by``, ``descending`` or ``nulls_first`` has an
unsupported type.
+ :raises ValueError: If ``by`` is empty, option lengths do not match, a
column does not
+ exist, or an expression already specifies ``asc`` or ``desc``.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records(
+ ... [(2, "b"), (1, "a")], schema=["id", "name"]
+ ... )
+ >>> ascending = df.sort("id")
+ >>> mixed = df.sort(["id", "name"], descending=[False, True])
+
+ .. versionadded:: 2.4.0
+ """
+ order_keys = _normalize_order_by(by, "by")
+ if order_keys is None:
+ raise TypeError("by must be a string, an expression, or a list or
tuple of them")
+ columns = self._table.get_resolved_schema().get_column_names()
+ for key in order_keys:
+ if isinstance(key, str) and key not in columns:
+ raise ValueError(
+ "by column '%s' does not exist, available columns: %s" %
(key, columns)
+ )
+ if isinstance(key, Expression) and
_contains_ordering_expression(key):
+ raise ValueError(
+ "sort() expressions must not specify asc or desc; use
descending instead"
+ )
+
+ descending_values = _normalize_descending(descending, len(order_keys))
+ nulls_values: List[Optional[bool]] = (
+ [None] * len(order_keys)
+ if nulls_first is None
+ else _normalize_nulls_first(nulls_first, len(order_keys))
+ )
+ if any(value is not None for value in nulls_values):
+ return DataFrame(
+ _build_sort_sql(self._table, order_keys, descending_values,
nulls_values)
+ )
+
+ order_expressions = []
+ for key, is_descending in zip(order_keys, descending_values):
+ expression = table_col(key) if isinstance(key, str) else key
+ order_expressions.append(expression.desc if is_descending else
expression.asc)
+ return DataFrame(self._table.order_by(*order_expressions))
+
# ======================== Slicing ========================
@PublicEvolving()
@@ -1090,6 +1165,7 @@ def _normalize_subset(subset: Union[str, List[str],
None]) -> Optional[List[str]
def _normalize_order_by(
order_by: Union[str, Expression, List[Union[str, Expression]], None],
+ parameter_name: str = "order_by",
) -> Optional[List[Union[str, Expression]]]:
if order_by is None:
return None
@@ -1101,15 +1177,30 @@ def _normalize_order_by(
keys.append(value)
else:
raise TypeError(
- "order_by must be a string, an expression, or a list or tuple
of them"
+ "%s must be a string, an expression, or a list or tuple of
them" % parameter_name
)
if not keys:
- raise ValueError("order_by must not be empty")
+ raise ValueError("%s must not be empty" % parameter_name)
return keys
+def _contains_ordering_expression(expression: Expression) -> bool:
+ gateway = get_gateway()
+ api_expression_utils =
gateway.jvm.org.apache.flink.table.expressions.ApiExpressionUtils
+ built_in_functions =
gateway.jvm.org.apache.flink.table.functions.BuiltInFunctionDefinitions
+
+ def contains_ordering(j_expression) -> bool:
+ if api_expression_utils.isFunction(
+ j_expression, built_in_functions.ORDER_ASC
+ ) or api_expression_utils.isFunction(j_expression,
built_in_functions.ORDER_DESC):
+ return True
+ return any(contains_ordering(child) for child in
j_expression.getChildren())
+
+ return contains_ordering(expression._j_expr.toExpr())
+
+
def _normalize_nulls_first(
nulls_first: Union[bool, List[bool], None], order_len: int
) -> Optional[List[bool]]:
@@ -1127,7 +1218,7 @@ def _normalize_nulls_first(
raise TypeError("nulls_first must be a boolean or a list of booleans")
if len(values) != order_len:
- raise ValueError("nulls_first must have the same length as order_by")
+ raise ValueError("nulls_first must have the same length as the sort
keys")
return values
@@ -1137,13 +1228,44 @@ def _normalize_descending(descending, order_len):
return [descending] * order_len
if isinstance(descending, (list, tuple)):
if len(descending) != order_len:
- raise ValueError("descending must have the same length as
order_by")
+ raise ValueError("descending must have the same length as the sort
keys")
if not all(isinstance(v, bool) for v in descending):
raise TypeError("descending must be a boolean or a list of
booleans")
return list(descending)
raise TypeError("descending must be a boolean or a list of booleans")
+def _build_sort_sql(table, order_keys, descending_flags, nulls) -> Table:
+ columns = table.get_resolved_schema().get_column_names()
+ taken = set(columns)
+ order_terms = []
+ for index, key in enumerate(order_keys):
+ direction = "DESC" if descending_flags[index] else "ASC"
+ if isinstance(key, str):
+ if key not in columns:
+ raise ValueError(
+ "by column '%s' does not exist, available columns: %s" %
(key, columns)
+ )
+ expression_sql = _quote_identifier(key)
+ else:
+ name = _unique_name("__pf_order_%d" % index, taken)
+ taken.add(name)
+ table = table.add_columns(key.alias(name))
+ expression_sql = _quote_identifier(name)
+ term = "%s %s" % (expression_sql, direction)
+ if nulls[index] is not None:
+ term += " NULLS FIRST" if nulls[index] else " NULLS LAST"
+ order_terms.append(term)
+
+ select_list = ", ".join(_quote_identifier(name) for name in columns)
+ query = "SELECT %s FROM %s ORDER BY %s" % (
+ select_list,
+ _quote_identifier(str(table)),
+ ", ".join(order_terms),
+ )
+ return table._t_env.sql_query(query)
+
+
def _build_rank_sql(table, partition_keys, order_keys, descending_flags,
nulls, n) -> Table:
columns = table.get_resolved_schema().get_column_names()
for name in partition_keys:
diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py
b/flink-python/pyflink/dataframe/tests/test_dataframe.py
index fd8f0dfe048..d51d82c857d 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -243,6 +243,104 @@ class DataFrameSlicingTests(unittest.TestCase):
self.table.execute.assert_not_called()
+class DataFrameSortingTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ self.dataframe = pf.from_records(
+ [(2, "b"), (1, "a")],
+ schema=["id", "name"],
+ )
+
+ def test_sort_is_lazy_and_returns_new_dataframe(self):
+ result = self.dataframe.sort("id")
+
+ self.assertIsInstance(result, pf.DataFrame)
+ self.assertIsNot(result, self.dataframe)
+ self.assertIsNot(result.to_table(), self.dataframe.to_table())
+ self.assertEqual(
+
result.to_table()._j_table.getQueryOperation().getOrder().toString(),
+ "[asc(id)]",
+ )
+
+ def test_sort_supports_descending_and_multiple_keys(self):
+ result = self.dataframe.sort(["id", "name"], descending=[True, False])
+ all_descending = self.dataframe.sort(["id", "name"], descending=True)
+
+ self.assertEqual(
+
result.to_table()._j_table.getQueryOperation().getOrder().toString(),
+ "[desc(id), asc(name)]",
+ )
+ self.assertEqual(
+
all_descending.to_table()._j_table.getQueryOperation().getOrder().toString(),
+ "[desc(id), desc(name)]",
+ )
+
+ def test_sort_supports_expression_keys(self):
+ result = self.dataframe.sort(pf.col("id") + 1, descending=True)
+
+ self.assertEqual(
+
result.to_table()._j_table.getQueryOperation().getOrder().toString(),
+ "[desc(plus(id, 1))]",
+ )
+
+ def test_sort_with_null_ordering_preserves_temporal_sort(self):
+ self.t_env.execute_sql(
+ """
+ CREATE TEMPORARY TABLE sort_source (
+ id INT,
+ ts TIMESTAMP(3),
+ WATERMARK FOR ts AS ts - INTERVAL '5' SECOND
+ ) WITH ('connector' = 'datagen', 'number-of-rows' = '1')
+ """
+ )
+ dataframe = pf.from_table(self.t_env.from_path("sort_source"))
+
+ plan = dataframe.sort("ts", nulls_first=True).to_table().explain()
+
+ self.assertIn("LogicalSort(sort0=[$1], dir0=[ASC-nulls-first])", plan)
+ self.assertIn("TemporalSort(orderBy=[ts ASC])", plan)
+
+ def test_sort_rejects_ordered_expressions(self):
+ for expression in (
+ pf.col("id").asc,
+ pf.col("id").desc,
+ pf.col("id").asc + 1,
+ ):
+ with self.subTest(expression=str(expression)):
+ with self.assertRaisesRegex(
+ ValueError,
+ "expressions must not specify asc or desc",
+ ):
+ self.dataframe.sort(expression)
+
+ def test_sort_rejects_invalid_keys_and_options(self):
+ with self.assertRaisesRegex(ValueError, "by must not be empty"):
+ self.dataframe.sort([])
+ with self.assertRaisesRegex(ValueError, "by column 'missing' does not
exist"):
+ self.dataframe.sort("missing")
+ with self.assertRaisesRegex(TypeError, "by must be a string, an
expression"):
+ self.dataframe.sort(1)
+ with self.assertRaisesRegex(TypeError, "by must be a string, an
expression"):
+ self.dataframe.sort(None)
+
+ with self.assertRaisesRegex(TypeError, "descending must be a boolean"):
+ self.dataframe.sort("id", descending=1)
+ with self.assertRaisesRegex(TypeError, "descending must be a boolean"):
+ self.dataframe.sort(["id", "name"], descending=[True, 1])
+ with self.assertRaisesRegex(
+ ValueError, "descending must have the same length as the sort keys"
+ ):
+ self.dataframe.sort(["id", "name"], descending=[True])
+ with self.assertRaisesRegex(TypeError, "nulls_first must be a
boolean"):
+ self.dataframe.sort("id", nulls_first=1)
+ with self.assertRaisesRegex(TypeError, "nulls_first must be a
boolean"):
+ self.dataframe.sort(["id", "name"], nulls_first=[True, 1])
+ with self.assertRaisesRegex(
+ ValueError, "nulls_first must have the same length as the sort
keys"
+ ):
+ self.dataframe.sort(["id", "name"], nulls_first=[True])
+
+
class DataFrameCreationTests(PyFlinkDataFrameUTTestCase):
def test_from_dict_uses_insertion_order_without_schema(self):
dataframe = pf.from_dict({"name": ["Alice"], "id": [1]})
@@ -1408,7 +1506,7 @@ class
DataFrameDropDuplicatesTests(PyFlinkDataFrameUTTestCase):
"id", order_by="score", nulls_first=[True, False]
)
self.assertEqual(
- str(error.exception), "nulls_first must have the same length as
order_by"
+ str(error.exception), "nulls_first must have the same length as
the sort keys"
)
def test_rejects_nulls_first_wrong_type(self):
@@ -1976,6 +2074,50 @@ class DataFrameBatchITTests(PyFlinkITTestCase):
)
return pf.from_table(table.order_by(table.id))
+ def _unsorted_dataframe(self):
+ table = self.t_env.sql_query(
+ "SELECT * FROM (VALUES (3, 'C'), (1, 'A'), (2, 'B')) AS T(id,
name)"
+ )
+ return pf.from_table(table)
+
+ def _nullable_dataframe(self):
+ table = self.t_env.sql_query(
+ "SELECT * FROM (VALUES (CAST(NULL AS INT), 'NULL'), (2, 'B'), (1,
'A')) "
+ "AS T(id, name)"
+ )
+ return pf.from_table(table)
+
+ def test_sort_returns_rows_in_ascending_order(self):
+ self.assertEqual(
+ self._unsorted_dataframe().sort("id").collect(),
+ [Row(1, "A"), Row(2, "B"), Row(3, "C")],
+ )
+
+ def test_sort_supports_per_key_descending_order(self):
+ dataframe = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (1, 10, 'A'), (1, 20, 'B'), (2, 5,
'C')) "
+ "AS T(group_id, score, name)"
+ )
+ )
+
+ self.assertEqual(
+ dataframe.sort(["group_id", "score"], descending=[False,
True]).collect(),
+ [Row(1, 20, "B"), Row(1, 10, "A"), Row(2, 5, "C")],
+ )
+
+ def test_sort_supports_explicit_null_ordering(self):
+ dataframe = self._nullable_dataframe()
+
+ self.assertEqual(
+ dataframe.sort("id", descending=True, nulls_first=False).collect(),
+ [Row(2, "B"), Row(1, "A"), Row(None, "NULL")],
+ )
+ self.assertEqual(
+ dataframe.sort(pf.col("id") + 1, nulls_first=True).collect(),
+ [Row(None, "NULL"), Row(1, "A"), Row(2, "B")],
+ )
+
def test_limit_returns_first_rows(self):
self.assertEqual(
self._ordered_dataframe().limit(2).collect(),
