feat: Add +, - as unary ops, ^ binary op (#724)

Author: TrevorBergeron

bigframes/core/compile/scalar_op_compiler.py

@@ -375,6 +375,16 @@ def abs_op_impl(x: ibis_types.Value):
     return typing.cast(ibis_types.NumericValue, x).abs()
 
 
+@scalar_op_compiler.register_unary_op(ops.pos_op)
+def pos_op_impl(x: ibis_types.Value):
+    return typing.cast(ibis_types.NumericValue, x)
+
+
+@scalar_op_compiler.register_unary_op(ops.neg_op)
+def neg_op_impl(x: ibis_types.Value):
+    return typing.cast(ibis_types.NumericValue, x).negate()
+
+
 @scalar_op_compiler.register_unary_op(ops.sqrt_op)
 def sqrt_op_impl(x: ibis_types.Value):
     numeric_value = typing.cast(ibis_types.NumericValue, x)
@@ -979,6 +989,16 @@ def or_op(
     )
 
 
+@scalar_op_compiler.register_binary_op(ops.xor_op)
+def xor_op(
+    x: ibis_types.Value,
+    y: ibis_types.Value,
+):
+    return typing.cast(ibis_types.BooleanValue, x) ^ typing.cast(
+        ibis_types.BooleanValue, y
+    )
+
+
 @scalar_op_compiler.register_binary_op(ops.add_op)
 @short_circuit_nulls()
 def add_op(

bigframes/dataframe.py

@@ -1112,6 +1112,33 @@ def __rpow__(self, other):
 
     __rpow__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__rpow__)
 
+    def __and__(self, other: bool | int | bigframes.series.Series) -> DataFrame:
+        return self._apply_binop(other, ops.and_op)
+
+    __and__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__and__)
+
+    __rand__ = __and__
+
+    def __or__(self, other: bool | int | bigframes.series.Series) -> DataFrame:
+        return self._apply_binop(other, ops.or_op)
+
+    __or__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__or__)
+
+    __ror__ = __or__
+
+    def __xor__(self, other: bool | int | bigframes.series.Series) -> DataFrame:
+        return self._apply_binop(other, ops.xor_op)
+
+    __xor__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__xor__)
+
+    __rxor__ = __xor__
+
+    def __pos__(self) -> DataFrame:
+        return self._apply_unary_op(ops.pos_op)
+
+    def __neg__(self) -> DataFrame:
+        return self._apply_unary_op(ops.neg_op)
+
     def align(
         self,
         other: typing.Union[DataFrame, bigframes.series.Series],

bigframes/operations/__init__.py

@@ -186,7 +186,7 @@ def create_binary_op(
         dtypes.is_binary_like,
         description="binary-like",
     ),
-)  # numeric
+)
 isnull_op = create_unary_op(
     name="isnull",
     type_signature=op_typing.FixedOutputType(
@@ -311,6 +311,8 @@ def create_binary_op(
 floor_op = create_unary_op(name="floor", type_signature=op_typing.UNARY_REAL_NUMERIC)
 ceil_op = create_unary_op(name="ceil", type_signature=op_typing.UNARY_REAL_NUMERIC)
 abs_op = create_unary_op(name="abs", type_signature=op_typing.UNARY_NUMERIC)
+pos_op = create_unary_op(name="pos", type_signature=op_typing.UNARY_NUMERIC)
+neg_op = create_unary_op(name="neg", type_signature=op_typing.UNARY_NUMERIC)
 exp_op = create_unary_op(name="exp", type_signature=op_typing.UNARY_REAL_NUMERIC)
 expm1_op = create_unary_op(name="expm1", type_signature=op_typing.UNARY_REAL_NUMERIC)
 ln_op = create_unary_op(name="log", type_signature=op_typing.UNARY_REAL_NUMERIC)
@@ -650,6 +652,7 @@ def output_type(self, *input_types):
 # Logical Ops
 and_op = create_binary_op(name="and", type_signature=op_typing.LOGICAL)
 or_op = create_binary_op(name="or", type_signature=op_typing.LOGICAL)
+xor_op = create_binary_op(name="xor", type_signature=op_typing.LOGICAL)
 
 ## Comparison Ops
 eq_op = create_binary_op(name="eq", type_signature=op_typing.COMPARISON)

bigframes/series.py

@@ -668,6 +668,13 @@ def __or__(self, other: bool | int | Series) -> Series:
 
     __ror__ = __or__
 
+    def __xor__(self, other: bool | int | Series) -> Series:
+        return self._apply_binary_op(other, ops.xor_op)
+
+    __or__.__doc__ = inspect.getdoc(vendored_pandas_series.Series.__xor__)
+
+    __rxor__ = __xor__
+
     def __add__(self, other: float | int | Series) -> Series:
         return self.add(other)
 
@@ -1036,6 +1043,12 @@ def __invert__(self) -> Series:
 
     __invert__.__doc__ = inspect.getdoc(vendored_pandas_series.Series.__invert__)
 
+    def __pos__(self) -> Series:
+        return self._apply_unary_op(ops.pos_op)
+
+    def __neg__(self) -> Series:
+        return self._apply_unary_op(ops.neg_op)
+
     def eq(self, other: object) -> Series:
         # TODO: enforce stricter alignment
         return self._apply_binary_op(other, ops.eq_op)

tests/system/small/test_dataframe.py

@@ -1699,6 +1699,22 @@ def test_df_abs(scalars_dfs):
     assert_pandas_df_equal(bf_result, pd_result)
 
 
+def test_df_pos(scalars_dfs):
+    scalars_df, scalars_pandas_df = scalars_dfs
+    bf_result = (+scalars_df[["int64_col", "numeric_col"]]).to_pandas()
+    pd_result = +scalars_pandas_df[["int64_col", "numeric_col"]]
+
+    assert_pandas_df_equal(pd_result, bf_result)
+
+
+def test_df_neg(scalars_dfs):
+    scalars_df, scalars_pandas_df = scalars_dfs
+    bf_result = (-scalars_df[["int64_col", "numeric_col"]]).to_pandas()
+    pd_result = -scalars_pandas_df[["int64_col", "numeric_col"]]
+
+    assert_pandas_df_equal(pd_result, bf_result)
+
+
 def test_df_invert(scalars_dfs):
     scalars_df, scalars_pandas_df = scalars_dfs
     columns = ["int64_col", "bool_col"]

tests/system/small/test_series.py

@@ -354,6 +354,36 @@ def test_abs(scalars_dfs, col_name):
     assert_series_equal(pd_result, bf_result)
 
 
+@pytest.mark.parametrize(
+    ("col_name",),
+    (
+        ("float64_col",),
+        ("int64_too",),
+    ),
+)
+def test_series_pos(scalars_dfs, col_name):
+    scalars_df, scalars_pandas_df = scalars_dfs
+    bf_result = (+scalars_df[col_name]).to_pandas()
+    pd_result = +scalars_pandas_df[col_name]
+
+    assert_series_equal(pd_result, bf_result)
+
+
+@pytest.mark.parametrize(
+    ("col_name",),
+    (
+        ("float64_col",),
+        ("int64_too",),
+    ),
+)
+def test_series_neg(scalars_dfs, col_name):
+    scalars_df, scalars_pandas_df = scalars_dfs
+    bf_result = (-scalars_df[col_name]).to_pandas()
+    pd_result = -scalars_pandas_df[col_name]
+
+    assert_series_equal(pd_result, bf_result)
+
+
 @pytest.mark.parametrize(
     ("col_name",),
     (
@@ -678,10 +708,12 @@ def test_series_pow_scalar_reverse(scalars_dfs):
     [
         (lambda x, y: x & y),
         (lambda x, y: x | y),
+        (lambda x, y: x ^ y),
     ],
     ids=[
         "and",
         "or",
+        "xor",
     ],
 )
 @pytest.mark.parametrize(("other_scalar"), [True, False, pd.NA])
@@ -714,6 +746,7 @@ def test_series_bool_bool_operators_scalar(
         (lambda x, y: x // y),
         (lambda x, y: x & y),
         (lambda x, y: x | y),
+        (lambda x, y: x ^ y),
     ],
     ids=[
         "add",
@@ -728,6 +761,7 @@ def test_series_bool_bool_operators_scalar(
         "floordivide",
         "bitwise_and",
         "bitwise_or",
+        "bitwise_xor",
     ],
 )
 def test_series_int_int_operators_series(scalars_dfs, operator):

third_party/bigframes_vendored/pandas/core/frame.py

@@ -3555,6 +3555,42 @@ def __rpow__(self, other):
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+    def __and__(self, other):
+        """Get bitwise AND of DataFrame and other, element-wise, using operator `&`.
+
+        Args:
+            other (scalar, Series or DataFrame):
+                Object to bitwise AND with the DataFrame.
+
+        Returns:
+            bigframes.dataframe.DataFrame: The result of the operation.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+    def __or__(self, other):
+        """Get bitwise OR of DataFrame and other, element-wise, using operator `|`.
+
+        Args:
+            other (scalar, Series or DataFrame):
+                Object to bitwise OR with the DataFrame.
+
+        Returns:
+            bigframes.dataframe.DataFrame: The result of the operation.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+    def __xor__(self, other):
+        """Get bitwise XOR of DataFrame and other, element-wise, using operator `^`.
+
+        Args:
+            other (scalar, Series or DataFrame):
+                Object to bitwise XOR with the DataFrame.
+
+        Returns:
+            bigframes.dataframe.DataFrame: The result of the operation.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
     def combine(
         self, other, func, fill_value=None, overwrite: bool = True
     ) -> DataFrame:

third_party/bigframes_vendored/pandas/core/series.py

@@ -4224,7 +4224,7 @@ def __and__(self, other):
                 Object to bitwise AND with the Series.
 
         Returns:
-            Series: The result of the operation.
+            bigframes.series.Series: The result of the operation.
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
@@ -4262,7 +4262,45 @@ def __or__(self, other):
                 Object to bitwise OR with the Series.
 
         Returns:
-            Series: The result of the operation.
+            bigframes.series.Series: The result of the operation.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+    def __xor__(self, other):
+        """Get bitwise XOR of Series and other, element-wise, using operator `^`.
+
+        **Examples:**
+
+            >>> import bigframes.pandas as bpd
+            >>> bpd.options.display.progress_bar = None
+
+            >>> s = bpd.Series([0, 1, 2, 3])
+
+        You can operate with a scalar.
+
+            >>> s ^ 6
+            0    6
+            1    7
+            2    4
+            3    5
+            dtype: Int64
+
+        You can operate with another Series.
+
+            >>> s1 = bpd.Series([5, 6, 7, 8])
+            >>> s ^ s1
+            0     5
+            1     7
+            2     5
+            3    11
+            dtype: Int64
+
+        Args:
+            other (scalar or Series):
+                Object to bitwise XOR with the Series.
+
+        Returns:
+            bigframes.series.Series: The result of the operation.
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)