dianfu commented on code in PR #29002:
URL: https://github.com/apache/flink/pull/29002#discussion_r3831086428
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -212,6 +228,177 @@ def with_column(
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"])
+ >>> result = df.with_columns(
Review Comment:
We can also add an example for named arguments
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -212,6 +228,177 @@ def with_column(
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"])
+ >>> result = df.with_columns(
+ ... (pf.col("left") + 1).alias("left"),
+ ... total=pf.col("left") + pf.col("right"),
+ ... )
+
+ .. 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.
+
+ Supply a mapping or callable as the sole positional argument, use the
``mapping`` keyword,
+ or provide positional old/new name pairs. Mapping entries for columns
that are not present
+ are ignored. A callable is applied to every current column name.
+
+ :param args: A mapping, a callable, or positional old/new name pairs.
+ :param mapping: A mapping from old names to new names, or a callable
that returns each new
+ name.
+ :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_callable = df.rename(str.upper)
Review Comment:
Could you also add an lambda function example?
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -212,6 +228,177 @@ def with_column(
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"])
+ >>> result = df.with_columns(
+ ... (pf.col("left") + 1).alias("left"),
+ ... total=pf.col("left") + pf.col("right"),
+ ... )
+
+ .. 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[
Review Comment:
It seems that the `args` isn't that necessary. Removing it will make the API
more clear.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
