Add tests

databricks · Feb 15, 2021 · 2e20719 · 2e20719
1 parent 87f5b18
commit 2e20719
Show file tree

Hide file tree

Showing 2 changed files with 733 additions and 0 deletions.
diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py
@@ -452,6 +452,10 @@ def spark_type(self):
     def add(self, other) -> "Series":
         return self + other
 
+    def __add__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__add__(self, other)
+
     add.__doc__ = _flex_doc_SERIES.format(
         desc="Addition",
         op_name="+",
@@ -463,6 +467,10 @@ def add(self, other) -> "Series":
     def radd(self, other) -> "Series":
         return other + self
 
+    def __radd__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__radd__(self, other)
+
     radd.__doc__ = _flex_doc_SERIES.format(
         desc="Reverse Addition",
         op_name="+",
@@ -498,6 +506,10 @@ def rdiv(self, other) -> "Series":
     def truediv(self, other) -> "Series":
         return self / other
 
+    def __truediv__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__truediv__(self, other)
+
     truediv.__doc__ = _flex_doc_SERIES.format(
         desc="Floating division",
         op_name="/",
@@ -509,6 +521,10 @@ def truediv(self, other) -> "Series":
     def rtruediv(self, other) -> "Series":
         return other / self
 
+    def __rtruediv__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__rtruediv__(self, other)
+
     rtruediv.__doc__ = _flex_doc_SERIES.format(
         desc="Reverse Floating division",
         op_name="/",
@@ -520,6 +536,10 @@ def rtruediv(self, other) -> "Series":
     def mul(self, other) -> "Series":
         return self * other
 
+    def __mul__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__mul__(self, other)
+
     mul.__doc__ = _flex_doc_SERIES.format(
         desc="Multiplication",
         op_name="*",
@@ -533,6 +553,10 @@ def mul(self, other) -> "Series":
     def rmul(self, other) -> "Series":
         return other * self
 
+    def __rmul__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__rmul__(self, other)
+
     rmul.__doc__ = _flex_doc_SERIES.format(
         desc="Reverse Multiplication",
         op_name="*",
@@ -544,6 +568,10 @@ def rmul(self, other) -> "Series":
     def sub(self, other) -> "Series":
         return self - other
 
+    def __sub__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__sub__(self, other)
+
     sub.__doc__ = _flex_doc_SERIES.format(
         desc="Subtraction",
         op_name="-",
@@ -557,6 +585,10 @@ def sub(self, other) -> "Series":
     def rsub(self, other) -> "Series":
         return other - self
 
+    def __rsub__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__rsub__(self, other)
+
     rsub.__doc__ = _flex_doc_SERIES.format(
         desc="Reverse Subtraction",
         op_name="-",
@@ -568,6 +600,10 @@ def rsub(self, other) -> "Series":
     def mod(self, other) -> "Series":
         return self % other
 
+    def __mod__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__mod__(self, other)
+
     mod.__doc__ = _flex_doc_SERIES.format(
         desc="Modulo",
         op_name="%",
@@ -579,6 +615,10 @@ def mod(self, other) -> "Series":
     def rmod(self, other) -> "Series":
         return other % self
 
+    def __rmod__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__rmod__(self, other)
+
     rmod.__doc__ = _flex_doc_SERIES.format(
         desc="Reverse Modulo",
         op_name="%",
@@ -590,6 +630,10 @@ def rmod(self, other) -> "Series":
     def pow(self, other) -> "Series":
         return self ** other
 
+    def __pow__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__pow__(self, other)
+
     pow.__doc__ = _flex_doc_SERIES.format(
         desc="Exponential power of series",
         op_name="**",
@@ -601,6 +645,10 @@ def pow(self, other) -> "Series":
     def rpow(self, other) -> "Series":
         return other ** self
 
+    def __rpow__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__rpow__(self, other)
+
     rpow.__doc__ = _flex_doc_SERIES.format(
         desc="Reverse Exponential power",
         op_name="**",
@@ -612,6 +660,10 @@ def rpow(self, other) -> "Series":
     def floordiv(self, other) -> "Series":
         return self // other
 
+    def __floordiv__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__floordiv__(self, other)
+
     floordiv.__doc__ = _flex_doc_SERIES.format(
         desc="Integer division",
         op_name="//",
@@ -623,6 +675,10 @@ def floordiv(self, other) -> "Series":
     def rfloordiv(self, other) -> "Series":
         return other // self
 
+    def __rfloordiv__(self, other):
+        other = ks.Index(other, name=self.name) if isinstance(other, (list, tuple)) else other
+        return IndexOpsMixin.__rfloordiv__(self, other)
+
     rfloordiv.__doc__ = _flex_doc_SERIES.format(
         desc="Reverse Integer division",
         op_name="//",