Skip to content

Commit 54050d5

Browse files
feat: Add BigQuery ML CREATE MODEL support
- Refactor `bigframes.core.sql` to a package. - Add `bigframes.core.sql.ml` for DDL generation. - Add `bigframes.bigquery.ml` module with `create_model` function. - Add unit tests for SQL generation. - Use `_start_query_ml_ddl` for execution. - Return the created model object using `read_gbq_model`.
1 parent 0356206 commit 54050d5

File tree

1 file changed

+4
-1
lines changed
  • bigframes/bigquery/_operations

1 file changed

+4
-1
lines changed

bigframes/bigquery/_operations/ml.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import bigframes.core.sql.ml
2121
import bigframes.core.log_adapter as log_adapter
2222
import bigframes.dataframe as dataframe
23+
import bigframes.ml.base
2324
import bigframes.session
2425

2526
@log_adapter.method_logger(custom_base_name="bigquery_ml")
@@ -37,7 +38,7 @@ def create_model(
3738
training_data: Optional[Union[dataframe.DataFrame, str]] = None,
3839
custom_holiday: Optional[Union[dataframe.DataFrame, str]] = None,
3940
session: Optional[bigframes.session.Session] = None,
40-
) -> None:
41+
) -> bigframes.ml.base.BaseModel:
4142
"""
4243
Creates a BigQuery ML model.
4344
"""
@@ -81,3 +82,5 @@ def _to_sql(df_or_sql: Union[dataframe.DataFrame, str]) -> str:
8182

8283
# Use _start_query_ml_ddl which is designed for this
8384
session._start_query_ml_ddl(sql)
85+
86+
return session.read_gbq_model(model_name)

0 commit comments

Comments
 (0)