MattBelle commented on code in PR #28979:
URL: https://github.com/apache/flink/pull/28979#discussion_r3856285259
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -412,96 +331,210 @@ def __getitem__(
return self.filter(key)
raise TypeError("key must be a string, list, tuple, or Expression")
- # ======================== Conversion ========================
+ def _validate_subset(self, subset: Optional[List[str]]) -> List[str]:
+ """
+ Validate and normalize the subset parameter.
+
+ :param subset: Column names to validate, or None for all columns.
+ :return: Validated list of column names.
+ :raises ValueError: If subset is empty or contains invalid column
names.
+ :raises TypeError: If subset is not a list of strings.
+ """
+ schema = self._table.get_schema()
+ all_columns = schema.get_field_names()
+
+ if subset is None:
+ return all_columns
+
+ if not isinstance(subset, list):
+ raise TypeError("subset must be a list of strings")
+
+ if not subset:
+ raise ValueError("subset cannot be empty")
+
+ # Validate all column names exist
+ all_columns_set = set(all_columns)
+ invalid_columns = set(subset) - all_columns_set
+ if invalid_columns:
+ raise ValueError(f"Columns not found in DataFrame:
{sorted(invalid_columns)}")
+
+ return subset
+
+ def _fill_values(
+ self,
+ value: Any,
+ subset: Optional[List[str]],
+ condition_fn: Callable[[Expression], Expression]
+ ) -> "DataFrame":
+ """
+ Helper method to fill values based on a condition.
+
+ :param value: The value to use as replacement.
+ :param subset: Column names to fill, or None for all columns.
+ :param condition_fn: Function that takes a column expression and
returns
+ a boolean expression indicating when to replace.
+ :return: A new DataFrame with values replaced.
+ """
+ subset = self._validate_subset(subset)
+ subset_set = set(subset)
+
+ schema = self._table.get_schema()
+ all_columns = schema.get_field_names()
+
+ expressions = []
+ for col_name in all_columns:
+ col_expr = table_col(col_name)
+ if col_name in subset_set:
+ col_type = schema.get_field_data_type(col_name)
+ typed_value = table_lit(value).cast(col_type)
+ filled_expr = if_then_else(
+ condition_fn(col_expr),
+ typed_value,
+ col_expr
+ ).alias(col_name)
+ expressions.append(filled_expr)
+ else:
+ expressions.append(col_expr)
+
+ return DataFrame(self._table.select(*expressions))
@PublicEvolving()
- def collect(self) -> List[Row]:
+ def drop_null(self, subset: Optional[List[str]] = None) -> "DataFrame":
"""
- Execute this DataFrame and return all rows.
+ Remove rows containing NULL values.
- The result iterator is always closed before this method returns or
propagates an error.
+ This method uses three-valued logic: NULL values in the specified
columns
+ will cause the row to be filtered out. Rows where all checked columns
are
+ non-NULL will be retained.
- :return: All result rows in collection order.
+ :param subset: Column names to check. If None, checks all columns.
+ :return: A new DataFrame with rows containing NULL values removed.
+ :raises ValueError: If subset is empty or contains invalid column
names.
+ :raises TypeError: If subset is not a list of strings.
Example::
>>> import pyflink.dataframe as pf
- >>> df = pf.from_records([{"id": 1}, {"id": 2}])
- >>> rows = df.collect()
+ >>> df = pf.from_records([
+ ... {"id": 1, "name": "Alice", "age": 30},
+ ... {"id": 2, "name": None, "age": 25},
+ ... {"id": 3, "name": "Bob", "age": None},
+ ... ])
+ >>> df.drop_null() # Drop rows with any NULL
+ >>> df.drop_null(subset=["age"]) # Drop rows where "age" is NULL
.. versionadded:: 2.4.0
"""
- with self._table.execute().collect() as rows:
- return list(rows)
+ subset = self._validate_subset(subset)
+ conditions = [table_col(col_name).is_not_null for col_name in subset]
+ condition = and_(*conditions) if len(conditions) > 1 else conditions[0]
+ return DataFrame(self._table.filter(condition))
+ @PublicEvolving()
+ def drop_nan(self, subset: Optional[List[str]] = None) -> "DataFrame":
+ """
+ Remove rows containing NaN values (for float/double columns).
-@PublicEvolving()
-class GroupedDataFrame:
- """
- A DataFrame grouped by one or more keys and ready for aggregation.
+ This method uses three-valued logic: NaN values in the specified
columns
+ will cause the row to be filtered out. NULL values are preserved (not
+ treated as NaN). Only applies to floating-point numeric types.
- Instances are created by :meth:`DataFrame.group_by`.
+ :param subset: Column names to check. If None, checks all columns.
+ :return: A new DataFrame with rows containing NaN values removed.
+ :raises ValueError: If subset is empty or contains invalid column
names.
+ :raises TypeError: If subset is not a list of strings.
- .. versionadded:: 2.4.0
- """
+ Example::
- def __init__(self, dataframe: DataFrame, grouping_keys: List[Expression]):
- self._dataframe = dataframe
- self._grouping_keys = grouping_keys
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([
+ ... {"id": 1, "score": 0.95},
+ ... {"id": 2, "score": float('nan')},
+ ... ])
+ >>> df.drop_nan() # Drop rows with any NaN
+ >>> df.drop_nan(subset=["score"]) # Drop rows where "score" is NaN
+
+ .. versionadded:: 2.4.0
+ """
+ subset = self._validate_subset(subset)
+ conditions = [table_col(col_name).is_not_nan for col_name in subset]
Review Comment:
Good catch! Wrote a new test for this scenario
(`test_drop_nan_preserves_null_values`). Added `OR col IS NULL` to the filter
condition in `drop_nan()` to explicitly preserve NULL values. All
DataFrameNullNanITTests are passing.
--
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]
