Skip to content

Commit

Permalink
feat: Add +, - as unary ops, ^ binary op (#724)
Browse files Browse the repository at this point in the history
  • Loading branch information
TrevorBergeron committed May 29, 2024
1 parent 6a78c89 commit 968d825
Show file tree
Hide file tree
Showing 8 changed files with 190 additions and 3 deletions.
20 changes: 20 additions & 0 deletions bigframes/core/compile/scalar_op_compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand Down
27 changes: 27 additions & 0 deletions bigframes/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down
5 changes: 4 additions & 1 deletion bigframes/operations/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
13 changes: 13 additions & 0 deletions bigframes/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down
16 changes: 16 additions & 0 deletions tests/system/small/test_dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
34 changes: 34 additions & 0 deletions tests/system/small/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",),
(
Expand Down Expand Up @@ -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])
Expand Down Expand Up @@ -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",
Expand All @@ -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):
Expand Down
36 changes: 36 additions & 0 deletions third_party/bigframes_vendored/pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
42 changes: 40 additions & 2 deletions third_party/bigframes_vendored/pandas/core/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit 968d825

Please sign in to comment.