diff --git a/pydeequ/scala_utils.py b/pydeequ/scala_utils.py index b6d3e83..030e43f 100644 --- a/pydeequ/scala_utils.py +++ b/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!") diff --git a/tests/test_checks.py b/tests/test_checks.py index 878257b..f462d4c 100644 --- a/tests/test_checks.py +++ b/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")])