feat: add bigframes.ml.compose.SQLScalarColumnTransformer to create custom SQL-based transformations

bigframes/ml/compose.py

@@ -21,7 +21,9 @@
 import re
 import types
 import typing
-from typing import cast, Iterable, List, Optional, Set, Tuple, Union
+from typing import cast, Iterable, List, Optional, Set, Tuple, Union, Dict, Type
+import abc
+import json
 
 from bigframes_vendored import constants
 import bigframes_vendored.sklearn.compose._column_transformer
@@ -46,6 +48,114 @@
 )
 
 
+CUSTOM_TRANSFORMER_SQL_RX = re.compile(
+    "^(?P<sql>.*)/[*]CT.(?P<id>[A-Z]+[A-Z0-9]*)[(](?P<config>[^*]*)[)][*]/$",
+    re.IGNORECASE,
+)
+
+
+class CustomTransformer(base.BaseTransformer):
+    _CTID = None
+    _custom_transformer_classes = {}
+
+    @classmethod
+    def register(cls, transformer_cls: Type[base.BaseTransformer]):
+        assert transformer_cls._CTID
+        cls._custom_transformer_classes[transformer_cls._CTID] = transformer_cls
+
+    @classmethod
+    def find_matching_transformer(
+        cls, transform_sql: str
+    ) -> Optional[Type[base.BaseTransformer]]:
+        for transform_cls in cls._custom_transformer_classes.values():
+            if transform_cls.understands(transform_sql):
+                return transform_cls
+        return None
+
+    @classmethod
+    def understands(cls, transform_sql: str) -> bool:
+        """
+        may be overwritten to have a more advanced matching, possibly without comments in SQL
+        """
+        m = CUSTOM_TRANSFORMER_SQL_RX.match(transform_sql)
+        if m and m.group("id").strip() == cls._CTID:
+            return True
+        return False
+
+    def __init__(self):
+        super().__init__()
+
+    def _compile_to_sql(
+        self, X: bpd.DataFrame, columns: Optional[Iterable[str]] = None
+    ) -> List[str]:
+        if columns is None:
+            columns = X.columns
+        return [
+            f"{self.custom_compile_to_sql(X, column)} {self._get_sql_comment(column)} AS {self.get_target_column_name(column)}"
+            for column in columns
+        ]
+
+    def get_target_column_name(self, column: str) -> str:
+        return f"{self._CTID.lower()}_{column}"
+
+    @classmethod
+    @abc.abstractclassmethod
+    def custom_compile_to_sql(cls, X: bpd.DataFrame, column: str) -> str:
+        pass
+
+    def get_persistent_config(self, column: str) -> Optional[Union[Dict, List]]:
+        """
+        return structure to be persisted in the comment of the sql
+        """
+        return None
+
+    def _get_pc_as_args(self, column: str) -> str:
+        config = self.get_persistent_config(column)
+        if not config:
+            return ""
+        return json.dumps(config)
+
+    def _get_sql_comment(self, column: str) -> str:
+        args = self._get_pc_as_args(column)
+        return f"/*CT.{self._CTID}({args})*/"
+
+    @classmethod
+    def _parse_from_sql(cls, transform_sql: str) -> Tuple[base.BaseTransformer, str]:
+        m = CUSTOM_TRANSFORMER_SQL_RX.match(transform_sql)
+        if m and m.group("id").strip() != cls._CTID:
+            raise ValueError("understand() does not match _parse_from_sql!")
+        args = m.group("config").strip()
+        if args != "":
+            config = json.loads(args)
+        else:
+            config = None
+        sql = m.group("sql").strip()
+        return cls.custom_parse_from_sql(config, sql)
+
+    @classmethod
+    @abc.abstractclassmethod
+    def custom_parse_from_sql(
+        cls, config: Optional[Union[Dict, List]], sql: str
+    ) -> Tuple[base.BaseTransformer, str]:
+        """
+        return transformer instance and column name
+        """
+        pass
+
+    def _keys(self):
+        return ()
+
+    # CustomTransformers are thought to be used inside a column transformer.
+    # So there is no need to implement fit() and transform() directly.
+    # ColumnTransformer.merge() takes care, that a single custom transformer
+    # is not returned as a standalone transformer.
+    def fit(self, y: Union[bpd.DataFrame, bpd.Series]) -> base.BaseTransformer:
+        raise NotImplementedError("Unsupported")
+
+    def transform(self, y: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
+        raise NotImplementedError("Unsupported")
+
+
 @log_adapter.class_logger
 class ColumnTransformer(
     base.Transformer,
@@ -130,10 +240,7 @@ def camel_to_snake(name):
             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"])
             found_transformer = False
             for prefix in _BQML_TRANSFROM_TYPE_MAPPING:
                 if transform_sql.startswith(prefix):
@@ -147,10 +254,30 @@ def camel_to_snake(name):
 
                     found_transformer = True
                     break
+
             if not found_transformer:
-                raise NotImplementedError(
-                    f"Unsupported transformer type. {constants.FEEDBACK_LINK}"
+                transformer_cls = CustomTransformer.find_matching_transformer(
+                    transform_sql
                 )
+                if transformer_cls:
+                    transformers_set.add(
+                        (
+                            camel_to_snake(transformer_cls.__name__),
+                            *transformer_cls._parse_from_sql(transform_sql),  # type: ignore
+                        )
+                    )
+                    found_transformer = True
+
+            if not found_transformer:
+                if not transform_sql.startswith("ML.") and "/*CT." not in transform_sql:
+                    continue  # ignore other patterns, only report unhandled known patterns
+                if transform_sql.startswith("ML."):
+                    raise NotImplementedError(
+                        f"Unsupported transformer type. {constants.FEEDBACK_LINK}"
+                    )
+                raise ValueError("Missing custom transformer")
+
+            output_names.append(transform_col_dict["name"])
 
         transformer = cls(transformers=list(transformers_set))
         transformer._output_names = output_names
@@ -167,6 +294,8 @@ def _merge(
 
         assert len(transformers) > 0
         _, transformer_0, column_0 = transformers[0]
+        if isinstance(transformer_0, CustomTransformer):
+            return self  # CustomTransformers only work inside ColumnTransformer
         feature_columns_sorted = sorted(
             [
                 cast(str, feature_column.name)
@@ -234,6 +363,26 @@ def fit(
         self._extract_output_names()
         return self
 
+    # Overwrite the implementation in BaseTransformer, as it only supports the "ML." transformers.
+    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."):
+                if not CustomTransformer.find_matching_transformer(transform_sql):
+                    continue
+
+            output_names.append(transform_col_dict["name"])
+
+        self._output_names = output_names
+
     def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
         if not self._bqml_model:
             raise RuntimeError("Must be fitted before transform")

tests/unit/ml/compose_custom_transformers.py

@@ -0,0 +1,91 @@
+import bigframes.pandas as bpd
+from bigframes.ml.compose import CustomTransformer
+from typing import List, Optional, Union, Dict
+import re
+
+
+class IdentityTransformer(CustomTransformer):
+    _CTID = "IDENT"
+    IDENT_BQSQL_RX = re.compile("^(?P<colname>[a-z][a-z0-9_]+)$", flags=re.IGNORECASE)
+
+    def custom_compile_to_sql(self, X: bpd.DataFrame, column: str) -> str:
+        return f"{column}"
+
+    @classmethod
+    def custom_parse_from_sql(
+        cls, config: Optional[Union[Dict, List]], sql: str
+    ) -> tuple[CustomTransformer, str]:
+        col_label = cls.IDENT_BQSQL_RX.match(sql).group("colname")
+        return cls(), col_label
+
+
+CustomTransformer.register(IdentityTransformer)
+
+
+class Length1Transformer(CustomTransformer):
+    _CTID = "LEN1"
+    _DEFAULT_VALUE_DEFAULT = -1
+    LEN1_BQSQL_RX = re.compile(
+        "^CASE WHEN (?P<colname>[a-z][a-z0-9_]*) IS NULL THEN (?P<defaultvalue>[-]?[0-9]+) ELSE LENGTH[(](?P=colname)[)] END$",
+        flags=re.IGNORECASE,
+    )
+
+    def __init__(self, default_value: Optional[int] = None):
+        self.default_value = default_value
+
+    def custom_compile_to_sql(self, X: bpd.DataFrame, column: str) -> str:
+        default_value = (
+            self.default_value
+            if self.default_value is not None
+            else Length1Transformer._DEFAULT_VALUE_DEFAULT
+        )
+        return (
+            f"CASE WHEN {column} IS NULL THEN {default_value} ELSE LENGTH({column}) END"
+        )
+
+    @classmethod
+    def custom_parse_from_sql(
+        cls, config: Optional[Union[Dict, List]], sql: str
+    ) -> tuple[CustomTransformer, str]:
+        m = cls.LEN1_BQSQL_RX.match(sql)
+        col_label = m.group("colname")
+        default_value = int(m.group("defaultvalue"))
+        return cls(default_value), col_label
+
+
+CustomTransformer.register(Length1Transformer)
+
+
+class Length2Transformer(CustomTransformer):
+    _CTID = "LEN2"
+    _DEFAULT_VALUE_DEFAULT = -1
+    LEN2_BQSQL_RX = re.compile(
+        "^CASE WHEN (?P<colname>[a-z][a-z0-9_]*) .*$", flags=re.IGNORECASE
+    )
+
+    def __init__(self, default_value: Optional[int] = None):
+        self.default_value = default_value
+
+    def get_persistent_config(self, column: str) -> Optional[Union[Dict, List]]:
+        return [self.default_value]
+
+    def custom_compile_to_sql(self, X: bpd.DataFrame, column: str) -> str:
+        default_value = (
+            self.default_value
+            if self.default_value is not None
+            else Length2Transformer._DEFAULT_VALUE_DEFAULT
+        )
+        return (
+            f"CASE WHEN {column} IS NULL THEN {default_value} ELSE LENGTH({column}) END"
+        )
+
+    @classmethod
+    def custom_parse_from_sql(
+        cls, config: Optional[Union[Dict, List]], sql: str
+    ) -> tuple[CustomTransformer, str]:
+        col_label = cls.LEN2_BQSQL_RX.match(sql).group("colname")
+        default_value = config[0]  # get default value from persistent_config
+        return cls(default_value), col_label
+
+
+CustomTransformer.register(Length2Transformer)