(fix) partition block overwriting (#1695)

* Fix blocks overwriting

* Add test case

Author: kukushking

awswrangler/distributed/ray/modin/s3/_write_dataset.py

@@ -138,7 +138,7 @@ def _to_partitions_distributed(  # pylint: disable=unused-argument
         func = engine.dispatch_func(func, PandasDataFrame)
 
         @ray_remote
-        def write_partitions(df: pd.DataFrame) -> Tuple[List[str], Dict[str, List[str]]]:
+        def write_partitions(df: pd.DataFrame, block_index: int) -> Tuple[List[str], Dict[str, List[str]]]:
             paths, partitions_values = _to_partitions_func(
                 # Passing a copy of the data frame because data in ray object store is immutable
                 # and that leads to "ValueError: buffer source array is read-only" during df.groupby()
@@ -154,7 +154,7 @@ def write_partitions(df: pd.DataFrame) -> Tuple[List[str], Dict[str, List[str]]]
                 table_type=table_type,
                 transaction_id=transaction_id,
                 bucketing_info=None,
-                filename_prefix=filename_prefix,
+                filename_prefix=f"{filename_prefix}_{block_index:05d}",
                 partition_cols=partition_cols,
                 partitions_types=partitions_types,
                 boto3_session=None,
@@ -163,7 +163,9 @@ def write_partitions(df: pd.DataFrame) -> Tuple[List[str], Dict[str, List[str]]]
             return paths, partitions_values
 
         block_object_refs = ray.data.from_modin(df).get_internal_block_refs()
-        result = ray_get([write_partitions(object_ref) for object_ref in block_object_refs])
+        result = ray_get(
+            [write_partitions(object_ref, block_index) for block_index, object_ref in enumerate(block_object_refs)]
+        )
         paths = [path for row in result for path in row[0]]
         partitions_values = {
             partition_key: partition_value for row in result for partition_key, partition_value in row[1].items()

tests/load/test_s3.py

@@ -1,4 +1,4 @@
-import pandas as pd
+import modin.pandas as pd
 import pytest
 import ray
 
@@ -31,6 +31,13 @@ def big_modin_df():
     return frame
 
 
+def _modin_repartition(df: pd.DataFrame, num_blocks: int) -> pd.DataFrame:
+    """Repartition modin dataframe into n blocks"""
+    dataset = ray.data.from_modin(df)
+    dataset = dataset.repartition(num_blocks)
+    return dataset.to_modin()
+
+
 @pytest.mark.repeat(1)
 @pytest.mark.parametrize("benchmark_time", [180])
 def test_s3_select(benchmark_time):
@@ -90,6 +97,20 @@ def test_s3_write_parquet_dataset(df_s, path, partition_cols, bucketing_info, be
     assert timer.elapsed_time < benchmark_time
 
 
+@pytest.mark.parametrize("benchmark_time", [200])
+@pytest.mark.parametrize("partition_cols", [None, ["payment_type"]])
+@pytest.mark.parametrize("num_blocks", [None, 1, 5])
+def test_s3_write_parquet_blocks(df_s, path, partition_cols, num_blocks, benchmark_time):
+    dataset = True if partition_cols else False
+    if num_blocks:
+        df_s = _modin_repartition(df_s, num_blocks)
+    with ExecutionTimer(f"elapsed time of wr.s3.to_parquet() with repartitioning into {num_blocks} blocks") as timer:
+        wr.s3.to_parquet(df_s, path=path, dataset=dataset, partition_cols=partition_cols)
+    df = wr.s3.read_parquet(path=path, dataset=dataset)
+    assert df.shape == df_s.shape
+    assert timer.elapsed_time < benchmark_time
+
+
 @pytest.mark.parametrize("benchmark_time", [5])
 def test_s3_delete_objects(path, path2, benchmark_time):
     df = pd.DataFrame({"id": [1, 2, 3]})