fix(ingestion/s3): groupby group-splitting issue

Author: eagle-25

metadata-ingestion/src/datahub/ingestion/source/s3/source.py

@@ -5,8 +5,8 @@
 import pathlib
 import re
 import time
+from collections import defaultdict
 from datetime import datetime
-from itertools import groupby
 from pathlib import PurePath
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 from urllib.parse import urlparse
@@ -139,6 +139,14 @@ def partitioned_folder_comparator(folder1: str, folder2: str) -> int:
             return 1 if folder1 > folder2 else -1
 
 
+def _group_s3_objects_by_dirname(s3_objects: Any) -> Dict[str, List[Any]]:
+    grouped_objects = defaultdict(list)
+    for obj in s3_objects:
+        dirname = os.path.dirname(obj.key)
+        grouped_objects[dirname].append(obj)
+    return grouped_objects
+
+
 @dataclasses.dataclass
 class Folder:
     creation_time: datetime
@@ -863,16 +871,9 @@ def get_folder_info(
         Returns:
         List[Folder]: A list of Folder objects representing the partitions found.
         """
-
-        prefix_to_list = prefix
-        files = list(
-            bucket.objects.filter(Prefix=f"{prefix_to_list}").page_size(PAGE_SIZE)
-        )
-        files = sorted(files, key=lambda a: a.last_modified)
-        grouped_files = groupby(files, lambda x: x.key.rsplit("/", 1)[0])
-
         partitions: List[Folder] = []
-        for key, group in grouped_files:
+        s3_objects = bucket.objects.filter(Prefix=prefix).page_size(PAGE_SIZE)
+        for key, group in _group_s3_objects_by_dirname(s3_objects).items():
             file_size = 0
             creation_time = None
             modification_time = None

metadata-ingestion/tests/unit/s3/test_s3_source.py

@@ -1,12 +1,16 @@
 from typing import List, Tuple
+from unittest.mock import Mock
 
 import pytest
 
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
 from datahub.ingestion.source.data_lake_common.path_spec import PathSpec
-from datahub.ingestion.source.s3.source import partitioned_folder_comparator
+from datahub.ingestion.source.s3.source import (
+    _group_s3_objects_by_dirname,
+    partitioned_folder_comparator,
+)
 
 
 def test_partition_comparator_numeric_folder_name():
@@ -240,3 +244,20 @@ def container_properties_filter(x: MetadataWorkUnit) -> bool:
         "folder_abs_path": "my-bucket/my-dir/my-dir2",
         "platform": "s3",
     }
+
+
+def test_group_s3_objects_by_dir_name():
+    # arrange
+    s3_objects = [
+        Mock(key="/dir1/file1.txt"),
+        Mock(key="/dir1/file2.txt"),
+        Mock(key="/dir2/file3.txt"),
+        Mock(key="/dir2/file4.txt"),
+    ]
+
+    # act
+    grouped_objects = _group_s3_objects_by_dirname(s3_objects)
+
+    # assert
+    assert grouped_objects["/dir1"] == [s3_objects[0], s3_objects[1]]
+    assert grouped_objects["/dir2"] == [s3_objects[2], s3_objects[3]]