Add s3.copy_objects() #186

Author: igorborgest

awswrangler/_utils.py

@@ -28,6 +28,13 @@ def client(service_name: str, session: Optional[boto3.Session] = None) -> boto3.
     )
 
 
+def resource(service_name: str, session: Optional[boto3.Session] = None) -> boto3.resource:
+    """Create a valid boto3.resource."""
+    return ensure_session(session=session).resource(
+        service_name=service_name, use_ssl=True, config=botocore.config.Config(retries={"max_attempts": 15})
+    )
+
+
 def parse_path(path: str) -> Tuple[str, str]:
     """Split a full S3 path in bucket and key strings.
 

awswrangler/s3.py

@@ -16,6 +16,7 @@
 import pyarrow.lib  # type: ignore
 import pyarrow.parquet  # type: ignore
 import s3fs  # type: ignore
+from boto3.s3.transfer import TransferConfig  # type: ignore
 from pandas.io.common import infer_compression  # type: ignore
 
 from awswrangler import _data_types, _utils, catalog, exceptions
@@ -2051,3 +2052,99 @@ def read_parquet_table(
         boto3_session=boto3_session,
         s3_additional_kwargs=s3_additional_kwargs,
     )
+
+
+def copy_objects(
+    paths: List[str],
+    source_path: str,
+    target_path: str,
+    mode: str = "append",
+    use_threads: bool = True,
+    boto3_session: Optional[boto3.Session] = None,
+) -> List[str]:
+    """Copy a list of S3 objects to another S3 directory.
+
+    Note
+    ----
+    In case of `use_threads=True` the number of process that will be spawned will be get from os.cpu_count().
+
+    Parameters
+    ----------
+    paths : List[str]
+        List of S3 objects paths (e.g. [s3://bucket/dir0/key0, s3://bucket/dir0/key1]).
+    source_path : str,
+        S3 Path for the source directory.
+    target_path : str,
+        S3 Path for the target directory.
+    mode: str, optional
+        ``append`` (Default), ``overwrite``, ``overwrite_partitions``.
+    use_threads : bool
+        True to enable concurrent requests, False to disable multiple threads.
+        If enabled os.cpu_count() will be used as the max number of threads.
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    List[str]
+        List of new objects paths.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> wr.s3.copy_objects(
+    ...     paths=["s3://bucket0/dir0/key0", "s3://bucket0/dir0/key1"])
+    ...     source_path="s3://bucket0/dir0/",
+    ...     target_path="s3://bucket1/dir1/",
+    ...     mode="append"
+    ... )
+    ["s3://bucket1/dir1/key0", "s3://bucket1/dir1/key1"]
+
+    """
+    _logger.debug(f"len(paths): {len(paths)}")
+    if len(paths) < 1:
+        return []
+    source_path = source_path[:-1] if source_path[-1] == "/" else source_path
+    target_path = target_path[:-1] if target_path[-1] == "/" else target_path
+    session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    if mode == "overwrite":
+        _logger.debug(f"Deleting to overwrite: {target_path}/")
+        delete_objects(path=f"{target_path}/", use_threads=use_threads, boto3_session=session)
+    elif mode == "overwrite_partitions":
+        paths_wo_prefix: List[str] = [x.replace(f"{source_path}/", "") for x in paths]
+        paths_wo_filename: List[str] = [f"{x.rpartition('/')[0]}/" for x in paths_wo_prefix]
+        partitions_paths: List[str] = list(set(paths_wo_filename))
+        target_partitions_paths = [f"{target_path}/{x}" for x in partitions_paths]
+        for path in target_partitions_paths:
+            _logger.debug(f"Deleting to overwrite_partitions: {path}")
+            delete_objects(path=path, use_threads=use_threads, boto3_session=session)
+    elif mode != "append":
+        raise exceptions.InvalidArgumentValue(f"{mode} is a invalid mode option.")
+
+    batch: List[Tuple[str, str]] = []
+    new_objects: List[str] = []
+    for path in paths:
+        path_wo_prefix: str = path.replace(f"{source_path}/", "")
+        path_final: str = f"{target_path}/{path_wo_prefix}"
+        new_objects.append(path_final)
+        batch.append((path, path_final))
+    _logger.debug(f"len(new_objects): {len(new_objects)}")
+    _copy_objects(batch=batch, use_threads=use_threads, boto3_session=session)
+    return new_objects
+
+
+def _copy_objects(batch: List[Tuple[str, str]], use_threads: bool, boto3_session: boto3.Session) -> None:
+    _logger.debug(f"len(batch): {len(batch)}")
+    client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session)
+    resource_s3: boto3.resource = _utils.resource(service_name="s3", session=boto3_session)
+    for source, target in batch:
+        source_bucket, source_key = _utils.parse_path(path=source)
+        copy_source: Dict[str, str] = {"Bucket": source_bucket, "Key": source_key}
+        target_bucket, target_key = _utils.parse_path(path=target)
+        resource_s3.meta.client.copy(
+            CopySource=copy_source,
+            Bucket=target_bucket,
+            Key=target_key,
+            SourceClient=client_s3,
+            Config=TransferConfig(num_download_attempts=15, use_threads=use_threads),
+        )

testing/test_awswrangler/test_data_lake.py

@@ -1010,3 +1010,50 @@ def test_parquet_char_length(bucket, database, external_schema):
 
     wr.s3.delete_objects(path=path)
     assert wr.catalog.delete_table_if_exists(database=database, table=table) is True
+
+
+def test_copy(bucket):
+    path = f"s3://{bucket}/test_copy/"
+    df = pd.DataFrame({"id": [1, 2, 3], "par": [1, 2, 3]})
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, partition_cols=["par"], mode="overwrite")["paths"]
+    wr.s3.wait_objects_exist(paths=paths)
+    df = wr.s3.read_parquet(path=path, dataset=True)
+    assert df.id.sum() == 6
+    assert df.par.astype("Int64").sum() == 6
+
+    path2 = f"s3://{bucket}/test_copy2/"
+    df = pd.DataFrame({"id": [1, 2, 3], "par": [1, 2, 3]})
+    paths = wr.s3.to_parquet(df=df, path=path2, dataset=True, partition_cols=["par"], mode="overwrite")["paths"]
+    wr.s3.wait_objects_exist(paths=paths)
+    paths = wr.s3.copy_objects(paths, source_path=path2, target_path=path, mode="append", use_threads=True)
+    wr.s3.wait_objects_exist(paths=paths, use_threads=False)
+    df = wr.s3.read_parquet(path=path, dataset=True)
+    assert df.id.sum() == 12
+    assert df.par.astype("Int64").sum() == 12
+
+    paths = wr.s3.copy_objects(
+        wr.s3.list_objects(path2), source_path=path2, target_path=path, mode="overwrite", use_threads=False
+    )
+    wr.s3.wait_objects_exist(paths=paths, use_threads=False)
+    df = wr.s3.read_parquet(path=path, dataset=True)
+    assert df.id.sum() == 6
+    assert df.par.astype("Int64").sum() == 6
+
+    df = pd.DataFrame({"id": [4], "par": [3]})
+    paths = wr.s3.to_parquet(df=df, path=path2, dataset=True, partition_cols=["par"], mode="overwrite")["paths"]
+    wr.s3.wait_objects_exist(paths=paths)
+    paths = wr.s3.copy_objects(
+        paths, source_path=path2, target_path=path, mode="overwrite_partitions", use_threads=True
+    )
+    wr.s3.wait_objects_exist(paths=paths, use_threads=False)
+    df = wr.s3.read_parquet(path=path, dataset=True)
+    assert df.id.sum() == 7
+    assert df.par.astype("Int64").sum() == 6
+
+    with pytest.raises(wr.exceptions.InvalidArgumentValue):
+        wr.s3.copy_objects(["foo"], source_path="boo", target_path="bar", mode="WRONG")
+
+    assert len(wr.s3.copy_objects([], source_path="boo", target_path="bar")) == 0
+
+    wr.s3.delete_objects(path=path)
+    wr.s3.delete_objects(path=path2)