acryl-datahub 0.15.0.3__py3-none-any.whl → 0.15.0.4__py3-none-any.whl

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -2,7 +2,6 @@ import contextlib
 import dataclasses
 import enum
 import functools
-import itertools
 import json
 import logging
 import os
@@ -63,6 +62,7 @@ from datahub.utilities.file_backed_collections import (
     FileBackedDict,
     FileBackedList,
 )
+from datahub.utilities.groupby import groupby_unsorted
 from datahub.utilities.lossy_collections import LossyDict, LossyList
 from datahub.utilities.ordered_set import OrderedSet
 from datahub.utilities.perf_timer import PerfTimer
@@ -1314,8 +1314,8 @@ class SqlParsingAggregator(Closeable):
         upstream_aspect.fineGrainedLineages = []
         for downstream_column, all_upstream_columns in cll.items():
             # Group by query ID.
-            for query_id, upstream_columns_for_query in itertools.groupby(
-                sorted(all_upstream_columns.items(), key=lambda x: x[1]),
+            for query_id, upstream_columns_for_query in groupby_unsorted(
+                all_upstream_columns.items(),
                 key=lambda x: x[1],
             ):
                 upstream_columns = [x[0] for x in upstream_columns_for_query]

datahub/utilities/groupby.py

@@ -0,0 +1,17 @@
+import collections
+from typing import Callable, Iterable, Tuple, TypeVar
+
+T = TypeVar("T")
+K = TypeVar("K")
+
+
+def groupby_unsorted(
+    iterable: Iterable[T], key: Callable[[T], K]
+) -> Iterable[Tuple[K, Iterable[T]]]:
+    """The default itertools.groupby() requires that the iterable is already sorted by the key.
+    This method is similar to groupby() but without the pre-sorted requirement."""
+
+    values = collections.defaultdict(list)
+    for v in iterable:
+        values[key(v)].append(v)
+    return values.items()

datahub/utilities/urns/_urn_base.py

@@ -6,6 +6,7 @@ from typing import ClassVar, Dict, List, Optional, Type, Union
 from deprecated import deprecated
 from typing_extensions import Self
 
+from datahub._codegen.aspect import _Aspect
 from datahub.utilities.urns.error import InvalidUrnError
 
 URN_TYPES: Dict[str, Type["_SpecificUrn"]] = {}
@@ -270,7 +271,7 @@ class Urn:
 
 
 class _SpecificUrn(Urn):
-    ENTITY_TYPE: str = ""
+    ENTITY_TYPE: ClassVar[str] = ""
 
     def __init_subclass__(cls) -> None:
         # Validate the subclass.
@@ -286,7 +287,10 @@ class _SpecificUrn(Urn):
         return super().__init_subclass__()
 
     @classmethod
-    def underlying_key_aspect_type(cls) -> Type:
+    def underlying_key_aspect_type(cls) -> Type[_Aspect]:
+        raise NotImplementedError()
+
+    def to_key_aspect(self) -> _Aspect:
         raise NotImplementedError()
 
     @classmethod