fix: bind py4j callback server to a dynamic port to avoid 25334 collision (#86, #19)

pydeequ/scala_utils.py

@@ -1,7 +1,21 @@
 # -*- coding: utf-8 -*-
 """A collection of utility functions and classes for manipulating with scala objects anc classes through py4j
 """
-from py4j.java_gateway import JavaObject
+from py4j.java_gateway import DEFAULT_PYTHON_PROXY_PORT, JavaObject
+
+
+def _ensure_dynamic_callback_port(gateway):
+    """Switch the gateway's callback server off the hardcoded default port (25334).
+
+    If it is still on ``DEFAULT_PYTHON_PROXY_PORT``, mutate the existing
+    ``callback_server_parameters`` to an OS-assigned free port (``0``) so
+    concurrent applications don't collide on 25334. Mutating the existing
+    parameters (rather than passing a fresh object) keeps PySpark's callback
+    wiring intact so shutdown stays clean.
+    """
+    params = getattr(gateway, "callback_server_parameters", None)
+    if params is not None and getattr(params, "port", 0) == DEFAULT_PYTHON_PROXY_PORT:
+        params.port = 0
 
 
 class PythonCallback:
@@ -12,16 +26,14 @@ def __init__(self, gateway):
         # P4j will return false if the callback server is already started
         # https://github.com/bartdag/py4j/blob/master/py4j-python/src/py4j/java_gateway.py
         callback_server = self.gateway.get_callback_server()
-        # TODO clean
         if callback_server is None:
+            _ensure_dynamic_callback_port(self.gateway)
             self.gateway.start_callback_server()
             print("Python Callback server started!")  # TODO Logging
         elif callback_server.is_shutdown:
             callback_server.close()
+            _ensure_dynamic_callback_port(self.gateway)
             self.gateway.restart_callback_server()
-            # Have you tried turning it off and on again?
-            # TODO why do we need to restart this every time?
-            # TODO Will this break during chained function calls?
             print("PythonCallback server restarted!")
 
 

tests/test_checks.py

@@ -452,6 +452,28 @@ def test_isUnique(self):
         self.assertEqual(self.isUnique("b", "All rows are unique"), [Row(constraint_status="Success")])
         self.assertEqual(self.isUnique("email", "All rows are unique"), [Row(constraint_status="Success")])
 
+    def test_lambda_check_uses_dynamic_callback_port(self):
+        """A lambda-based Check should bind the callback server to a dynamic port, not 25334."""
+        gateway = self.spark.sparkContext._gateway
+
+        result = self.hasSize(lambda x: x == 3.0)
+        self.assertEqual(result, [Row(constraint_status="Success")])
+
+        callback_server = gateway.get_callback_server()
+        self.assertIsNotNone(callback_server, "a lambda Check should have started the callback server")
+        listening_port = callback_server.get_listening_port()
+        self.assertNotEqual(
+            listening_port,
+            25334,
+            "Callback server must not use the hardcoded default port 25334; "
+            f"got {listening_port}",
+        )
+        self.assertGreater(listening_port, 0, "callback server should be bound to a real port")
+
+        # second lambda Check: the JVM callback client must reach the dynamic port
+        result2 = self.hasSize(lambda x: x >= 2.0 and x < 5.0)
+        self.assertEqual(result2, [Row(constraint_status="Success")])
+
     def test_fail_isUnique(self):
         self.assertEqual(self.isUnique("d"), [Row(constraint_status="Failure")])
         self.assertEqual(self.isUnique("f", "All rows are unique"), [Row(constraint_status="Failure")])