auroflow commented on code in PR #29002:
URL: https://github.com/apache/flink/pull/29002#discussion_r3844252933
##########
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:
The `*args` is primarily used to support passing a list of alternating old
and new name pairs. I have updated the docstring to explain the usage more
clearly.
--
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]
