dianfu commented on code in PR #28979:
URL: https://github.com/apache/flink/pull/28979#discussion_r3794041302
##########
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:
For NULL input, IS_NOT_NAN(NULL) returns NULL, and Table.filter only retains
rows for which the predicate is TRUE. Therefore, the current implementation
incorrectly drops NULL rows which violates the doc: NULL values are preserved.
##########
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:
With subset=None, the current implementation applies IS_NAN/IS_NOT_NAN to
every column for drop_nan/fill_nan. A common mixed schema containing STRING or
BOOLEAN columns will therefore fail during type inference.
I think we could select only FLOAT/DOUBLE columns from the schema and return
an equivalent DataFrame when no floating-point columns exist. I checked that
Polars, Daft follow this behavior which seem reasonable for me.
##########
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]
+ condition = and_(*conditions) if len(conditions) > 1 else conditions[0]
+ return DataFrame(self._table.filter(condition))
@PublicEvolving()
- def agg(self, *aggs: Expression, **named_aggs: Expression) -> DataFrame:
+ def fill_null(self, value: Any, subset: Optional[List[str]] = None) ->
"DataFrame":
"""
- Aggregate the rows in each group.
+ Replace NULL values with a specified value.
- Grouping keys are included first in their supplied order, followed by
positional
- aggregation expressions and then named aggregations. Each named
aggregation is aliased to
- its keyword name.
+ This method uses three-valued logic: NULL values in the specified
columns
+ are replaced with the provided value, while non-NULL values are
preserved.
+ The replacement value is automatically cast to match each column's
data type.
- :param aggs: Aggregation expressions.
- :param named_aggs: Aggregation expressions keyed by their result
column names.
- :return: A DataFrame containing the grouping keys and aggregation
results.
- :raises TypeError: If an aggregation is not an expression.
- :raises ValueError: If no aggregations are provided.
+ :param value: The value to replace NULL with.
+ :param subset: Column names to fill. If None, fills all columns.
+ :return: A new DataFrame with NULL values replaced.
+ :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([
- ... ("engineering", 10),
- ... ("engineering", 20),
- ... ("sales", 5),
- ... ], schema=["department", "amount"])
- >>> totals = df.group_by("department").agg(
- ... pf.col("amount").sum.alias("total_amount"),
- ... row_count=pf.col("amount").count,
- ... )
- >>> # totals schema: [department: STRING, total_amount: BIGINT,
- >>> # row_count: BIGINT NOT NULL]
+ ... {"id": 1, "name": "Alice", "quantity": 10},
+ ... {"id": 2, "name": None, "quantity": None},
+ ... ])
+ >>> df.fill_null(0, subset=["quantity"])
+ >>> df.fill_null("unknown", subset=["name"])
.. versionadded:: 2.4.0
"""
- aggregations = _normalize_aggregations(aggs, named_aggs)
- grouped_table = self._dataframe._table.group_by(*self._grouping_keys)
- return DataFrame(grouped_table.select(*self._grouping_keys,
*aggregations))
-
-
-# ======================== Internal Helpers ========================
-
-
-def _normalize_aggregations(
- aggs: Tuple[Expression, ...], named_aggs: Dict[str, Expression]
-) -> List[Expression]:
- if not aggs and not named_aggs:
- raise ValueError("agg() requires at least one aggregation")
-
- aggregations: List[Expression] = []
- for aggregation in aggs:
- if not isinstance(aggregation, Expression):
- raise TypeError("agg() aggregations must be expressions")
- aggregations.append(aggregation)
- for name, aggregation in named_aggs.items():
- if not isinstance(aggregation, Expression):
- raise TypeError("agg() aggregations must be expressions")
- aggregations.append(aggregation.alias(name))
- return aggregations
+ return self._fill_values(value, subset, lambda col: col.is_null)
+
+ @PublicEvolving()
+ def fill_nan(self, value: Any, subset: Optional[List[str]] = None) ->
"DataFrame":
Review Comment:
For drop_nan/fill_nan API, we should define clearly the behavior of
non-number columns, eg. String, Boolean, etc. I suggest validating column
existence first and then ignoring non-floating-point columns, consistent with
Daft, and Polars’ default behavior.
##########
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)
Review Comment:
The current implementation casts the replacement to every target column
type. For example, fill_null(0) may replace a STRING NULL with "0" and may fail
for ARRAY, ROW, or TIMESTAMP columns.
I checked that it only handle the columns which supports the cast in Spark
and Daft.
##########
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")
Review Comment:
Nit: Polars and Pandas treat empty subset as a no-op instead of raising
exceptions. I have no preference. Just comment here for your reference.
##########
flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/functions/scalar/IsNotNanFunction.java:
##########
@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.table.runtime.functions.scalar;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.table.functions.BuiltInFunctionDefinitions;
+import org.apache.flink.table.functions.SpecializedFunction.SpecializedContext;
+
+import javax.annotation.Nullable;
+
+import java.math.BigDecimal;
+
+/** Implementation of {@link BuiltInFunctionDefinitions#IS_NOT_NAN}. */
+@Internal
+public final class IsNotNanFunction extends BuiltInScalarFunction {
Review Comment:
Do we really need IsNotNan?
##########
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]
+ condition = and_(*conditions) if len(conditions) > 1 else conditions[0]
+ return DataFrame(self._table.filter(condition))
@PublicEvolving()
- def agg(self, *aggs: Expression, **named_aggs: Expression) -> DataFrame:
+ def fill_null(self, value: Any, subset: Optional[List[str]] = None) ->
"DataFrame":
"""
- Aggregate the rows in each group.
+ Replace NULL values with a specified value.
- Grouping keys are included first in their supplied order, followed by
positional
- aggregation expressions and then named aggregations. Each named
aggregation is aliased to
- its keyword name.
+ This method uses three-valued logic: NULL values in the specified
columns
+ are replaced with the provided value, while non-NULL values are
preserved.
+ The replacement value is automatically cast to match each column's
data type.
- :param aggs: Aggregation expressions.
- :param named_aggs: Aggregation expressions keyed by their result
column names.
- :return: A DataFrame containing the grouping keys and aggregation
results.
- :raises TypeError: If an aggregation is not an expression.
- :raises ValueError: If no aggregations are provided.
+ :param value: The value to replace NULL with.
+ :param subset: Column names to fill. If None, fills all columns.
+ :return: A new DataFrame with NULL values replaced.
+ :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([
- ... ("engineering", 10),
- ... ("engineering", 20),
- ... ("sales", 5),
- ... ], schema=["department", "amount"])
- >>> totals = df.group_by("department").agg(
- ... pf.col("amount").sum.alias("total_amount"),
- ... row_count=pf.col("amount").count,
- ... )
- >>> # totals schema: [department: STRING, total_amount: BIGINT,
- >>> # row_count: BIGINT NOT NULL]
+ ... {"id": 1, "name": "Alice", "quantity": 10},
+ ... {"id": 2, "name": None, "quantity": None},
+ ... ])
+ >>> df.fill_null(0, subset=["quantity"])
+ >>> df.fill_null("unknown", subset=["name"])
.. versionadded:: 2.4.0
"""
- aggregations = _normalize_aggregations(aggs, named_aggs)
- grouped_table = self._dataframe._table.group_by(*self._grouping_keys)
- return DataFrame(grouped_table.select(*self._grouping_keys,
*aggregations))
-
-
-# ======================== Internal Helpers ========================
-
-
-def _normalize_aggregations(
- aggs: Tuple[Expression, ...], named_aggs: Dict[str, Expression]
-) -> List[Expression]:
- if not aggs and not named_aggs:
- raise ValueError("agg() requires at least one aggregation")
-
- aggregations: List[Expression] = []
- for aggregation in aggs:
- if not isinstance(aggregation, Expression):
- raise TypeError("agg() aggregations must be expressions")
- aggregations.append(aggregation)
- for name, aggregation in named_aggs.items():
- if not isinstance(aggregation, Expression):
- raise TypeError("agg() aggregations must be expressions")
- aggregations.append(aggregation.alias(name))
- return aggregations
+ return self._fill_values(value, subset, lambda col: col.is_null)
+
+ @PublicEvolving()
+ def fill_nan(self, value: Any, subset: Optional[List[str]] = None) ->
"DataFrame":
+ """
+ Replace NaN values with a specified value (for float/double columns).
Review Comment:
The current implementation appears to support fill_nan(None) which converts
NaN values to NULL. I think this is reasonable. What about documenting this
behavior explicitly?
--
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]
