Fix index recovery combined with columns filter. #408

Author: igorborgest

awswrangler/s3/_read_parquet.py

@@ -170,19 +170,24 @@ def _read_parquet_metadata(
 
 def _apply_index(df: pd.DataFrame, metadata: Dict[str, Any]) -> pd.DataFrame:
     index_columns: List[Any] = metadata["index_columns"]
+    ignore_index: bool = True
+    _logger.debug("df.columns: %s", df.columns)
+
     if index_columns:
         if isinstance(index_columns[0], str):
-            df = df.set_index(keys=index_columns, drop=True, inplace=False, verify_integrity=False)
+            indexes: List[str] = [i for i in index_columns if i in df.columns]
+            if indexes:
+                df = df.set_index(keys=indexes, drop=True, inplace=False, verify_integrity=False)
+                ignore_index = False
         elif isinstance(index_columns[0], dict) and index_columns[0]["kind"] == "range":
             col = index_columns[0]
             if col["kind"] == "range":
                 df.index = pd.RangeIndex(start=col["start"], stop=col["stop"], step=col["step"])
+                ignore_index = False
                 if col["name"] is not None and col["name"].startswith("__index_level_") is False:
                     df.index.name = col["name"]
         df.index.names = [None if n is not None and n.startswith("__index_level_") else n for n in df.index.names]
-        ignore_index: bool = False
-    else:
-        ignore_index = True
+
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", category=UserWarning)
         df._awswrangler_ignore_index = ignore_index  # pylint: disable=protected-access

tests/test_s3_parquet.py

@@ -325,6 +325,43 @@ def test_multi_index_recovery_nameless(path, use_threads):
     assert df.reset_index().equals(df2.reset_index())
 
 
+@pytest.mark.parametrize("use_threads", [True, False])
+@pytest.mark.parametrize("name", [None, "foo"])
+@pytest.mark.parametrize("pandas", [True, False])
+def test_index_columns(path, use_threads, name, pandas):
+    df = pd.DataFrame({"c0": [0, 1], "c1": [2, 3]}, dtype="Int64")
+    df.index.name = name
+    path_file = f"{path}0.parquet"
+    if pandas:
+        df.to_parquet(path_file, index=True)
+    else:
+        wr.s3.to_parquet(df, path_file, index=True)
+    wr.s3.wait_objects_exist(paths=[path_file], use_threads=use_threads)
+    df2 = wr.s3.read_parquet([path_file], columns=["c0"], use_threads=use_threads)
+    assert df[["c0"]].equals(df2)
+
+
+@pytest.mark.parametrize("use_threads", [True, False])
+@pytest.mark.parametrize("name", [None, "foo"])
+@pytest.mark.parametrize("pandas", [True, False])
+@pytest.mark.parametrize("drop", [True, False])
+def test_range_index_columns(path, use_threads, name, pandas, drop):
+    df = pd.DataFrame({"c0": [0, 1], "c1": [2, 3]}, dtype="Int64", index=pd.RangeIndex(start=5, stop=7, step=1))
+    df.index.name = name
+    path_file = f"{path}0.parquet"
+    if pandas:
+        df.to_parquet(path_file, index=True)
+    else:
+        wr.s3.to_parquet(df, path_file, index=True)
+    wr.s3.wait_objects_exist(paths=[path_file], use_threads=use_threads)
+
+    name = "__index_level_0__" if name is None else name
+    columns = ["c0"] if drop else [name, "c0"]
+    df2 = wr.s3.read_parquet([path_file], columns=columns, use_threads=use_threads)
+
+    assert df[["c0"]].reset_index(level=0, drop=drop).equals(df2.reset_index(level=0, drop=drop))
+
+
 def test_to_parquet_dataset_sanitize(path):
     df = pd.DataFrame({"C0": [0, 1], "camelCase": [2, 3], "c**--2": [4, 5], "Par": ["a", "b"]})