fix: Fall back when UDF globals are missing

Signed-off-by: chloe <823083779@qq.com>
zqzq0422 · Sep 29, 2026 · Sep 29, 2026 · a3fd7bdd6069befa0c60045713ef4e101fcbe8b3
commit a3fd7bdd6069befa0c60045713ef4e101fcbe8b3
@@ -9,14 +9,39 @@
 
 from __future__ import annotations
 
+import builtins
+import dis
 import logging
+import types
 from typing import Callable, Optional
 
 import dill
 
 logger = logging.getLogger(__name__)
 
 
+def _has_unresolved_globals(code: types.CodeType, namespace: dict) -> bool:
+    """Return whether ``code`` reads a global absent from its namespace/builtins."""
+    builtin_namespace = namespace.get("__builtins__", builtins)
+    if not isinstance(builtin_namespace, dict):
+        builtin_namespace = vars(builtin_namespace)
+
+    for instruction in dis.get_instructions(code):
+        # Function bodies use LOAD_GLOBAL for module names. LOAD_NAME also
+        # appears in class bodies, where names may be provided by that body's
+        # local namespace and cannot be validated against the exec namespace.
+        if instruction.opname == "LOAD_GLOBAL":
+            name = instruction.argval
+            if name not in namespace and name not in builtin_namespace:
+                return True
+
+    return any(
+        _has_unresolved_globals(value, namespace)
+        for value in code.co_consts
+        if isinstance(value, types.CodeType)
+    )
+
+
 def _strip_leading_decorators(udf_string: str) -> str:
     """Remove leading ``@decorator`` lines before ``def`` / ``async def``.
 
@@ -141,7 +166,11 @@ def rehydrate_udf_from_source(
         return None
 
     if preferred_name and preferred_name in ns and callable(ns[preferred_name]):
-        return ns[preferred_name]
+        value = ns[preferred_name]
+        code = getattr(value, "__code__", None)
+        if isinstance(code, types.CodeType) and _has_unresolved_globals(code, ns):
+            return None
+        return value
 
     for value in ns.values():
         if not callable(value):
@@ -152,6 +181,9 @@ def rehydrate_udf_from_source(
         # Skip imported modules / classes we seeded
         if name in ("DataFrame",):
             continue
+        code = getattr(value, "__code__", None)
+        if isinstance(code, types.CodeType) and _has_unresolved_globals(code, ns):
+            continue
         return value
 
     return None

@@ -16,6 +16,13 @@ def _sample_udf(df: pd.DataFrame) -> pd.DataFrame:
     return out
 
 
+def _rehydrate_fallback_udf(value):
+    return value + _rehydrate_fallback_constant
+
+
+_rehydrate_fallback_constant = 40
+
+
 _SAMPLE_SRC = """
 def _sample_udf(df):
     import pandas as pd
@@ -50,6 +57,43 @@ def test_resolve_udf_falls_back_to_dill_when_source_empty():
     assert result["doubled"].iloc[0] == 8
 
 
+def test_resolve_udf_falls_back_to_dill_for_missing_module_global():
+    source = """
+def _rehydrate_fallback_udf(value):
+    return value + _rehydrate_fallback_constant
+"""
+    body = dill.dumps(_rehydrate_fallback_udf, recurse=True)
+
+    udf = resolve_udf(
+        udf_string=source,
+        body=body,
+        preferred_name="_rehydrate_fallback_udf",
+    )
+
+    assert udf(2) == 42
+
+
+def test_rehydrate_udf_allows_attribute_names():
+    source = """
+def copy_frame(frame):
+    return frame.copy()
+"""
+
+    udf = rehydrate_udf_from_source(source, preferred_name="copy_frame")
+
+    assert udf is not None
+    assert udf(pd.DataFrame({"x": [1]})).equals(pd.DataFrame({"x": [1]}))
+
+
+def test_rehydrate_udf_checks_globals_in_nested_code():
+    source = """
+def make_callback(value):
+    return lambda: missing_module_global + value
+"""
+
+    assert rehydrate_udf_from_source(source, preferred_name="make_callback") is None
+
+
 def test_rehydrate_strips_on_demand_decorator():
     src = """@on_demand_feature_view(
     sources=[feature_view_1],