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 1138d45143d [FLINK-40417][python] Add
with_columns/drop_columns/rename_columns to DataFrame API (#29002)
1138d45143d is described below
commit 1138d45143d4a5b0295c5cf5bfe06f7a31fe7cd8
Author: Liu Liu <[email protected]>
AuthorDate: Tue Aug 25 09:51:26 2026 +0800
[FLINK-40417][python] Add with_columns/drop_columns/rename_columns to
DataFrame API (#29002)
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 27 ++
flink-python/pyflink/dataframe/dataframe.py | 276 ++++++++++++++++++++-
.../pyflink/dataframe/tests/test_dataframe.py | 261 ++++++++++++++++++-
3 files changed, 558 insertions(+), 6 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index 29c09bee435..4c78cb5d857 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -51,7 +51,13 @@ Transformations
DataFrame.select
DataFrame.with_column
+ DataFrame.with_columns
+ DataFrame.drop_columns
+ DataFrame.drop
+ DataFrame.rename_columns
+ DataFrame.rename
DataFrame.filter
+ DataFrame.where
DataFrame.__getitem__
Aggregations
@@ -67,6 +73,27 @@ Aggregations
GroupedDataFrame
GroupedDataFrame.agg
+Composition
+-----------
+
+.. currentmodule:: pyflink.dataframe
+
+.. autosummary::
+ :toctree: api/
+
+ DataFrame.pipe
+
+Properties
+----------
+
+.. currentmodule:: pyflink.dataframe
+
+.. autosummary::
+ :toctree: api/
+
+ DataFrame.schema
+ DataFrame.columns
+
Results
-------
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index 793f074f8da..ac4c12d24d5 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -16,10 +16,22 @@
# limitations under the License.
################################################################################
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple,
Union, overload
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Callable,
+ Dict,
+ List,
+ Optional,
+ Tuple,
+ TypeVar,
+ Union,
+ overload,
+)
if TYPE_CHECKING:
import pandas
+ from pyflink.table.table_schema import TableSchema
from pyflink.common import Row
from pyflink.dataframe.datatype import DataType
@@ -35,6 +47,8 @@ from pyflink.util.api_stability_decorators import
PublicEvolving
__all__ = ["DataFrame", "GroupedDataFrame", "col", "lit"]
+T = TypeVar("T")
+
@PublicEvolving()
def col(name: str) -> Expression:
@@ -173,6 +187,8 @@ class DataFrame:
condition = conditions[0] if len(conditions) == 1 else
and_(*conditions)
return DataFrame(self._table.filter(condition))
+ where = filter
+
@PublicEvolving()
def with_column(
self,
@@ -212,6 +228,189 @@ class DataFrame:
raise TypeError("expr must be an Expression")
return
DataFrame(self._table.add_or_replace_columns(expression.alias(name)))
+ @PublicEvolving()
+ def with_columns(
+ self,
+ *exprs: Expression,
+ **named_exprs: Expression,
+ ) -> "DataFrame":
+ """
+ Add or replace multiple columns in one call.
+
+ Positional expressions are applied first and must carry their desired
output names. Named
+ expressions are appended afterward and are aliased to their keyword
names.
+
+ :param exprs: Expressions to add or replace.
+ :param named_exprs: Expressions keyed by their output column names.
+ :return: A new DataFrame with the requested columns.
+ :raises TypeError: If a positional or named value is not an expression.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([(2, 3)], schema=["left", "right"])
+ >>> df.with_columns(
+ ... (pf.col("left") + 1).alias("left_plus_one"),
+ ... (pf.col("right") + 1).alias("right_plus_one"),
+ ... )
+ >>> df.with_columns(
+ ... left_plus_one=pf.col("left") + 1,
+ ... right_plus_one=pf.col("right") + 1,
+ ... )
+
+ .. versionadded:: 2.4.0
+ """
+ expressions: List[Expression] = []
+ for expression in exprs:
+ if not isinstance(expression, Expression):
+ raise TypeError("exprs must be expressions")
+ expressions.append(expression)
+
+ for name, expression in named_exprs.items():
+ if not isinstance(expression, Expression):
+ raise TypeError("named_exprs must be expressions")
+ expressions.append(expression.alias(name))
+
+ return DataFrame(self._table.add_or_replace_columns(*expressions))
+
+ @PublicEvolving()
+ def drop_columns(
+ self,
+ *columns: Union[str, Expression],
+ strict: bool = True,
+ ) -> "DataFrame":
+ """
+ Remove columns from this DataFrame.
+
+ String column names are checked against the current schema. When
``strict`` is ``False``,
+ names that are not present are ignored. Expression arguments are
validated by the Table
+ API.
+
+ :param columns: Column names or expressions to remove.
+ :param strict: Whether a missing column name raises an error.
+ :return: A new DataFrame without the requested columns, or this
DataFrame if no columns
+ remain to be dropped.
+ :raises TypeError: If ``strict`` is not a boolean or a column has an
unsupported type.
+ :raises ValueError: If a named column is missing in strict mode.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([(1, "debug")], schema=["id",
"temporary"])
+ >>> result = df.drop_columns("temporary")
+ >>> unchanged = df.drop("missing", strict=False)
+
+ .. versionadded:: 2.4.0
+ """
+ if not isinstance(strict, bool):
+ raise TypeError("strict must be a boolean")
+
+ existing_columns = set(self.columns)
+ expressions: List[Expression] = []
+ for column in columns:
+ if isinstance(column, str):
+ if column not in existing_columns:
+ if strict:
+ raise ValueError(f"Column '{column}' not found in
schema")
+ continue
+ expressions.append(table_col(column))
+ elif isinstance(column, Expression):
+ expressions.append(column)
+ else:
+ raise TypeError("columns must be strings or expressions")
+
+ if not expressions:
+ return self
+ return DataFrame(self._table.drop_columns(*expressions))
+
+ drop = drop_columns
+
+ @PublicEvolving()
+ def rename_columns(
+ self,
+ *args: Any,
+ mapping: Optional[
+ Union[Dict[str, str], Callable[[str], str]]
+ ] = None,
+ ) -> "DataFrame":
+ """
+ Rename one or more columns.
+
+ Use exactly one of the following forms:
+
+ * A dictionary defines mappings from existing column names to new
names. It can be supplied
+ as the only positional argument or through the keyword-only
``mapping`` parameter.
+ Entries whose existing column name is not present are ignored.
+ * An even number of positional string arguments is interpreted as
alternating old and new
+ column name pairs.
+ * A function or lambda expression is applied to every current column
name and must return
+ the new name as a string. It can be supplied as the only positional
argument or through
+ ``mapping``.
+
+ :param args: One dictionary or callable, or an even number of
alternating old/new names.
+ :param mapping: Keyword-only alternative for passing a dictionary or
callable.
+ :return: A new DataFrame with renamed columns, or this DataFrame if no
names change.
+ :raises TypeError: If the mapping, a name, or a callable result has an
unsupported type.
+ :raises ValueError: If positional pairs are incomplete or ``mapping``
is combined with
+ positional arguments.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([(1, "Alice")], schema=["id", "name"])
+ >>> by_mapping = df.rename_columns({"id": "user_id"})
+ >>> by_keyword = df.rename_columns(mapping={"id": "user_id"})
+ >>> by_pairs = df.rename("id", "user_id", "name", "user_name")
+ >>> by_function = df.rename(str.upper)
+ >>> by_lambda = df.rename(lambda name: name.upper())
+
+ .. versionadded:: 2.4.0
+ """
+ if args and mapping is not None:
+ raise ValueError(
+ "rename_columns() accepts either positional arguments or
mapping, not both"
+ )
+
+ rename_spec: Any = mapping
+ if len(args) == 1:
+ rename_spec = args[0]
+ elif args:
+ if len(args) % 2 != 0:
+ raise ValueError(
+ "rename_columns() positional arguments must be old/new
name pairs"
+ )
+ positional_mapping: Dict[str, str] = {}
+ for index in range(0, len(args), 2):
+ old_name, new_name = args[index], args[index + 1]
+ if not isinstance(old_name, str) or not isinstance(new_name,
str):
+ raise TypeError("column names must be strings")
+ positional_mapping[old_name] = new_name
+ rename_spec = positional_mapping
+
+ current_columns = self.columns
+ rename_expressions: List[Expression] = []
+ if isinstance(rename_spec, dict):
+ for old_name, new_name in rename_spec.items():
+ if not isinstance(old_name, str) or not isinstance(new_name,
str):
+ raise TypeError("mapping keys and values must be strings")
+ if old_name in current_columns and new_name != old_name:
+
rename_expressions.append(table_col(old_name).alias(new_name))
+ elif callable(rename_spec):
+ for old_name in current_columns:
+ new_name = rename_spec(old_name)
+ if not isinstance(new_name, str):
+ raise TypeError("rename_columns() callable must return a
string")
+ if new_name != old_name:
+
rename_expressions.append(table_col(old_name).alias(new_name))
+ else:
+ raise TypeError("mapping must be a dictionary or callable")
+
+ if not rename_expressions:
+ return self
+ return DataFrame(self._table.rename_columns(*rename_expressions))
+
+ rename = rename_columns
+
@PublicEvolving()
def select(
self,
@@ -404,6 +603,41 @@ class DataFrame:
return self.filter(key)
raise TypeError("key must be a string, list, tuple, or Expression")
+ # ======================== Composition ========================
+
+ @PublicEvolving()
+ def pipe(
+ self,
+ func: Callable[..., T],
+ *args: Any,
+ **kwargs: Any,
+ ) -> T:
+ """
+ Apply a function to this DataFrame for reusable functional composition.
+
+ This DataFrame is passed as the first argument, followed by ``args``
and ``kwargs``. The
+ function's return value is returned unchanged.
+
+ :param func: Function whose first argument receives this DataFrame.
+ :param args: Additional positional arguments passed to ``func``.
+ :param kwargs: Additional keyword arguments passed to ``func``.
+ :return: The value returned by ``func``.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([(1, 2)], schema=["left", "right"])
+ >>> result = df.pipe(
+ ... lambda current, name: current.with_column(
+ ... name, pf.col("left") + pf.col("right")
+ ... ),
+ ... "total",
+ ... )
+
+ .. versionadded:: 2.4.0
+ """
+ return func(self, *args, **kwargs)
+
# ======================== Conversion ========================
@PublicEvolving()
@@ -466,6 +700,46 @@ class DataFrame:
"""
return self._table.to_pandas()
+ # ======================== Properties ========================
+
+ @property
+ @PublicEvolving()
+ def schema(self) -> "TableSchema":
+ """
+ Return this DataFrame's schema.
+
+ :return: The TableSchema exposed by the underlying Table.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([(1, "Alice")], schema=["id", "name"])
+ >>> df.schema.get_field_names()
+ ['id', 'name']
+
+ .. versionadded:: 2.4.0
+ """
+ return self._table.get_schema()
+
+ @property
+ @PublicEvolving()
+ def columns(self) -> List[str]:
+ """
+ Return this DataFrame's column names in schema order.
+
+ :return: A new list containing the column names.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([(1, "Alice")], schema=["id", "name"])
+ >>> df.columns
+ ['id', 'name']
+
+ .. versionadded:: 2.4.0
+ """
+ return list(self._table.get_resolved_schema().get_column_names())
+
# ======================== I/O ========================
@PublicEvolving()
diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py
b/flink-python/pyflink/dataframe/tests/test_dataframe.py
index 0e6651e44d7..0531b0fc6e9 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -31,6 +31,7 @@ from pyflink.table import (
DataTypes as TableDataTypes,
EnvironmentSettings,
TableEnvironment,
+ TableSchema,
)
from pyflink.table.expression import Expression
from pyflink.table.types import LocalZonedTimestampType, TimestampType
@@ -130,6 +131,25 @@ class DataFrameConversionTests(unittest.TestCase):
).to_pandas()
+class DataFrameCompositionTests(unittest.TestCase):
+ def test_pipe_forwards_dataframe_arguments_and_return_value(self):
+ dataframe = pf.DataFrame(object())
+ expected = object()
+
+ def transform(current, value, *, label):
+ self.assertIs(current, dataframe)
+ self.assertEqual(value, 42)
+ self.assertEqual(label, "answer")
+ return expected
+
+ self.assertIs(dataframe.pipe(transform, 42, label="answer"), expected)
+
+ def test_aliases_reference_the_original_methods(self):
+ self.assertIs(pf.DataFrame.where, pf.DataFrame.filter)
+ self.assertIs(pf.DataFrame.drop, pf.DataFrame.drop_columns)
+ self.assertIs(pf.DataFrame.rename, pf.DataFrame.rename_columns)
+
+
class DataFrameCreationTests(PyFlinkDataFrameUTTestCase):
def test_from_dict_uses_insertion_order_without_schema(self):
dataframe = pf.from_dict({"name": ["Alice"], "id": [1]})
@@ -521,6 +541,210 @@ class
DataFrameWithColumnTests(PyFlinkDataFrameUTTestCase):
with self.assertRaisesRegex(TypeError, "name must be a string"):
self.dataframe.with_column(42, object())
+ def test_with_columns_adds_and_replaces_positional_and_named_columns(self):
+ result = self.dataframe.with_columns(
+ (pf.col("id") + 1).alias("id"),
+ (pf.col("age") + 2).alias("age_in_two_years"),
+ age_next_year=pf.col("age") + 1,
+ doubled_age=pf.col("age") * 2,
+ )
+
+ self.assert_dataframe_schema(
+ result,
+ [
+ "id",
+ "name",
+ "age",
+ "age_in_two_years",
+ "age_next_year",
+ "doubled_age",
+ ],
+ [
+ TableDataTypes.BIGINT(),
+ TableDataTypes.STRING(),
+ TableDataTypes.BIGINT(),
+ TableDataTypes.BIGINT(),
+ TableDataTypes.BIGINT(),
+ TableDataTypes.BIGINT(),
+ ],
+ )
+
+ def test_with_columns_rejects_non_expressions(self):
+ invalid_calls = [
+ ("positional", lambda: self.dataframe.with_columns(42), "exprs"),
+ (
+ "named",
+ lambda: self.dataframe.with_columns(answer=42),
+ "named_exprs",
+ ),
+ ]
+ for name, invalid_call, message in invalid_calls:
+ with self.subTest(name=name):
+ with self.assertRaisesRegex(TypeError, message):
+ invalid_call()
+
+
+class DataFrameDropColumnsTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ self.dataframe = pf.from_records(
+ [(1, "Alice", 30)],
+ schema=["id", "name", "age"],
+ )
+
+ def test_drop_alias_accepts_names_and_expressions(self):
+ result = self.dataframe.drop("name", pf.col("age"))
+
+ self.assert_dataframe_schema(
+ result,
+ ["id"],
+ [TableDataTypes.BIGINT()],
+ )
+
+ def test_drop_columns_handles_missing_names_and_no_op(self):
+ with self.assertRaisesRegex(ValueError, "Column 'missing' not found"):
+ self.dataframe.drop_columns("missing")
+
+ self.assertIs(
+ self.dataframe.drop_columns("missing", strict=False),
+ self.dataframe,
+ )
+ self.assertIs(self.dataframe.drop_columns(), self.dataframe)
+
+ def test_drop_columns_rejects_invalid_arguments(self):
+ invalid_calls = [
+ (
+ "column",
+ lambda: self.dataframe.drop_columns(42),
+ "columns must be strings or expressions",
+ ),
+ (
+ "strict",
+ lambda: self.dataframe.drop_columns("id", strict="yes"),
+ "strict must be a boolean",
+ ),
+ ]
+ for name, invalid_call, message in invalid_calls:
+ with self.subTest(name=name):
+ with self.assertRaisesRegex(TypeError, message):
+ invalid_call()
+
+
+class DataFrameRenameColumnsTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ self.dataframe = pf.from_records(
+ [(1, "Alice", 30)],
+ schema=["id", "name", "age"],
+ )
+
+ def test_rename_columns_supports_all_input_forms(self):
+ cases = [
+ (
+ "mapping_alias",
+ lambda: self.dataframe.rename(
+ {"id": "identifier", "missing": "ignored"}
+ ),
+ ["identifier", "name", "age"],
+ ),
+ (
+ "mapping_keyword",
+ lambda: self.dataframe.rename_columns(
+ mapping={"name": "customer"}
+ ),
+ ["id", "customer", "age"],
+ ),
+ (
+ "callable",
+ lambda: self.dataframe.rename_columns(str.upper),
+ ["ID", "NAME", "AGE"],
+ ),
+ (
+ "pairs",
+ lambda: self.dataframe.rename_columns(
+ "id", "identifier", "age", "years"
+ ),
+ ["identifier", "name", "years"],
+ ),
+ ]
+ for name, rename, expected_columns in cases:
+ with self.subTest(name=name):
+ self.assert_dataframe_schema(rename(), expected_columns)
+
+ def test_rename_columns_returns_self_when_nothing_changes(self):
+ self.assertIs(
+ self.dataframe.rename_columns({"missing": "ignored"}),
+ self.dataframe,
+ )
+
+ def test_rename_columns_rejects_invalid_arguments(self):
+ invalid_calls = [
+ (
+ "missing_mapping",
+ lambda: self.dataframe.rename_columns(),
+ TypeError,
+ "mapping must be a dictionary or callable",
+ ),
+ (
+ "odd_pairs",
+ lambda: self.dataframe.rename_columns("id", "identifier",
"age"),
+ ValueError,
+ "must be old/new name pairs",
+ ),
+ (
+ "non_string_pair",
+ lambda: self.dataframe.rename_columns("id", 42),
+ TypeError,
+ "column names must be strings",
+ ),
+ (
+ "non_string_mapping",
+ lambda: self.dataframe.rename_columns({"id": 42}),
+ TypeError,
+ "mapping keys and values must be strings",
+ ),
+ (
+ "invalid_callable_result",
+ lambda: self.dataframe.rename_columns(lambda name: 42),
+ TypeError,
+ "callable must return a string",
+ ),
+ (
+ "ambiguous_mapping",
+ lambda: self.dataframe.rename_columns(
+ {"id": "identifier"}, mapping={"name": "customer"}
+ ),
+ ValueError,
+ "either positional arguments or mapping",
+ ),
+ ]
+ for name, invalid_call, error, message in invalid_calls:
+ with self.subTest(name=name):
+ with self.assertRaisesRegex(error, message):
+ invalid_call()
+
+
+class DataFramePropertyTests(PyFlinkDataFrameUTTestCase):
+ def test_schema_exposes_ordered_metadata(self):
+ dataframe = pf.from_records(
+ [(1, "Alice")],
+ schema=["id", "name"],
+ )
+
+ self.assertIsInstance(dataframe.schema, TableSchema)
+ self.assertEqual(dataframe.schema.get_field_names(), ["id", "name"])
+
+ def test_columns_returns_defensive_ordered_list(self):
+ dataframe = pf.from_records(
+ [(1, "Alice")],
+ schema=["id", "name"],
+ )
+
+ columns = dataframe.columns
+ self.assertEqual(columns, ["id", "name"])
+ columns.append("mutated")
+ self.assertEqual(dataframe.columns, ["id", "name"])
+
class DataFrameFilterTests(PyFlinkDataFrameUTTestCase):
def setUp(self):
@@ -1069,7 +1293,7 @@ class DataFrameITTests(PyFlinkStreamDataFrameTestCase):
result = (
df[df["id"] > 0]
- .filter(
+ .where(
"score >= 0.9",
lambda current: current["id"] < 6,
city="SF",
@@ -1080,11 +1304,19 @@ class DataFrameITTests(PyFlinkStreamDataFrameTestCase):
lambda current: current["age"] + 1,
)
.with_column("age", pf.col("age") + 1)
+ .with_columns(
+ (pf.col("age_next_year") + 1).alias("age_in_two_years"),
+ score_percent=pf.col("score") * 100,
+ )
+ .drop("score", "city", "destination")
+ .rename({"name": "customer_name"})
.select(
"id",
- "name",
+ "customer_name",
"age",
- age_next_year=pf.col("age_next_year"),
+ "age_next_year",
+ "age_in_two_years",
+ "score_percent",
inferred_int=pf.lit(2),
inferred_string=pf.lit("x"),
explicit_int=pf.lit(3, pf.DataType.int64()),
@@ -1098,10 +1330,12 @@ class DataFrameITTests(PyFlinkStreamDataFrameTestCase):
),
)[
(
- "name",
+ "customer_name",
"id",
"age",
"age_next_year",
+ "age_in_two_years",
+ "score_percent",
"inferred_int",
"inferred_string",
"explicit_int",
@@ -1116,7 +1350,24 @@ class DataFrameITTests(PyFlinkStreamDataFrameTestCase):
self.assertEqual(
result.collect(),
- [Row("Alice", 1, 31, 31, 2, "x", 3, 1 << 40, "y", None, None, 3)],
+ [
+ Row(
+ "Alice",
+ 1,
+ 31,
+ 31,
+ 32,
+ 95.0,
+ 2,
+ "x",
+ 3,
+ 1 << 40,
+ "y",
+ None,
+ None,
+ 3,
+ )
+ ],
)
