feat(ingest/s3): type aware directory sorting (#8089)

Author: treff7es

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

@@ -1,4 +1,5 @@
 import dataclasses
+import functools
 import logging
 import os
 import pathlib
@@ -182,6 +183,29 @@ def get_column_type(
 # )
 
 
+def partitioned_folder_comparator(folder1: str, folder2: str) -> int:
+    # Try to convert to number and compare if the folder name is a number
+    try:
+        # Stripping = from the folder names as it most probably partition name part like year=2021
+        if "=" in folder1 and "=" in folder2:
+            if folder1.split("=", 1)[0] == folder2.split("=", 1)[0]:
+                folder1 = folder1.split("=", 1)[1]
+                folder2 = folder2.split("=", 1)[1]
+
+        num_folder1 = int(folder1)
+        num_folder2 = int(folder2)
+        if num_folder1 == num_folder2:
+            return 0
+        else:
+            return 1 if num_folder1 > num_folder2 else -1
+    except Exception:
+        # If folder name is not a number then do string comparison
+        if folder1 == folder2:
+            return 0
+        else:
+            return 1 if folder1 > folder2 else -1
+
+
 @dataclasses.dataclass
 class TableData:
     display_name: str
@@ -700,7 +724,12 @@ def get_dir_to_process(self, bucket_name: str, folder: str) -> str:
         )
         iterator = peekable(iterator)
         if iterator:
-            sorted_dirs = sorted(iterator, reverse=True)
+            sorted_dirs = sorted(
+                iterator,
+                key=functools.cmp_to_key(partitioned_folder_comparator),
+                reverse=True,
+            )
+
             return self.get_dir_to_process(
                 bucket_name=bucket_name, folder=sorted_dirs[0] + "/"
             )
@@ -786,7 +815,8 @@ def local_browser(self, path_spec: PathSpec) -> Iterable[Tuple[str, datetime, in
         else:
             logger.debug(f"Scanning files under local folder: {prefix}")
             for root, dirs, files in os.walk(prefix):
-                dirs.sort()
+                dirs.sort(key=functools.cmp_to_key(partitioned_folder_comparator))
+
                 for file in sorted(files):
                     full_path = os.path.join(root, file)
                     yield full_path, datetime.utcfromtimestamp(

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

@@ -0,0 +1,49 @@
+from datahub.ingestion.source.s3.source import partitioned_folder_comparator
+
+
+def test_partition_comparator_numeric_folder_name():
+    folder1 = "3"
+    folder2 = "12"
+    assert partitioned_folder_comparator(folder1, folder2) == -1
+
+
+def test_partition_comparator_numeric_folder_name2():
+    folder1 = "12"
+    folder2 = "3"
+    assert partitioned_folder_comparator(folder1, folder2) == 1
+
+
+def test_partition_comparator_string_folder():
+    folder1 = "myfolderB"
+    folder2 = "myFolderA"
+    assert partitioned_folder_comparator(folder1, folder2) == 1
+
+
+def test_partition_comparator_string_same_folder():
+    folder1 = "myFolderA"
+    folder2 = "myFolderA"
+    assert partitioned_folder_comparator(folder1, folder2) == 0
+
+
+def test_partition_comparator_with_numeric_partition():
+    folder1 = "year=3"
+    folder2 = "year=12"
+    assert partitioned_folder_comparator(folder1, folder2) == -1
+
+
+def test_partition_comparator_with_padded_numeric_partition():
+    folder1 = "year=03"
+    folder2 = "year=12"
+    assert partitioned_folder_comparator(folder1, folder2) == -1
+
+
+def test_partition_comparator_with_equal_sign_in_name():
+    folder1 = "month=12"
+    folder2 = "year=0"
+    assert partitioned_folder_comparator(folder1, folder2) == -1
+
+
+def test_partition_comparator_with_string_partition():
+    folder1 = "year=year2020"
+    folder2 = "year=year2021"
+    assert partitioned_folder_comparator(folder1, folder2) == -1