refactor: use core.convert for series conversions under the ml packages (#1178)

* refactor: use core.convert for series conversions under the ml packages

* update method name

* fetch global session lazily

Author: sycai

bigframes/core/convert.py

@@ -17,6 +17,7 @@
 
 import pandas as pd
 
+from bigframes.core import global_session
 import bigframes.core.indexes as index
 import bigframes.series as series
 
@@ -36,7 +37,9 @@ def is_series_convertible(obj) -> bool:
         return False
 
 
-def to_bf_series(obj, default_index: Optional[index.Index], session) -> series.Series:
+def to_bf_series(
+    obj, default_index: Optional[index.Index], session=None
+) -> series.Series:
     """
     Convert a an object to a bigframes series
 
@@ -51,6 +54,10 @@ def to_bf_series(obj, default_index: Optional[index.Index], session) -> series.S
     """
     if isinstance(obj, series.Series):
         return obj
+
+    if session is None:
+        session = global_session.get_global_session()
+
     if isinstance(obj, pd.Series):
         return series.Series(obj, session=session)
     if isinstance(obj, index.Index):

bigframes/ml/metrics/_metrics.py

@@ -37,7 +37,7 @@ def r2_score(
     *,
     force_finite=True,
 ) -> float:
-    y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred)
+    y_true_series, y_pred_series = utils.batch_convert_to_series(y_true, y_pred)
 
     # total sum of squares
     # (dataframe, scalar) binops
@@ -66,7 +66,7 @@ def accuracy_score(
     normalize=True,
 ) -> float:
     # TODO(ashleyxu): support sample_weight as the parameter
-    y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred)
+    y_true_series, y_pred_series = utils.batch_convert_to_series(y_true, y_pred)
 
     # Compute accuracy for each possible representation
     # TODO(ashleyxu): add multilabel classification support where y_type
@@ -97,7 +97,7 @@ def roc_curve(
             f"drop_intermediate is not yet implemented. {constants.FEEDBACK_LINK}"
         )
 
-    y_true_series, y_score_series = utils.convert_to_series(y_true, y_score)
+    y_true_series, y_score_series = utils.batch_convert_to_series(y_true, y_score)
 
     session = y_true_series._block.expr.session
 
@@ -157,7 +157,7 @@ def roc_auc_score(
 ) -> float:
     # TODO(bmil): Add multi-class support
     # TODO(bmil): Add multi-label support
-    y_true_series, y_score_series = utils.convert_to_series(y_true, y_score)
+    y_true_series, y_score_series = utils.batch_convert_to_series(y_true, y_score)
 
     fpr, tpr, _ = roc_curve(y_true_series, y_score_series, drop_intermediate=False)
 
@@ -174,7 +174,7 @@ def auc(
     x: Union[bpd.DataFrame, bpd.Series],
     y: Union[bpd.DataFrame, bpd.Series],
 ) -> float:
-    x_series, y_series = utils.convert_to_series(x, y)
+    x_series, y_series = utils.batch_convert_to_series(x, y)
 
     # TODO(b/286410053) Support ML exceptions and error handling.
     auc = sklearn_metrics.auc(x_series.to_pandas(), y_series.to_pandas())
@@ -189,7 +189,7 @@ def confusion_matrix(
     y_pred: Union[bpd.DataFrame, bpd.Series],
 ) -> pd.DataFrame:
     # TODO(ashleyxu): support labels and sample_weight parameters
-    y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred)
+    y_true_series, y_pred_series = utils.batch_convert_to_series(y_true, y_pred)
 
     y_true_series = y_true_series.rename("y_true")
     confusion_df = y_true_series.to_frame().assign(y_pred=y_pred_series)
@@ -235,7 +235,7 @@ def recall_score(
             f"Only average=None is supported. {constants.FEEDBACK_LINK}"
         )
 
-    y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred)
+    y_true_series, y_pred_series = utils.batch_convert_to_series(y_true, y_pred)
 
     is_accurate = y_true_series == y_pred_series
     unique_labels = (
@@ -272,7 +272,7 @@ def precision_score(
             f"Only average=None is supported. {constants.FEEDBACK_LINK}"
         )
 
-    y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred)
+    y_true_series, y_pred_series = utils.batch_convert_to_series(y_true, y_pred)
 
     is_accurate = y_true_series == y_pred_series
     unique_labels = (
@@ -306,7 +306,7 @@ def f1_score(
     average: typing.Optional[str] = "binary",
 ) -> pd.Series:
     # TODO(ashleyxu): support more average type, default to "binary"
-    y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred)
+    y_true_series, y_pred_series = utils.batch_convert_to_series(y_true, y_pred)
 
     if average is not None:
         raise NotImplementedError(
@@ -337,7 +337,7 @@ def mean_squared_error(
     y_true: Union[bpd.DataFrame, bpd.Series],
     y_pred: Union[bpd.DataFrame, bpd.Series],
 ) -> float:
-    y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred)
+    y_true_series, y_pred_series = utils.batch_convert_to_series(y_true, y_pred)
 
     return (y_pred_series - y_true_series).pow(2).sum() / len(y_true_series)
 

bigframes/ml/metrics/pairwise.py

@@ -25,7 +25,7 @@
 def paired_cosine_distances(
     X: Union[bpd.DataFrame, bpd.Series], Y: Union[bpd.DataFrame, bpd.Series]
 ) -> bpd.DataFrame:
-    X, Y = utils.convert_to_series(X, Y)
+    X, Y = utils.batch_convert_to_series(X, Y)
     joined_block, _ = X._block.join(Y._block, how="outer")
 
     result_block, _ = joined_block.project_expr(
@@ -45,7 +45,7 @@ def paired_cosine_distances(
 def paired_manhattan_distance(
     X: Union[bpd.DataFrame, bpd.Series], Y: Union[bpd.DataFrame, bpd.Series]
 ) -> bpd.DataFrame:
-    X, Y = utils.convert_to_series(X, Y)
+    X, Y = utils.batch_convert_to_series(X, Y)
     joined_block, _ = X._block.join(Y._block, how="outer")
 
     result_block, _ = joined_block.project_expr(
@@ -65,7 +65,7 @@ def paired_manhattan_distance(
 def paired_euclidean_distances(
     X: Union[bpd.DataFrame, bpd.Series], Y: Union[bpd.DataFrame, bpd.Series]
 ) -> bpd.DataFrame:
-    X, Y = utils.convert_to_series(X, Y)
+    X, Y = utils.batch_convert_to_series(X, Y)
     joined_block, _ = X._block.join(Y._block, how="outer")
 
     result_block, _ = joined_block.project_expr(

bigframes/ml/model_selection.py

@@ -162,7 +162,7 @@ def _convert_to_bf_type(
         if isinstance(type_instance, pd.Series) or isinstance(
             type_instance, bpd.Series
         ):
-            return next(utils.convert_to_series(input))
+            return next(utils.batch_convert_to_series(input))
 
         if isinstance(type_instance, pd.DataFrame) or isinstance(
             type_instance, bpd.DataFrame

bigframes/ml/utils.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 
 import typing
-from typing import Any, Generator, Literal, Mapping, Optional, Tuple, Union
+from typing import Any, Generator, Hashable, Literal, Mapping, Optional, Tuple, Union
 
 import bigframes_vendored.constants as constants
 from google.cloud import bigquery
 import pandas as pd
 
-from bigframes.core import blocks, guid
+from bigframes.core import convert, guid
 import bigframes.pandas as bpd
 from bigframes.session import Session
 
@@ -65,7 +65,7 @@ def _convert_to_dataframe(
     )
 
 
-def convert_to_series(
+def batch_convert_to_series(
     *input: ArrayType, session: Optional[Session] = None
 ) -> Generator[bpd.Series, None, None]:
     """Converts the input to BigFrames Series.
@@ -76,37 +76,29 @@ def convert_to_series(
             It is not used if the input itself is already a BigFrame data frame or series.
 
     """
-    return (_convert_to_series(frame, session) for frame in input)
+    return (
+        convert.to_bf_series(
+            _get_only_column(frame), default_index=None, session=session
+        )
+        for frame in input
+    )
 
 
-def _convert_to_series(
-    frame: ArrayType, session: Optional[Session] = None
-) -> bpd.Series:
-    if isinstance(frame, bpd.DataFrame):
-        if len(frame.columns) != 1:
-            raise ValueError(
-                "To convert into Series, DataFrames can only contain one column. "
-                f"Try input with only one column. {constants.FEEDBACK_LINK}"
-            )
-
-        label = typing.cast(blocks.Label, frame.columns.tolist()[0])
-        return typing.cast(bpd.Series, frame[label])
-    if isinstance(frame, bpd.Series):
-        return frame
-    if isinstance(frame, pd.DataFrame):
-        # Recursively call this method to re-use the length-checking logic
-        if session is None:
-            return _convert_to_series(bpd.read_pandas(frame))
-        else:
-            return _convert_to_series(session.read_pandas(frame), session)
-    if isinstance(frame, pd.Series):
-        if session is None:
-            return bpd.read_pandas(frame)
-        else:
-            return session.read_pandas(frame)
-    raise ValueError(
-        f"Unsupported type {type(frame)} to convert to Series. {constants.FEEDBACK_LINK}"
-    )
+def _get_only_column(input: ArrayType) -> Union[pd.Series, bpd.Series]:
+    if isinstance(input, pd.Series) or isinstance(input, bpd.Series):
+        return input
+
+    if len(input.columns) != 1:
+        raise ValueError(
+            "To convert into Series, DataFrames can only contain one column. "
+            f"Try input with only one column. {constants.FEEDBACK_LINK}"
+        )
+
+    label = typing.cast(Hashable, input.columns.tolist()[0])
+    if isinstance(input, pd.DataFrame):
+        return typing.cast(pd.Series, input[label])
+
+    return typing.cast(bpd.Series, input[label])
 
 
 def parse_model_endpoint(model_endpoint: str) -> tuple[str, Optional[str]]:

tests/system/small/ml/test_utils.py

@@ -61,7 +61,7 @@ def test_convert_pandas_to_dataframe(data, session):
 def test_convert_to_series(session, data):
     bf_data = session.read_pandas(data)
 
-    (actual_result,) = utils.convert_to_series(bf_data)
+    (actual_result,) = utils.batch_convert_to_series(bf_data)
 
     pandas.testing.assert_series_equal(
         actual_result.to_pandas(), _SERIES, check_index_type=False, check_dtype=False
@@ -73,7 +73,7 @@ def test_convert_to_series(session, data):
     [pytest.param(_DATA_FRAME, id="dataframe"), pytest.param(_SERIES, id="series")],
 )
 def test_convert_pandas_to_series(data, session):
-    (actual_result,) = utils.convert_to_series(data, session=session)
+    (actual_result,) = utils.batch_convert_to_series(data, session=session)
 
     pandas.testing.assert_series_equal(
         actual_result.to_pandas(), _SERIES, check_index_type=False, check_dtype=False