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

[eagle-25]

Summary

Fix the group-splitting issue caused by itertools.groupby when the input data is not ordered by the group key.

Changes

As-is

• Use itertools.groupby for grouping. (group_key: s3_object.key)
• Sort s3_objects before grouping. (sort_key: s3_object.last_modified)

To-be

• Use defaultdict for grouping without pre-sorting.

Background

Behavior of itertools.groupby

• itertools.groupby() generates a new group every time the value of the key function changes, which requires the input data to be sorted.

• This means the input data should be sorted and the group key must be equal to the sort key for correct grouping.

  from itertools import groupby
  
  data = [1, 1, 2, 1]
  
  # Without sorting
  for key, val in groupby(data, key=lambda x: x):
      print(key, list(val))
  
  # --- output --- (Group Splitting)
  # 1 [1, 1] 
  # 2 [2]
  # 1 [1]
  
  # With sorting
  sorted_data = sorted(data, key=lambda x: x)
  
  for key, val in groupby(sorted_data, key=lambda x: x):
      print(key, list(val))
  
  # --- output ---
  # 1 [1, 1, 1]
  # 2 [2]

Problem part

• in our code, the sorting order is changed to file.last_modified even though the group key is file.key. This mismatch causes group-splitting issues.

  # Current code
  files = sorted(files, key=lambda a: a.last_modified) 
  grouped_files = groupby(files, lambda x: x.key.rsplit("/", 1)[0])

Checklist

• The PR conforms to DataHub's Contributing Guideline (particularly Commit Message Format)
• Links to related issues (if applicable)
• Tests for the changes have been added/updated (if applicable)
• Docs related to the changes have been added/updated (if applicable). If a new feature has been added a Usage Guide has been added for the same.
• For any breaking change/potential downtime/deprecation/big changes an entry has been made in Updating DataHub

[eagle-25]

I think there is no need to sort the s3_objects by last_modified, as the logic below iterates through all files within the same folder and determines the maximum value

[sgomezvillamor]

Sounds fair to me. @treff7es , since you were involved in the development here, are you okay with this?

[eagle-25]

I used defaultdict for grouping to eliminate the need to consider the order of the collection.

[eagle-25]

Unfixed code returns

len(res): 3
res[0].sample_file: s3://my-bucket/my-folder/dir1/0001.csv
res[1].sample_file: s3://my-bucket/my-folder/dir2/0001.csv
res[2].sample_file: s3://my-bucket/my-folder/dir1/0002.csv

[sgomezvillamor]

yes, nice catch!

[sgomezvillamor]

The _group_s3_objects_by_dirname function could be placed in either s3_util.py or s3_boto_utils.py within metadata-ingestion/src/datahub/ingestion/source/aws.

To keep things as generic as possible, we should consider what happens if there's no directory (i.e., no /) in the key. In such cases, we could include the file in the group None or /. WDYT?

[sgomezvillamor]

Also, regarding s3_objects: Any, could you specify a more precise type?

[sgomezvillamor]

Similar functionality existed in the past
https://github.com/datahub-project/datahub/pull/4011/files

def groupby_unsorted(
    iterable: Iterable[T], key: Callable[[T], K]
) -> Iterable[Tuple[K, Iterable[T]]]:

Yet another example to keep this logic as generic as possible

[eagle-25]

@sgomezvillamor

Thank you for the detailed review. I’ll make the changes you suggested and request your feedback again.

Regarding your question, here’s my response:

To keep things as generic as possible, we should consider what happens if there's no directory (i.e., no /) in the key. In such cases, we could include the file in the group None or /. WDYT?

-> I think / is the better choice because it intuitively represents the root directory, whereas it's harder to understand the meaning of None in this context. I'll add the necessary logic and tests for this case. Thanks for pointing this out. 👍

The _group_s3_objects_by_dirname function could be placed in either s3_util.py or s3_boto_utils.py within metadata-ingestion/src/datahub/ingestion/source/aws.

-> Sounds good. I'll move it to s3_util.py

Also, regarding s3_objects: Any, could you specify a more precise type?

-> Okay, I'll specify the types with mypy_boto3_s3.

[sgomezvillamor]

LGTM
Let's double confirm with @treff7es 's review, if possible

[sgomezvillamor]

yes, nice catch!

[eagle-25]

Changes

• change a function name (group_s3_objects_by_dirname)

[sgomezvillamor]

LGTM
Thanks for the contrib!

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Files with missing lines	Coverage Δ
...estion/src/datahub/ingestion/source/aws/s3_util.py	91.66% <100.00%> (+1.92%)	⬆️
...ngestion/src/datahub/ingestion/source/s3/source.py	86.82% <100.00%> (-0.11%)	⬇️

... and 636 files with indirect coverage changes

---

Continue to review full report in Codecov by Sentry.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 3b827f3...2412ab9. Read the comment docs.

[eagle-25]

@sgomezvillamor Could you please merge this PR? 🙏