Skip to content

Commit

Permalink
Implement a mixin for binops (#10360)
Browse files Browse the repository at this point in the history 
This PR builds on the framework introduced in #9925 to implement scans. I plan to apply this mixin to ColumnBase as well, but that will require more work to clean up binary operations for column types and it is large enough to merit a separate PR.

Contributes to #10177.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Michael Wang (https://github.com/isVoid)
  - Ashwin Srinath (https://github.com/shwina)

URL: #10360
  • Loading branch information
@vyasr
vyasr authored Mar 2, 2022
1 parent 1217f24 commit 7120694
Show file tree
Hide file tree
Showing 11 changed files with 212 additions and 228 deletions.
131 changes: 27 additions & 104 deletions python/cudf/cudf/core/frame.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
)
from cudf.core.column_accessor import ColumnAccessor
from cudf.core.join import Merge, MergeSemi
from cudf.core.mixins import BinaryOperand
from cudf.core.window import Rolling
from cudf.utils import ioutils
from cudf.utils.docutils import copy_docstring
Expand Down Expand Up @@ -97,7 +98,7 @@
}


class Frame:
class Frame(BinaryOperand):
"""A collection of Column objects with an optional index.
Parameters
Expand All @@ -114,6 +115,8 @@ class Frame:
_index: Optional[cudf.core.index.BaseIndex]
_names: Optional[List]

_VALID_BINARY_OPERATIONS = BinaryOperand._SUPPORTED_BINARY_OPERATIONS

def __init__(self, data=None, index=None):
if data is None:
data = {}
Expand Down Expand Up @@ -3555,29 +3558,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 +3610,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 +3893,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 +5108,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 +5188,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 +5269,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 +5355,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 +5438,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 +5525,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 +5596,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 +5679,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 +5759,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 +5839,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 +5919,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 +6016,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 +6101,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 +6195,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 +6273,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 +6349,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 +6425,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 +6501,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 +6577,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 +6653,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

0 comments on commit 7120694

Please sign in to comment.