refactor: push down SQL generate logic in core.BqmlModelFactory (#62)

Author: GarrettWu

bigframes/ml/core.py

@@ -251,17 +251,10 @@ def create_model(
 
         session = X_train._session
 
-        source_sql = input_data.sql
-        options_sql = self._model_creation_sql_generator.options(**options)
-        transform_sql = (
-            self._model_creation_sql_generator.transform(*transforms)
-            if transforms is not None
-            else None
-        )
         sql = self._model_creation_sql_generator.create_model(
-            source_sql=source_sql,
-            transform_sql=transform_sql,
-            options_sql=options_sql,
+            source=input_data,
+            transforms=transforms,
+            options=options,
         )
 
         return self._create_model_with_sql(session=session, sql=sql)
@@ -287,18 +280,10 @@ def create_time_series_model(
 
         session = X_train._session
 
-        source_sql = input_data.sql
-        options_sql = self._model_creation_sql_generator.options(**options)
-
-        transform_sql = (
-            self._model_creation_sql_generator.transform(*transforms)
-            if transforms is not None
-            else None
-        )
         sql = self._model_creation_sql_generator.create_model(
-            source_sql=source_sql,
-            transform_sql=transform_sql,
-            options_sql=options_sql,
+            source=input_data,
+            transforms=transforms,
+            options=options,
         )
 
         return self._create_model_with_sql(session=session, sql=sql)
@@ -320,10 +305,9 @@ def create_remote_model(
         Returns:
             BqmlModel: a BqmlModel wrapping a trained model in BigQuery
         """
-        options_sql = self._model_creation_sql_generator.options(**options)
         sql = self._model_creation_sql_generator.create_remote_model(
             connection_name=connection_name,
-            options_sql=options_sql,
+            options=options,
         )
 
         return self._create_model_with_sql(session=session, sql=sql)
@@ -341,9 +325,8 @@ def create_imported_model(
 
         Returns: a BqmlModel, wrapping a trained model in BigQuery
         """
-        options_sql = self._model_creation_sql_generator.options(**options)
         sql = self._model_creation_sql_generator.create_imported_model(
-            options_sql=options_sql,
+            options=options,
         )
 
         return self._create_model_with_sql(session=session, sql=sql)

bigframes/ml/sql.py

@@ -16,9 +16,10 @@
 Generates SQL queries needed for BigQuery DataFrames ML
 """
 
-from typing import Iterable, Optional, Union
+from typing import Iterable, Mapping, Optional, Union
 
 import bigframes.constants as constants
+import bigframes.pandas as bpd
 
 
 class BaseSqlGenerator:
@@ -113,11 +114,15 @@ def __init__(self, model_id: str):
     # Model create and alter
     def create_model(
         self,
-        source_sql: str,
-        transform_sql: Optional[str] = None,
-        options_sql: Optional[str] = None,
+        source: bpd.DataFrame,
+        options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
+        transforms: Optional[Iterable[str]] = None,
     ) -> str:
         """Encode the CREATE TEMP MODEL statement for BQML"""
+        source_sql = source.sql
+        transform_sql = self.transform(*transforms) if transforms is not None else None
+        options_sql = self.options(**options)
+
         parts = [f"CREATE TEMP MODEL `{self._model_id}`"]
         if transform_sql:
             parts.append(transform_sql)
@@ -129,9 +134,11 @@ def create_model(
     def create_remote_model(
         self,
         connection_name: str,
-        options_sql: Optional[str] = None,
+        options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> str:
         """Encode the CREATE TEMP MODEL statement for BQML remote model."""
+        options_sql = self.options(**options)
+
         parts = [f"CREATE TEMP MODEL `{self._model_id}`"]
         parts.append(self.connection(connection_name))
         if options_sql:
@@ -140,17 +147,19 @@ def create_remote_model(
 
     def create_imported_model(
         self,
-        options_sql: Optional[str] = None,
+        options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> str:
         """Encode the CREATE TEMP MODEL statement for BQML remote model."""
+        options_sql = self.options(**options)
+
         parts = [f"CREATE TEMP MODEL `{self._model_id}`"]
         if options_sql:
             parts.append(options_sql)
         return "\n".join(parts)
 
 
 class ModelManipulationSqlGenerator(BaseSqlGenerator):
-    """Sql generator for manipulating a model entity. Model name is the fully model path of project_id.dataset_id.model_id."""
+    """Sql generator for manipulating a model entity. Model name is the full model path of project_id.dataset_id.model_id."""
 
     def __init__(self, model_name: str):
         self._model_name = model_name

tests/unit/ml/test_sql.py

@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest import mock
+
 import pytest
 
 import bigframes.ml.sql as ml_sql
+import bigframes.pandas as bpd
 
 
 @pytest.fixture(scope="session")
@@ -34,6 +37,14 @@ def model_manipulation_sql_generator() -> ml_sql.ModelManipulationSqlGenerator:
     )
 
 
+@pytest.fixture(scope="session")
+def mock_df():
+    mock_df = mock.create_autospec(spec=bpd.DataFrame)
+    mock_df.sql = "input_X_y_sql"
+
+    return mock_df
+
+
 def test_options_produces_correct_sql(base_sql_generator: ml_sql.BaseSqlGenerator):
     sql = base_sql_generator.options(
         model_type="lin_reg", input_label_cols=["col_a"], l1_reg=0.6
@@ -96,33 +107,44 @@ def test_label_encoder_produces_correct_sql(
 
 def test_create_model_produces_correct_sql(
     model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator,
+    mock_df: bpd.DataFrame,
 ):
     sql = model_creation_sql_generator.create_model(
-        source_sql="my_source_sql",
-        options_sql="my_options_sql",
+        source=mock_df,
+        options={"option_key1": "option_value1", "option_key2": 2},
     )
     assert (
         sql
         == """CREATE TEMP MODEL `my_model_id`
-my_options_sql
-AS my_source_sql"""
+OPTIONS(
+  option_key1="option_value1",
+  option_key2=2)
+AS input_X_y_sql"""
     )
 
 
 def test_create_model_transform_produces_correct_sql(
     model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator,
+    mock_df: bpd.DataFrame,
 ):
     sql = model_creation_sql_generator.create_model(
-        source_sql="my_source_sql",
-        options_sql="my_options_sql",
-        transform_sql="my_transform_sql",
+        source=mock_df,
+        options={"option_key1": "option_value1", "option_key2": 2},
+        transforms=[
+            "ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a",
+            "ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b",
+        ],
     )
     assert (
         sql
         == """CREATE TEMP MODEL `my_model_id`
-my_transform_sql
-my_options_sql
-AS my_source_sql"""
+TRANSFORM(
+  ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a,
+  ML.ONE_HOT_ENCODER(col_b) OVER(col_b) AS encoded_col_b)
+OPTIONS(
+  option_key1="option_value1",
+  option_key2=2)
+AS input_X_y_sql"""
     )
 
 
@@ -131,26 +153,30 @@ def test_create_remote_model_produces_correct_sql(
 ):
     sql = model_creation_sql_generator.create_remote_model(
         connection_name="my_project.us.my_connection",
-        options_sql="my_options_sql",
+        options={"option_key1": "option_value1", "option_key2": 2},
     )
     assert (
         sql
         == """CREATE TEMP MODEL `my_model_id`
 REMOTE WITH CONNECTION `my_project.us.my_connection`
-my_options_sql"""
+OPTIONS(
+  option_key1="option_value1",
+  option_key2=2)"""
     )
 
 
 def test_create_imported_model_produces_correct_sql(
     model_creation_sql_generator: ml_sql.ModelCreationSqlGenerator,
 ):
     sql = model_creation_sql_generator.create_imported_model(
-        options_sql="my_options_sql",
+        options={"option_key1": "option_value1", "option_key2": 2},
     )
     assert (
         sql
         == """CREATE TEMP MODEL `my_model_id`
-my_options_sql"""
+OPTIONS(
+  option_key1="option_value1",
+  option_key2=2)"""
     )