feat: Support partitioning by index cols (#2528)

Signed-off-by: Anton Kukushkin <[email protected]>

Author: kukushking

awswrangler/_data_types.py

@@ -554,9 +554,11 @@ def pyarrow_types_from_pandas(  # pylint: disable=too-many-branches,too-many-sta
             fields = pa.Schema.from_pandas(df=df.reset_index().drop(columns=cols), preserve_index=False)
         for field in fields:
             name = str(field.name)
-            _logger.debug("Inferring PyArrow type from index: %s", name)
-            cols_dtypes[name] = field.type
-            indexes.append(name)
+            # Check if any of the index columns must be ignored
+            if name not in ignore_cols:
+                _logger.debug("Inferring PyArrow type from index: %s", name)
+                cols_dtypes[name] = field.type
+                indexes.append(name)
 
     # Merging Index
     sorted_cols: List[str] = indexes + list(df.columns) if index_left is True else list(df.columns) + indexes

awswrangler/s3/_write_dataset.py

@@ -150,7 +150,13 @@ def _to_partitions(
     for keys, subgroup in df.groupby(by=partition_cols, observed=True):
         # Keys are either a primitive type or a tuple if partitioning by multiple cols
         keys = (keys,) if not isinstance(keys, tuple) else keys  # ruff: noqa: PLW2901
-        subgroup = subgroup.drop(partition_cols, axis="columns")  # ruff: noqa: PLW2901
+        # Drop partition columns from df
+        subgroup.drop(
+            columns=[col for col in partition_cols if col in subgroup.columns],
+            inplace=True,
+        )  # ruff: noqa: PLW2901
+        # Drop index levels if partitioning by index columns
+        subgroup = subgroup.droplevel(level=[col for col in partition_cols if col in subgroup.index.names])
         prefix = _delete_objects(
             keys=keys,
             path_root=path_root,