Merge pull request #34 from usc-isi-i2/dongyul

allow use to choose if to keep the dataframe index

rltk/io/reader/dataframe_reader.py

@@ -9,11 +9,19 @@ class DataFrameReader(Reader):
     
     Args:
         df (pandas.DataFrame): DataFrame.
+        keep_dataframe_default_index (bool, optional): if True, add a key "dataframe_default_index" holding the \
+                                                       original index in df. [default value: False]
+
     """
 
-    def __init__(self, df: pd.DataFrame):
+    def __init__(self, df: pd.DataFrame, keep_dataframe_default_index: bool=False):
         self._df = df
+        self._keep_dataframe_default_index = keep_dataframe_default_index
 
     def __next__(self):
-        for i, item in self._df.iterrows():
-            yield dict(item.to_dict(), dataframe_default_index=i)
+        if self._keep_dataframe_default_index:
+            for i, item in self._df.iterrows():
+                yield dict(item.to_dict(), dataframe_default_index=i)
+        else:
+            for _, item in self._df.iterrows():
+                yield item.to_dict()

rltk/tests/test_io_reader.py

@@ -17,6 +17,12 @@ def test_array_reader():
 def test_dataframe_reader():
     df = pd.DataFrame(arr)
     for idx, obj in enumerate(DataFrameReader(df)):
+        assert obj == arr[idx]
+
+
+def test_dataframe_reader_keep_index():
+    df = pd.DataFrame(arr)
+    for idx, obj in enumerate(DataFrameReader(df, True)):
         assert obj == dict(**arr[idx], dataframe_default_index=idx)