Fix Series.div when divide by zero

databricks/koalas/base.py

@@ -184,8 +184,14 @@ def __sub__(self, other):
             return _column_op(spark.Column.__sub__)(self, other)
 
     __mul__ = _column_op(spark.Column.__mul__)
-    __div__ = _numpy_column_op(spark.Column.__div__)
-    __truediv__ = _numpy_column_op(spark.Column.__truediv__)
+
+    def __truediv__(self, other):
+        def truediv(left, right):
+            return F.when(F.lit(right) == 0, F.lit("Infinity").__div__(F.lit(left))).otherwise(
+                F.lit(left).__truediv__(F.lit(right))
+            )
+
+        return self._with_new_scol(_numpy_column_op(truediv)(self, other)._scol)
 
     def __mod__(self, other):
         def mod(left, right):
@@ -202,18 +208,30 @@ def __radd__(self, other):
 
     __rsub__ = _column_op(spark.Column.__rsub__)
     __rmul__ = _column_op(spark.Column.__rmul__)
-    __rdiv__ = _numpy_column_op(spark.Column.__rdiv__)
-    __rtruediv__ = _numpy_column_op(spark.Column.__rtruediv__)
+
+    def __rtruediv__(self, other):
+        def rtruediv(left, right):
+            return F.when(F.lit(left) == 0, F.lit("Infinity").__div__(F.lit(right))).otherwise(
+                F.lit(right).__truediv__(F.lit(left))
+            )
+
+        return self._with_new_scol(_numpy_column_op(rtruediv)(self, other)._scol)
 
     def __floordiv__(self, other):
-        return self._with_new_scol(
-            F.floor(_numpy_column_op(spark.Column.__div__)(self, other)._scol)
-        )
+        def floordiv(left, right):
+            return F.when(F.lit(right) == 0, F.lit("Infinity").__div__(F.lit(left))).otherwise(
+                F.floor(F.lit(left).__div__(F.lit(right)))
+            )
+
+        return self._with_new_scol(_numpy_column_op(floordiv)(self, other)._scol)
 
     def __rfloordiv__(self, other):
-        return self._with_new_scol(
-            F.floor(_numpy_column_op(spark.Column.__rdiv__)(self, other)._scol)
-        )
+        def rfloordiv(left, right):
+            return F.when(F.lit(left) == 0, F.lit("Infinity").__div__(F.lit(right))).otherwise(
+                F.floor(F.lit(right).__div__(F.lit(left)))
+            )
+
+        return self._with_new_scol(_numpy_column_op(rfloordiv)(self, other)._scol)
 
     def __rmod__(self, other):
         def rmod(left, right):

databricks/koalas/tests/test_series.py

@@ -1446,3 +1446,13 @@ def test_squeeze(self):
         kser = ks.Series([90, 91, 85], index=midx)
         pser = kser.to_pandas()
         self.assert_eq(kser.squeeze(), pser.squeeze())
+
+    def test_div_zero(self):
+        pser = pd.Series([100, None, -300, None, 500, -700], name="Koalas")
+        kser = ks.from_pandas(pser)
+
+        self.assert_eq(repr(pser.div(0)), repr(kser.div(0)))
+        self.assert_eq(repr(pser.truediv(0)), repr(kser.truediv(0)))
+        self.assert_eq(repr(pser.floordiv(0)), repr(kser.floordiv(0)))
+        self.assert_eq(repr(pser / 0), repr(kser / 0))
+        self.assert_eq(repr(pser // 0), repr(kser // 0))