Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement a mixin for binops #10360

Merged
merged 10 commits into from
Mar 2, 2022
Merged

Implement a mixin for binops #10360

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
26b6a84
Add BinaryOperand mixin and use it for Frame.
vyasr Feb 24, 2022
4339a72
Make RangeIndex a BinaryOperand.
vyasr Feb 24, 2022
adbf831
Add sentinel to simplify indicating that all supported operations are…
vyasr Feb 24, 2022
1168cd6
Stop preprocessing op names for Frame and change all calls appropriat…
vyasr Feb 24, 2022
ada6b45
Remove reflection from the general binop API and just add a helper to…
vyasr Feb 24, 2022
c58570c
Make Scalar a BinaryOperand.
vyasr Feb 24, 2022
70a13a0
Include matmul in list of supported ops.
vyasr Feb 25, 2022
32c22d2
Remove unnecessary sentinel value in favor of just using the full lis…
vyasr Mar 1, 2022
00d16ea
Merge remote-tracking branch 'origin/branch-22.04' into refactor/bino…
vyasr Mar 2, 2022
4ba4db9
Update comment.
vyasr Mar 2, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 28 additions & 104 deletions python/cudf/cudf/core/frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@
)
from cudf.core.column_accessor import ColumnAccessor
from cudf.core.join import Merge, MergeSemi
from cudf.core.mixins import BinaryOperand
from cudf.core.mixins.mixin_factory import AllOperations
from cudf.core.window import Rolling
from cudf.utils import ioutils
from cudf.utils.docutils import copy_docstring
Expand Down Expand Up @@ -97,7 +99,7 @@
}


class Frame:
class Frame(BinaryOperand):
"""A collection of Column objects with an optional index.

Parameters
Expand All @@ -114,6 +116,8 @@ class Frame:
_index: Optional[cudf.core.index.BaseIndex]
_names: Optional[List]

_VALID_BINARY_OPERATIONS = AllOperations
vyasr marked this conversation as resolved.
Show resolved Hide resolved

def __init__(self, data=None, index=None):
if data is None:
data = {}
Expand Down Expand Up @@ -3555,29 +3559,19 @@ def _unaryop(self, op):
)

def _binaryop(
self,
other: T,
fn: str,
fill_value: Any = None,
reflect: bool = False,
*args,
**kwargs,
self, other: T, op: str, fill_value: Any = None, *args, **kwargs,
) -> Frame:
"""Perform a binary operation between two frames.

Parameters
----------
other : Frame
The second operand.
fn : str
op : str
The operation to perform.
fill_value : Any, default None
The value to replace null values with. If ``None``, nulls are not
filled before the operation.
reflect : bool, default False
If ``True``, swap the order of the operands. See
https://docs.python.org/3/reference/datamodel.html#object.__ror__
for more information on when this is necessary.

Returns
-------
Expand Down Expand Up @@ -3617,6 +3611,7 @@ def _colwise_binop(
A dict of columns constructed from the result of performing the
requested operation on the operands.
"""
fn = fn[2:-2]

# Now actually perform the binop on the columns in left and right.
output = {}
Expand Down Expand Up @@ -3899,83 +3894,12 @@ def dot(self, other, reflect=False):
return cudf.DataFrame(result)
return result.item()

# Binary arithmetic operations.
def __add__(self, other):
return self._binaryop(other, "add")

def __radd__(self, other):
return self._binaryop(other, "add", reflect=True)

def __sub__(self, other):
return self._binaryop(other, "sub")

def __rsub__(self, other):
return self._binaryop(other, "sub", reflect=True)

def __matmul__(self, other):
return self.dot(other)

def __rmatmul__(self, other):
return self.dot(other, reflect=True)

def __mul__(self, other):
return self._binaryop(other, "mul")

def __rmul__(self, other):
return self._binaryop(other, "mul", reflect=True)

def __mod__(self, other):
return self._binaryop(other, "mod")

def __rmod__(self, other):
return self._binaryop(other, "mod", reflect=True)

def __pow__(self, other):
return self._binaryop(other, "pow")

def __rpow__(self, other):
return self._binaryop(other, "pow", reflect=True)

def __floordiv__(self, other):
return self._binaryop(other, "floordiv")

def __rfloordiv__(self, other):
return self._binaryop(other, "floordiv", reflect=True)

def __truediv__(self, other):
return self._binaryop(other, "truediv")

def __rtruediv__(self, other):
return self._binaryop(other, "truediv", reflect=True)

def __and__(self, other):
return self._binaryop(other, "and")

def __or__(self, other):
return self._binaryop(other, "or")

def __xor__(self, other):
return self._binaryop(other, "xor")

# Binary rich comparison operations.
def __eq__(self, other):
return self._binaryop(other, "eq")

def __ne__(self, other):
return self._binaryop(other, "ne")

def __lt__(self, other):
return self._binaryop(other, "lt")

def __le__(self, other):
return self._binaryop(other, "le")

def __gt__(self, other):
return self._binaryop(other, "gt")

def __ge__(self, other):
return self._binaryop(other, "ge")

# Unary logical operators
def __neg__(self):
return -1 * self
Expand Down Expand Up @@ -5185,7 +5109,7 @@ def add(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "add", fill_value)
return self._binaryop(other, "__add__", fill_value)

@annotate("FRAME_RADD", color="green", domain="cudf_python")
def radd(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -5265,7 +5189,7 @@ def radd(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "add", fill_value, reflect=True)
return self._binaryop(other, "__radd__", fill_value)

@annotate("FRAME_SUBTRACT", color="green", domain="cudf_python")
def subtract(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -5346,7 +5270,7 @@ def subtract(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "sub", fill_value)
return self._binaryop(other, "__sub__", fill_value)

sub = subtract

Expand Down Expand Up @@ -5432,7 +5356,7 @@ def rsub(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "sub", fill_value, reflect=True)
return self._binaryop(other, "__rsub__", fill_value)

@annotate("FRAME_MULTIPLY", color="green", domain="cudf_python")
def multiply(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -5515,7 +5439,7 @@ def multiply(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "mul", fill_value)
return self._binaryop(other, "__mul__", fill_value)

mul = multiply

Expand Down Expand Up @@ -5602,7 +5526,7 @@ def rmul(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "mul", fill_value, reflect=True)
return self._binaryop(other, "__rmul__", fill_value)

@annotate("FRAME_MOD", color="green", domain="cudf_python")
def mod(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -5673,7 +5597,7 @@ def mod(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "mod", fill_value)
return self._binaryop(other, "__mod__", fill_value)

@annotate("FRAME_RMOD", color="green", domain="cudf_python")
def rmod(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -5756,7 +5680,7 @@ def rmod(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "mod", fill_value, reflect=True)
return self._binaryop(other, "__rmod__", fill_value)

@annotate("FRAME_POW", color="green", domain="cudf_python")
def pow(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -5836,7 +5760,7 @@ def pow(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "pow", fill_value)
return self._binaryop(other, "__pow__", fill_value)

@annotate("FRAME_RPOW", color="green", domain="cudf_python")
def rpow(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -5916,7 +5840,7 @@ def rpow(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "pow", fill_value, reflect=True)
return self._binaryop(other, "__rpow__", fill_value)

@annotate("FRAME_FLOORDIV", color="green", domain="cudf_python")
def floordiv(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -5996,7 +5920,7 @@ def floordiv(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "floordiv", fill_value)
return self._binaryop(other, "__floordiv__", fill_value)

@annotate("FRAME_RFLOORDIV", color="green", domain="cudf_python")
def rfloordiv(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -6093,7 +6017,7 @@ def rfloordiv(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "floordiv", fill_value, reflect=True)
return self._binaryop(other, "__rfloordiv__", fill_value)

@annotate("FRAME_TRUEDIV", color="green", domain="cudf_python")
def truediv(self, other, axis, level=None, fill_value=None):
Expand Down Expand Up @@ -6178,7 +6102,7 @@ def truediv(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "truediv", fill_value)
return self._binaryop(other, "__truediv__", fill_value)

# Alias for truediv
div = truediv
Expand Down Expand Up @@ -6272,7 +6196,7 @@ def rtruediv(self, other, axis, level=None, fill_value=None):
if level is not None:
raise NotImplementedError("level parameter is not supported yet.")

return self._binaryop(other, "truediv", fill_value, reflect=True)
return self._binaryop(other, "__rtruediv__", fill_value)

# Alias for rtruediv
rdiv = rtruediv
Expand Down Expand Up @@ -6350,7 +6274,7 @@ def eq(self, other, axis="columns", level=None, fill_value=None):
dtype: bool
"""
return self._binaryop(
other=other, fn="eq", fill_value=fill_value, can_reindex=True
other=other, op="__eq__", fill_value=fill_value, can_reindex=True
)

@annotate("FRAME_NE", color="green", domain="cudf_python")
Expand Down Expand Up @@ -6426,7 +6350,7 @@ def ne(self, other, axis="columns", level=None, fill_value=None):
dtype: bool
""" # noqa: E501
return self._binaryop(
other=other, fn="ne", fill_value=fill_value, can_reindex=True
other=other, op="__ne__", fill_value=fill_value, can_reindex=True
)

@annotate("FRAME_LT", color="green", domain="cudf_python")
Expand Down Expand Up @@ -6502,7 +6426,7 @@ def lt(self, other, axis="columns", level=None, fill_value=None):
dtype: bool
""" # noqa: E501
return self._binaryop(
other=other, fn="lt", fill_value=fill_value, can_reindex=True
other=other, op="__lt__", fill_value=fill_value, can_reindex=True
)

@annotate("FRAME_LE", color="green", domain="cudf_python")
Expand Down Expand Up @@ -6578,7 +6502,7 @@ def le(self, other, axis="columns", level=None, fill_value=None):
dtype: bool
""" # noqa: E501
return self._binaryop(
other=other, fn="le", fill_value=fill_value, can_reindex=True
other=other, op="__le__", fill_value=fill_value, can_reindex=True
)

@annotate("FRAME_GT", color="green", domain="cudf_python")
Expand Down Expand Up @@ -6654,7 +6578,7 @@ def gt(self, other, axis="columns", level=None, fill_value=None):
dtype: bool
""" # noqa: E501
return self._binaryop(
other=other, fn="gt", fill_value=fill_value, can_reindex=True
other=other, op="__gt__", fill_value=fill_value, can_reindex=True
)

@annotate("FRAME_GE", color="green", domain="cudf_python")
Expand Down Expand Up @@ -6730,7 +6654,7 @@ def ge(self, other, axis="columns", level=None, fill_value=None):
dtype: bool
""" # noqa: E501
return self._binaryop(
other=other, fn="ge", fill_value=fill_value, can_reindex=True
other=other, op="__ge__", fill_value=fill_value, can_reindex=True
)

def nunique(self, method: builtins.str = "sort", dropna: bool = True):
Expand Down
Loading