Skip to content

[FLINK-40196][python] Add missing-value handling to DataFrame API #28979

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MattBelle wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[FLINK-40196][python] Add missing-value handling to DataFrame API #28979

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
190 changes: 190 additions & 0 deletions flink-python/pyflink/dataframe/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
and_,
call_sql,
col as table_col,
if_then_else,
is_nan,
lit as table_lit,
)
from pyflink.table.table import Table
Expand Down Expand Up @@ -329,6 +331,194 @@ def __getitem__(
return self.filter(key)
raise TypeError("key must be a string, list, tuple, or Expression")

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")

@dianfu dianfu Aug 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


# 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)

@dianfu dianfu Aug 17, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

It only handles the columns which supports the type cast in Spark and Daft.

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 drop_null(self, subset: Optional[List[str]] = None) -> "DataFrame":
"""
Remove rows containing NULL values.

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.

: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, "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
"""
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).

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.

: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.

Example::

>>> 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]

@dianfu dianfu Aug 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@dianfu dianfu Aug 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

condition = and_(*conditions) if len(conditions) > 1 else conditions[0]
return DataFrame(self._table.filter(condition))

@PublicEvolving()
def fill_null(self, value: Any, subset: Optional[List[str]] = None) -> "DataFrame":
"""
Replace NULL values with a specified value.

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 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([
... {"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
"""
return self._fill_values(value, subset, lambda col: col.is_null)

@PublicEvolving()
def fill_nan(self, value: Any, subset: Optional[List[str]] = None) -> "DataFrame":

@dianfu dianfu Aug 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

"""
Replace NaN values with a specified value (for float/double columns).

@dianfu dianfu Aug 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 method uses three-valued logic: NaN values in the specified columns
are replaced with the provided value, while non-NaN values are preserved.
NULL values are preserved (not treated as NaN). The replacement value is
automatically cast to match each column's data type.

:param value: The value to replace NaN with.
:param subset: Column names to fill. If None, fills all columns.
:return: A new DataFrame with NaN 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([
... {"id": 1, "score": 0.95},
... {"id": 2, "score": float('nan')},
... ])
>>> df.fill_nan(0.0, subset=["score"])

.. versionadded:: 2.4.0
"""
return self._fill_values(value, subset, lambda col: is_nan(col))

@PublicEvolving()
def collect(self) -> List[Row]:
"""
Expand Down
Loading