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
+ )