diff --git a/sdk/python/feast/transformation/udf_rehydrate.py b/sdk/python/feast/transformation/udf_rehydrate.py index 0111bebe6d1..8f86ee08758 100644 --- a/sdk/python/feast/transformation/udf_rehydrate.py +++ b/sdk/python/feast/transformation/udf_rehydrate.py @@ -9,8 +9,11 @@ from __future__ import annotations +import builtins +import dis import logging -from typing import Callable, Optional +from types import CodeType +from typing import Callable, Optional, Set import dill @@ -119,6 +122,46 @@ def _exec_namespace() -> dict: return ns +def _unresolved_global_names(func: Callable, namespace: dict) -> Set[str]: + """Global names *func* loads that neither *namespace* nor builtins provide. + + ``exec``-ing a function definition only binds the function; the globals its + body reads are looked up when it is *called*. Source text alone therefore + cannot tell us whether the callable actually works — a UDF referencing a + helper or constant from its defining module execs cleanly and then raises + ``NameError`` mid-retrieval. + + Only ``LOAD_GLOBAL`` operands count. ``co_names`` also holds attribute names + (``df.columns``), which would flag working UDFs and push them onto the dill + path for no reason. Nested code objects (inner functions, comprehensions) + are walked too. + """ + missing: Set[str] = set() + code = getattr(func, "__code__", None) + if code is None: + return missing + + seen: Set[int] = set() + pending = [code] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + for instruction in dis.get_instructions(current): + if instruction.opname != "LOAD_GLOBAL": + continue + name = instruction.argval + if not isinstance(name, str): + continue + if name not in namespace and not hasattr(builtins, name): + missing.add(name) + for const in current.co_consts: + if isinstance(const, CodeType): + pending.append(const) + return missing + + def rehydrate_udf_from_source( udf_string: str, *, @@ -141,7 +184,7 @@ def rehydrate_udf_from_source( return None if preferred_name and preferred_name in ns and callable(ns[preferred_name]): - return ns[preferred_name] + return _accept_if_self_contained(ns[preferred_name], ns) for value in ns.values(): if not callable(value): @@ -152,11 +195,29 @@ def rehydrate_udf_from_source( # Skip imported modules / classes we seeded if name in ("DataFrame",): continue - return value + return _accept_if_self_contained(value, ns) return None +def _accept_if_self_contained(func: Callable, ns: dict) -> Optional[Callable]: + """Return *func* only if every global it reads is available, else ``None``. + + Returning ``None`` lets :func:`resolve_udf` fall back to the dill body, which + carries the UDF's captured globals and can still run. + """ + missing = _unresolved_global_names(func, ns) + if missing: + logger.debug( + "udf source rehydrate skipped for %s: unresolved global names %s; " + "falling back to the serialized body", + getattr(func, "__name__", func), + sorted(missing), + ) + return None + return func + + def resolve_udf( *, udf_string: str = "", diff --git a/sdk/python/tests/unit/transformation/test_udf_rehydrate.py b/sdk/python/tests/unit/transformation/test_udf_rehydrate.py index def9ba4c401..7237d2a34a3 100644 --- a/sdk/python/tests/unit/transformation/test_udf_rehydrate.py +++ b/sdk/python/tests/unit/transformation/test_udf_rehydrate.py @@ -92,3 +92,66 @@ def test_strip_leading_decorators_adversarial_at_spam_is_linear(): # rehydrate should fail closed so dill fallback can run. assert rehydrate_udf_from_source(spam) is None assert isinstance(stripped, str) + + +_SCALE = 10 + + +def _scale(value): + return value * _SCALE + + +def _udf_using_module_globals(df: pd.DataFrame) -> pd.DataFrame: + out = pd.DataFrame() + out["scaled"] = _scale(df["x"]) + return out + + +_MODULE_GLOBALS_SRC = """def _udf_using_module_globals(df): + out = __import__("pandas").DataFrame() + out["scaled"] = _scale(df["x"]) + return out +""" + + +def test_rehydrate_rejects_udf_with_unresolved_globals(): + """Source alone cannot supply names the UDF reads from its defining module.""" + assert ( + rehydrate_udf_from_source( + _MODULE_GLOBALS_SRC, preferred_name="_udf_using_module_globals" + ) + is None + ) + + +def test_resolve_udf_falls_back_to_dill_when_source_misses_globals(): + """The serialized body carries the captured globals, so it must win here.""" + udf = resolve_udf( + udf_string=_MODULE_GLOBALS_SRC, + body=dill.dumps(_udf_using_module_globals, recurse=True), + preferred_name="_udf_using_module_globals", + ) + result = udf(pd.DataFrame({"x": [1, 2]})) + assert list(result["scaled"]) == [10, 20] + + +def test_attribute_access_is_not_mistaken_for_a_missing_global(): + """Attribute names live in co_names too; flagging them would strand good UDFs.""" + src = """def _attr_udf(df): + out = __import__("pandas").DataFrame() + out["cols"] = len(df.columns) + out["upper"] = df["s"].str.upper() + return out +""" + fn = rehydrate_udf_from_source(src, preferred_name="_attr_udf") + assert fn is not None + result = fn(pd.DataFrame({"s": ["a"]})) + assert result["upper"].iloc[0] == "A" + + +def test_nested_scopes_are_inspected_for_missing_globals(): + """A comprehension or inner function has its own code object.""" + src = """def _nested_udf(df): + return [_scale(v) for v in df["x"]] +""" + assert rehydrate_udf_from_source(src, preferred_name="_nested_udf") is None