feat: Add Series.combine (#680)

Author: TrevorBergeron

bigframes/core/compile/scalar_op_compiler.py

@@ -1298,22 +1298,36 @@ def coalesce_impl(
         return ibis.coalesce(x, y)
 
 
-@scalar_op_compiler.register_binary_op(ops.cliplower_op)
-def clip_lower(
+@scalar_op_compiler.register_binary_op(ops.maximum_op)
+def maximum_impl(
     value: ibis_types.Value,
     lower: ibis_types.Value,
 ):
+    # Note: propagates nulls
     return ibis.case().when(lower.isnull() | (value < lower), lower).else_(value).end()
 
 
-@scalar_op_compiler.register_binary_op(ops.clipupper_op)
-def clip_upper(
+@scalar_op_compiler.register_binary_op(ops.minimum_op)
+def minimum_impl(
     value: ibis_types.Value,
     upper: ibis_types.Value,
 ):
+    # Note: propagates nulls
     return ibis.case().when(upper.isnull() | (value > upper), upper).else_(value).end()
 
 
+@scalar_op_compiler.register_binary_op(ops.BinaryRemoteFunctionOp, pass_op=True)
+def binary_remote_function_op_impl(
+    x: ibis_types.Value, y: ibis_types.Value, op: ops.BinaryRemoteFunctionOp
+):
+    if not hasattr(op.func, "bigframes_remote_function"):
+        raise TypeError(
+            f"only a bigframes remote function is supported as a callable. {constants.FEEDBACK_LINK}"
+        )
+    x_transformed = op.func(x, y)
+    return x_transformed
+
+
 # Ternary Operations
 @scalar_op_compiler.register_ternary_op(ops.where_op)
 def where_op(

bigframes/operations/__init__.py

@@ -545,8 +545,8 @@ def output_type(self, *input_types):
 
 # Binary Ops
 fillna_op = create_binary_op(name="fillna", type_signature=op_typing.COERCE)
-cliplower_op = create_binary_op(name="clip_lower", type_signature=op_typing.COERCE)
-clipupper_op = create_binary_op(name="clip_upper", type_signature=op_typing.COERCE)
+maximum_op = create_binary_op(name="maximum", type_signature=op_typing.COERCE)
+minimum_op = create_binary_op(name="minimum", type_signature=op_typing.COERCE)
 coalesce_op = create_binary_op(name="coalesce", type_signature=op_typing.COERCE)
 
 
@@ -587,6 +587,16 @@ def output_type(self, *input_types):
         raise TypeError(f"Cannot subtract dtypes {left_type} and {right_type}")
 
 
+@dataclasses.dataclass(frozen=True)
+class BinaryRemoteFunctionOp(BinaryOp):
+    name: typing.ClassVar[str] = "binary_remote_function"
+    func: typing.Callable
+
+    def output_type(self, *input_types):
+        # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
+        return self.func.output_dtype
+
+
 add_op = AddOp()
 sub_op = SubOp()
 mul_op = create_binary_op(name="mul", type_signature=op_typing.BINARY_NUMERIC)
@@ -713,4 +723,6 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
     np.divide: div_op,
     np.power: pow_op,
     np.arctan2: arctan2_op,
+    np.maximum: maximum_op,
+    np.minimum: minimum_op,
 }

bigframes/series.py

@@ -1031,9 +1031,9 @@ def clip(self, lower, upper):
         if lower is None and upper is None:
             return self
         if lower is None:
-            return self._apply_binary_op(upper, ops.clipupper_op, alignment="left")
+            return self._apply_binary_op(upper, ops.minimum_op, alignment="left")
         if upper is None:
-            return self._apply_binary_op(lower, ops.cliplower_op, alignment="left")
+            return self._apply_binary_op(lower, ops.maximum_op, alignment="left")
         value_id, lower_id, upper_id, block = self._align3(lower, upper)
         block, result_id = block.apply_ternary_op(
             value_id, lower_id, upper_id, ops.clip_op
@@ -1374,6 +1374,38 @@ def apply(
         materialized_series = result_series._cached()
         return materialized_series
 
+    def combine(
+        self,
+        other,
+        func,
+    ) -> Series:
+        if not callable(func):
+            raise ValueError(
+                "Only a ufunc (a function that applies to the entire Series) or a remote function that only works on single values are supported."
+            )
+
+        if not hasattr(func, "bigframes_remote_function"):
+            # Keep this in sync with .apply
+            try:
+                return func(self, other)
+            except Exception as ex:
+                # This could happen if any of the operators in func is not
+                # supported on a Series. Let's guide the customer to use a
+                # remote function instead
+                if hasattr(ex, "message"):
+                    ex.message += f"\n{_remote_function_recommendation_message}"
+                raise
+
+        reprojected_series = Series(self._block._force_reproject())
+        result_series = reprojected_series._apply_binary_op(
+            other, ops.BinaryRemoteFunctionOp(func=func)
+        )
+
+        # return Series with materialized result so that any error in the remote
+        # function is caught early
+        materialized_series = result_series._cached()
+        return materialized_series
+
     def add_prefix(self, prefix: str, axis: int | str | None = None) -> Series:
         return Series(self._get_block().add_prefix(prefix))
 

tests/system/large/test_remote_function.py

@@ -221,6 +221,41 @@ def stringify(x):
         )
 
 
+# @pytest.mark.flaky(retries=2, delay=120)
+def test_remote_function_binop(session, scalars_dfs, dataset_id, bq_cf_connection):
+    try:
+
+        def func(x, y):
+            return x * abs(y % 4)
+
+        remote_func = session.remote_function(
+            [str, int],
+            str,
+            dataset_id,
+            bq_cf_connection,
+            reuse=False,
+        )(func)
+
+        scalars_df, scalars_pandas_df = scalars_dfs
+
+        scalars_df = scalars_df.dropna()
+        scalars_pandas_df = scalars_pandas_df.dropna()
+        bf_result = (
+            scalars_df["string_col"]
+            .combine(scalars_df["int64_col"], remote_func)
+            .to_pandas()
+        )
+        pd_result = scalars_pandas_df["string_col"].combine(
+            scalars_pandas_df["int64_col"], func
+        )
+        pandas.testing.assert_series_equal(bf_result, pd_result)
+    finally:
+        # clean up the gcp assets created for the remote function
+        cleanup_remote_function_assets(
+            session.bqclient, session.cloudfunctionsclient, remote_func
+        )
+
+
 @pytest.mark.flaky(retries=2, delay=120)
 def test_remote_function_decorator_with_bigframes_series(
     session, scalars_dfs, dataset_id, bq_cf_connection

tests/system/small/test_numpy.py

@@ -73,27 +73,6 @@ def test_df_ufuncs(scalars_dfs, opname):
     pd.testing.assert_frame_equal(bf_result, pd_result)
 
 
-@pytest.mark.parametrize(
-    ("opname",),
-    [
-        ("add",),
-        ("subtract",),
-        ("multiply",),
-        ("divide",),
-        ("power",),
-        ("arctan2",),
-    ],
-)
-def test_series_binary_ufuncs(floats_product_pd, floats_product_bf, opname):
-    bf_result = getattr(np, opname)(
-        floats_product_bf.float64_col_x, floats_product_bf.float64_col_y
-    ).to_pandas()
-    pd_result = getattr(np, opname)(
-        floats_product_pd.float64_col_x, floats_product_pd.float64_col_y
-    )
-    pd.testing.assert_series_equal(bf_result, pd_result)
-
-
 @pytest.mark.parametrize(
     ("opname",),
     [
@@ -106,30 +85,42 @@ def test_series_binary_ufuncs(floats_product_pd, floats_product_bf, opname):
 )
 def test_df_binary_ufuncs(scalars_dfs, opname):
     scalars_df, scalars_pandas_df = scalars_dfs
+    op = getattr(np, opname)
 
-    bf_result = getattr(np, opname)(
-        scalars_df[["float64_col", "int64_col"]], 5.1
-    ).to_pandas()
-    pd_result = getattr(np, opname)(
-        scalars_pandas_df[["float64_col", "int64_col"]], 5.1
-    )
+    bf_result = op(scalars_df[["float64_col", "int64_col"]], 5.1).to_pandas()
+    pd_result = op(scalars_pandas_df[["float64_col", "int64_col"]], 5.1)
 
     pd.testing.assert_frame_equal(bf_result, pd_result)
 
 
+# Operations tested here don't work on full dataframe in numpy+pandas
+# Maybe because of nullable dtypes?
 @pytest.mark.parametrize(
     ("x", "y"),
     [
         ("int64_col", "int64_col"),
         ("float64_col", "int64_col"),
     ],
 )
-def test_series_atan2(scalars_dfs, x, y):
-    # Test atan2 separately as pandas errors when passing entire df as input, so pass only series
+@pytest.mark.parametrize(
+    ("opname",),
+    [
+        ("add",),
+        ("subtract",),
+        ("multiply",),
+        ("divide",),
+        ("arctan2",),
+        ("minimum",),
+        ("maximum",),
+    ],
+)
+def test_series_binary_ufuncs(scalars_dfs, x, y, opname):
     scalars_df, scalars_pandas_df = scalars_dfs
 
-    bf_result = np.arctan2(scalars_df[x], scalars_df[y]).to_pandas()
-    pd_result = np.arctan2(scalars_pandas_df[x], scalars_pandas_df[y])
+    op = getattr(np, opname)
+
+    bf_result = op(scalars_df[x], scalars_df[y]).to_pandas()
+    pd_result = op(scalars_pandas_df[x], scalars_pandas_df[y])
 
     pd.testing.assert_series_equal(bf_result, pd_result)
 

tests/system/small/test_series.py

@@ -3509,6 +3509,41 @@ def test_apply_numpy_ufunc(scalars_dfs, ufunc):
     assert_series_equal(bf_result, pd_result)
 
 
+@pytest.mark.parametrize(
+    ("ufunc",),
+    [
+        pytest.param(numpy.add),
+        pytest.param(numpy.divide),
+    ],
+    ids=[
+        "add",
+        "divide",
+    ],
+)
+def test_combine_series_ufunc(scalars_dfs, ufunc):
+    scalars_df, scalars_pandas_df = scalars_dfs
+
+    bf_col = scalars_df["int64_col"].dropna()
+    bf_result = bf_col.combine(bf_col, ufunc).to_pandas()
+
+    pd_col = scalars_pandas_df["int64_col"].dropna()
+    pd_result = pd_col.combine(pd_col, ufunc)
+
+    assert_series_equal(bf_result, pd_result, check_dtype=False)
+
+
+def test_combine_scalar_ufunc(scalars_dfs):
+    scalars_df, scalars_pandas_df = scalars_dfs
+
+    bf_col = scalars_df["int64_col"].dropna()
+    bf_result = bf_col.combine(2.5, numpy.add).to_pandas()
+
+    pd_col = scalars_pandas_df["int64_col"].dropna()
+    pd_result = pd_col.combine(2.5, numpy.add)
+
+    assert_series_equal(bf_result, pd_result, check_dtype=False)
+
+
 def test_apply_simple_udf(scalars_dfs):
     scalars_df, scalars_pandas_df = scalars_dfs
 

third_party/bigframes_vendored/pandas/core/series.py

@@ -1279,6 +1279,62 @@ def apply(
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+    def combine(
+        self,
+        other: Series | Hashable,
+        func,
+    ) -> Series:
+        """
+        Combine the Series with a Series or scalar according to `func`.
+
+        Combine the Series and `other` using `func` to perform elementwise
+        selection for combined Series.
+        `fill_value` is assumed when value is missing at some index
+        from one of the two objects being combined.
+
+        **Examples:**
+
+            >>> import bigframes.pandas as bpd
+            >>> import numpy as np
+            >>> bpd.options.display.progress_bar = None
+
+            Consider 2 Datasets ``s1`` and ``s2`` containing
+            highest clocked speeds of different birds.
+
+            >>> s1 = bpd.Series({'falcon': 330.0, 'eagle': 160.0})
+            >>> s1
+            falcon    330.0
+            eagle     160.0
+            dtype: Float64
+            >>> s2 = bpd.Series({'falcon': 345.0, 'eagle': 200.0, 'duck': 30.0})
+            >>> s2
+            falcon    345.0
+            eagle     200.0
+            duck       30.0
+            dtype: Float64
+
+            Now, to combine the two datasets and view the highest speeds
+            of the birds across the two datasets
+
+            >>> s1.combine(s2, np.maximum)
+            falcon    345.0
+            eagle     200.0
+            duck       <NA>
+            dtype: Float64
+
+        Args:
+            other (Series or scalar):
+                The value(s) to be combined with the `Series`.
+            func (function):
+                BigFrames DataFrames ``remote_function`` to apply.
+                Takes two scalars as inputs and returns an element.
+                Also accepts some numpy binary functions.
+
+        Returns:
+            Series: The result of combining the Series with the other object.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
     def groupby(
         self,
         by=None,