From 3dd102598cfc72cf6e9835a29f7c9da96ebe64be Mon Sep 17 00:00:00 2001
From: Devin Petersohn <devin.petersohn@gmail.com>
Date: Fri, 17 Jan 2020 11:34:12 -0800
Subject: [PATCH] Implement `tz_localize` for `Series` and `DataFrame`

* Resolves #1012
* Add implementation for `tz_localize` using pandas to dermine the new index
* Update and improve tests
---
 docs/UsingPandasonRay/dataframe_supported.rst |  2 +-
 modin/pandas/base.py                          | 21 ++++++++++++-------
 modin/pandas/test/test_dataframe.py           | 16 +++++++++-----
 modin/pandas/test/test_series.py              | 15 +++++++++----
 4 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/docs/UsingPandasonRay/dataframe_supported.rst b/docs/UsingPandasonRay/dataframe_supported.rst
index 9563587197b..79e19d0dd91 100644
--- a/docs/UsingPandasonRay/dataframe_supported.rst
+++ b/docs/UsingPandasonRay/dataframe_supported.rst
@@ -443,7 +443,7 @@ default to pandas.
 +---------------------------+---------------------------------+----------------------------------------------------+
 | `tz_convert`_             | D                               |                                                    |
 +---------------------------+---------------------------------+----------------------------------------------------+
-| `tz_localize`_            | D                               |                                                    |
+| `tz_localize`_            | Y                               |                                                    |
 +---------------------------+---------------------------------+----------------------------------------------------+
 | `unstack`_                | D                               |                                                    |
 +---------------------------+---------------------------------+----------------------------------------------------+
diff --git a/modin/pandas/base.py b/modin/pandas/base.py
index 92e6bd73119..4ad3371bb86 100644
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -3161,15 +3161,20 @@ def tz_convert(self, tz, axis=0, level=None, copy=True):
     def tz_localize(
         self, tz, axis=0, level=None, copy=True, ambiguous="raise", nonexistent="raise"
     ):
-        return self._default_to_pandas(
-            "tz_localize",
-            tz,
-            axis=axis,
-            level=level,
-            copy=copy,
-            ambiguous=ambiguous,
-            nonexistent=nonexistent,
+        axis = self._get_axis_number(axis)
+        new_labels = (
+            pandas.Series(index=self.axes[axis])
+            .tz_localize(
+                tz,
+                axis=axis,
+                level=level,
+                copy=False,
+                ambiguous=ambiguous,
+                nonexistent=nonexistent,
+            )
+            .index
         )
+        return self.set_axis(labels=new_labels, axis=axis, inplace=not copy)
 
     def unstack(self, level=-1, fill_value=None):
         return self._default_to_pandas("unstack", level=level, fill_value=fill_value)
diff --git a/modin/pandas/test/test_dataframe.py b/modin/pandas/test/test_dataframe.py
index d8ee3d6a810..f82d39e36c7 100644
--- a/modin/pandas/test/test_dataframe.py
+++ b/modin/pandas/test/test_dataframe.py
@@ -2622,11 +2622,17 @@ def test_tz_convert(self):
             df.tz_localize("America/Los_Angeles").tz_convert("America/Los_Angeles")
 
     def test_tz_localize(self):
-        idx = pd.date_range("1/1/2012", periods=5, freq="M")
-        df = pd.DataFrame(np.random.randint(0, 100, size=(len(idx), 4)), index=idx)
-
-        with pytest.warns(UserWarning):
-            df.tz_localize("America/Los_Angeles")
+        idx = pd.date_range("1/1/2012", periods=400, freq="2D")
+        data = np.random.randint(0, 100, size=(len(idx), 4))
+        modin_df = pd.DataFrame(data, index=idx)
+        pandas_df = pandas.DataFrame(data, index=idx)
+        df_equals(
+            modin_df.tz_localize("UTC", axis=0), pandas_df.tz_localize("UTC", axis=0)
+        )
+        df_equals(
+            modin_df.tz_localize("America/Los_Angeles", axis=0),
+            pandas_df.tz_localize("America/Los_Angeles", axis=0),
+        )
 
     def test_unstack(self):
         data = test_data_values[0]
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
index f2188c29795..1ac2b52c11f 100644
--- a/modin/pandas/test/test_series.py
+++ b/modin/pandas/test/test_series.py
@@ -2607,10 +2607,17 @@ def test_tz_convert():
 
 
 def test_tz_localize():
-    idx = pd.date_range("1/1/2012", periods=5, freq="M")
-    modin_series = pd.Series(np.random.randint(0, 100, size=len(idx)), index=idx)
-    with pytest.warns(UserWarning):
-        modin_series.tz_localize("America/Los_Angeles")
+    idx = pd.date_range("1/1/2012", periods=400, freq="2D")
+    data = np.random.randint(0, 100, size=len(idx))
+    modin_series = pd.Series(data, index=idx)
+    pandas_series = pandas.Series(data, index=idx)
+    df_equals(
+        modin_series.tz_localize("America/Los_Angeles"),
+        pandas_series.tz_localize("America/Los_Angeles"),
+    )
+    df_equals(
+        modin_series.tz_localize("UTC"), pandas_series.tz_localize("UTC"),
+    )
 
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)