refactor: push down SQL generate logic in core.BqmlModel (#66)

Author: GarrettWu

bigframes/ml/core.py

@@ -58,7 +58,7 @@ def model(self) -> bigquery.Model:
     def _apply_sql(
         self,
         input_data: bpd.DataFrame,
-        func: Callable[[str], str],
+        func: Callable[[bpd.DataFrame], str],
     ) -> bpd.DataFrame:
         """Helper to wrap a dataframe in a SQL query, keeping the index intact.
 
@@ -74,11 +74,9 @@ def _apply_sql(
                 string from which to construct the output dataframe. It must
                 include the index columns of the input SQL.
         """
-        source_sql, index_col_ids, index_labels = input_data._to_sql_query(
-            include_index=True
-        )
+        _, index_col_ids, index_labels = input_data._to_sql_query(include_index=True)
 
-        sql = func(source_sql)
+        sql = func(input_data)
         df = self._session.read_gbq(sql, index_col=index_col_ids)
         df.index.names = index_labels
 
@@ -106,11 +104,9 @@ def generate_text(
         # TODO: validate input data schema
         return self._apply_sql(
             input_data,
-            lambda source_sql: self._model_manipulation_sql_generator.ml_generate_text(
-                source_sql=source_sql,
-                struct_options=self._model_manipulation_sql_generator.struct_options(
-                    **options
-                ),
+            lambda source_df: self._model_manipulation_sql_generator.ml_generate_text(
+                source_df=source_df,
+                struct_options=options,
             ),
         )
 
@@ -122,11 +118,9 @@ def generate_text_embedding(
         # TODO: validate input data schema
         return self._apply_sql(
             input_data,
-            lambda source_sql: self._model_manipulation_sql_generator.ml_generate_text_embedding(
-                source_sql=source_sql,
-                struct_options=self._model_manipulation_sql_generator.struct_options(
-                    **options
-                ),
+            lambda source_df: self._model_manipulation_sql_generator.ml_generate_text_embedding(
+                source_df=source_df,
+                struct_options=options,
             ),
         )
 
@@ -136,13 +130,7 @@ def forecast(self) -> bpd.DataFrame:
 
     def evaluate(self, input_data: Optional[bpd.DataFrame] = None):
         # TODO: validate input data schema
-        # Note: don't need index as evaluate returns a new table
-        source_sql, _, _ = (
-            input_data._to_sql_query(include_index=False)
-            if (input_data is not None)
-            else (None, None, None)
-        )
-        sql = self._model_manipulation_sql_generator.ml_evaluate(source_sql)
+        sql = self._model_manipulation_sql_generator.ml_evaluate(input_data)
 
         return self._session.read_gbq(sql)
 
@@ -188,11 +176,8 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel:
         # truncate as Vertex ID only accepts 63 characters, easily exceeding the limit for temp models.
         # The possibility of conflicts should be low.
         vertex_ai_model_id = vertex_ai_model_id[:63]
-        options_sql = self._model_manipulation_sql_generator.options(
-            **{"vertex_ai_model_id": vertex_ai_model_id}
-        )
         sql = self._model_manipulation_sql_generator.alter_model(
-            options_sql=options_sql
+            options={"vertex_ai_model_id": vertex_ai_model_id}
         )
         # Register the model and wait it to finish
         self._session._start_query(sql)
@@ -252,7 +237,7 @@ def create_model(
         session = X_train._session
 
         sql = self._model_creation_sql_generator.create_model(
-            source=input_data,
+            source_df=input_data,
             transforms=transforms,
             options=options,
         )
@@ -281,7 +266,7 @@ def create_time_series_model(
         session = X_train._session
 
         sql = self._model_creation_sql_generator.create_model(
-            source=input_data,
+            source_df=input_data,
             transforms=transforms,
             options=options,
         )

bigframes/ml/sql.py

@@ -118,12 +118,12 @@ def __init__(self, model_id: str):
     # Model create and alter
     def create_model(
         self,
-        source: bpd.DataFrame,
+        source_df: 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
+        source_sql = source_df.sql
         transform_sql = self.transform(*transforms) if transforms is not None else None
         options_sql = self.options(**options)
 
@@ -168,39 +168,58 @@ class ModelManipulationSqlGenerator(BaseSqlGenerator):
     def __init__(self, model_name: str):
         self._model_name = model_name
 
+    def _source_sql(self, source_df: bpd.DataFrame) -> str:
+        """Return DataFrame sql with index columns."""
+        _source_sql, _, _ = source_df._to_sql_query(include_index=True)
+        return _source_sql
+
     # Alter model
     def alter_model(
         self,
-        options_sql: str,
+        options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> str:
         """Encode the ALTER MODEL statement for BQML"""
+        options_sql = self.options(**options)
+
         parts = [f"ALTER MODEL `{self._model_name}`"]
         parts.append(f"SET {options_sql}")
         return "\n".join(parts)
 
     # ML prediction TVFs
-    def ml_predict(self, source_sql: str) -> str:
+    def ml_predict(self, source_df: bpd.DataFrame) -> str:
         """Encode ML.PREDICT for BQML"""
         return f"""SELECT * FROM ML.PREDICT(MODEL `{self._model_name}`,
-  ({source_sql}))"""
+  ({self._source_sql(source_df)}))"""
 
     def ml_forecast(self) -> str:
         """Encode ML.FORECAST for BQML"""
         return f"""SELECT * FROM ML.FORECAST(MODEL `{self._model_name}`)"""
 
-    def ml_generate_text(self, source_sql: str, struct_options: str) -> str:
+    def ml_generate_text(
+        self, source_df: bpd.DataFrame, struct_options: Mapping[str, Union[int, float]]
+    ) -> str:
         """Encode ML.GENERATE_TEXT for BQML"""
+        struct_options_sql = self.struct_options(**struct_options)
         return f"""SELECT * FROM ML.GENERATE_TEXT(MODEL `{self._model_name}`,
-  ({source_sql}), {struct_options})"""
+  ({self._source_sql(source_df)}), {struct_options_sql})"""
 
-    def ml_generate_text_embedding(self, source_sql: str, struct_options: str) -> str:
+    def ml_generate_text_embedding(
+        self, source_df: bpd.DataFrame, struct_options: Mapping[str, Union[int, float]]
+    ) -> str:
         """Encode ML.GENERATE_TEXT_EMBEDDING for BQML"""
+        struct_options_sql = self.struct_options(**struct_options)
         return f"""SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(MODEL `{self._model_name}`,
-  ({source_sql}), {struct_options})"""
+  ({self._source_sql(source_df)}), {struct_options_sql})"""
 
     # ML evaluation TVFs
-    def ml_evaluate(self, source_sql: Optional[str] = None) -> str:
+    def ml_evaluate(self, source_df: Optional[bpd.DataFrame] = None) -> str:
         """Encode ML.EVALUATE for BQML"""
+        if source_df is None:
+            source_sql = None
+        else:
+            # Note: don't need index as evaluate returns a new table
+            source_sql, _, _ = source_df._to_sql_query(include_index=False)
+
         if source_sql is None:
             return f"""SELECT * FROM ML.EVALUATE(MODEL `{self._model_name}`)"""
         else:
@@ -222,7 +241,7 @@ def ml_principal_component_info(self) -> str:
         )
 
     # ML transform TVF, that require a transform_only type model
-    def ml_transform(self, source_sql: str) -> str:
+    def ml_transform(self, source_df: bpd.DataFrame) -> str:
         """Encode ML.TRANSFORM for BQML"""
         return f"""SELECT * FROM ML.TRANSFORM(MODEL `{self._model_name}`,
-  ({source_sql}))"""
+  ({self._source_sql(source_df)}))"""

tests/unit/ml/test_sql.py

@@ -41,6 +41,7 @@ def model_manipulation_sql_generator() -> ml_sql.ModelManipulationSqlGenerator:
 def mock_df():
     mock_df = mock.create_autospec(spec=bpd.DataFrame)
     mock_df.sql = "input_X_y_sql"
+    mock_df._to_sql_query.return_value = "input_X_sql", None, None
 
     return mock_df
 
@@ -117,7 +118,7 @@ def test_create_model_produces_correct_sql(
     mock_df: bpd.DataFrame,
 ):
     sql = model_creation_sql_generator.create_model(
-        source=mock_df,
+        source_df=mock_df,
         options={"option_key1": "option_value1", "option_key2": 2},
     )
     assert (
@@ -135,7 +136,7 @@ def test_create_model_transform_produces_correct_sql(
     mock_df: bpd.DataFrame,
 ):
     sql = model_creation_sql_generator.create_model(
-        source=mock_df,
+        source_df=mock_df,
         options={"option_key1": "option_value1", "option_key2": 2},
         transforms=[
             "ML.STANDARD_SCALER(col_a) OVER(col_a) AS scaled_col_a",
@@ -191,38 +192,38 @@ def test_alter_model_correct_sql(
     model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
 ):
     sql = model_manipulation_sql_generator.alter_model(
-        options_sql="my_options_sql",
+        options={"option_key1": "option_value1", "option_key2": 2},
     )
     assert (
         sql
         == """ALTER MODEL `my_project_id.my_dataset_id.my_model_id`
-SET my_options_sql"""
+SET OPTIONS(
+  option_key1="option_value1",
+  option_key2=2)"""
     )
 
 
 def test_ml_predict_produces_correct_sql(
     model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
+    mock_df: bpd.DataFrame,
 ):
-    sql = model_manipulation_sql_generator.ml_predict(
-        source_sql="SELECT * FROM my_table"
-    )
+    sql = model_manipulation_sql_generator.ml_predict(source_df=mock_df)
     assert (
         sql
         == """SELECT * FROM ML.PREDICT(MODEL `my_project_id.my_dataset_id.my_model_id`,
-  (SELECT * FROM my_table))"""
+  (input_X_sql))"""
     )
 
 
 def test_ml_evaluate_produces_correct_sql(
     model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
+    mock_df: bpd.DataFrame,
 ):
-    sql = model_manipulation_sql_generator.ml_evaluate(
-        source_sql="SELECT * FROM my_table"
-    )
+    sql = model_manipulation_sql_generator.ml_evaluate(source_df=mock_df)
     assert (
         sql
         == """SELECT * FROM ML.EVALUATE(MODEL `my_project_id.my_dataset_id.my_model_id`,
-  (SELECT * FROM my_table))"""
+  (input_X_sql))"""
     )
 
 
@@ -248,15 +249,35 @@ def test_ml_centroids_produces_correct_sql(
 
 def test_ml_generate_text_produces_correct_sql(
     model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
+    mock_df: bpd.DataFrame,
 ):
     sql = model_manipulation_sql_generator.ml_generate_text(
-        source_sql="SELECT * FROM my_table",
-        struct_options="STRUCT(value AS item)",
+        source_df=mock_df,
+        struct_options={"option_key1": 1, "option_key2": 2.2},
     )
     assert (
         sql
         == """SELECT * FROM ML.GENERATE_TEXT(MODEL `my_project_id.my_dataset_id.my_model_id`,
-  (SELECT * FROM my_table), STRUCT(value AS item))"""
+  (input_X_sql), STRUCT(
+  1 AS option_key1,
+  2.2 AS option_key2))"""
+    )
+
+
+def test_ml_generate_text_embedding_produces_correct_sql(
+    model_manipulation_sql_generator: ml_sql.ModelManipulationSqlGenerator,
+    mock_df: bpd.DataFrame,
+):
+    sql = model_manipulation_sql_generator.ml_generate_text_embedding(
+        source_df=mock_df,
+        struct_options={"option_key1": 1, "option_key2": 2.2},
+    )
+    assert (
+        sql
+        == """SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(MODEL `my_project_id.my_dataset_id.my_model_id`,
+  (input_X_sql), STRUCT(
+  1 AS option_key1,
+  2.2 AS option_key2))"""
     )