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 52659086b59 [FLINK-40194][python] Add drop_duplicates to the DataFrame
API (#29022)
52659086b59 is described below
commit 52659086b597bc23705d80c237d3bc0938410e76
Author: Timo Theusner <[email protected]>
AuthorDate: Thu Aug 27 14:11:46 2026 +0200
[FLINK-40194][python] Add drop_duplicates to the DataFrame API (#29022)
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 3 +
flink-python/pyflink/dataframe/dataframe.py | 204 ++++++++++
.../pyflink/dataframe/tests/test_dataframe.py | 431 ++++++++++++++++++++-
flink-python/pyflink/testing/test_case_utils.py | 22 +-
4 files changed, 657 insertions(+), 3 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index 4c78cb5d857..dc197a55e9a 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -58,6 +58,9 @@ Transformations
DataFrame.rename
DataFrame.filter
DataFrame.where
+ DataFrame.drop_duplicates
+ DataFrame.distinct
+ DataFrame.unique
DataFrame.__getitem__
Aggregations
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index ac4c12d24d5..0e542d0ab29 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -23,6 +23,7 @@ from typing import (
Dict,
List,
Optional,
+ Set,
Tuple,
TypeVar,
Union,
@@ -464,6 +465,83 @@ class DataFrame:
return DataFrame(self._table.select(*expressions))
+ @PublicEvolving()
+ def drop_duplicates(
+ self,
+ subset: Union[str, List[str]] = None,
+ *,
+ keep: str = "first",
+ order_by: Union[str, Expression, List[Union[str, Expression]]] = None,
+ nulls_first: Union[bool, List[bool]] = None,
+ ) -> "DataFrame":
+ """
+ Remove duplicate rows.
+
+ When ``subset`` is omitted, fully identical rows are dropped,
equivalent to a whole-row
+ ``DISTINCT``. When ``subset`` is given, rows are deduplicated by those
key columns, keeping
+ one row per key; ``order_by`` together with ``keep`` decides which row
survives. When
+ ``order_by`` is omitted, processing time is used, so ``keep`` keeps
the first or last row to
+ arrive.
+
+ :param subset: Column name or list of column names that define a
duplicate. When omitted,
+ all columns are considered.
+ :param keep: ``"first"`` keeps the earliest row, ``"last"`` the
latest, in ``order_by``
+ order. Ignored when ``subset`` is omitted.
+ :param order_by: Column name or expression (or a list of them)
defining the order in which
+ ``keep`` selects the surviving row. When omitted, processing time
is used.
+ :param nulls_first: Where NULLs rank in ``order_by``: a single boolean
applied to every key,
+ or a list with one boolean per key. When omitted, the engine
default applies. Requires
+ ``order_by``.
+ :return: A new DataFrame with duplicate rows removed.
+ :raises ValueError: If ``keep`` is not ``"first"`` or ``"last"``, if
``order_by`` or
+ ``nulls_first`` is combined with an omitted ``subset``, if
``nulls_first`` is given
+ without ``order_by`` or with a mismatched length, or if a named
column does not exist.
+ :raises TypeError: If ``subset``, ``order_by`` or ``nulls_first`` has
an unsupported type.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([{"id": 1, "ts": 1}, {"id": 1, "ts": 2}])
+ >>> unique_rows = df.drop_duplicates()
+ >>> latest_per_id = df.drop_duplicates("id", order_by="ts",
keep="last")
+
+ .. versionadded:: 2.4.0
+ """
+ if keep not in ("first", "last"):
+ raise ValueError('keep must be "first" or "last"')
+
+ subset_keys = _normalize_subset(subset)
+ order_keys = _normalize_order_by(order_by)
+
+ if nulls_first is not None and order_keys is None:
+ raise ValueError("nulls_first requires order_by")
+ nulls = _normalize_nulls_first(
+ nulls_first, len(order_keys) if order_keys else 0
+ )
+
+ if subset_keys is None:
+ if order_keys is not None:
+ raise ValueError(
+ "order_by requires subset; whole-row duplicates cannot be
ordered"
+ )
+ return DataFrame(self._table.distinct())
+
+ columns = self._table.get_resolved_schema().get_column_names()
+ for name in subset_keys:
+ if name not in columns:
+ raise ValueError(
+ "subset column '%s' does not exist, available columns: %s"
% (name, columns)
+ )
+
+ return DataFrame(
+ _build_deduplication_query(
+ self._table, columns, subset_keys, order_keys, keep, nulls
+ )
+ )
+
+ distinct = drop_duplicates
+ unique = drop_duplicates
+
# ======================== Aggregation ========================
@PublicEvolving()
@@ -836,6 +914,132 @@ class GroupedDataFrame:
# ======================== Internal Helpers ========================
+def _normalize_subset(subset: Union[str, List[str], None]) ->
Optional[List[str]]:
+ if subset is None:
+ return None
+ if isinstance(subset, str):
+ return [subset]
+ if isinstance(subset, (list, tuple)):
+ if not subset:
+ raise ValueError("subset must not be empty")
+ for name in subset:
+ if not isinstance(name, str):
+ raise TypeError("subset must be a string or a list of strings")
+ return list(subset)
+ raise TypeError("subset must be a string or a list of strings")
+
+
+def _normalize_order_by(
+ order_by: Union[str, Expression, List[Union[str, Expression]], None],
+) -> Optional[List[Union[str, Expression]]]:
+ if order_by is None:
+ return None
+
+ values = order_by if isinstance(order_by, (list, tuple)) else [order_by]
+ keys: List[Union[str, Expression]] = []
+ for value in values:
+ if isinstance(value, (str, Expression)):
+ keys.append(value)
+ else:
+ raise TypeError(
+ "order_by must be a string, an expression, or a list or tuple
of them"
+ )
+
+ if not keys:
+ raise ValueError("order_by must not be empty")
+
+ return keys
+
+
+def _normalize_nulls_first(
+ nulls_first: Union[bool, List[bool], None], order_len: int
+) -> Optional[List[bool]]:
+ if nulls_first is None:
+ return None
+
+ if isinstance(nulls_first, bool):
+ values = [nulls_first] * order_len
+ elif isinstance(nulls_first, (list, tuple)):
+ for value in nulls_first:
+ if not isinstance(value, bool):
+ raise TypeError("nulls_first must be a boolean or a list of
booleans")
+ values = list(nulls_first)
+ else:
+ 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")
+
+ return values
+
+
+def _build_deduplication_query(
+ table: Table,
+ columns: List[str],
+ subset_keys: List[str],
+ order_keys: Optional[List[Union[str, Expression]]],
+ keep: str,
+ nulls: Optional[List[bool]],
+) -> Table:
+ direction = "DESC" if keep == "last" else "ASC"
+ taken = set(columns)
+
+ if order_keys is None:
+ # Arrival order (processing time); keep decides its direction.
+ order_terms = ["PROCTIME() " + direction]
+ else:
+ order_terms = []
+ for index, key in enumerate(order_keys):
+ if isinstance(key, str):
+ if key not in columns:
+ raise ValueError(
+ "order_by column '%s' does not exist, available
columns: %s"
+ % (key, columns)
+ )
+ name = key
+ else:
+ # An Expression cannot be rendered to SQL text, so materialize
it as a helper
+ # column and reference it by name.
+ name = _unique_name("__pf_order_%d" % index, taken)
+ taken.add(name)
+ table = table.add_columns(key.alias(name))
+ term = _quote_identifier(name) + " " + direction
+ if nulls is not None:
+ term += " NULLS FIRST" if nulls[index] else " NULLS LAST"
+ order_terms.append(term)
+
+ rank_column = _quote_identifier(_unique_name("__pf_row_number", taken))
+ source = _quote_identifier(str(table))
+ select_list = ", ".join(_quote_identifier(name) for name in columns)
+ partition_by = ", ".join(_quote_identifier(name) for name in subset_keys)
+ query = (
+ "SELECT %s FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY %s ORDER BY %s) AS %s\n"
+ " FROM %s\n"
+ ") WHERE %s = 1"
+ % (
+ select_list,
+ partition_by,
+ ", ".join(order_terms),
+ rank_column,
+ source,
+ rank_column,
+ )
+ )
+ return table._t_env.sql_query(query)
+
+
+def _unique_name(base: str, taken: Set[str]) -> str:
+ name = base
+ while name in taken:
+ name += "_"
+ return name
+
+
+def _quote_identifier(name: str) -> str:
+ return "`" + name.replace("`", "``") + "`"
+
+
def _normalize_aggregations(
aggs: Tuple[Expression, ...], named_aggs: Dict[str, Expression]
) -> List[Expression]:
diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py
b/flink-python/pyflink/dataframe/tests/test_dataframe.py
index 0531b0fc6e9..dc34f28229e 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -16,17 +16,18 @@
# limitations under the License.
################################################################################
+import inspect
import array
import decimal
import unittest
+from py4j.protocol import Py4JJavaError
from datetime import date, datetime, time, timedelta, 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
+from pyflink.common import Row, RowKind
from pyflink.table import (
DataTypes as TableDataTypes,
EnvironmentSettings,
@@ -1149,6 +1150,432 @@ class
DataFrameAggregationTests(PyFlinkDataFrameUTTestCase):
)
+class DataFrameDropDuplicatesTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ self.dataframe = pf.from_records(
+ [
+ {"id": 1, "name": "a", "score": 10},
+ {"id": 1, "name": "b", "score": 20},
+ {"id": 2, "name": "c", "score": 30},
+ ]
+ )
+
+ def test_sql_proctime_default_keep_first(self):
+ self.assert_dataframe_sql(
+ self.dataframe,
+ "SELECT `id`, `name`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id` ORDER BY
PROCTIME() ASC)"
+ " AS `__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: self.dataframe.drop_duplicates(subset="id"),
+ )
+
+ def test_sql_proctime_default_keep_last(self):
+ self.assert_dataframe_sql(
+ self.dataframe,
+ "SELECT `id`, `name`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id` ORDER BY
PROCTIME() DESC)"
+ " AS `__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: self.dataframe.drop_duplicates(subset="id", keep="last"),
+ )
+
+ def test_sql_multi_column_subset(self):
+ self.assert_dataframe_sql(
+ self.dataframe,
+ "SELECT `id`, `name`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id`, `name` ORDER BY
PROCTIME() ASC)"
+ " AS `__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: self.dataframe.drop_duplicates(subset=["id", "name"]),
+ )
+
+ def test_sql_order_by_column_name(self):
+ self.assert_dataframe_sql(
+ self.dataframe,
+ "SELECT `id`, `name`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id` ORDER BY `score`
DESC)"
+ " AS `__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: self.dataframe.drop_duplicates(subset="id",
order_by="score", keep="last"),
+ )
+
+ def test_sql_nulls_first_per_key(self):
+ self.assert_dataframe_sql(
+ self.dataframe,
+ "SELECT `id`, `name`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id`"
+ " ORDER BY `score` ASC NULLS FIRST, `name` ASC NULLS LAST) AS
`__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: self.dataframe.drop_duplicates(
+ subset="id", order_by=["score", "name"], nulls_first=[True,
False]
+ ),
+ )
+
+ def test_sql_order_by_expression_is_materialized(self):
+ self.assert_dataframe_sql(
+ self.dataframe,
+ "SELECT `id`, `name`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id` ORDER BY
`__pf_order_0` ASC)"
+ " AS `__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: self.dataframe.drop_duplicates(subset="id",
order_by=pf.col("score")),
+ )
+
+ def test_sql_after_select_lists_only_source_columns(self):
+ source = self.dataframe.select("id", "score")
+ self.assert_dataframe_sql(
+ source,
+ "SELECT `id`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id` ORDER BY
PROCTIME() ASC)"
+ " AS `__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: source.drop_duplicates(subset="id"),
+ )
+
+ def test_sql_rank_column_avoids_collision(self):
+ dataframe = pf.from_records([{"__pf_row_number": 1, "value": 2}])
+ self.assert_dataframe_sql(
+ dataframe,
+ "SELECT `__pf_row_number`, `value` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `value` ORDER BY
PROCTIME() ASC)"
+ " AS `__pf_row_number_`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number_` = 1",
+ lambda: dataframe.drop_duplicates(subset="value"),
+ )
+
+ def test_whole_row_produces_no_sql(self):
+ # Whole-row deduplication uses distinct(), not a generated query.
+ self.assert_dataframe_sql(
+ self.dataframe, None, lambda: self.dataframe.drop_duplicates()
+ )
+
+ def test_schema_preserved(self):
+ self.assert_dataframe_schema(
+ self.dataframe.drop_duplicates("id", order_by="score",
keep="last"),
+ ["id", "name", "score"],
+ )
+
+ def test_whole_row_schema_preserved(self):
+ self.assert_dataframe_schema(
+ self.dataframe.drop_duplicates(), ["id", "name", "score"]
+ )
+
+ def test_rejects_invalid_keep(self):
+ with self.assertRaises(ValueError) as error:
+ self.dataframe.drop_duplicates("id", keep="middle")
+ self.assertEqual(str(error.exception), 'keep must be "first" or
"last"')
+
+ def test_rejects_empty_subset_list(self):
+ with self.assertRaises(ValueError) as error:
+ self.dataframe.drop_duplicates([])
+ self.assertEqual(str(error.exception), "subset must not be empty")
+
+ def test_rejects_non_string_subset(self):
+ with self.assertRaises(TypeError) as error:
+ self.dataframe.drop_duplicates([1])
+ self.assertEqual(
+ str(error.exception), "subset must be a string or a list of
strings"
+ )
+
+ def test_rejects_unknown_subset_column(self):
+ # The message embeds the (py4j) column list, so match only the stable
prefix.
+ with self.assertRaisesRegex(ValueError, "subset column 'nope' does not
exist"):
+ self.dataframe.drop_duplicates("nope")
+
+ def test_rejects_unknown_order_column(self):
+ # The message embeds the (py4j) column list, so match only the stable
prefix.
+ with self.assertRaisesRegex(ValueError, "order_by column 'nope' does
not exist"):
+ self.dataframe.drop_duplicates("id", order_by="nope")
+
+ def test_rejects_order_by_without_subset(self):
+ with self.assertRaises(ValueError) as error:
+ self.dataframe.drop_duplicates(order_by="score")
+ self.assertEqual(
+ str(error.exception),
+ "order_by requires subset; whole-row duplicates cannot be ordered",
+ )
+
+ def test_rejects_nulls_first_without_order_by(self):
+ with self.assertRaises(ValueError) as error:
+ self.dataframe.drop_duplicates("id", nulls_first=True)
+ self.assertEqual(str(error.exception), "nulls_first requires order_by")
+
+ def test_rejects_nulls_first_length_mismatch(self):
+ with self.assertRaises(ValueError) as error:
+ self.dataframe.drop_duplicates(
+ "id", order_by="score", nulls_first=[True, False]
+ )
+ self.assertEqual(
+ str(error.exception), "nulls_first must have the same length as
order_by"
+ )
+
+ def test_rejects_nulls_first_wrong_type(self):
+ with self.assertRaises(TypeError) as error:
+ self.dataframe.drop_duplicates("id", order_by="score",
nulls_first=["x"])
+ self.assertEqual(
+ str(error.exception), "nulls_first must be a boolean or a list of
booleans"
+ )
+
+
+class DataFrameDistinctTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ self.dataframe = pf.from_records(
+ [{"id": 1, "score": 10}, {"id": 1, "score": 20}]
+ )
+
+ def test_shares_drop_duplicates_signature(self):
+ self.assertEqual(
+ inspect.signature(pf.DataFrame.distinct),
+ inspect.signature(pf.DataFrame.drop_duplicates),
+ )
+
+ def test_produces_the_same_sql_as_drop_duplicates(self):
+ self.assert_dataframe_sql(
+ self.dataframe,
+ "SELECT `id`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id` ORDER BY `score`
ASC)"
+ " AS `__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: self.dataframe.distinct(subset="id", order_by="score"),
+ )
+
+
+class DataFrameUniqueTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ self.dataframe = pf.from_records(
+ [{"id": 1, "score": 10}, {"id": 1, "score": 20}]
+ )
+
+ def test_shares_drop_duplicates_signature(self):
+ self.assertEqual(
+ inspect.signature(pf.DataFrame.unique),
+ inspect.signature(pf.DataFrame.drop_duplicates),
+ )
+
+ def test_produces_the_same_sql_as_drop_duplicates(self):
+ self.assert_dataframe_sql(
+ self.dataframe,
+ "SELECT `id`, `score` FROM (\n"
+ " SELECT *, ROW_NUMBER() OVER (PARTITION BY `id` ORDER BY `score`
ASC)"
+ " AS `__pf_row_number`\n"
+ " FROM `SRC`\n"
+ ") WHERE `__pf_row_number` = 1",
+ lambda: self.dataframe.unique(subset="id", order_by="score"),
+ )
+
+
+class DataFrameDropDuplicatesITTests(PyFlinkStreamDataFrameTestCase):
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+ cls.t_env.get_config().set("table.exec.resource.default-parallelism",
"1")
+
+ @staticmethod
+ def _materialize(dataframe, key=None):
+ # Fold the collected changelog into the final table so assertions read
as the
+ # resulting rows rather than the raw +I/-U/+U events.
+ columns = dataframe._table.get_resolved_schema().get_column_names()
+ if key is None:
+ indices = list(range(len(columns)))
+ else:
+ indices = [columns.index(name) for name in key]
+
+ state = {}
+ for row in dataframe.collect():
+ key_value = tuple(row[index] for index in indices)
+ if row.get_row_kind() in (RowKind.INSERT, RowKind.UPDATE_AFTER):
+ state[key_value] = tuple(row)
+ else:
+ state.pop(key_value, None)
+ return sorted(state.values())
+
+ def test_whole_row_removes_identical_rows(self):
+ dataframe = pf.from_records(
+ [(1, "a"), (1, "a"), (2, "b")],
+ schema=["id", "name"],
+ )
+
+ self.assertEqual(
+ self._materialize(dataframe.drop_duplicates()),
+ [(1, "a"), (2, "b")],
+ )
+
+ def test_whole_row_keeps_rows_differing_in_any_column(self):
+ dataframe = pf.from_records(
+ [(1, "a"), (1, "b"), (1, "a")],
+ schema=["id", "name"],
+ )
+
+ self.assertEqual(
+ self._materialize(dataframe.drop_duplicates()),
+ [(1, "a"), (1, "b")],
+ )
+
+ def test_distinct_alias_removes_identical_rows(self):
+ dataframe = pf.from_records(
+ [(1, "a"), (1, "a"), (2, "b")],
+ schema=["id", "name"],
+ )
+
+ self.assertEqual(
+ self._materialize(dataframe.distinct()),
+ [(1, "a"), (2, "b")],
+ )
+
+ def test_unique_alias_removes_identical_rows(self):
+ dataframe = pf.from_records(
+ [(1, "a"), (1, "a"), (2, "b")],
+ schema=["id", "name"],
+ )
+
+ self.assertEqual(
+ self._materialize(dataframe.unique()),
+ [(1, "a"), (2, "b")],
+ )
+
+ def test_from_dict_input(self):
+ dataframe = pf.from_dict({"id": [1, 1, 2], "name": ["a", "a", "b"]})
+
+ self.assertEqual(
+ self._materialize(dataframe.drop_duplicates()),
+ [(1, "a"), (2, "b")],
+ )
+
+ def test_subset_keep_first_by_order_column(self):
+ dataframe = pf.from_records(
+ [
+ (1, "a", 10),
+ (1, "b", 20),
+ (2, "c", 30),
+ (2, "d", 5),
+ (3, "e", 7),
+ ],
+ schema=["id", "name", "score"],
+ )
+
+ result = dataframe.drop_duplicates("id", order_by="score",
keep="first")
+
+ self.assertEqual(
+ self._materialize(result, key=["id"]),
+ [(1, "a", 10), (2, "d", 5), (3, "e", 7)],
+ )
+
+ def test_subset_keep_last_by_order_column(self):
+ dataframe = pf.from_records(
+ [
+ (1, "a", 10),
+ (1, "b", 20),
+ (2, "c", 30),
+ (2, "d", 5),
+ (3, "e", 7),
+ ],
+ schema=["id", "name", "score"],
+ )
+
+ result = dataframe.drop_duplicates("id", order_by="score", keep="last")
+
+ self.assertEqual(
+ self._materialize(result, key=["id"]),
+ [(1, "b", 20), (2, "c", 30), (3, "e", 7)],
+ )
+
+ def test_multi_column_subset_keep_first_by_order_column(self):
+ dataframe = pf.from_records(
+ [
+ (1, "a", 10),
+ (1, "a", 20),
+ (1, "b", 30),
+ (2, "a", 40),
+ ],
+ schema=["id", "name", "score"],
+ )
+
+ result = dataframe.drop_duplicates(["id", "name"], order_by="score",
keep="first")
+
+ self.assertEqual(
+ self._materialize(result, key=["id", "name"]),
+ [(1, "a", 10), (1, "b", 30), (2, "a", 40)],
+ )
+
+ def test_subset_keep_first_by_arrival(self):
+ dataframe = pf.from_records(
+ [(1, "a"), (1, "b"), (2, "c")],
+ schema=["id", "name"],
+ )
+
+ self.assertEqual(
+ self._materialize(dataframe.drop_duplicates("id"), key=["id"]),
+ [(1, "a"), (2, "c")],
+ )
+
+ def test_subset_keep_last_by_arrival(self):
+ dataframe = pf.from_records(
+ [(1, "a"), (1, "b"), (2, "c")],
+ schema=["id", "name"],
+ )
+
+ result = dataframe.drop_duplicates("id", keep="last")
+
+ self.assertEqual(
+ self._materialize(result, key=["id"]),
+ [(1, "b"), (2, "c")],
+ )
+
+ def test_dedup_after_select_and_filter(self):
+ dataframe = pf.from_records(
+ [
+ (1, "a", 10),
+ (1, "b", 20),
+ (2, "c", 30),
+ (3, "d", 40),
+ ],
+ schema=["id", "name", "score"],
+ )
+
+ result = (
+ dataframe.filter(pf.col("id") > 1)
+ .select("id", "score")
+ .drop_duplicates("id", order_by="score", keep="last")
+ )
+
+ self.assertEqual(
+ self._materialize(result, key=["id"]),
+ [(2, 30), (3, 40)],
+ )
+
+ def test_filter_after_dedup(self):
+ dataframe = pf.from_records(
+ [
+ (1, "a", 10),
+ (1, "b", 20),
+ (2, "c", 30),
+ ],
+ schema=["id", "name", "score"],
+ )
+
+ result = dataframe.drop_duplicates("id", order_by="score",
keep="last").filter(
+ pf.col("id") > 1
+ )
+
+ self.assertEqual(
+ self._materialize(result, key=["id"]),
+ [(2, "c", 30)],
+ )
+
+
class DataFrameITTests(PyFlinkStreamDataFrameTestCase):
def test_from_records(self):
dataframe = pf.from_records(
diff --git a/flink-python/pyflink/testing/test_case_utils.py
b/flink-python/pyflink/testing/test_case_utils.py
index bc547bd8d36..095438b927d 100644
--- a/flink-python/pyflink/testing/test_case_utils.py
+++ b/flink-python/pyflink/testing/test_case_utils.py
@@ -29,7 +29,6 @@ import unittest
from abc import abstractmethod
from decimal import Decimal
from functools import wraps
-
from py4j.java_gateway import JavaObject
from pyflink.common import JobExecutionResult, Time, Instant, Row
@@ -183,6 +182,27 @@ class PyFlinkDataFrameUTTestCase(PyFlinkUTTestCase):
expected_column_data_types,
)
+ def assert_dataframe_sql(self, dataframe, expected_sql, invoke):
+ environment = dataframe._table._t_env
+ original = environment.sql_query
+ captured = {}
+
+ def capture(query):
+ captured["sql"] = query
+ return original(query)
+
+ environment.sql_query = capture
+ try:
+ invoke()
+ finally:
+ del environment.sql_query
+
+ if "sql" in captured:
+ actual_sql = re.sub(r"UnnamedTable\$\d+", "SRC", captured["sql"])
+ else:
+ actual_sql = None
+ self.assertEqual(actual_sql, expected_sql)
+
class PyFlinkStreamTableTestCase(PyFlinkITTestCase):
"""
