Skip to content

Commit dced460

Browse files
authored
fix: update session metrics in read_gbq_query (#1084)
* fix: update read_gbq_query query metric and fix ML transform metric missing issue. * update test
1 parent d73fe9d commit dced460

File tree

2 files changed

+10
-2
lines changed

2 files changed

+10
-2
lines changed

bigframes/session/loader.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -529,6 +529,9 @@ def read_gbq_query(
529529
configuration=configuration,
530530
)
531531

532+
if self._metrics is not None:
533+
self._metrics.count_job_stats(query_job)
534+
532535
# If there was no destination table, that means the query must have
533536
# been DDL or DML. Return some job metadata, instead.
534537
if not destination:

tests/system/large/ml/test_core.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -146,10 +146,15 @@ def test_bqml_standalone_transform(penguins_df_default_index, new_penguins_df):
146146
"ML.ONE_HOT_ENCODER(species, 'none', 1000000, 0) OVER() AS onehotencoded_species",
147147
],
148148
)
149+
start_execution_count = model.session._metrics.execution_count
150+
151+
transformed = model.transform(new_penguins_df)
152+
153+
end_execution_count = model.session._metrics.execution_count
154+
assert end_execution_count - start_execution_count == 1
149155

150-
transformed = model.transform(new_penguins_df).to_pandas()
151156
utils.check_pandas_df_schema_and_index(
152-
transformed,
157+
transformed.to_pandas(),
153158
columns=["scaled_culmen_length_mm", "onehotencoded_species"],
154159
index=[1633, 1672, 1690],
155160
col_exact=False,

0 commit comments

Comments
 (0)