Skip to content

Commit 65e2e70

Browse files
authored
chore: add exec_seconds attribute to ExecutionMetrics (#1104)
* chore: score query count assertion fix. * chore: add exec_seconds attribute to ExecutionMetrics * update naming * update logic. * update logic. * update logic. * update logic.
1 parent e7fa913 commit 65e2e70

File tree

1 file changed

+15
-16
lines changed

1 file changed

+15
-16
lines changed

bigframes/session/metrics.py

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -29,48 +29,47 @@ class ExecutionMetrics:
2929
execution_count: int = 0
3030
slot_millis: int = 0
3131
bytes_processed: int = 0
32+
execution_secs: float = 0
3233

3334
def count_job_stats(self, query_job: bq_job.QueryJob):
3435
stats = get_performance_stats(query_job)
3536
if stats is not None:
36-
bytes_processed, slot_millis, exec_seconds = stats
37+
bytes_processed, slot_millis, execution_secs = stats
3738
self.execution_count += 1
3839
self.bytes_processed += bytes_processed
3940
self.slot_millis += slot_millis
41+
self.execution_secs += execution_secs
4042
if LOGGING_NAME_ENV_VAR in os.environ:
4143
# when running notebooks via pytest nbmake
42-
write_stats_to_disk(bytes_processed, slot_millis, exec_seconds)
44+
write_stats_to_disk(bytes_processed, slot_millis, execution_secs)
4345

4446

4547
def get_performance_stats(
4648
query_job: bigquery.QueryJob,
47-
) -> Optional[Tuple[int, int, Optional[float]]]:
49+
) -> Optional[Tuple[int, int, float]]:
4850
"""Parse the query job for performance stats.
4951
5052
Return None if the stats do not reflect real work done in bigquery.
5153
"""
54+
55+
if (
56+
query_job.configuration.dry_run
57+
or query_job.created is None
58+
or query_job.ended is None
59+
):
60+
return None
61+
5262
bytes_processed = query_job.total_bytes_processed
5363
if not isinstance(bytes_processed, int):
5464
return None # filter out mocks
55-
if query_job.configuration.dry_run:
56-
# dry run stats are just predictions of the real run
57-
bytes_processed = 0
5865

5966
slot_millis = query_job.slot_millis
6067
if not isinstance(slot_millis, int):
6168
return None # filter out mocks
6269

63-
if query_job.configuration.dry_run:
64-
# dry run stats are just predictions of the real run
65-
slot_millis = 0
66-
67-
exec_seconds = (
68-
(query_job.ended - query_job.created).total_seconds()
69-
if query_job.created is not None and query_job.ended is not None
70-
else None
71-
)
70+
execution_secs = (query_job.ended - query_job.created).total_seconds()
7271

73-
return bytes_processed, slot_millis, exec_seconds
72+
return bytes_processed, slot_millis, execution_secs
7473

7574

7675
def write_stats_to_disk(

0 commit comments

Comments
 (0)