[SPARK-58582][PS][FOLLOWUP] Keep integral operands in integer space for NumPy fmod

python/pyspark/pandas/numpy_compat.py

@@ -152,18 +152,28 @@ def _copysign_func(c1: Column, c2: Column) -> Column:
 def _fmod_func(c1: Column, c2: Column) -> Column:
     c1_double = c1.cast("double")
     c2_double = c2.cast("double")
+    integral_types = ["tinyint", "smallint", "int", "bigint"]
 
-    return F.when(
-        F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float", "double"),
-        F.when(c1.isNull() | F.isnan(c1), c1_double)
-        .when(c2.isNull() | F.isnan(c2), c2_double)
-        .when(c2_double == 0, F.lit(float("nan")))
-        .otherwise(F.try_mod(c1_double, c2_double)),
-    ).otherwise(
+    return (
+        # A null or a nan propagates for every operand type.
         F.when(c1.isNull() | F.isnan(c1), c1_double)
         .when(c2.isNull() | F.isnan(c2), c2_double)
-        .when(c2_double == 0, F.lit(0.0))
-        .otherwise(F.try_mod(c1_double, c2_double))
+        # Integral operands take the remainder as longs: a double cannot hold 9007199254740993.
+        .when(
+            F.typeof(c1).isin(integral_types) & F.typeof(c2).isin(integral_types),
+            F.when(c2_double == 0, F.lit(0.0)).otherwise(
+                F.try_mod(c1.cast("long"), c2.cast("long")).cast("double")
+            ),
+        )
+        # Floating operands, where a zero divisor is nan rather than the 0 above.
+        .when(
+            F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float", "double"),
+            F.when(c2_double == 0, F.lit(float("nan"))).otherwise(F.try_mod(c1_double, c2_double)),
+        )
+        # A decimal falls through here, since typeof carries its precision, as in decimal(10,2).
+        # np.fmod raises on Decimal objects, so there is no NumPy behavior to match: a zero
+        # divisor returns 0 and any other divisor takes the remainder in double.
+        .otherwise(F.when(c2_double == 0, F.lit(0.0)).otherwise(F.try_mod(c1_double, c2_double)))
     )
 
 

python/pyspark/pandas/tests/test_numpy_compat.py

@@ -313,6 +313,27 @@ def test_np_fmod(self):
 
             self.assert_eq(np.fmod(psdf.x1, psdf.x2), np.fmod(pdf.x1, pdf.x2), almost=True)
 
+        # Integral operands above 2**53, where casting an operand to double would drop its low
+        # bits: fmod(9007199254740993, 2) is 1, not 0. Compared exactly, since almost=True would
+        # accept an off-by-one at these magnitudes.
+        pdf = pd.DataFrame(
+            {
+                "x1": [9007199254740993, 9007199254740995, -9007199254740993, 4611686018427387905],
+                "x2": [2, 4, 2, 1000000000],
+            }
+        )
+        psdf = ps.from_pandas(pdf)
+        self.assert_eq(np.fmod(psdf.x1, psdf.x2), np.fmod(pdf.x1, pdf.x2).astype("float64"))
+
+        # almost=True treats -0.0 and 0.0 as equal, so check the sign of zero explicitly:
+        # an integral remainder is never negative zero, unlike a double one.
+        pdf = pd.DataFrame({"x1": [-64, -2, 0, 64], "x2": [2, 2, 3, 2]})
+        psdf = ps.from_pandas(pdf)
+
+        result = np.fmod(psdf.x1, psdf.x2)
+        expected = np.fmod(pdf.x1, pdf.x2)
+        self.assert_eq(np.signbit(result.to_pandas()), np.signbit(expected))
+
     def test_np_modf(self):
         # np.modf(x) returns a tuple (fractional part, integral part).
         for pdf in (