adding fillna (#335)

* adding fillna

* update nested logic for value matching

* minor type change

Author: Graciaaa3

docs/reference/nestedframe.rst

@@ -43,6 +43,7 @@ Extended Pandas.DataFrame Interface
     NestedFrame.max
     NestedFrame.describe
     NestedFrame.explode
+    NestedFrame.fillna
 
 I/O
 ~~~~~~~~~

src/nested_pandas/nestedframe/core.py

@@ -10,7 +10,7 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 from pandas._libs import lib
-from pandas._typing import Any, AnyAll, Axis, IndexLabel
+from pandas._typing import Any, AnyAll, Axis, Hashable, IndexLabel, Mapping
 from pandas.api.extensions import no_default
 from pandas.core.computation.eval import Expr, ensure_scope
 from pandas.core.dtypes.inference import is_list_like
@@ -1114,6 +1114,91 @@ def explode(self, column: IndexLabel, ignore_index: bool = False):
             result.index.name = None
         return result
 
+    def fillna(
+        self,
+        value: Hashable | Mapping | pd.Series | pd.DataFrame | None = None,
+        *,
+        axis: Axis | None = None,
+        inplace: bool = False,
+        limit: int | None = None,
+    ) -> NestedFrame | None:
+        """
+        Fill NA/NaN values using the specified method for base and nested columns.
+
+        Parameters
+        ----------
+        value : scalar, dict, Series, or DataFrame
+            Value to use to fill holes (e.g. 0), alternately a
+            dict/Series/DataFrame of values specifying which value to use for
+            each column.  Values not in the dict/Series/DataFrame will not be filled.
+            This value cannot be a list.
+        axis : {axes_single_arg}, default None
+            Axis along which to fill missing values.
+        inplace : bool, default False
+            If True, fill in-place. Note: this will modify any
+            other views on this object (e.g., a no-copy slice for a column in a
+            NestedFrame).
+        limit : int, default None
+            The maximum number of entries along the entire axis where NaNs will be
+            filled. Must be greater than 0 if not None. Currently, limit on nested
+            columns is not supported, meaning that all Nans will be filled (if there
+            is a value specified) regardless of the input.
+
+        Returns
+        -------
+        NestedFrame or None
+            NestedFrame with missing values filled or None if ``inplace=True``.
+
+        See Also
+        --------
+        :meth:`pandas.DataFrame.fillna`
+
+        Examples
+        --------
+        >>> import nested_pandas as npd
+        >>> nf = npd.NestedFrame(
+        ...     data={"a": [np.nan, 20, np.nan], "b": [np.nan, np.nan, 30], "c": [10, np.nan, np.nan]},
+        ...     index=[0, 1, 2]
+        ... )
+        >>> nested = pd.DataFrame(
+        ...     data={"d": [np.nan, np.nan, np.nan], "e": [np.nan, 1, np.nan]},
+        ...     index=[0, 1, 2]
+        ... )
+        >>> nf = nf.add_nested(nested, "nested")
+
+        >>> nf.fillna(0)
+              a     b     c              nested
+        0   0.0   0.0  10.0  [{d: 0.0, e: 0.0}]
+        1  20.0   0.0   0.0  [{d: 0.0, e: 1.0}]
+        2   0.0  30.0   0.0  [{d: 0.0, e: 0.0}]
+
+        """
+
+        if not self.nested_columns:
+            return super().fillna(value=value, axis=axis, inplace=inplace, limit=limit)
+
+        base_cols = [col for col in self.columns if col not in self.nested_columns]
+        filled_df = super().__getitem__(base_cols).fillna(value=value, axis=axis, inplace=False, limit=limit)
+
+        for nest_col in self.nested_columns:
+            nested_df = self[nest_col].nest.to_flat()
+            nested_value: Any
+            if isinstance(value, Mapping):
+                nested_value = {}
+                for k, v in value.items():
+                    if k.startswith(f"{nest_col}."):
+                        subcol = k.split(".", 1)[1]  # strip prefix
+                        nested_value[subcol] = v
+            else:
+                nested_value = value
+            nested_df = nested_df.fillna(value=nested_value, axis=axis, inplace=False, limit=None)
+            filled_df = filled_df.add_nested(nested_df, nest_col)
+
+        if inplace:
+            self._update_inplace(filled_df)
+            return None
+        return filled_df
+
     def eval(self, expr: str, *, inplace: bool = False, **kwargs) -> Any | None:
         """Evaluate a string describing operations on NestedFrame columns.
 

tests/nested_pandas/nestedframe/test_nestedframe.py

@@ -1787,6 +1787,53 @@ def test_explode_non_unique_index():
         nf.explode(["nested", "XXX", "AAA"])
 
 
+def test_fillna():
+    """Test NestedFrame.fillna give correct result with NA/NaN filled"""
+
+    base = NestedFrame(data={"a": [np.nan, np.nan, 3], "b": [2, np.nan, 6]}, index=[0, 1, 2])
+    nested = pd.DataFrame(
+        data={"c": [0, 2, np.nan, 1, np.nan, 3, 1, 4, 1], "d": [np.nan, 4, np.nan, 5, 3, 1, np.nan, 3, 4]},
+        index=[0, 0, 0, 1, 1, 1, 2, 2, 2],
+    )
+    nested2 = pd.DataFrame(
+        data={"e": [np.nan, np.nan, np.nan, 1, 4, np.nan, 4, 1], "f": [5, 4, 7, 5, 1, 9, 3, 4]},
+        index=[0, 0, 0, 1, 1, 2, 2, 2],
+    )
+
+    # only base columns
+    r0 = base.fillna(0)
+    assert (r0["a"] == pd.Series([0, 0, 3])).all()
+    assert (r0["b"] == pd.Series([2, 0, 6])).all()
+
+    # 1 nested column
+    base = base.add_nested(nested, "nested")
+    r1 = base.fillna(0)
+    expected1 = pd.Series([0, 4, 0, 5, 3, 1, 0, 3, 4], index=[0, 0, 0, 1, 1, 1, 2, 2, 2])
+    assert (r1["nested.d"] == expected1).all()
+
+    r2 = base.fillna({"a": 0, "b": 1, "nested.c": 2, "nested.d": 3})
+    expected2 = pd.Series([0, 2, 2, 1, 2, 3, 1, 4, 1], index=[0, 0, 0, 1, 1, 1, 2, 2, 2])
+    assert (r2["nested.c"] == expected2).all()
+    assert (r2["a"] == pd.Series([0, 0, 3])).all()
+
+    r3 = base.fillna(0, limit=1)
+    assert np.isnan(r3["a"][1])
+
+    # 2 nested columns
+    base = base.add_nested(nested2, "nested2")
+    r4 = base.fillna(0)
+    expected4 = pd.Series([0, 0, 0, 1, 4, 0, 4, 1], index=[0, 0, 0, 1, 1, 2, 2, 2])
+    assert (r4["nested2.e"] == expected4).all()
+
+    # inplace check
+    expected5 = pd.Series([0, 2, np.nan, 1, np.nan, 3, 1, 4, 1], index=[0, 0, 0, 1, 1, 1, 2, 2, 2])
+    assert (base["nested.c"] == expected5).all()
+
+    base.fillna({"nested.c": 0}, inplace=True)
+    expected6 = pd.Series([0, 2, 0, 1, 0, 3, 1, 4, 1], index=[0, 0, 0, 1, 1, 1, 2, 2, 2])
+    assert (base["nested.c"] == expected6).all()
+
+
 def test_eval():
     """
     Test basic behavior of NestedFrame.eval, and that it can handle nested references