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.

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

Supports list-like Python objects for Series comparison. #2022

Open
wants to merge 13 commits into
base: master
Choose a base branch
from
14 changes: 13 additions & 1 deletion databricks/koalas/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -598,7 +598,19 @@ def rpow_func(left, right):
__abs__ = column_op(F.abs)

# comparison operators
__eq__ = column_op(Column.__eq__)
def __eq__(self, other) -> Union["Series", "Index"]: # type: ignore[override]
if isinstance(other, (list, tuple)):
with ks.option_context("compute.ordered_head", True):
pindex_ops = self.head(len(other) + 1)._to_internal_pandas() # type: ignore
if len(pindex_ops) != len(other):
raise ValueError("Lengths must be equal")
return ks.from_pandas(pindex_ops == other) # type: ignore
# pandas always returns False for all items with dict and set.
elif isinstance(other, (dict, set)):
return self != self
else:
return column_op(Column.__eq__)(self, other)

__ne__ = column_op(Column.__ne__)
__lt__ = column_op(Column.__lt__)
__le__ = column_op(Column.__le__)
Expand Down
18 changes: 17 additions & 1 deletion databricks/koalas/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -636,7 +636,7 @@ def rfloordiv(self, other) -> "Series":
koalas = CachedAccessor("koalas", KoalasSeriesMethods)

# Comparison Operators
def eq(self, other) -> bool:
def eq(self, other) -> "Series":
itholic marked this conversation as resolved.
Show resolved Hide resolved
"""
Compare if the current value is equal to the other.

Expand All @@ -657,6 +657,22 @@ def eq(self, other) -> bool:
c True
d False
Name: b, dtype: bool

Support for list-like Python object with same length

>>> df.a == [1, 3, 2, 4]
a True
b False
c False
d True
Name: a, dtype: bool

>>> df.a.eq([1, 3, 2, 4])
a True
b False
c False
d True
Name: a, dtype: bool
"""
return self == other

Expand Down
69 changes: 69 additions & 0 deletions databricks/koalas/tests/test_ops_on_diff_frames.py
Original file line number Diff line number Diff line change
Expand Up @@ -1651,6 +1651,61 @@ def test_index_ops(self):
else:
self.assert_eq(kidx1 * 10 + kidx3, (pidx1 * 10 + pidx3).rename(None))

def test_series_eq(self):
pser = pd.Series([1, 2, 3, 4, 5, 6], name="x")
kser = ks.from_pandas(pser)

# other = Series
pandas_other = pd.Series([np.nan, 1, 3, 4, np.nan, 6], name="x")
koalas_other = ks.from_pandas(pandas_other)
self.assert_eq(pser.eq(pandas_other), kser.eq(koalas_other).sort_index())
self.assert_eq(pser == pandas_other, (kser == koalas_other).sort_index())

# other = Series with different Index
pandas_other = pd.Series(
[np.nan, 1, 3, 4, np.nan, 6], index=[10, 20, 30, 40, 50, 60], name="x"
)
koalas_other = ks.from_pandas(pandas_other)
self.assert_eq(pser.eq(pandas_other), kser.eq(koalas_other).sort_index())

# other = Index
pandas_other = pd.Index([np.nan, 1, 3, 4, np.nan, 6], name="x")
koalas_other = ks.from_pandas(pandas_other)
self.assert_eq(pser.eq(pandas_other), kser.eq(koalas_other).sort_index())
self.assert_eq(pser == pandas_other, (kser == koalas_other).sort_index())

# other = list
other = [np.nan, 1, 3, 4, np.nan, 6]
if LooseVersion(pd.__version__) >= LooseVersion("1.2"):
self.assert_eq(pser.eq(other), kser.eq(other).sort_index())
self.assert_eq(pser == other, (kser == other).sort_index())
else:
self.assert_eq(pser.eq(other).rename("x"), kser.eq(other).sort_index())
self.assert_eq((pser == other).rename("x"), (kser == other).sort_index())

# other = tuple
other = (np.nan, 1, 3, 4, np.nan, 6)
if LooseVersion(pd.__version__) >= LooseVersion("1.2"):
self.assert_eq(pser.eq(other), kser.eq(other).sort_index())
self.assert_eq(pser == other, (kser == other).sort_index())
else:
self.assert_eq(pser.eq(other).rename("x"), kser.eq(other).sort_index())
self.assert_eq((pser == other).rename("x"), (kser == other).sort_index())

# other = list with the different length
other = [np.nan, 1, 3, 4, np.nan]
with self.assertRaisesRegex(ValueError, "Lengths must be equal"):
kser.eq(other)
with self.assertRaisesRegex(ValueError, "Lengths must be equal"):
kser == other

# other = tuple with the different length
other = (np.nan, 1, 3, 4, np.nan)
with self.assertRaisesRegex(ValueError, "Lengths must be equal"):
kser.eq(other)
with self.assertRaisesRegex(ValueError, "Lengths must be equal"):
kser == other

def test_align(self):
pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}, index=[10, 20, 30])
pdf2 = pd.DataFrame({"a": [4, 5, 6], "c": ["d", "e", "f"]}, index=[10, 11, 12])
Expand Down Expand Up @@ -1869,6 +1924,20 @@ def test_mask(self):
with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"):
kdf1.mask(kdf2 > -250)

def test_series_eq(self):
pser = pd.Series([1, 2, 3, 4, 5, 6], name="x")
kser = ks.from_pandas(pser)

others = (
ks.Series([np.nan, 1, 3, 4, np.nan, 6], name="x"),
ks.Index([np.nan, 1, 3, 4, np.nan, 6], name="x"),
)
for other in others:
with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"):
kser.eq(other)
with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"):
kser == other

def test_align(self):
pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}, index=[10, 20, 30])
pdf2 = pd.DataFrame({"a": [4, 5, 6], "c": ["d", "e", "f"]}, index=[10, 11, 12])
Expand Down
36 changes: 36 additions & 0 deletions databricks/koalas/tests/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -2819,6 +2819,42 @@ def test_backfill(self):
kser.backfill(inplace=True)
self.assert_eq(expected, kser)

def test_eq(self):
pser = pd.Series([1, 2, 3, 4, 5, 6], name="x")
kser = ks.from_pandas(pser)

# other = Series
self.assert_eq(pser.eq(pser), kser.eq(kser))
self.assert_eq(pser == pser, kser == kser)

# other = list
other = [np.nan, 1, 3, 4, np.nan, 6]
if LooseVersion(pd.__version__) >= LooseVersion("1.2"):
self.assert_eq(pser.eq(other), kser.eq(other))
self.assert_eq(pser == other, kser == other)
else:
self.assert_eq(pser.eq(other).rename("x"), kser.eq(other))
self.assert_eq((pser == other).rename("x"), kser == other)

# other = tuple
other = (np.nan, 1, 3, 4, np.nan, 6)
if LooseVersion(pd.__version__) >= LooseVersion("1.2"):
self.assert_eq(pser.eq(other), kser.eq(other))
self.assert_eq(pser == other, kser == other)
else:
self.assert_eq(pser.eq(other).rename("x"), kser.eq(other))
self.assert_eq((pser == other).rename("x"), kser == other)

# other = dict
other = {1: None, 2: None, 3: None, 4: None, np.nan: None, 6: None}
self.assert_eq(pser.eq(other), kser.eq(other))
self.assert_eq(pser == other, kser == other)

# other = set
other = {1, 2, 3, 4, np.nan, 6}
self.assert_eq(pser.eq(other), kser.eq(other))
self.assert_eq(pser == other, kser == other)

def test_align(self):
pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]})
kdf = ks.from_pandas(pdf)
Expand Down