Skip to content

Commit 8b08869

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`. - Remove `query` argument, simplify SQL generation logic. - Fix linting and mypy errors.
1 parent c1adfd9 commit 8b08869

File tree

2 files changed

+8
-4
lines changed

2 files changed

+8
-4
lines changed

bigframes/bigquery/_operations/ml.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,15 @@
1515
from __future__ import annotations
1616

1717
import typing
18-
from typing import Mapping, Optional, Union
18+
from typing import Mapping, Optional, TYPE_CHECKING, Union
1919

2020
import bigframes.core.log_adapter as log_adapter
2121
import bigframes.core.sql.ml
2222
import bigframes.dataframe as dataframe
23-
import bigframes.ml.base
24-
import bigframes.session
23+
24+
if TYPE_CHECKING:
25+
import bigframes.ml.base
26+
import bigframes.session
2527

2628

2729
# Helper to convert DataFrame to SQL string
@@ -47,7 +49,7 @@ def create_model(
4749
training_data: Optional[Union[dataframe.DataFrame, str]] = None,
4850
custom_holiday: Optional[Union[dataframe.DataFrame, str]] = None,
4951
session: Optional[bigframes.session.Session] = None,
50-
) -> bigframes.ml.base.BaseModel:
52+
) -> bigframes.ml.base.BaseEstimator:
5153
"""
5254
Creates a BigQuery ML model.
5355
"""

tests/unit/core/sql/test_ml.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,10 @@
1313
# limitations under the License.
1414

1515
import pytest
16+
1617
import bigframes.core.sql.ml
1718

19+
1820
def test_create_model_basic(snapshot):
1921
sql = bigframes.core.sql.ml.create_model_ddl(
2022
model_name="my_project.my_dataset.my_model",

0 commit comments

Comments
 (0)