github-actions[bot] commented on code in PR #69255:
URL: https://github.com/apache/doris/pull/69255#discussion_r4396320091
##########
be/src/udf/python/python_server.py:
##########
@@ -958,63 +937,130 @@ def load(self) -> AdaptivePythonUDF:
return AdaptivePythonUDF(self.python_udf_meta, func)
+class UDFModuleAliasLoader(importlib.abc.Loader):
+ """Makes "pkg" the same module object as
"__doris_pyudf_<function_id>.pkg"."""
+
+ def __init__(self, udf_module_name: str):
+ self.udf_module_name = udf_module_name
+ self.udf_module_spec = None
+
+ def create_module(self, spec) -> ModuleType:
+ module = importlib.import_module(self.udf_module_name)
+ self.udf_module_spec = module.__spec__
+ return module
+
+ def exec_module(self, module: ModuleType) -> None:
+ # module_from_spec() replaced the spec; keep it consistent with
__name__.
+ module.__spec__ = self.udf_module_spec
+
+
+class UDFModuleFinder(importlib.abc.MetaPathFinder):
+ """
+ Loads user modules under the root package "__doris_pyudf_<function_id>",
so a
+ user module named like a server module (e.g. types, json) does not resolve
to
+ the server module. The name depends only on the function id, which keeps
+ pickled UDAF state classes loadable in every worker.
+
+ Absolute imports in user code look in the UDF directory first, e.g.
+ "from json import helper" in the user package json loads
+ "__doris_pyudf_<function_id>.json.helper".
+ """
+
+ def __init__(self, function_id: int, location: str):
+ self.root_package = self.root_package_name(function_id)
+ self.root_package_prefix = self.root_package + "."
+ self.location = location
+ self.original_import = builtins.__import__
+ # Top-level name -> whether the UDF directory provides it.
+ self.udf_top_level_module_cache: Dict[str, bool] = {}
+ # importlib.import_module bypasses builtins.__import__, so user code
gets
+ # an importlib whose import_module adds the prefix.
+ self.udf_importlib = ModuleType("importlib")
+ self.udf_importlib.__dict__.update(vars(importlib))
+ self.udf_importlib.import_module = self.import_module_for_udf_code
+ self.udf_importlib.__getattr__ = lambda name: getattr(importlib, name)
+
+ @staticmethod
+ def root_package_name(function_id: int) -> str:
+ return f"__doris_pyudf_{function_id}"
+
+ def install(self) -> None:
+ sys.meta_path.insert(0, self)
+ builtins.__import__ = self.import_for_udf_code
+
+ def find_spec(self, fullname, path=None, target=None):
+ if fullname == self.root_package:
+ spec = importlib.machinery.ModuleSpec(fullname, None,
is_package=True)
+ spec.submodule_search_locations = [self.location]
+ return spec
+
+ if fullname.startswith(self.root_package_prefix):
+ return None
+
+ parent_name = fullname.rpartition(".")[0]
+ if parent_name:
+ # Only children of a user package, so "json.decoder" stays
standard.
+ parent = sys.modules.get(parent_name)
+ if getattr(parent, "__name__", None) != self.root_package_prefix +
parent_name:
+ return None
+ if importlib.machinery.PathFinder.find_spec(
+ self.root_package_prefix + fullname, path
+ ) is None:
+ return None
+ elif not self.is_udf_top_level_module(fullname):
+ return None
+
+ return importlib.machinery.ModuleSpec(
+ fullname, UDFModuleAliasLoader(self.root_package_prefix + fullname)
+ )
+
+ def is_udf_top_level_module(self, name: str) -> bool:
+ is_udf_module = self.udf_top_level_module_cache.get(name)
+ if is_udf_module is None:
+ is_udf_module = (
+ importlib.machinery.PathFinder.find_spec(name,
[self.location]) is not None
+ )
+ self.udf_top_level_module_cache[name] = is_udf_module
+ return is_udf_module
+
+ def add_root_package_prefix(self, name: str) -> str:
+ if self.is_udf_top_level_module(name.partition(".")[0]):
+ return self.root_package_prefix + name
+ return name
+
+ def import_for_udf_code(self, name, globals=None, locals=None,
fromlist=(), level=0):
+ """builtins.__import__ that adds the prefix to absolute imports in
user code."""
+ if (
+ level != 0
+ or globals is None
+ or not globals.get("__name__",
"").startswith(self.root_package_prefix)
+ ):
+ return self.original_import(name, globals, locals, fromlist, level)
+
+ udf_module_name = self.add_root_package_prefix(name)
Review Comment:
[P2] Preserve UDAF state compatibility across rolling BE upgrades
This namespace also becomes the module name stored in pickled aggregate
state. Take an existing ZIP where `entry.py` imports `State` from
`state_helper.py` and returns it from `aggregate_state`. Previously the helper
remained registered as `state_helper`, so old workers could exchange these
states. On an upgraded BE it is imported as `__doris_pyudf_<id>.state_helper`;
`UDAFStateManager.serialize()` pickles that name, and
`AggregatePythonUDAFData::write/read/merge()` forwards the bytes unchanged.
During the supported BE-first rolling upgrade, a merge fragment on an older BE
calls its unchanged `pickle.loads()`, which [imports the recorded module
name](https://github.com/python/cpython/blob/3.10/Lib/pickle.py#L1447-L1459),
cannot find this new namespace, and fails with `ModuleNotFoundError`. Keep
serialized user-class references readable by pre-upgrade workers, or prevent
incompatible state exchanges, and cover a mixed-version custom-state case.
##########
be/src/udf/python/python_server.py:
##########
@@ -958,63 +937,130 @@ def load(self) -> AdaptivePythonUDF:
return AdaptivePythonUDF(self.python_udf_meta, func)
+class UDFModuleAliasLoader(importlib.abc.Loader):
+ """Makes "pkg" the same module object as
"__doris_pyudf_<function_id>.pkg"."""
+
+ def __init__(self, udf_module_name: str):
+ self.udf_module_name = udf_module_name
+ self.udf_module_spec = None
+
+ def create_module(self, spec) -> ModuleType:
+ module = importlib.import_module(self.udf_module_name)
+ self.udf_module_spec = module.__spec__
+ return module
+
+ def exec_module(self, module: ModuleType) -> None:
+ # module_from_spec() replaced the spec; keep it consistent with
__name__.
+ module.__spec__ = self.udf_module_spec
+
+
+class UDFModuleFinder(importlib.abc.MetaPathFinder):
+ """
+ Loads user modules under the root package "__doris_pyudf_<function_id>",
so a
+ user module named like a server module (e.g. types, json) does not resolve
to
+ the server module. The name depends only on the function id, which keeps
+ pickled UDAF state classes loadable in every worker.
+
+ Absolute imports in user code look in the UDF directory first, e.g.
+ "from json import helper" in the user package json loads
+ "__doris_pyudf_<function_id>.json.helper".
+ """
+
+ def __init__(self, function_id: int, location: str):
+ self.root_package = self.root_package_name(function_id)
+ self.root_package_prefix = self.root_package + "."
+ self.location = location
+ self.original_import = builtins.__import__
+ # Top-level name -> whether the UDF directory provides it.
+ self.udf_top_level_module_cache: Dict[str, bool] = {}
+ # importlib.import_module bypasses builtins.__import__, so user code
gets
+ # an importlib whose import_module adds the prefix.
+ self.udf_importlib = ModuleType("importlib")
+ self.udf_importlib.__dict__.update(vars(importlib))
+ self.udf_importlib.import_module = self.import_module_for_udf_code
+ self.udf_importlib.__getattr__ = lambda name: getattr(importlib, name)
+
+ @staticmethod
+ def root_package_name(function_id: int) -> str:
+ return f"__doris_pyudf_{function_id}"
+
+ def install(self) -> None:
+ sys.meta_path.insert(0, self)
+ builtins.__import__ = self.import_for_udf_code
+
+ def find_spec(self, fullname, path=None, target=None):
+ if fullname == self.root_package:
+ spec = importlib.machinery.ModuleSpec(fullname, None,
is_package=True)
+ spec.submodule_search_locations = [self.location]
+ return spec
+
+ if fullname.startswith(self.root_package_prefix):
+ return None
+
+ parent_name = fullname.rpartition(".")[0]
+ if parent_name:
+ # Only children of a user package, so "json.decoder" stays
standard.
+ parent = sys.modules.get(parent_name)
+ if getattr(parent, "__name__", None) != self.root_package_prefix +
parent_name:
+ return None
+ if importlib.machinery.PathFinder.find_spec(
+ self.root_package_prefix + fullname, path
+ ) is None:
+ return None
+ elif not self.is_udf_top_level_module(fullname):
+ return None
+
+ return importlib.machinery.ModuleSpec(
+ fullname, UDFModuleAliasLoader(self.root_package_prefix + fullname)
Review Comment:
[P2] Preserve resource loading through module aliases
For an uploaded ZIP containing `mypkg/__init__.py`, `mypkg/udf.py` and
`mypkg/data.txt`, an initializer such as `DATA = pkgutil.get_data("mypkg",
"data.txt").decode()` now fails before `mypkg.udf.evaluate` can load. The
package initially exists only under `__doris_pyudf_<id>.mypkg`, so `pkgutil`'s
bare-name `find_spec` gets this alias spec. `UDFModuleAliasLoader` has no
`get_data`; [CPython returns `None`
immediately](https://github.com/python/cpython/blob/3.10/Lib/pkgutil.py#L569-L574),
and `.decode()` raises `AttributeError`. Restoring the real spec in
`exec_module` happens after that capability check. At the merge base, the
ordinary package was already in `sys.modules` during its initializer and its
`SourceFileLoader` read the resource. Preserve the underlying loader's resource
support and package metadata in aliases, and cover resource access during
package initialization.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
