feat: add ml.metrics.pairwise.cosine_similarity function (#374)

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

Author: GarrettWu

bigframes/ml/core.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import datetime
-from typing import Callable, cast, Iterable, Mapping, Optional, Union
+from typing import Callable, cast, Iterable, Literal, Mapping, Optional, Union
 import uuid
 
 from google.cloud import bigquery
@@ -28,34 +28,12 @@
 import bigframes.pandas as bpd
 
 
-class BqmlModel:
-    """Represents an existing BQML model in BigQuery.
-
-    Wraps the BQML API and SQL interface to expose the functionality needed for
-    BigQuery DataFrames ML.
-    """
+class BaseBqml:
+    """Base class for BQML functionalities."""
 
-    def __init__(self, session: bigframes.Session, model: bigquery.Model):
+    def __init__(self, session: bigframes.Session):
         self._session = session
-        self._model = model
-        self._model_manipulation_sql_generator = ml_sql.ModelManipulationSqlGenerator(
-            self.model_name
-        )
-
-    @property
-    def session(self) -> bigframes.Session:
-        """Get the BigQuery DataFrames session that this BQML model wrapper is tied to"""
-        return self._session
-
-    @property
-    def model_name(self) -> str:
-        """Get the fully qualified name of the model, i.e. project_id.dataset_id.model_id"""
-        return f"{self._model.project}.{self._model.dataset_id}.{self._model.model_id}"
-
-    @property
-    def model(self) -> bigquery.Model:
-        """Get the BQML model associated with this wrapper"""
-        return self._model
+        self._base_sql_generator = ml_sql.BaseSqlGenerator()
 
     def _apply_sql(
         self,
@@ -84,6 +62,71 @@ def _apply_sql(
 
         return df
 
+    def distance(
+        self,
+        x: bpd.DataFrame,
+        y: bpd.DataFrame,
+        type: Literal["EUCLIDEAN", "MANHATTAN", "COSINE"],
+        name: str,
+    ) -> bpd.DataFrame:
+        """Calculate ML.DISTANCE from DataFrame inputs.
+
+        Args:
+            x:
+                input DataFrame
+            y:
+                input DataFrame
+            type:
+                Distance types, accept values are  "EUCLIDEAN", "MANHATTAN", "COSINE".
+            name:
+                name of the output result column
+        """
+        assert len(x.columns) == 1 and len(y.columns) == 1
+
+        input_data = x._cached().join(y._cached(), how="outer")
+        x_column_id, y_column_id = x._block.value_columns[0], y._block.value_columns[0]
+
+        return self._apply_sql(
+            input_data,
+            lambda source_df: self._base_sql_generator.ml_distance(
+                x_column_id,
+                y_column_id,
+                type=type,
+                source_df=source_df,
+                name=name,
+            ),
+        )
+
+
+class BqmlModel(BaseBqml):
+    """Represents an existing BQML model in BigQuery.
+
+    Wraps the BQML API and SQL interface to expose the functionality needed for
+    BigQuery DataFrames ML.
+    """
+
+    def __init__(self, session: bigframes.Session, model: bigquery.Model):
+        self._session = session
+        self._model = model
+        self._model_manipulation_sql_generator = ml_sql.ModelManipulationSqlGenerator(
+            self.model_name
+        )
+
+    @property
+    def session(self) -> bigframes.Session:
+        """Get the BigQuery DataFrames session that this BQML model wrapper is tied to"""
+        return self._session
+
+    @property
+    def model_name(self) -> str:
+        """Get the fully qualified name of the model, i.e. project_id.dataset_id.model_id"""
+        return f"{self._model.project}.{self._model.dataset_id}.{self._model.model_id}"
+
+    @property
+    def model(self) -> bigquery.Model:
+        """Get the BQML model associated with this wrapper"""
+        return self._model
+
     def predict(self, input_data: bpd.DataFrame) -> bpd.DataFrame:
         # TODO: validate input data schema
         return self._apply_sql(

bigframes/ml/metrics/__init__.py

@@ -0,0 +1,39 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from bigframes.ml.metrics import pairwise
+from bigframes.ml.metrics._metrics import (
+    accuracy_score,
+    auc,
+    confusion_matrix,
+    f1_score,
+    precision_score,
+    r2_score,
+    recall_score,
+    roc_auc_score,
+    roc_curve,
+)
+
+__all__ = [
+    "r2_score",
+    "recall_score",
+    "accuracy_score",
+    "roc_curve",
+    "roc_auc_score",
+    "auc",
+    "confusion_matrix",
+    "precision_score",
+    "f1_score",
+    "pairwise",
+]

bigframes/ml/metrics.py renamed to bigframes/ml/metrics/_metrics.py

@@ -0,0 +1,34 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Union
+
+from bigframes.ml import core, utils
+import bigframes.pandas as bpd
+import third_party.bigframes_vendored.sklearn.metrics.pairwise as vendored_metrics_pairwise
+
+
+def cosine_similarity(
+    X: Union[bpd.DataFrame, bpd.Series], Y: Union[bpd.DataFrame, bpd.Series]
+) -> bpd.DataFrame:
+    X, Y = utils.convert_to_dataframe(X, Y)
+    if len(X.columns) != 1 or len(Y.columns) != 1:
+        raise ValueError("Inputs X and Y can only contain 1 column.")
+
+    base_bqml = core.BaseBqml(session=X._session)
+    return base_bqml.distance(X, Y, type="COSINE", name="cosine_similarity")
+
+
+cosine_similarity.__doc__ = inspect.getdoc(vendored_metrics_pairwise.cosine_similarity)

bigframes/ml/metrics/pairwise.py

@@ -16,7 +16,7 @@
 Generates SQL queries needed for BigQuery DataFrames ML
 """
 
-from typing import Iterable, Mapping, Optional, Union
+from typing import Iterable, Literal, Mapping, Optional, Union
 
 import google.cloud.bigquery
 
@@ -133,6 +133,19 @@ def ml_label_encoder(
         https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-label-encoder for params."""
         return f"""ML.LABEL_ENCODER({numeric_expr_sql}, {top_k}, {frequency_threshold}) OVER() AS {name}"""
 
+    def ml_distance(
+        self,
+        col_x: str,
+        col_y: str,
+        type: Literal["EUCLIDEAN", "MANHATTAN", "COSINE"],
+        source_df: bpd.DataFrame,
+        name: str,
+    ) -> str:
+        """Encode ML.DISTANCE for BQML.
+        https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-distance"""
+        source_sql, _, _ = source_df._to_sql_query(include_index=True)
+        return f"""SELECT *, ML.DISTANCE({col_x}, {col_y}, '{type}') AS {name} FROM ({source_sql})"""
+
 
 class ModelCreationSqlGenerator(BaseSqlGenerator):
     """Sql generator for creating a model entity. Model id is the standalone id without project id and dataset id."""

bigframes/ml/sql.py

@@ -116,6 +116,10 @@
       - name: metrics
         uid: bigframes.ml.metrics
       name: metrics
+    - items:
+      - name: metrics.pairwise
+        uid: bigframes.ml.metrics.pairwise
+      name: metrics.pairwise
     - items:
       - name: model_selection
         uid: bigframes.ml.model_selection

docs/templates/toc.yml

@@ -0,0 +1,35 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pandas as pd
+
+from bigframes.ml import metrics
+import bigframes.pandas as bpd
+
+
+def test_cosine_similarity():
+    x_col = [np.array([4.1, 0.5, 1.0])]
+    y_col = [np.array([3.0, 0.0, 2.5])]
+    X = bpd.read_pandas(pd.DataFrame({"X": x_col}))
+    Y = bpd.read_pandas(pd.DataFrame({"Y": y_col}))
+
+    result = metrics.pairwise.cosine_similarity(X, Y)
+    expected_pd_df = pd.DataFrame(
+        {"X": x_col, "Y": y_col, "cosine_similarity": [0.108199]}
+    )
+
+    pd.testing.assert_frame_equal(
+        result.to_pandas(), expected_pd_df, check_dtype=False, check_index_type=False
+    )