refactor: read transformer output columns from model entity (#817)

* refactor: read transformer output columns from model entity

* fix tests and docs

* remove dup code

* fix comment

Author: GarrettWu

bigframes/ml/base.py

@@ -188,6 +188,24 @@ def __init__(self):
     def _keys(self):
         pass
 
+    def _extract_output_names(self):
+        """Extract transform output column names. Save the results to self._output_names."""
+        assert self._bqml_model is not None
+
+        output_names = []
+        for transform_col in self._bqml_model._model._properties["transformColumns"]:
+            transform_col_dict = cast(dict, transform_col)
+            # pass the columns that are not transformed
+            if "transformSql" not in transform_col_dict:
+                continue
+            transform_sql: str = transform_col_dict["transformSql"]
+            if not transform_sql.startswith("ML."):
+                continue
+
+            output_names.append(transform_col_dict["name"])
+
+        self._output_names = output_names
+
     def __eq__(self, other) -> bool:
         return type(self) is type(other) and self._keys() == other._keys()
 

bigframes/ml/compose.py

@@ -201,25 +201,20 @@ def _merge(
 
     def _compile_to_sql(
         self,
-        columns: List[str],
         X: bpd.DataFrame,
-    ) -> List[Tuple[str, str]]:
+    ) -> List[str]:
         """Compile this transformer to a list of SQL expressions that can be included in
         a BQML TRANSFORM clause
 
         Args:
-            columns (List[str]):
-                a list of column names to transform
-            X (bpd.DataFrame):
-                The Dataframe with training data.
+            X: DataFrame to transform.
 
-        Returns:
-            a list of tuples of (sql_expression, output_name)"""
+        Returns: a list of sql_expr."""
         result = []
         for _, transformer, target_columns in self.transformers:
             if isinstance(target_columns, str):
                 target_columns = [target_columns]
-            result += transformer._compile_to_sql(target_columns, X=X)
+            result += transformer._compile_to_sql(X, target_columns)
         return result
 
     def fit(
@@ -229,17 +224,14 @@ def fit(
     ) -> ColumnTransformer:
         (X,) = utils.convert_to_dataframe(X)
 
-        compiled_transforms = self._compile_to_sql(X.columns.tolist(), X)
-        transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms]
-
+        transform_sqls = self._compile_to_sql(X)
         self._bqml_model = self._bqml_model_factory.create_model(
             X,
             options={"model_type": "transform_only"},
             transforms=transform_sqls,
         )
 
-        # The schema of TRANSFORM output is not available in the model API, so save it during fitting
-        self._output_names = [name for _, name in compiled_transforms]
+        self._extract_output_names()
         return self
 
     def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:

bigframes/ml/impute.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Iterable, List, Literal, Optional, Tuple, Union
+from typing import Iterable, List, Literal, Optional, Union
 
 import bigframes_vendored.sklearn.impute._base
 
@@ -49,25 +49,22 @@ def _keys(self):
 
     def _compile_to_sql(
         self,
-        columns: Iterable[str],
-        X=None,
-    ) -> List[Tuple[str, str]]:
+        X: bpd.DataFrame,
+        columns: Optional[Iterable[str]] = None,
+    ) -> List[str]:
         """Compile this transformer to a list of SQL expressions that can be included in
         a BQML TRANSFORM clause
 
         Args:
-            columns:
-                A list of column names to transform.
-            X:
-                The Dataframe with training data.
+            X: DataFrame to transform.
+            columns: transform columns. If None, transform all columns in X.
 
-        Returns: a list of tuples of (sql_expression, output_name)"""
+        Returns: a list of tuples sql_expr."""
+        if columns is None:
+            columns = X.columns
         return [
-            (
-                self._base_sql_generator.ml_imputer(
-                    column, self.strategy, f"imputer_{column}"
-                ),
-                f"imputer_{column}",
+            self._base_sql_generator.ml_imputer(
+                column, self.strategy, f"imputer_{column}"
             )
             for column in columns
         ]
@@ -92,17 +89,14 @@ def fit(
     ) -> SimpleImputer:
         (X,) = utils.convert_to_dataframe(X)
 
-        compiled_transforms = self._compile_to_sql(X.columns.tolist(), X)
-        transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms]
-
+        transform_sqls = self._compile_to_sql(X)
         self._bqml_model = self._bqml_model_factory.create_model(
             X,
             options={"model_type": "transform_only"},
             transforms=transform_sqls,
         )
 
-        # The schema of TRANSFORM output is not available in the model API, so save it during fitting
-        self._output_names = [name for _, name in compiled_transforms]
+        self._extract_output_names()
         return self
 
     def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:

bigframes/ml/pipeline.py

@@ -106,9 +106,7 @@ def fit(
     ) -> Pipeline:
         (X,) = utils.convert_to_dataframe(X)
 
-        compiled_transforms = self._transform._compile_to_sql(X.columns.tolist(), X=X)
-        transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms]
-
+        transform_sqls = self._transform._compile_to_sql(X)
         if y is not None:
             # If labels columns are present, they should pass through un-transformed
             (y,) = utils.convert_to_dataframe(y)