yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-09-26 01:23:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(ingest): followup on bigquery queries v2 ordering (#11353)

Browse Source
This commit is contained in:
Mayuri Nehate 2024-09-11 23:17:21 +05:30 committed by GitHub
parent fd6d4c88ec
commit 325a608cc5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/queries_extractor.py
View File

@ -247,32 +247,21 @@ class BigQueryQueriesExtractor:
self.report.num_queries_by_project[project.id] += 1 self.report.num_queries_by_project[project.id] += 1
queries.append(entry) queries.append(entry)
self.report.num_total_queries = len(queries) self.report.num_total_queries = len(queries)
logger.info(f"Found {self.report.num_total_queries} total queries")
with self.report.audit_log_preprocessing_timer: with self.report.audit_log_preprocessing_timer:
# Preprocessing stage that deduplicates the queries using query hash per usage bucket # Preprocessing stage that deduplicates the queries using query hash per usage bucket
# Using regular dictionary with # Note: FileBackedDict is an ordered dictionary, so the order of execution of
# key: usage bucket # queries is inherently maintained
# value: File backed dictionary with query hash as key and observed query as value queries_deduped: FileBackedDict[Dict[int, ObservedQuery]]
# This structure is chosen in order to maintain order of execution of queries
queries_deduped: Dict[int, FileBackedDict[ObservedQuery]]
queries_deduped = self.deduplicate_queries(queries) queries_deduped = self.deduplicate_queries(queries)
self.report.num_unique_queries = len( self.report.num_unique_queries = len(queries_deduped)
set( logger.info(f"Found {self.report.num_unique_queries} unique queries")
query_hash
for bucket in queries_deduped.values()
for query_hash in bucket
)
)
with self.report.audit_log_load_timer: with self.report.audit_log_load_timer:
i = 0 i = 0
for queries_in_bucket in queries_deduped.values(): for _, query_instances in queries_deduped.items():
# Ordering is essential for column-level lineage via temporary table for query in query_instances.values():
for row in queries_in_bucket.sql_query_iterator(
"select value from data order by last_query_timestamp asc",
):
query = queries_in_bucket.deserializer(row["value"])
if i > 0 and i % 10000 == 0: if i > 0 and i % 10000 == 0:
logger.info(f"Added {i} query log entries to SQL aggregator") logger.info(f"Added {i} query log entries to SQL aggregator")
@ -283,7 +272,7 @@ class BigQueryQueriesExtractor:
def deduplicate_queries( def deduplicate_queries(
self, queries: FileBackedList[ObservedQuery] self, queries: FileBackedList[ObservedQuery]
) -> Dict[int, FileBackedDict[ObservedQuery]]: ) -> FileBackedDict[Dict[int, ObservedQuery]]:
# This fingerprint based deduplication is done here to reduce performance hit due to # This fingerprint based deduplication is done here to reduce performance hit due to
# repetitive sql parsing while adding observed query to aggregator that would otherwise # repetitive sql parsing while adding observed query to aggregator that would otherwise
@ -291,7 +280,7 @@ class BigQueryQueriesExtractor:
# With current implementation, it is possible that "Operation"(e.g. INSERT) is reported # With current implementation, it is possible that "Operation"(e.g. INSERT) is reported
# only once per day, although it may have happened multiple times throughout the day. # only once per day, although it may have happened multiple times throughout the day.
queries_deduped: Dict[int, FileBackedDict[ObservedQuery]] = dict() queries_deduped: FileBackedDict[Dict[int, ObservedQuery]] = FileBackedDict()
for i, query in enumerate(queries): for i, query in enumerate(queries):
if i > 0 and i % 10000 == 0: if i > 0 and i % 10000 == 0:
@ -310,20 +299,14 @@ class BigQueryQueriesExtractor:
query.query, self.identifiers.platform, fast=True query.query, self.identifiers.platform, fast=True
) )
if time_bucket not in queries_deduped: query_instances = queries_deduped.setdefault(query.query_hash, {})
# TODO: Cleanup, etc as required for file backed dicts after use
queries_deduped[time_bucket] = FileBackedDict[ObservedQuery](
extra_columns={"last_query_timestamp": lambda e: e.timestamp}
)
observed_query = queries_deduped[time_bucket].get(query.query_hash) observed_query = query_instances.setdefault(time_bucket, query)
# If the query already exists for this time bucket, update its attributes # If the query already exists for this time bucket, update its attributes
if observed_query is not None: if observed_query is not query:
observed_query.usage_multiplier += 1 observed_query.usage_multiplier += 1
observed_query.timestamp = query.timestamp observed_query.timestamp = query.timestamp
else:
queries_deduped[time_bucket][query.query_hash] = query
return queries_deduped return queries_deduped