ARROW-7914: [Python] Allow pandas datetime as index for feather

Salonijain27 · salonijain27 · jorisvandenbossche · commit 20ec0fda708b · 2022-04-21T18:48:14.000+02:00
Closes apache#12821 from Salonijain27/ARROW-7914_fetch_update Lead-authored-by: Salonijain27 <salj7856@gmail.com> Co-authored-by: salonijain27 <salonijain@Salonis-MacBook-Pro.local> Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
diff --git a/python/pyarrow/feather.py b/python/pyarrow/feather.py
@@ -151,7 +151,17 @@ def write_feather(df, dest, compression=None, compression_level=None,
             df = df.to_dense()
 
     if _pandas_api.is_data_frame(df):
-        table = Table.from_pandas(df, preserve_index=False)
+        # Feather v1 creates a new column in the resultant Table to
+        # store index information if index type is not RangeIndex
+
+        if version == 1:
+            preserve_index = False
+        elif version == 2:
+            preserve_index = None
+        else:
+            raise ValueError("Version value should either be 1 or 2")
+
+        table = Table.from_pandas(df, preserve_index=preserve_index)
 
         if version == 1:
             # Version 1 does not chunking
diff --git a/python/pyarrow/tests/test_feather.py b/python/pyarrow/tests/test_feather.py
@@ -30,7 +30,6 @@
 from pyarrow.feather import (read_feather, write_feather, read_table,
                              FeatherDataset)
 
-
 try:
     from pandas.testing import assert_frame_equal
     import pandas as pd
@@ -90,13 +89,18 @@ def _check_pandas_roundtrip(df, expected=None, path=None,
     if path is None:
         path = random_path()
 
+    if version is None:
+        version = 2
+
     TEST_FILES.append(path)
     write_feather(df, path, compression=compression,
                   compression_level=compression_level, version=version)
+
     if not os.path.exists(path):
         raise Exception('file not written')
 
     result = read_feather(path, columns, use_threads=use_threads)
+
     if expected is None:
         expected = df
 
@@ -504,8 +508,10 @@ def test_out_of_float64_timestamp_with_nulls(version):
 def test_non_string_columns(version):
     df = pd.DataFrame({0: [1, 2, 3, 4],
                        1: [True, False, True, False]})
+    expected = df
 
-    expected = df.rename(columns=str)
+    if version == 1:
+        expected = df.rename(columns=str)
     _check_pandas_roundtrip(df, expected, version=version)
 
 
@@ -820,3 +826,15 @@ def test_feather_v017_experimental_compression_backward_compatibility(datadir):
     expected = pa.table({'a': range(5)})
     result = read_table(datadir / "v0.17.0.version.2-compression.lz4.feather")
     assert result.equals(expected)
+
+
+@pytest.mark.pandas
+def test_preserve_index_pandas(version):
+    df = pd.DataFrame({'a': [1, 2, 3]}, index=['a', 'b', 'c'])
+
+    if version == 1:
+        expected = df.reset_index(drop=True).rename(columns=str)
+    else:
+        expected = df
+
+    _check_pandas_roundtrip(df, expected, version=version)
