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

Enabling binary operations with list-like Python objects. #2054

Open
wants to merge 6 commits into
base: master
Choose a base branch
from
Open

Enabling binary operations with list-like Python objects. #2054

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
2e20719
Add tests
itholic Feb 15, 2021
922ba6a
Fix rsub
itholic Feb 15, 2021
02c3334
Add Index
itholic Feb 18, 2021
8df4ff9
Fix test for mod and rmod
itholic Feb 18, 2021
5d2d1c5
Use pandas
itholic Feb 18, 2021
0fd3666
Fix test
itholic Feb 18, 2021
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
51 changes: 50 additions & 1 deletion databricks/koalas/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
scol_for,
validate_axis,
ERROR_MESSAGE_CANNOT_COMBINE,
check_same_length,
)
from databricks.koalas.frame import DataFrame

Expand Down Expand Up @@ -321,6 +322,9 @@ def spark_column(self) -> Column:
__neg__ = column_op(Column.__neg__)

def __add__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(pindex_ops + other) # type: ignore
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we avoid using # type: ignore as possible? We can use cast instead.

if not isinstance(self.spark.data_type, StringType) and (
(isinstance(other, IndexOpsMixin) and isinstance(other.spark.data_type, StringType))
or isinstance(other, str)
Expand All @@ -339,6 +343,9 @@ def __add__(self, other) -> Union["Series", "Index"]:
return column_op(Column.__add__)(self, other)

def __sub__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(pindex_ops - other) # type: ignore
if (
isinstance(self.spark.data_type, StringType)
or (isinstance(other, IndexOpsMixin) and isinstance(other.spark.data_type, StringType))
Expand Down Expand Up @@ -383,6 +390,9 @@ def __sub__(self, other) -> Union["Series", "Index"]:
return column_op(Column.__sub__)(self, other)

def __mul__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(pindex_ops * other) # type: ignore
if isinstance(other, str):
raise TypeError("multiplication can not be applied to a string literal.")

Expand Down Expand Up @@ -422,6 +432,9 @@ def __truediv__(self, other) -> Union["Series", "Index"]:
| -10 | null | -np.inf |
+-----------------------|---------|---------+
"""
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(pindex_ops / other) # type: ignore

if (
isinstance(self.spark.data_type, StringType)
Expand All @@ -440,6 +453,9 @@ def truediv(left, right):
return numpy_column_op(truediv)(self, other)

def __mod__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(pindex_ops % other) # type: ignore
if (
isinstance(self.spark.data_type, StringType)
or (isinstance(other, IndexOpsMixin) and isinstance(other.spark.data_type, StringType))
Expand All @@ -453,6 +469,11 @@ def mod(left, right):
return column_op(mod)(self, other)

def __radd__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(other + pindex_ops) # type: ignore
if isinstance(other, (list, tuple)):
other = ks.Index(other, name=self.name) # type: ignore
Comment on lines +475 to +476
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not needed?

# Handle 'literal' + df['col']
if not isinstance(self.spark.data_type, StringType) and isinstance(other, str):
raise TypeError("string addition can only be applied to string series or literals.")
Expand All @@ -466,6 +487,9 @@ def __radd__(self, other) -> Union["Series", "Index"]:
return column_op(Column.__radd__)(self, other)

def __rsub__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(other - pindex_ops) # type: ignore
if isinstance(self.spark.data_type, StringType) or isinstance(other, str):
raise TypeError("substraction can not be applied to string series or literals.")

Expand Down Expand Up @@ -495,9 +519,12 @@ def __rsub__(self, other) -> Union["Series", "Index"]:
return -column_op(F.datediff)(self, F.lit(other)).astype("long")
else:
raise TypeError("date subtraction can only be applied to date series.")
return column_op(Column.__rsub__)(self, other)
return column_op(lambda left, right: right - left)(self, other)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI: Column.__rsub__ doesn't support pyspark.sql.column.Column for second parameter.

>>> kdf = ks.DataFrame({"A": [1, 2, 3, 4], "B": [10, 20, 30, 40]})
>>> sdf = kdf.to_spark()
>>> col1 = sdf.A
>>> col2 = sdf.B
>>> Column.__rsub__(col1, col2)
Traceback (most recent call last):
...
TypeError: Column is not iterable

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It does support:

>>> Column.__rsub__(df.id, 1)
Column<'(1 - id)'>

It doesn't work in your case above because the instance is Spark column. In practice, that wouldn't happen because it will only be called when the first operand doesn't know how to handle Spark column e.g.) 1 - df.id.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does it cause any exception?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we use column_op(Column.__rsub__)(self, other) as it is, it raises TypeError: Column is not iterable for the case below.

>>> kser = ks.Series([1, 2, 3, 4])
>>> [10, 20, 30, 40] - kser
Traceback (most recent call last):
...
TypeError: Column is not iterable

Copy link
Collaborator

@ueshin ueshin Feb 18, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not that this case must be handled in lines 490-492. We can move back to Column.__rsub__.


def __rmul__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(other * pindex_ops) # type: ignore
if isinstance(other, str):
raise TypeError("multiplication can not be applied to a string literal.")

Expand All @@ -512,6 +539,9 @@ def __rmul__(self, other) -> Union["Series", "Index"]:
return column_op(Column.__rmul__)(self, other)

def __rtruediv__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(other / pindex_ops) # type: ignore
if isinstance(self.spark.data_type, StringType) or isinstance(other, str):
raise TypeError("division can not be applied on string series or literals.")

Expand Down Expand Up @@ -539,6 +569,9 @@ def __floordiv__(self, other) -> Union["Series", "Index"]:
| -10 | null | -np.inf |
+-----------------------|---------|---------+
"""
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(pindex_ops // other) # type: ignore
if (
isinstance(self.spark.data_type, StringType)
or (isinstance(other, IndexOpsMixin) and isinstance(other.spark.data_type, StringType))
Expand All @@ -560,6 +593,11 @@ def floordiv(left, right):
return numpy_column_op(floordiv)(self, other)

def __rfloordiv__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(other // pindex_ops) # type: ignore
if isinstance(other, (list, tuple)):
other = ks.Index(other, name=self.name) # type: ignore
Comment on lines +599 to +600
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not needed?

if isinstance(self.spark.data_type, StringType) or isinstance(other, str):
raise TypeError("division can not be applied on string series or literals.")

Expand All @@ -571,6 +609,9 @@ def rfloordiv(left, right):
return numpy_column_op(rfloordiv)(self, other)

def __rmod__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(other % pindex_ops) # type: ignore
if isinstance(self.spark.data_type, StringType) or isinstance(other, str):
raise TypeError("modulo can not be applied on string series or literals.")

Expand All @@ -580,12 +621,20 @@ def rmod(left, right):
return column_op(rmod)(self, other)

def __pow__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(pindex_ops ** other) # type: ignore

def pow_func(left, right):
return F.when(left == 1, left).otherwise(Column.__pow__(left, right))

return column_op(pow_func)(self, other)

def __rpow__(self, other) -> Union["Series", "Index"]:
if isinstance(other, (list, tuple)):
pindex_ops, other = check_same_length(self, other)
return ks.from_pandas(other ** pindex_ops) # type: ignore

def rpow_func(left, right):
return F.when(F.lit(right == 1), right).otherwise(Column.__rpow__(left, right))

Expand Down
2 changes: 1 addition & 1 deletion databricks/koalas/indexes/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ class Index(IndexOpsMixin):
Index(['a', 'b', 'c'], dtype='object')
"""

def __new__(cls, data: Union[DataFrame, list], dtype=None, name=None, names=None):
def __new__(cls, data: Union[DataFrame, list, tuple], dtype=None, name=None, names=None):
from databricks.koalas.indexes.datetimes import DatetimeIndex
from databricks.koalas.indexes.multi import MultiIndex
from databricks.koalas.indexes.numeric import Float64Index, Int64Index
Expand Down
Loading