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.

Author: google-labs-jules[bot]

bigframes/bigquery/_operations/ml.py

@@ -23,6 +23,16 @@
 import bigframes.ml.base
 import bigframes.session
 
+
+# Helper to convert DataFrame to SQL string
+def _to_sql(df_or_sql: Union[dataframe.DataFrame, str]) -> str:
+    if isinstance(df_or_sql, str):
+        return df_or_sql
+    # It's a DataFrame
+    sql, _, _ = df_or_sql._to_sql_query(include_index=False)
+    return sql
+
+
 @log_adapter.method_logger(custom_base_name="bigquery_ml")
 def create_model(
     model_name: str,
@@ -34,7 +44,6 @@ def create_model(
     output_schema: Optional[Mapping[str, str]] = None,
     connection_name: Optional[str] = None,
     options: Optional[Mapping[str, Union[str, int, float, bool, list]]] = None,
-    query: Optional[Union[dataframe.DataFrame, str]] = None,
     training_data: Optional[Union[dataframe.DataFrame, str]] = None,
     custom_holiday: Optional[Union[dataframe.DataFrame, str]] = None,
     session: Optional[bigframes.session.Session] = None,
@@ -44,22 +53,13 @@ def create_model(
     """
     import bigframes.pandas as bpd
 
-    # Helper to convert DataFrame to SQL string
-    def _to_sql(df_or_sql: Union[dataframe.DataFrame, str]) -> str:
-        if isinstance(df_or_sql, str):
-            return df_or_sql
-        # It's a DataFrame
-        sql, _, _ = df_or_sql._to_sql_query(include_index=True)
-        return sql
-
-    query_statement = _to_sql(query) if query is not None else None
     training_data_sql = _to_sql(training_data) if training_data is not None else None
     custom_holiday_sql = _to_sql(custom_holiday) if custom_holiday is not None else None
 
     # Determine session from DataFrames if not provided
     if session is None:
         # Try to get session from inputs
-        dfs = [obj for obj in [query, training_data, custom_holiday] if hasattr(obj, "_session")]
+        dfs = [obj for obj in [training_data, custom_holiday] if hasattr(obj, "_session")]
         if dfs:
             session = dfs[0]._session
 
@@ -72,7 +72,6 @@ def _to_sql(df_or_sql: Union[dataframe.DataFrame, str]) -> str:
         output_schema=output_schema,
         connection_name=connection_name,
         options=options,
-        query_statement=query_statement,
         training_data=training_data_sql,
         custom_holiday=custom_holiday_sql,
     )

bigframes/core/sql/ml.py

@@ -30,11 +30,13 @@ def create_model_ddl(
     output_schema: Optional[Mapping[str, str]] = None,
     connection_name: Optional[str] = None,
     options: Optional[Mapping[str, Union[str, int, float, bool, list]]] = None,
-    query_statement: Optional[str] = None,
     training_data: Optional[str] = None,
     custom_holiday: Optional[str] = None,
 ) -> str:
-    """Encode the CREATE MODEL statement."""
+    """Encode the CREATE MODEL statement.
+
+    See https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create for reference.
+    """
 
     if replace:
         create = "CREATE OR REPLACE MODEL "
@@ -83,18 +85,15 @@ def create_model_ddl(
 
     # [AS {query_statement | ( training_data AS (query_statement), custom_holiday AS (holiday_statement) )}]
 
-    if query_statement and (training_data or custom_holiday):
-        raise ValueError("Cannot specify both `query_statement` and (`training_data` or `custom_holiday`).")
-
-    if query_statement:
-        ddl += f"AS {query_statement}"
-    elif training_data:
-        # specialized AS clause
-        parts = []
-        parts.append(f"training_data AS ({training_data})")
+    if training_data:
         if custom_holiday:
+            # When custom_holiday is present, we need named clauses
+            parts = []
+            parts.append(f"training_data AS ({training_data})")
             parts.append(f"custom_holiday AS ({custom_holiday})")
-
-        ddl += f"AS (\n  {', '.join(parts)}\n)"
+            ddl += f"AS (\n  {', '.join(parts)}\n)"
+        else:
+             # Just training_data is treated as the query_statement
+             ddl += f"AS {training_data}"
 
     return ddl

tests/unit/core/sql/test_ml.py

@@ -19,7 +19,7 @@ def test_create_model_basic(snapshot):
     sql = bigframes.core.sql.ml.create_model_ddl(
         model_name="my_project.my_dataset.my_model",
         options={"model_type": "LINEAR_REG", "input_label_cols": ["label"]},
-        query_statement="SELECT * FROM my_table",
+        training_data="SELECT * FROM my_table",
     )
     snapshot.assert_match(sql, "create_model_basic.sql")
 
@@ -28,7 +28,7 @@ def test_create_model_replace(snapshot):
         model_name="my_model",
         replace=True,
         options={"model_type": "LOGISTIC_REG"},
-        query_statement="SELECT * FROM t",
+        training_data="SELECT * FROM t",
     )
     snapshot.assert_match(sql, "create_model_replace.sql")
 
@@ -37,7 +37,7 @@ def test_create_model_if_not_exists(snapshot):
         model_name="my_model",
         if_not_exists=True,
         options={"model_type": "KMEANS"},
-        query_statement="SELECT * FROM t",
+        training_data="SELECT * FROM t",
     )
     snapshot.assert_match(sql, "create_model_if_not_exists.sql")
 
@@ -46,7 +46,7 @@ def test_create_model_transform(snapshot):
         model_name="my_model",
         transform=["ML.STANDARD_SCALER(c1) OVER() AS c1_scaled", "c2"],
         options={"model_type": "LINEAR_REG"},
-        query_statement="SELECT c1, c2, label FROM t",
+        training_data="SELECT c1, c2, label FROM t",
     )
     snapshot.assert_match(sql, "create_model_transform.sql")
 
@@ -81,6 +81,6 @@ def test_create_model_list_option(snapshot):
     sql = bigframes.core.sql.ml.create_model_ddl(
         model_name="my_model",
         options={"hidden_units": [32, 16], "dropout": 0.2},
-        query_statement="SELECT * FROM t",
+        training_data="SELECT * FROM t",
     )
     snapshot.assert_match(sql, "create_model_list_option.sql")