fix

Author: xinrong-meng

python/pyspark/pandas/data_type_ops/num_ops.py

@@ -271,14 +271,23 @@ def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
         _sanitize_list_like(right)
         if not is_valid_operand_for_numeric_arithmetic(right):
             raise TypeError("Floor division can not be applied to given types.")
+        spark_session = left._internal.spark_frame.sparkSession
 
         def floordiv(left: PySparkColumn, right: Any) -> PySparkColumn:
-            return F.when(F.lit(right is np.nan), np.nan).otherwise(
-                F.when(
-                    F.lit(right != 0) | F.lit(right).isNull(),
-                    F.floor(left.__div__(right)),
-                ).otherwise(F.lit(np.inf).__div__(left))
-            )
+            if is_ansi_mode_enabled(spark_session):
+                return F.when(F.lit(right is np.nan), np.nan).otherwise(
+                    F.when(
+                        F.lit(right != 0) | F.lit(right).isNull(),
+                        F.floor(left.__div__(right)),
+                    ).otherwise(F.try_divide(F.lit(np.inf), left))
+                )
+            else:
+                return F.when(F.lit(right is np.nan), np.nan).otherwise(
+                    F.when(
+                        F.lit(right != 0) | F.lit(right).isNull(),
+                        F.floor(left.__div__(right)),
+                    ).otherwise(F.lit(np.inf).__div__(left))
+                )
 
         right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)
         return numpy_column_op(floordiv)(left, right)
@@ -373,11 +382,14 @@ def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
 
         def floordiv(left: PySparkColumn, right: Any) -> PySparkColumn:
             if is_ansi_mode_enabled(spark_session):
-                return F.when(F.lit(right is np.nan), F.lit(np.nan)).otherwise(
-                    F.when(F.lit(right != 0), F.floor(left / F.lit(right))).otherwise(
-                        F.when(
-                            F.lit(left == np.inf) | F.lit(left == -np.inf), F.lit(left)
-                        ).otherwise(F.lit(np.inf).__div__(left))
+                return F.when(F.lit(right is np.nan), np.nan).otherwise(
+                    F.when(
+                        F.lit(right != 0) | F.lit(right).isNull(),
+                        F.floor(left.__div__(right)),
+                    ).otherwise(
+                        F.when(F.lit(left == np.inf) | F.lit(left == -np.inf), left).otherwise(
+                            F.try_divide(F.lit(np.inf), left)
+                        )
                     )
                 )
             else:

python/pyspark/pandas/tests/computation/test_binary_ops.py

@@ -208,9 +208,10 @@ def test_binary_operator_truediv(self):
         self.assertRaisesRegex(TypeError, ks_err_msg, lambda: 1 / psdf["a"])
 
     def test_binary_operator_floordiv(self):
-        pdf = pd.DataFrame({"a": ["x"], "b": [1]})
+        pdf = pd.DataFrame({"a": ["x"], "b": [1], "c": [1.0]})
         psdf = ps.from_pandas(pdf)
-        self.assert_eq(pdf["b"] / 0, psdf["b"] / 0)
+        self.assert_eq(pdf["b"] // 0, psdf["b"] // 0)
+        self.assert_eq(pdf["c"] // 0, psdf["c"] // 0)
 
         ks_err_msg = "Floor division can not be applied to strings"
         self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] // psdf["b"])