[SNOW-1524894]: Add support for reindex (#1893)

snowflakedb · Jul 15, 2024 · d1f5e16 · d1f5e16
1 parent 1a734c9
commit d1f5e16
Show file tree

Hide file tree

Showing 11 changed files with 1,833 additions and 33 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -25,6 +25,7 @@
 - Added support for `DataFrame.corr`.
 - Added support for `limit` parameter when `method` parameter is used in `fillna`.
 - Added support for `DataFrame.equals` and `Series.equals`.
+- Added support for `DataFrame.reindex` and `Series.reindex`.
 
 #### Bug Fixes
 - Fixed an issue when using np.where and df.where when the scalar 'other' is the literal 0.

diff --git a/docs/source/modin/supported/dataframe_supported.rst b/docs/source/modin/supported/dataframe_supported.rst
@@ -330,7 +330,8 @@ Methods
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``rdiv``                    | P                               | ``level``                        |                                                    |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
-| ``reindex``                 | N                               |                                  |                                                    |
+| ``reindex``                 | P                               |                                  | ``N`` if axis is MultiIndex or method is           |
+|                             |                                 |                                  | ``nearest``.                                       |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``reindex_like``            | N                               |                                  |                                                    |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+

diff --git a/docs/source/modin/supported/series_supported.rst b/docs/source/modin/supported/series_supported.rst
@@ -322,7 +322,8 @@ Methods
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``rdivmod``                 | N                               |                                  |                                                    |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
-| ``reindex``                 | N                               |                                  |                                                    |
+| ``reindex``                 | P                               |                                  | ``N`` if the series has MultiIndex, or method      |
+|                             |                                 |                                  | is ``nearest``.                                    |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``reindex_like``            | N                               |                                  |                                                    |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+

diff --git a/src/snowflake/snowpark/modin/pandas/base.py b/src/snowflake/snowpark/modin/pandas/base.py
@@ -2584,7 +2584,6 @@ def _ensure_index(self, index_like, axis=0):  # noqa: PR01, RT01, D200
                 pass
         return ensure_index(index_like)
 
-    @base_not_implemented()
     def reindex(
         self,
         index=None,
@@ -2596,7 +2595,10 @@ def reindex(
         Conform `BasePandasDataset` to new index with optional filling logic.
         """
         # TODO: SNOW-1119855: Modin upgrade - modin.pandas.base.BasePandasDataset
-
+        if kwargs.get("limit", None) is not None and kwargs.get("method", None) is None:
+            raise ValueError(
+                "limit argument only valid if doing pad, backfill or nearest reindexing"
+            )
         new_query_compiler = None
         if index is not None:
             if not isinstance(index, pandas.Index) or not index.equals(self.index):

diff --git a/src/snowflake/snowpark/modin/pandas/dataframe.py b/src/snowflake/snowpark/modin/pandas/dataframe.py
@@ -2073,7 +2073,6 @@ def rename(
             new_query_compiler=new_qc, inplace=inplace
         )
 
-    @dataframe_not_implemented()
     def reindex(
         self,
         labels=None,

diff --git a/src/snowflake/snowpark/modin/pandas/series.py b/src/snowflake/snowpark/modin/pandas/series.py
@@ -1672,10 +1672,12 @@ def ravel(self, order="C"):  # noqa: PR01, RT01, D200
     def reindex(self, *args, **kwargs):
         if args:
             if len(args) > 1:
-                raise TypeError("Only one positional argument ('index') is allowed")
+                raise TypeError(
+                    "Series.reindex() takes from 1 to 2 positional arguments but 3 were given"
+                )
             if "index" in kwargs:
                 raise TypeError(
-                    "'index' passed as both positional and keyword argument"
+                    "Series.reindex() got multiple values for argument 'index'"
                 )
             kwargs.update({"index": args[0]})
         index = kwargs.pop("index", None)
@@ -1685,10 +1687,11 @@ def reindex(self, *args, **kwargs):
         limit = kwargs.pop("limit", None)
         tolerance = kwargs.pop("tolerance", None)
         fill_value = kwargs.pop("fill_value", None)
+        kwargs.pop("axis", None)
         if kwargs:
             raise TypeError(
-                "reindex() got an unexpected keyword "
-                + f'argument "{list(kwargs.keys())[0]}"'
+                "Series.reindex() got an unexpected keyword "
+                + f"argument '{list(kwargs.keys())[0]}'"
             )
         return super().reindex(
             index=index,