From dee6ef046a8e6e87c210c4cbb2e0ee2425818736 Mon Sep 17 00:00:00 2001 From: Daksha1611 Date: Sun, 6 Sep 2026 22:56:44 +0530 Subject: [PATCH] fix: Fall back to the serialized body when UDF source misses globals Source-first rehydration execs only the UDF's own source into a namespace seeded with pandas/numpy. A UDF that reads a helper, constant, or aliased import from its defining module has no way to resolve those names, but the exec still succeeds, because a function body's free variables are only looked up when it is called. resolve_udf therefore treated rehydration as successful and never fell back to the dill body, even when that body carried the captured globals and would have run. The failure surfaced later as a NameError during retrieval, so an on-demand feature view that served correctly before started raising once it was loaded from the registry. Validate the rebuilt callable before accepting it: walk its LOAD_GLOBAL operands, and the operands of nested code objects, and return None when a name resolves in neither the exec namespace nor builtins. resolve_udf then falls back as it did before. Only LOAD_GLOBAL is inspected. co_names also holds attribute names, so checking it would flag working UDFs and push them onto the dill path for no reason. Self-contained UDFs are unaffected and still take the source path, keeping the Spark and cross-Python-version benefits intact. Signed-off-by: Daksha1611 --- .../feast/transformation/udf_rehydrate.py | 67 ++++++++++++++++++- .../unit/transformation/test_udf_rehydrate.py | 63 +++++++++++++++++ 2 files changed, 127 insertions(+), 3 deletions(-) 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