Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings
This repository was archived by the owner on May 7, 2026. It is now read-only.

Commit d29a609

Browse filesBrowse files
GarrettWuGarrettWu
authored
feat: add bigquery.ml.get_insights function (#2493)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent ce5de57 commit d29a609
Copy full SHA for d29a609

8 files changed

+88-5Lines changed: 88 additions & 5 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

‎GEMINI.md‎

Copy file name to clipboardExpand all lines: GEMINI.md
-5Lines changed: 0 additions & 5 deletions
This file was deleted.
Collapse file

‎bigframes/bigquery/_operations/ml.py‎

Copy file name to clipboardExpand all lines: bigframes/bigquery/_operations/ml.py
+33Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -480,6 +480,39 @@ def generate_text(
480480
return session.read_gbq_query(sql)
481481

482482

483+
@log_adapter.method_logger(custom_base_name="bigquery_ml")
484+
def get_insights(
485+
model: Union[bigframes.ml.base.BaseEstimator, str, pd.Series],
486+
) -> dataframe.DataFrame:
487+
"""
488+
Gets insights from a BigQuery ML model.
489+
490+
See the `BigQuery ML GET_INSIGHTS function syntax
491+
<https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-get-insights>`_
492+
for additional reference.
493+
494+
Args:
495+
model (bigframes.ml.base.BaseEstimator, str, or pd.Series):
496+
The model to get insights from.
497+
498+
Returns:
499+
bigframes.pandas.DataFrame:
500+
The insights.
501+
"""
502+
import bigframes.pandas as bpd
503+
504+
model_name, session = utils.get_model_name_and_session(model)
505+
506+
sql = bigframes.core.sql.ml.get_insights(
507+
model_name=model_name,
508+
)
509+
510+
if session is None:
511+
return bpd.read_gbq_query(sql)
512+
else:
513+
return session.read_gbq_query(sql)
514+
515+
483516
@log_adapter.method_logger(custom_base_name="bigquery_ml")
484517
def generate_embedding(
485518
model: Union[bigframes.ml.base.BaseEstimator, str, pd.Series],
Collapse file

‎bigframes/bigquery/ml.py‎

Copy file name to clipboardExpand all lines: bigframes/bigquery/ml.py
+2Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
explain_predict,
2626
generate_embedding,
2727
generate_text,
28+
get_insights,
2829
global_explain,
2930
predict,
3031
transform,
@@ -39,4 +40,5 @@
3940
"transform",
4041
"generate_text",
4142
"generate_embedding",
43+
"get_insights",
4244
]
Collapse file

‎bigframes/core/sql/ml.py‎

Copy file name to clipboardExpand all lines: bigframes/core/sql/ml.py
+10Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,16 @@ def generate_text(
266266
return sql
267267

268268

269+
def get_insights(
270+
model_name: str,
271+
) -> str:
272+
"""Encode the ML.GET_INSIGHTS statement.
273+
See https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-get-insights for reference.
274+
"""
275+
sql = f"SELECT * FROM ML.GET_INSIGHTS(MODEL {sg_sql.to_sql(sg_sql.identifier(model_name))})\n"
276+
return sql
277+
278+
269279
def generate_embedding(
270280
model_name: str,
271281
table: str,
Collapse file

‎tests/system/large/bigquery/test_ml.py‎

Copy file name to clipboardExpand all lines: tests/system/large/bigquery/test_ml.py
+26Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,32 @@ def test_generate_embedding_with_options(embedding_model):
6464
assert len(embedding[0]) == 256
6565

6666

67+
def test_get_insights(dataset_id):
68+
df = bpd.DataFrame(
69+
{
70+
"dim1": ["a", "a", "b", "b", "a", "a", "b", "b"],
71+
"dim2": ["x", "y", "x", "y", "x", "y", "x", "y"],
72+
"metric": [10, 20, 30, 40, 12, 25, 35, 45],
73+
"is_test": [False, False, False, False, True, True, True, True],
74+
}
75+
)
76+
model_name = f"{dataset_id}.contribution_analysis_model"
77+
78+
ml.create_model(
79+
model_name=model_name,
80+
options={
81+
"model_type": "CONTRIBUTION_ANALYSIS",
82+
"contribution_metric": "SUM(metric)",
83+
"is_test_col": "is_test",
84+
},
85+
training_data=df,
86+
)
87+
88+
result = ml.get_insights(model_name)
89+
assert len(result) > 0
90+
assert "contributors" in result.columns
91+
92+
6793
def test_create_model_linear_regression(dataset_id):
6894
df = bpd.DataFrame({"x": [1, 2, 3], "y": [2, 4, 6]})
6995
model_name = f"{dataset_id}.linear_regression_model"
Collapse file

‎tests/unit/bigquery/test_ml.py‎

Copy file name to clipboardExpand all lines: tests/unit/bigquery/test_ml.py
+9Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,15 @@ def test_generate_text_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mo
177177
assert "'TYPE' AS request_type" in generated_sql
178178

179179

180+
@mock.patch("bigframes.pandas.read_gbq_query")
181+
def test_get_insights(read_gbq_query_mock):
182+
ml_ops.get_insights(MODEL_SERIES)
183+
read_gbq_query_mock.assert_called_once()
184+
generated_sql = read_gbq_query_mock.call_args[0][0]
185+
assert "ML.GET_INSIGHTS" in generated_sql
186+
assert f"MODEL `{MODEL_NAME}`" in generated_sql
187+
188+
180189
@mock.patch("bigframes.pandas.read_gbq_query")
181190
@mock.patch("bigframes.pandas.read_pandas")
182191
def test_generate_embedding_with_pandas_dataframe(
Collapse file

‎tests/unit/core/sql/snapshots/test_ml/test_get_insights_model_basic/get_insights_model_basic.sql‎

Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
SELECT * FROM ML.GET_INSIGHTS(MODEL `my_project.my_dataset.my_model`)
Collapse file

‎tests/unit/core/sql/test_ml.py‎

Copy file name to clipboardExpand all lines: tests/unit/core/sql/test_ml.py
+7Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,13 @@ def test_generate_text_model_with_options(snapshot):
203203
snapshot.assert_match(sql, "generate_text_model_with_options.sql")
204204

205205

206+
def test_get_insights_model_basic(snapshot):
207+
sql = bigframes.core.sql.ml.get_insights(
208+
model_name="my_project.my_dataset.my_model",
209+
)
210+
snapshot.assert_match(sql, "get_insights_model_basic.sql")
211+
212+
206213
def test_generate_embedding_model_basic(snapshot):
207214
sql = bigframes.core.sql.ml.generate_embedding(
208215
model_name="my_project.my_dataset.my_model",

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.