From 055e48bd4f4779d7ef0053502a4c9b92c763e0bb Mon Sep 17 00:00:00 2001
From: Andong Zhan <andong.zhan@snowflake.com>
Date: Fri, 26 Jul 2024 10:26:20 -0700
Subject: [PATCH] SNOW-1555437 Fix Index docstrings (#1985)

---
 CHANGELOG.md                                  |  1 +
 .../snowpark/modin/plugin/extensions/index.py | 25 +++++++++++--------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3c65fe30333..5673eae03b1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -42,6 +42,7 @@
 - Removed axis labels and callable names from error messages and telemetry about unsupported aggregations.
 - Fixed AssertionError in `Series.drop_duplicates` and `DataFrame.drop_duplicates` when called after `sort_values`.
 - Fixed a bug in `Index.to_frame` where the result frame's column name may be wrong where name is unspecified.  
+- Fixed a bug where some Index docstrings are ignored. 
 
 
 ## 1.20.0 (2024-07-17)
diff --git a/src/snowflake/snowpark/modin/plugin/extensions/index.py b/src/snowflake/snowpark/modin/plugin/extensions/index.py
index 65da5c6c1f1..0c0457b4679 100644
--- a/src/snowflake/snowpark/modin/plugin/extensions/index.py
+++ b/src/snowflake/snowpark/modin/plugin/extensions/index.py
@@ -23,8 +23,10 @@
 
 from __future__ import annotations
 
+from functools import wraps
 from typing import Any, Callable, Hashable, Iterator, Literal
 
+import modin
 import numpy as np
 import pandas as native_pd
 from pandas._libs import lib
@@ -50,6 +52,7 @@ def is_lazy_check(func: Callable) -> Callable:
     Decorator method for separating function calls for lazy indexes and non-lazy (column) indexes
     """
 
+    @wraps(func)
     def check_lazy(*args: Any, **kwargs: Any) -> Any:
         func_name = func.__name__
 
@@ -1601,9 +1604,9 @@ def to_series(
     @is_lazy_check
     def to_frame(
         self, index: bool = True, name: Hashable | None = lib.no_default
-    ) -> DataFrame:
+    ) -> modin.pandas.DataFrame:
         """
-        Create a DataFrame with a column containing the Index.
+        Create a :class:`DataFrame` with a column containing the Index.
 
         Parameters
         ----------
@@ -1616,8 +1619,8 @@ def to_frame(
 
         Returns
         -------
-        DataFrame
-            DataFrame containing the original Index data.
+        :class:`DataFrame`
+            :class:`DataFrame` containing the original Index data.
 
         See Also
         --------
@@ -1627,7 +1630,7 @@ def to_frame(
         Examples
         --------
         >>> idx = pd.Index(['Ant', 'Bear', 'Cow'], name='animal')
-        >>> idx.to_frame()
+        >>> idx.to_frame()   # doctest: +NORMALIZE_WHITESPACE
                animal
         animal
         Ant       Ant
@@ -1637,10 +1640,10 @@ def to_frame(
         By default, the original Index is reused. To enforce a new Index:
 
         >>> idx.to_frame(index=False)
-            animal
-        0   Ant
-        1  Bear
-        2   Cow
+          animal
+        0    Ant
+        1   Bear
+        2    Cow
 
         To override the name of the resulting column, specify `name`:
 
@@ -2099,9 +2102,9 @@ def get_indexer_for(self, target: Any) -> Any:
 
         Examples
         --------
-        # Snowpark pandas converts np.nan, pd.NA, pd.NaT to None
+        Note Snowpark pandas converts np.nan, pd.NA, pd.NaT to None
         >>> idx = pd.Index([np.nan, 'var1', np.nan])
-        >>> idx.get_indexer_for([np.nan])
+        >>> idx.get_indexer_for([None])
         array([0, 2])
         """
         WarningMessage.index_to_pandas_warning("get_indexer_for")