Refactoring _delete_objects

Author: malachi-constant

awswrangler/s3/_delete.py

@@ -19,34 +19,33 @@
 def _split_paths_by_bucket(paths: List[str]) -> Dict[str, List[str]]:
     buckets: Dict[str, List[str]] = {}
     bucket: str
-    key: str
     for path in paths:
-        bucket, key = _utils.parse_path(path=path)
+        bucket = _utils.parse_path(path=path)[0]
         if bucket not in buckets:
             buckets[bucket] = []
-        buckets[bucket].append(key)
+        buckets[bucket].append(path)
     return buckets
 
 
 @ray_remote
 def _delete_objects(
     boto3_session: Optional[boto3.Session],
-    bucket: str,
-    keys: List[str],
+    paths: List[str],
     s3_additional_kwargs: Optional[Dict[str, Any]],
 ) -> None:
     client_s3: boto3.client = _utils.client(
         service_name="s3",
         session=boto3_session,
     )
-    _logger.debug("len(keys): %s", len(keys))
-    batch: List[Dict[str, str]] = [{"Key": key} for key in keys]
+    _logger.debug("len(paths): %s", len(paths))
     if s3_additional_kwargs:
         extra_kwargs: Dict[str, Any] = get_botocore_valid_kwargs(
             function_name="list_objects_v2", s3_additional_kwargs=s3_additional_kwargs
         )
     else:
         extra_kwargs = {}
+    bucket = _utils.parse_path(path=paths[0])[0]
+    batch: List[Dict[str, str]] = [{"Key": _utils.parse_path(path)[1]} for path in paths]
     res = client_s3.delete_objects(Bucket=bucket, Delete={"Objects": batch}, **extra_kwargs)
     deleted: List[Dict[str, Any]] = res.get("Deleted", [])
     for obj in deleted:
@@ -118,11 +117,16 @@ def delete_objects(
         last_modified_end=last_modified_end,
         s3_additional_kwargs=s3_additional_kwargs,
     )
-
-    buckets: Dict[str, List[str]] = _split_paths_by_bucket(paths=paths)
-    for bucket, keys in buckets.items():
-        chunks: List[List[str]] = _utils.chunkify(lst=keys, max_length=1_000)
-        executor = _get_executor(use_threads=use_threads)
-        executor.map(
-            _delete_objects, boto3_session, itertools.repeat(bucket), chunks, itertools.repeat(s3_additional_kwargs)
-        )
+    paths_by_bucket: Dict[List[str, List[str]]] = _split_paths_by_bucket(paths)
+
+    chunks = []
+    for bucket in paths_by_bucket:
+        chunks += _utils.chunkify(lst=paths_by_bucket[bucket], max_length=5)
+
+    executor = _get_executor(use_threads=use_threads)
+    executor.map(
+        _delete_objects,
+        boto3_session,
+        chunks,
+        itertools.repeat(s3_additional_kwargs),
+    )

tests/test_s3.py

@@ -69,12 +69,15 @@ def test_list_by_last_modified_date(path):
     assert len(wr.s3.read_json(path, last_modified_begin=begin_utc, last_modified_end=end_utc).index) == 6
 
 
-def test_s3_delete_objects(path):
+def test_s3_delete_objects(path, path2):
     df = pd.DataFrame({"id": [1, 2, 3]})
-    for i in range(10):
-        wr.s3.to_json(df, f"s3://{path}delete-test{i}.json")
-
-    wr.s3.delete_objects(path=f"{path}delete-test*")
+    objects_per_bucket = 10
+    paths = [f"s3://{path}delete-test{i}.json" for i in range(objects_per_bucket)] + [
+        f"s3://{path2}delete-test{i}.json" for i in range(objects_per_bucket)
+    ]
+    for path in paths:
+        wr.s3.to_json(df, path)
+    wr.s3.delete_objects(path=paths)
     time.sleep(5)  # s3 read-consistency
 
     assert len(wr.s3.list_objects(f"{path}delete-test*")) == 0