feat: add DataFrame.fill_nan

Replace NaN in floating-point columns, optionally limited to a subset.
Mirrors fill_null and wraps upstream DataFrame::fill_nan.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: timsaucer

crates/core/src/dataframe.rs

@@ -1320,6 +1320,26 @@ impl PyDataFrame {
         let df = self.df.as_ref().fill_null(&scalar_value.0, &cols)?;
         Ok(Self::new(df))
     }
+
+    /// Fill NaN values with a specified value for specific floating-point columns
+    #[pyo3(signature = (value, columns=None))]
+    fn fill_nan(
+        &self,
+        value: Py<PyAny>,
+        columns: Option<Vec<PyBackedStr>>,
+        py: Python,
+    ) -> PyDataFusionResult<Self> {
+        let scalar_value: PyScalarValue = value.extract(py)?;
+
+        let cols = match columns {
+            Some(col_names) => col_names.iter().map(|c| c.to_string()).collect(),
+            None => Vec::new(), // Empty vector means fill NaN for all columns
+        };
+
+        let cols = cols.iter().map(String::as_str).collect::<Vec<_>>();
+        let df = self.df.as_ref().fill_nan(&scalar_value.0, &cols)?;
+        Ok(Self::new(df))
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]

python/datafusion/dataframe.py

@@ -1875,6 +1875,34 @@ def fill_null(self, value: Any, subset: list[str] | None = None) -> DataFrame:
         """
         return DataFrame(self.df.fill_null(value, subset))
 
+    def fill_nan(self, value: float, subset: list[str] | None = None) -> DataFrame:
+        """Fill NaN values in floating-point columns with a value.
+
+        Only floating-point columns are changed; others are kept unchanged, as is
+        any column ``value`` cannot be cast to. NaN is distinct from null, which
+        :py:meth:`fill_null` handles.
+
+        Args:
+            value: Value to replace NaN with. Will be cast to match column type.
+            subset: Optional list of column names to fill. If None, fills all
+                floating-point columns.
+
+        Returns:
+            DataFrame with NaN values replaced.
+
+        Examples:
+            >>> from datafusion import SessionContext
+            >>> ctx = SessionContext()
+            >>> nan = float("nan")
+            >>> df = ctx.from_pydict({"a": [1.0, nan, None], "b": [nan, 2.0, 3.0]})
+            >>> df.fill_nan(0.0).to_pydict()
+            {'a': [1.0, 0.0, None], 'b': [0.0, 2.0, 3.0]}
+
+            >>> df.fill_nan(0.0, subset=["a"]).collect_column("b")[0].as_py()
+            nan
+        """
+        return DataFrame(self.df.fill_nan(value, subset))
+
 
 class InsertOp(Enum):
     """Insert operation mode.

python/tests/test_dataframe.py

@@ -3450,6 +3450,50 @@ def test_fill_null_all_null_column(ctx):
     assert result.column(1).to_pylist() == ["filled", "filled", "filled"]
 
 
+def _nan_df(ctx):
+    nan = float("nan")
+    batch = pa.RecordBatch.from_arrays(
+        [
+            pa.array([1.0, nan, None], type=pa.float64()),
+            pa.array([nan, 2.0, 3.0], type=pa.float32()),
+            pa.array([1, 2, 3]),
+            pa.array(["x", "nan", None]),
+        ],
+        names=["f64", "f32", "i", "s"],
+    )
+    return ctx.create_dataframe([[batch]])
+
+
+def _is_nan(v):
+    return v is not None and v != v  # noqa: PLR0124
+
+
+def test_fill_nan_all_columns(ctx):
+    result = _nan_df(ctx).fill_nan(0.0).to_pydict()
+    # NaN replaced in both float widths; null is not NaN and stays null.
+    assert result["f64"] == [1.0, 0.0, None]
+    assert result["f32"] == [0.0, 2.0, 3.0]
+    # Non-float columns are untouched.
+    assert result["i"] == [1, 2, 3]
+    assert result["s"] == ["x", "nan", None]
+
+
+def test_fill_nan_subset(ctx):
+    result = _nan_df(ctx).fill_nan(-1.0, subset=["f32"]).to_pydict()
+    assert result["f32"] == [-1.0, 2.0, 3.0]
+    assert _is_nan(result["f64"][1])
+
+
+def test_fill_nan_preserves_schema(ctx):
+    df = _nan_df(ctx)
+    assert df.fill_nan(0.0).schema() == df.schema()
+
+
+def test_fill_nan_unknown_column_raises(ctx):
+    with pytest.raises(Exception, match="missing"):
+        _nan_df(ctx).fill_nan(0.0, subset=["missing"]).collect()
+
+
 _slow_udf_started = threading.Event()