From 3af4d54e53e1772b696d76a9f6b75b9ba42824f3 Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Thu, 5 Mar 2026 20:21:12 -0800 Subject: [PATCH] Fix UDF logging tests under coverage --- python/test_coverage/sitecustomize.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/test_coverage/sitecustomize.py b/python/test_coverage/sitecustomize.py index b84e62a20c330..f1b8e494accb1 100644 --- a/python/test_coverage/sitecustomize.py +++ b/python/test_coverage/sitecustomize.py @@ -63,6 +63,11 @@ def wrapper(*args, **kwargs): return wrapper frame.f_globals["worker"] = save_when_exit(frame.f_globals["worker"]) + # Pretend that this module has the same name as the worker module. + # UDF logging checks where pyspark code firstly calls into user code, and if + # the module name is "sitecustomize", it will confuse UDF logging and make + # it believe this is user code, which will result in a wrong context. + globals()["__name__"] = frame.f_globals.get("__name__", globals()["__name__"]) os.register_at_fork(after_in_child=patch_worker)