Move index replacement logic in fillna to IndexedFrame.

rapidsai · Mar 17, 2022 · 25b5df9 · 25b5df9
1 parent f2bf9aa
commit 25b5df9
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/python/cudf/cudf/core/frame.py b/python/cudf/cudf/core/frame.py
@@ -1037,7 +1037,7 @@ def fillna(
 
         Returns
         -------
-        result : DataFrame
+        result : DataFrame, Series, or Index
             Copy with nulls filled.
 
         Examples
@@ -1165,18 +1165,9 @@ def fillna(
             else:
                 filled_data[col_name] = col.copy(deep=True)
 
-        # TODO: This logic needs to move into the IndexedFrame class.
-        old_index = self._index
-        ret = self._mimic_inplace(
+        return self._mimic_inplace(
             self._from_data(data=filled_data), inplace=inplace,
         )
-        # TODO: Split this logic into the IndexedFrame class.
-        if isinstance(self, cudf.core.indexed_frame.IndexedFrame):
-            if inplace:
-                self._index = old_index
-            else:
-                ret._index = old_index
-        return ret
 
     @_cudf_nvtx_annotate
     def _drop_column(self, name):

diff --git a/python/cudf/cudf/core/indexed_frame.py b/python/cudf/cudf/core/indexed_frame.py
@@ -904,6 +904,18 @@ def _split(self, splits, keep_index=True):
             for i in range(len(splits) + 1)
         ]
 
+    @_cudf_nvtx_annotate
+    def fillna(
+        self, value=None, method=None, axis=None, inplace=False, limit=None
+    ):  # noqa: D102
+        old_index = self._index
+        ret = super().fillna(value, method, axis, inplace, limit)
+        if inplace:
+            self._index = old_index
+        else:
+            ret._index = old_index
+        return ret
+
     def add_prefix(self, prefix):
         """
         Prefix labels with string `prefix`.