add S3.copy_listed_objects feature.

Author: igorborgest

awswrangler/s3.py

@@ -57,6 +57,10 @@ def parse_path(path):
                 path += "/"
         return bucket, path
 
+    @staticmethod
+    def parse_object_path(path):
+        return path.replace("s3://", "").split("/", 1)
+
     def delete_objects(self, path):
         bucket, path = self.parse_path(path=path)
         client = self._session.boto3_session.client(
@@ -127,11 +131,13 @@ def delete_listed_objects(self, objects_paths, procs_io_bound=None):
                     proc.daemon = False
                     proc.start()
                     procs.append(proc)
+                for proc in procs:
+                    proc.join()
             else:
-                self.delete_objects_batch(self._session.primitives, bucket,
-                                          batch)
-            for proc in procs:
-                proc.join()
+                self.delete_objects_batch(
+                    session_primitives=self._session.primitives,
+                    bucket=bucket,
+                    batch=batch)
 
     def delete_not_listed_objects(self, objects_paths, procs_io_bound=None):
         if not procs_io_bound:
@@ -274,3 +280,75 @@ def get_objects_sizes(self, objects_paths, procs_io_bound=None):
             LOGGER.debug(f"Closing proc number: {i}")
             receive_pipes[i].close()
         return objects_sizes
+
+    def copy_listed_objects(self,
+                            objects_paths,
+                            source_path,
+                            target_path,
+                            mode="append",
+                            procs_io_bound=None):
+        if not procs_io_bound:
+            procs_io_bound = self._session.procs_io_bound
+        LOGGER.debug(f"procs_io_bound: {procs_io_bound}")
+        LOGGER.debug(f"len(objects_paths): {len(objects_paths)}")
+        if source_path[-1] == "/":
+            source_path = source_path[:-1]
+        if target_path[-1] == "/":
+            target_path = target_path[:-1]
+
+        if mode == "overwrite":
+            LOGGER.debug(f"Deleting to overwrite: {target_path}")
+            self._session.s3.delete_objects(path=target_path)
+        elif mode == "overwrite_partitions":
+            objects_wo_prefix = [
+                o.replace(f"{source_path}/", "") for o in objects_paths
+            ]
+            objects_wo_filename = [
+                f"{o.rpartition('/')[0]}/" for o in objects_wo_prefix
+            ]
+            partitions_paths = list(set(objects_wo_filename))
+            target_partitions_paths = [
+                f"{target_path}/{p}" for p in partitions_paths
+            ]
+            for path in target_partitions_paths:
+                LOGGER.debug(f"Deleting to overwrite_partitions: {path}")
+                self._session.s3.delete_objects(path=path)
+
+        batch = []
+        for obj in objects_paths:
+            object_wo_prefix = obj.replace(f"{source_path}/", "")
+            target_object = f"{target_path}/{object_wo_prefix}"
+            batch.append((obj, target_object))
+
+        if procs_io_bound > 1:
+            bounders = calculate_bounders(len(objects_paths), procs_io_bound)
+            LOGGER.debug(f"bounders: {bounders}")
+            procs = []
+            for bounder in bounders:
+                proc = mp.Process(
+                    target=self.copy_objects_batch,
+                    args=(
+                        self._session.primitives,
+                        batch[bounder[0]:bounder[1]],
+                    ),
+                )
+                proc.daemon = False
+                proc.start()
+                procs.append(proc)
+            for proc in procs:
+                proc.join()
+        else:
+            self.copy_objects_batch(
+                session_primitives=self._session.primitives, batch=batch)
+
+    @staticmethod
+    def copy_objects_batch(session_primitives, batch):
+        session = session_primitives.session
+        resource = session.boto3_session.resource(
+            service_name="s3", config=session.botocore_config)
+        LOGGER.debug(f"len(batch): {len(batch)}")
+        for source_obj, target_obj in batch:
+            source_bucket, source_key = S3.parse_object_path(path=source_obj)
+            copy_source = {"Bucket": source_bucket, "Key": source_key}
+            target_bucket, target_key = S3.parse_object_path(path=target_obj)
+            resource.meta.client.copy(copy_source, target_bucket, target_key)

testing/test_awswrangler/test_s3.py

@@ -4,6 +4,7 @@
 
 import pytest
 import boto3
+import pandas
 
 from awswrangler import Session
 
@@ -81,6 +82,15 @@ def bucket(session, cloudformation_outputs):
     session.s3.delete_objects(path=f"s3://{bucket}/")
 
 
+@pytest.fixture(scope="module")
+def database(cloudformation_outputs):
+    if "GlueDatabaseName" in cloudformation_outputs:
+        database = cloudformation_outputs.get("GlueDatabaseName")
+    else:
+        raise Exception("You must deploy the test infrastructure using SAM!")
+    yield database
+
+
 @pytest.mark.parametrize("objects_num", [1, 10, 1001, 2001, 3001])
 def test_delete_objects(session, bucket, objects_num):
     print("Starting writes...")
@@ -147,3 +157,53 @@ def test_get_objects_sizes(session, bucket, objects_num):
     session.s3.delete_objects(path=path)
     for _, object_size in objects_sizes.items():
         assert object_size == 10
+
+
+@pytest.mark.parametrize("mode, procs_io_bound", [
+    ("append", 1),
+    ("overwrite", 1),
+    ("overwrite_partitions", 1),
+    ("append", 8),
+    ("overwrite", 8),
+    ("overwrite_partitions", 8),
+])
+def test_copy_listed_objects(session, bucket, database, mode, procs_io_bound):
+    path0 = f"s3://{bucket}/test_move_objects_0/"
+    path1 = f"s3://{bucket}/test_move_objects_1/"
+    print("Starting deletes...")
+    session.s3.delete_objects(path=path0)
+    session.s3.delete_objects(path=path1)
+    dataframe = pandas.read_csv("data_samples/micro.csv")
+    print("Starting writing path0...")
+    session.pandas.to_parquet(
+        dataframe=dataframe,
+        database=database,
+        path=path0,
+        preserve_index=False,
+        mode="overwrite",
+        partition_cols=["name", "date"],
+    )
+    print("Starting writing path0...")
+    objects_paths = session.pandas.to_parquet(
+        dataframe=dataframe,
+        path=path1,
+        preserve_index=False,
+        mode="overwrite",
+        partition_cols=["name", "date"],
+    )
+    print("Starting move...")
+    session.s3.copy_listed_objects(
+        objects_paths=objects_paths,
+        source_path=path1,
+        target_path=path0,
+        mode=mode,
+        procs_io_bound=procs_io_bound,
+    )
+    print("Asserting...")
+    sleep(1)
+    dataframe2 = session.pandas.read_sql_athena(
+        sql="select * from test_move_objects_0", database=database)
+    if mode == "append":
+        assert 2 * len(dataframe.index) == len(dataframe2.index)
+    else:
+        assert len(dataframe.index) == len(dataframe2.index)