diff --git a/be/src/agent/task_worker_pool.cpp b/be/src/agent/task_worker_pool.cpp index ecf25c4f291095..790976601c8d6c 100644 --- a/be/src/agent/task_worker_pool.cpp +++ b/be/src/agent/task_worker_pool.cpp @@ -2599,8 +2599,13 @@ void clean_udf_cache_callback(const TAgentTaskRequest& req) { fmt::format("failed to clean Java UDF cache, function_signature={}, function_id={}", clean_req.function_signature, clean_req.function_id)); } + if (drop_by_function_id) { - UserFunctionCache::instance()->drop_function_cache(clean_req.function_id); + auto status = UserFunctionCache::instance()->drop_function_cache(clean_req.function_id); + if (!status.ok()) { + LOG(WARNING) << "failed to drop function cache for function_id=" + << clean_req.function_id << ": " << status.to_string(); + } PythonServerManager::instance().clear_udaf_state_cache(clean_req.function_id); } diff --git a/be/src/runtime/user_function_cache.cpp b/be/src/runtime/user_function_cache.cpp index 7199cc7e1ff407..e8b0738a610702 100644 --- a/be/src/runtime/user_function_cache.cpp +++ b/be/src/runtime/user_function_cache.cpp @@ -567,29 +567,36 @@ Status UserFunctionCache::_check_and_return_default_java_udf_url(const std::stri return Status::OK(); } -void UserFunctionCache::drop_function_cache(int64_t fid) { +Status UserFunctionCache::drop_function_cache(int64_t fid) { std::shared_ptr entry = nullptr; { std::lock_guard l(_cache_lock); auto it = _entry_map.find(fid); if (it == _entry_map.end()) { - return; + return Status::OK(); } entry = it->second; - _entry_map.erase(it); } + // lib_file changes from the downloaded zip path to the extracted directory + // while an entry is loaded. Wait for that transition before clearing Python. + std::unique_lock load_lock(entry->load_lock); + // For Python UDF, clear module cache in Python server before deleting files if (entry->type == LibType::PY_ZIP && !entry->lib_file.empty()) { - auto status = PythonServerManager::instance().clear_module_cache(entry->lib_file); - if (!status.ok()) [[unlikely]] { - LOG(WARNING) << "drop_function_cache: failed to clear Python module cache for " - << entry->lib_file << ": " << status.to_string(); - } + RETURN_IF_ERROR(PythonServerManager::instance().clear_module_cache(entry->lib_file)); } + { + std::lock_guard l(_cache_lock); + auto it = _entry_map.find(fid); + if (it != _entry_map.end() && it->second == entry) { + _entry_map.erase(it); + } + } // Mark for deletion, destructor will delete the files entry->should_delete_library.store(true); + return Status::OK(); } } // namespace doris diff --git a/be/src/runtime/user_function_cache.h b/be/src/runtime/user_function_cache.h index e606899f969d4d..3d414a807a3de3 100644 --- a/be/src/runtime/user_function_cache.h +++ b/be/src/runtime/user_function_cache.h @@ -63,7 +63,7 @@ class UserFunctionCache { std::string* libpath); // Drop the cached function library by function id. - void drop_function_cache(int64_t fid); + Status drop_function_cache(int64_t fid); #ifndef BE_TEST private: diff --git a/be/src/udf/python/python_server.cpp b/be/src/udf/python/python_server.cpp index cf78591a2ef537..2bc8ef0b9dbc90 100644 --- a/be/src/udf/python/python_server.cpp +++ b/be/src/udf/python/python_server.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -707,6 +708,7 @@ Status PythonServerManager::_broadcast_action_to_processes(const std::string& ac int success_count = 0; int fail_count = 0; bool has_active_process = false; + std::string failure_details; for (auto& [version, versioned_pool] : _snapshot_process_pools()) { std::lock_guard lock(versioned_pool->mutex); @@ -741,12 +743,39 @@ Status PythonServerManager::_broadcast_action_to_processes(const std::string& ac } auto result = (*result_stream)->Next(); - if (result.ok() && *result) { - success_count++; - } else { + if (!result.ok() || !*result || !(*result)->body) { + fail_count++; + continue; + } + + std::string result_body = (*result)->body->ToString(); + rapidjson::Document result_json; + result_json.Parse(result_body.data(), result_body.size()); + bool action_succeeded = false; + if (!result_json.HasParseError() && result_json.IsObject()) { + auto success = result_json.FindMember("success"); + action_succeeded = success != result_json.MemberEnd() && + success->value.IsBool() && success->value.GetBool(); + } + if (!action_succeeded) { fail_count++; + std::string error = "invalid action result"; + if (!result_json.HasParseError() && result_json.IsObject()) { + auto error_member = result_json.FindMember("error"); + if (error_member != result_json.MemberEnd() && + error_member->value.IsString()) { + error.assign(error_member->value.GetString(), + error_member->value.GetStringLength()); + } + } + if (!failure_details.empty()) { + failure_details.append("; "); + } + failure_details.append(fmt::format("{}: {}", process->get_uri(), error)); + continue; } + success_count++; } catch (...) { fail_count++; } @@ -761,8 +790,9 @@ Status PythonServerManager::_broadcast_action_to_processes(const std::string& ac << ", failed=" << fail_count; if (fail_count > 0) { - return Status::InternalError("{} failed for {}, success={}, failed={}", action_type, - log_name, success_count, fail_count); + return Status::InternalError("{} failed for {}, success={}, failed={}, errors=[{}]", + action_type, log_name, success_count, fail_count, + failure_details); } return Status::OK(); diff --git a/be/src/udf/python/python_server.py b/be/src/udf/python/python_server.py index c0683149ea2d9b..22c8b6097942b9 100644 --- a/be/src/udf/python/python_server.py +++ b/be/src/udf/python/python_server.py @@ -17,6 +17,8 @@ import argparse import base64 +import builtins +import contextvars import gc import importlib import inspect @@ -29,6 +31,7 @@ import time import threading import pickle +import io from abc import ABC, abstractmethod from contextlib import contextmanager from typing import Any, Callable, Optional, Tuple, get_origin, Dict @@ -42,6 +45,94 @@ from pyarrow import flight +ModuleContext = Tuple[str, Dict[str, Any]] + +# UDF imports may temporarily replace sys.modules entries. Keep an immutable +# view of modules loaded by the Python server itself for safe lock-free reuse. +_SERVER_MODULES = dict(sys.modules) + + +class _ModuleImportOperation: + """Identifies one root import and the child imports it creates.""" + + def __init__(self, kind: str, module_context: Optional[ModuleContext] = None): + self.kind = kind + self.module_context = module_context + self.active = True + self.readers = 0 + + +_current_module_context: contextvars.ContextVar[Optional[ModuleContext]] = ( + contextvars.ContextVar("current_module_context", default=None) +) +_current_module_import_operation: contextvars.ContextVar[ + Optional[_ModuleImportOperation] +] = contextvars.ContextVar("current_module_import_operation", default=None) +_module_import_in_progress: contextvars.ContextVar[bool] = contextvars.ContextVar( + "module_import_in_progress", default=False +) + + +_thread_start = threading.Thread.start + + +def _start_thread_with_module_import_context(thread, *args, **kwargs): + """Propagate only an active Doris import operation to a child thread.""" + operation = _current_module_import_operation.get() + module_context = _current_module_context.get() + thread_target = getattr(thread, "_target", None) + # A direct import target has no UDF caller frame from which the import + # wrapper could recover the context. Its thread ends with that one import, + # unlike a long-lived worker thread which must not retain its creator's UDF. + direct_import_context = ( + module_context + if module_context is not None + and ( + thread_target is builtins.__import__ + or thread_target is importlib.import_module + ) + else None + ) + should_clear_parent_context = ( + module_context is not None + or _module_import_in_progress.get() + ) + if operation is not None or should_clear_parent_context: + thread_run = thread.run + had_instance_run = "run" in thread.__dict__ + + def restore_thread_run(): + if had_instance_run: + thread.run = thread_run + else: + thread.__dict__.pop("run", None) + + def run_with_module_import_context(): + operation_token = _current_module_import_operation.set(operation) + module_context_token = _current_module_context.set( + direct_import_context + ) + import_token = _module_import_in_progress.set(False) + try: + return thread_run() + finally: + restore_thread_run() + _module_import_in_progress.reset(import_token) + _current_module_context.reset(module_context_token) + _current_module_import_operation.reset(operation_token) + + thread.run = run_with_module_import_context + try: + return _thread_start(thread, *args, **kwargs) + except BaseException: + restore_thread_run() + raise + return _thread_start(thread, *args, **kwargs) + + +threading.Thread.start = _start_thread_with_module_import_context + + class ServerState: """Global server state container.""" @@ -877,54 +968,249 @@ class ModuleUDFLoader(UDFLoader): # with one of these names would overwrite the entry in sys.modules and # could break the server itself. _FORBIDDEN_MODULE_NAMES: frozenset = frozenset({ - "argparse", "base64", "gc", "importlib", "inspect", "ipaddress", + "argparse", "base64", "builtins", "contextvars", "gc", "importlib", + "inspect", "ipaddress", "json", "sys", "os", "traceback", "logging", "time", "threading", - "pickle", "abc", "contextlib", "typing", "datetime", "enum", - "pathlib", "pandas", "pd", "pyarrow", "pa", "flight", + "pickle", "io", "abc", "contextlib", "typing", "datetime", "enum", + "pathlib", "pandas", "pyarrow", "logging.handlers", }) - # Class-level lock dictionary for thread-safe module imports - # Using RLock allows the same thread to acquire the lock multiple times + # sys.path and sys.modules are process-global. Ordinary imports may share + # the stable environment, while UDF environment changes require exclusivity. + _module_import_condition = threading.Condition() + _active_module_import_readers = 0 + _module_import_writer_active = False + _waiting_module_import_writers = 0 + # Same-context child threads may borrow the currently restored environment. + # Cleanup waits until every borrower (including nested borrowers) exits. + _active_module_context: Optional[ModuleContext] = None + _active_module_context_borrowers = 0 + _active_module_import_operation: Optional[_ModuleImportOperation] = None + + # {location: top_module}; location already contains a unique function_id. + _module_cache: Dict[str, Any] = {} - # Key for _import_locks: module_name only (not location) - # sys.modules is a global dict keyed by module name. - # we need to ensure that imports with the same module name - # do not interfere with each other across different threads, - # even if they come from different file paths. - _import_locks: Dict[str, threading.Lock] = {} - _import_locks_lock = threading.Lock() + @staticmethod + def _is_path_from_location(path: Any, location: str) -> bool: + """Return whether a path belongs to a UDF location.""" + try: + normalized_location = os.path.realpath(location) + normalized_path = os.path.realpath(os.fspath(path)) + return ( + os.path.commonpath((normalized_location, normalized_path)) + == normalized_location + ) + except (TypeError, ValueError): + return False - # Key for _module_cache: location only - # since location already contains a unique function_id - _module_cache: Dict[str, Any] = {} - _module_cache_lock = threading.Lock() + @classmethod + def _is_module_from_location( + cls, module: Any, location: str + ) -> bool: + """Return whether a module was loaded from a UDF location.""" + module_paths = [] + module_file = getattr(module, "__file__", None) + if module_file: + module_paths.append(module_file) + module_path = getattr(module, "__path__", None) + if module_path: + module_paths.extend(module_path) + + return any(cls._is_path_from_location(path, location) for path in module_paths) @classmethod - def _get_import_lock(cls, module_name: str) -> threading.Lock: - """ - Get or create an import lock for the given module namespace. + def _collect_modules_from_location(cls, location: str) -> Dict[str, Any]: + """Collect loaded modules whose files belong to a UDF location.""" + normalized_location = os.path.realpath(location) + return { + name: module + for name, module in list(sys.modules.items()) + if cls._is_module_from_location(module, normalized_location) + } - Uses double-checked locking pattern for optimal performance: - - Fast path: return existing lock without acquiring global lock - - Slow path: create new lock under global lock protection - """ - # Lock by top-level package to avoid concurrent imports mutating shared - # parent entries in sys.modules. If we lock by full module name instead, - # pkg.mod.func1 and pkg.mod.func2 can import in parallel and race while - # initializing pkg/pkg.mod, causing flaky import failures (for example KeyError). - cache_key = module_name.split(".", 1)[0] + @staticmethod + def _bind_module_context(module_context: ModuleContext) -> None: + """Bind the owning UDF context to every module loaded from its location.""" + _, udf_modules = module_context + module_type = type(sys) + for module in udf_modules.values(): + if isinstance(module, module_type): + module_type.__getattribute__(module, "__dict__")[ + "_doris_module_context" + ] = module_context - # Fast path: check without lock (read-only, safe for most cases) - if cache_key in cls._import_locks: - return cls._import_locks[cache_key] + @classmethod + def _find_active_context_for_module_globals( + cls, module_globals: Dict[str, Any] + ) -> Optional[ModuleContext]: + """Find the active context for a module that is still initializing.""" + module_file = module_globals.get("__file__") + if not module_file: + return None + with cls._module_import_condition: + module_context = cls._active_module_context + if module_context is None: + return None + location, _ = module_context + if cls._is_path_from_location(module_file, location): + return module_context + return None - # Slow path: create lock under protection - with cls._import_locks_lock: - # Double-check: another thread might have created it while we waited - if cache_key not in cls._import_locks: - cls._import_locks[cache_key] = threading.Lock() - return cls._import_locks[cache_key] + @classmethod + @contextmanager + def _shared_module_import(cls): + """Run an ordinary import while the process import view is stable.""" + operation = _current_module_import_operation.get() + with cls._module_import_condition: + if not ( + operation is not None + and operation.kind == "reader" + and operation.active + ): + operation = None + while ( + cls._module_import_writer_active + or cls._waiting_module_import_writers + ): + cls._module_import_condition.wait() + operation = _ModuleImportOperation("reader") + operation.readers += 1 + cls._active_module_import_readers += 1 + + operation_token = _current_module_import_operation.set(operation) + import_token = _module_import_in_progress.set(True) + try: + yield + finally: + _module_import_in_progress.reset(import_token) + _current_module_import_operation.reset(operation_token) + with cls._module_import_condition: + operation.readers -= 1 + if operation.readers == 0: + operation.active = False + cls._active_module_import_readers -= 1 + if cls._active_module_import_readers == 0: + cls._module_import_condition.notify_all() + + @classmethod + @contextmanager + def _exclusive_module_import( + cls, module_context: Optional[ModuleContext] = None + ): + """Run an import that temporarily changes the process import view.""" + borrowed_active_context = False + operation = _current_module_import_operation.get() + with cls._module_import_condition: + cls._waiting_module_import_writers += 1 + try: + while True: + if ( + module_context is not None + and cls._active_module_context is module_context + and cls._active_module_import_operation is operation + and operation is not None + and operation.active + ): + cls._active_module_context_borrowers += 1 + borrowed_active_context = True + break + if ( + not cls._module_import_writer_active + and cls._active_module_import_readers == 0 + ): + cls._module_import_writer_active = True + operation = _ModuleImportOperation( + "writer", module_context + ) + cls._active_module_import_operation = operation + break + cls._module_import_condition.wait() + finally: + cls._waiting_module_import_writers -= 1 + + operation_token = _current_module_import_operation.set(operation) + import_token = _module_import_in_progress.set(True) + try: + yield borrowed_active_context + finally: + _module_import_in_progress.reset(import_token) + _current_module_import_operation.reset(operation_token) + with cls._module_import_condition: + if borrowed_active_context: + cls._active_module_context_borrowers -= 1 + if cls._active_module_context_borrowers == 0: + cls._module_import_condition.notify_all() + else: + operation.active = False + operation.module_context = None + cls._active_module_import_operation = None + cls._module_import_writer_active = False + cls._module_import_condition.notify_all() + + @classmethod + @contextmanager + def temporarily_restore_udf_modules( + cls, module_context: Optional[ModuleContext] + ): + """Temporarily register one UDF's modules for an import operation.""" + if module_context is None: + yield + return + + with cls._exclusive_module_import(module_context) as borrowed_active_context: + if borrowed_active_context: + yield + return + + location, udf_modules = module_context + missing = object() + previous = { + name: sys.modules.get(name, missing) for name in udf_modules + } + try: + # UDAF state conversion and runtime imports resolve user classes by + # name, so this UDF's isolated modules must briefly be visible here. + sys.modules.update(udf_modules) + with temporary_sys_path(location): + with cls._module_import_condition: + cls._active_module_context = module_context + cls._module_import_condition.notify_all() + try: + yield + finally: + with cls._module_import_condition: + # Keep the context visible while a borrower may create + # another importing thread, then close it atomically. + while cls._active_module_context_borrowers: + cls._module_import_condition.wait() + cls._active_module_context = None + finally: + try: + udf_modules.update( + cls._collect_modules_from_location(location) + ) + cls._bind_module_context(module_context) + finally: + for name in udf_modules: + sys.modules.pop(name, None) + for name, module in previous.items(): + if module is not missing: + sys.modules[name] = module + + @staticmethod + @contextmanager + def use_module_context(module_context: Optional[ModuleContext]): + """Bind a module context without locking normal UDF execution.""" + if module_context is None: + yield + return + + token = _current_module_context.set(module_context) + try: + yield + finally: + _current_module_context.reset(token) def load(self) -> AdaptivePythonUDF: """ @@ -1007,6 +1293,25 @@ def _clear_modules_from_sys(full_module_name: str) -> None: ancestor = ".".join(parts[: i + 1]) sys.modules.pop(ancestor, None) + @classmethod + def _remove_modules_from_location(cls, location: str) -> None: + """Remove modules imported from a UDF location from sys.modules.""" + for module_name in cls._collect_modules_from_location(location): + sys.modules.pop(module_name, None) + + def _find_cached_top_module(self, location: str) -> Optional[Any]: + """Return a valid cached top module and restore this loader's context.""" + cached_module = self._module_cache.get(location) + if cached_module is None or not ( + hasattr(cached_module, "__file__") + or hasattr(cached_module, "__path__") + ): + return None + self.module_context = getattr( + cached_module, "_doris_module_context", None + ) + return cached_module + def _get_or_import_module(self, location: str, full_module_name: str) -> Any: """ Get module from cache or import it (thread-safe). @@ -1016,7 +1321,7 @@ def _get_or_import_module(self, location: str, full_module_name: str) -> Any: """ # Reject module names that would shadow server-critical modules top_level_name = full_module_name.split(".")[0] - if top_level_name in ModuleUDFLoader._FORBIDDEN_MODULE_NAMES: + if top_level_name in self._FORBIDDEN_MODULE_NAMES: raise ImportError( f"Module name '{full_module_name}' is not allowed for UDFs " f"because it conflicts with a module used by the server. " @@ -1026,37 +1331,30 @@ def _get_or_import_module(self, location: str, full_module_name: str) -> Any: cache_key = location - # Use a per-module lock to prevent race conditions during import - import_lock = ModuleUDFLoader._get_import_lock(full_module_name) + # Repeated SQL executions only read the cache and may proceed together. + with self._shared_module_import(): + cached_module = self._find_cached_top_module(cache_key) + if cached_module is not None: + return cached_module - with import_lock: - # Fast path: check cache first - if cache_key in ModuleUDFLoader._module_cache: - cached_module = ModuleUDFLoader._module_cache[cache_key] - if cached_module is not None and ( - hasattr(cached_module, "__file__") - or hasattr(cached_module, "__path__") - ): - return cached_module - else: - del ModuleUDFLoader._module_cache[cache_key] + module_context: ModuleContext = (location, {}) + with self.temporarily_restore_udf_modules(module_context): + # Another loader may have populated the cache while this one waited. + cached_module = self._find_cached_top_module(cache_key) + if cached_module is not None: + return cached_module + self._module_cache.pop(cache_key, None) self._clear_modules_from_sys(full_module_name) - with temporary_sys_path(location): - try: - module = importlib.import_module(full_module_name) - ModuleUDFLoader._module_cache[cache_key] = module - # Evict from sys.modules so future imports from a - # different location are not poisoned by this one. - self._clear_modules_from_sys(full_module_name) - return module - except Exception: - # Clean up any partially-imported modules - self._clear_modules_from_sys(full_module_name) - if cache_key in ModuleUDFLoader._module_cache: - del ModuleUDFLoader._module_cache[cache_key] - raise + try: + module = importlib.import_module(full_module_name) + self.module_context = module_context + self._module_cache[cache_key] = module + return module + except Exception: + self._module_cache.pop(cache_key, None) + raise def _extract_function( self, module: Any, func_name: str, module_name: str @@ -1207,6 +1505,143 @@ def load_udf_from_module( ) +def _find_cached_module_for_builtin_import( + module_context: ModuleContext, + name: str, + globals: Optional[Dict[str, Any]] = None, + locals: Optional[Dict[str, Any]] = None, + fromlist: Any = (), + level: int = 0, +) -> Optional[Any]: + """Return a fully loaded module that can satisfy __import__ directly.""" + del locals + _, udf_modules = module_context + original_name = name + try: + if level < 0: + return None + if level: + package = globals.get("__package__") if globals else None + if not package: + return None + name = importlib.util.resolve_name("." * level + name, package) + except (AttributeError, ImportError, TypeError, ValueError): + return None + + available_modules = udf_modules + if ( + name.partition(".")[0] in ModuleUDFLoader._FORBIDDEN_MODULE_NAMES + and name in _SERVER_MODULES + ): + # This fixed mapping is not affected when another UDF temporarily + # replaces a same-named sys.modules entry. + available_modules = _SERVER_MODULES + + imported_module = available_modules.get(name) + if imported_module is None: + return None + + if fromlist: + module_attributes = getattr(imported_module, "__dict__", {}) + if not isinstance(fromlist, (tuple, list)) or any( + not isinstance(item, str) + or item == "*" + or item not in module_attributes + for item in fromlist + ): + return None + return imported_module + + if not level: + return available_modules.get(name.partition(".")[0]) + if not original_name: + return imported_module + + relative_root_length = len(name) - len(original_name) + len( + original_name.partition(".")[0] + ) + return available_modules.get(name[:relative_root_length]) + + +def _find_cached_module_for_importlib( + module_context: ModuleContext, name: str, package: Optional[str] = None +) -> Optional[Any]: + """Return a cached module for importlib.import_module, including relatives.""" + _, udf_modules = module_context + try: + absolute_name = ( + importlib.util.resolve_name(name, package) + if name.startswith(".") + else name + ) + except (AttributeError, ImportError, TypeError, ValueError): + return None + cached_module = udf_modules.get(absolute_name) + if cached_module is not None: + return cached_module + if ( + absolute_name.partition(".")[0] + in ModuleUDFLoader._FORBIDDEN_MODULE_NAMES + ): + return _SERVER_MODULES.get(absolute_name) + return None + + +def _wrap_import_with_module_context( + import_func: Callable, find_cached_module: Callable +) -> Callable: + """Use isolated cached modules first and take the global lock only on miss.""" + def wrapped_import(*args, **kwargs): + if _module_import_in_progress.get(): + return import_func(*args, **kwargs) + + # Prefer the calling UDF module over an inherited execution context. + # This lets long-lived worker threads execute callbacks from another UDF. + try: + caller_globals = sys._getframe(1).f_globals + module_context = caller_globals.get("_doris_module_context") + except (AttributeError, ValueError): + caller_globals = None + module_context = None + if module_context is None: + module_context = _current_module_context.get() + if module_context is None: + operation = _current_module_import_operation.get() + with ModuleUDFLoader._module_import_condition: + if ( + operation is not None + and operation.active + and ModuleUDFLoader._active_module_import_operation is operation + ): + module_context = operation.module_context + if module_context is None and caller_globals is not None: + module_context = ModuleUDFLoader._find_active_context_for_module_globals( + caller_globals + ) + if module_context is None: + with ModuleUDFLoader._shared_module_import(): + return import_func(*args, **kwargs) + + cached_module = find_cached_module(module_context, *args, **kwargs) + if cached_module is not None: + return cached_module + + with ModuleUDFLoader.temporarily_restore_udf_modules(module_context): + return import_func(*args, **kwargs) + + return wrapped_import + + +# Execution only binds a ContextVar; the process-wide lock is taken when user +# code actually imports a module. +builtins.__import__ = _wrap_import_with_module_context( + builtins.__import__, _find_cached_module_for_builtin_import +) +importlib.import_module = _wrap_import_with_module_context( + importlib.import_module, _find_cached_module_for_importlib +) + + class UDFLoaderFactory: """Factory to select the appropriate loader based on UDF location.""" @@ -1377,6 +1812,9 @@ def load_from_module( if not inspect.isclass(udaf_class): raise ValueError(f"'{symbol}' is not a class (type: {type(udaf_class)})") + # Instances are created after module loading has cleaned sys.modules, so + # keep the owning UDAF context on the class for later lifecycle calls. + udaf_class._doris_module_context = loader.module_context UDAFClassLoader.validate_udaf_class(udaf_class) return udaf_class @@ -1420,6 +1858,59 @@ def validate_udaf_class(udaf_class: type): ) +_PICKLE_BUILTIN_SCALAR_TYPES = ( + type(None), + bool, + int, + float, + complex, + str, + bytes, + bytearray, +) + + +def _is_pickle_builtin_value(value: Any, visited: Optional[set] = None) -> bool: + """Return whether pickle can encode a value without resolving modules.""" + value_type = type(value) + if value_type in _PICKLE_BUILTIN_SCALAR_TYPES: + return True + if value_type not in (list, tuple, set, frozenset, dict): + return False + + if visited is None: + visited = set() + value_id = id(value) + if value_id in visited: + return True + visited.add(value_id) + + if value_type is dict: + return all( + _is_pickle_builtin_value(key, visited) + and _is_pickle_builtin_value(item, visited) + for key, item in value.items() + ) + return all(_is_pickle_builtin_value(item, visited) for item in value) + + +class _ModuleContextUnpickler(pickle.Unpickler): + """Resolve UDAF classes from its isolated module cache.""" + + def __init__(self, serialized_state: bytes, module_context: ModuleContext): + super().__init__(io.BytesIO(serialized_state)) + self._udf_modules = module_context[1] + + def find_class(self, module_name: str, class_name: str) -> Any: + module = self._udf_modules.get(module_name) + if module is None: + return super().find_class(module_name, class_name) + value = module + for attribute in class_name.split("."): + value = getattr(value, attribute) + return value + + class UDAFStateManager: """ Manages UDAF aggregate states for Python UDAF execution. @@ -1437,6 +1928,8 @@ def __init__(self): """Initialize the state manager.""" self.states: Dict[int, Any] = {} # place_id -> UDAF instance self.udaf_class = None # UDAF class to instantiate + # Module UDAFs need this context for imports in later lifecycle calls. + self.module_context: Optional[ModuleContext] = None self._destroy_counter = 0 # Track number of destroys since last GC self._gc_threshold = 100 # Trigger GC every N destroys @@ -1451,6 +1944,9 @@ def set_udaf_class(self, udaf_class: type): Validation is performed by UDAFClassLoader before calling this method. """ self.udaf_class = udaf_class + self.module_context = getattr( + udaf_class, "_doris_module_context", None + ) def create_state(self, place_id: int) -> None: """ @@ -1517,9 +2013,19 @@ def serialize(self, place_id: int) -> bytes: """ state = self.states[place_id] try: - aggregate_state = state.aggregate_state - serialized = pickle.dumps(aggregate_state) - return serialized + with ModuleUDFLoader.use_module_context(self.module_context): + aggregate_state = state.aggregate_state + if ( + self.module_context is None + or _is_pickle_builtin_value(aggregate_state) + ): + return pickle.dumps(aggregate_state) + # Pickle resolves user-defined state classes by module name. Restore + # the isolated modules before running user serialization code once. + with ModuleUDFLoader.temporarily_restore_udf_modules( + self.module_context + ): + return pickle.dumps(aggregate_state) except Exception as e: logging.error( "Error serializing state for place_id %s: %s", @@ -1536,23 +2042,28 @@ def merge(self, place_id: int, other_state_bytes: bytes) -> None: place_id: Unique identifier for the aggregate state other_state_bytes: Serialized state to merge (pickle bytes) """ - try: - other_state = pickle.loads(other_state_bytes) - except Exception as e: - logging.error("Error deserializing state bytes: %s", e) - raise RuntimeError(f"Error deserializing state: {e}") from e - - state = self.states[place_id] + with ModuleUDFLoader.use_module_context(self.module_context): + try: + if self.module_context is None: + other_state = pickle.loads(other_state_bytes) + else: + other_state = _ModuleContextUnpickler( + other_state_bytes, self.module_context + ).load() + except Exception as e: + logging.error("Error deserializing state bytes: %s", e) + raise RuntimeError(f"Error deserializing state: {e}") from e - try: - state.merge(other_state) - except Exception as e: - logging.error( - "Error in merge for place_id %s: %s", - place_id, - e, - ) - raise RuntimeError(f"Error in merge: {e}") from e + state = self.states[place_id] + try: + state.merge(other_state) + except Exception as e: + logging.error( + "Error in merge for place_id %s: %s", + place_id, + e, + ) + raise RuntimeError(f"Error in merge: {e}") from e def finalize(self, place_id: int) -> Any: """ @@ -2110,6 +2621,7 @@ def _handle_exchange_udf( """Handle bidirectional streaming for UDF execution.""" loader = UDFLoaderFactory.get_loader(python_udf_meta) udf = loader.load() + module_context = getattr(loader, "module_context", None) logging.info("Loaded UDF: %s", udf) started = False @@ -2125,7 +2637,8 @@ def _handle_exchange_udf( logging.error("Schema mismatch: %s", error_msg) raise ValueError(f"Schema mismatch: {error_msg}") - result_array = udf(chunk.data) + with ModuleUDFLoader.use_module_context(module_context): + result_array = udf(chunk.data) if not python_udf_meta.output_type.equals(result_array.type): logging.error( @@ -2258,7 +2771,12 @@ def _handle_exchange_udaf( # Handle different operations and convert to unified format try: if operation_type == UDAFOperationType.CREATE: - result_batch = self._handle_udaf_create(place_id, state_manager) + with ModuleUDFLoader.use_module_context( + state_manager.module_context + ): + result_batch = self._handle_udaf_create( + place_id, state_manager + ) success = result_batch.column(0)[0].as_py() result_batch = self._create_unified_response( success=success, rows_processed=0, data=b"" @@ -2271,14 +2789,17 @@ def _handle_exchange_udaf( [batch.schema.field(i) for i in range(num_data_cols)] ), ) - result_batch_accumulate = self._handle_udaf_accumulate( - place_id, - is_single_place, - row_start, - row_end, - data_batch, - state_manager, - ) + with ModuleUDFLoader.use_module_context( + state_manager.module_context + ): + result_batch_accumulate = self._handle_udaf_accumulate( + place_id, + is_single_place, + row_start, + row_end, + data_batch, + state_manager, + ) rows_processed = result_batch_accumulate.column(0)[0].as_py() result_batch = self._create_unified_response( success=(rows_processed > 0), @@ -2308,9 +2829,12 @@ def _handle_exchange_udaf( success=success, rows_processed=0, data=b"" ) elif operation_type == UDAFOperationType.FINALIZE: - result_batch_finalize = self._handle_udaf_finalize( - place_id, python_udaf_meta.output_type, state_manager - ) + with ModuleUDFLoader.use_module_context( + state_manager.module_context + ): + result_batch_finalize = self._handle_udaf_finalize( + place_id, python_udaf_meta.output_type, state_manager + ) # Serialize the result to binary (including NULL results) # NULL is a valid aggregation result, not an error sink = pa.BufferOutputStream() @@ -2324,9 +2848,12 @@ def _handle_exchange_udaf( data=result_data, ) elif operation_type == UDAFOperationType.RESET: - result_batch_reset = self._handle_udaf_reset( - place_id, state_manager - ) + with ModuleUDFLoader.use_module_context( + state_manager.module_context + ): + result_batch_reset = self._handle_udaf_reset( + place_id, state_manager + ) success = result_batch_reset.column(0)[0].as_py() result_batch = self._create_unified_response( success=success, rows_processed=0, data=b"" @@ -2424,6 +2951,7 @@ def _handle_exchange_udtf( loader = UDFLoaderFactory.get_loader(python_udtf_meta) adaptive_udtf = loader.load() udtf_func = adaptive_udtf._eval_func + module_context = getattr(loader, "module_context", None) started = False for chunk in reader: @@ -2443,9 +2971,10 @@ def _handle_exchange_udtf( # Process all input rows and build ListArray try: - response_batch = self._process_udtf_with_list_array( - udtf_func, input_batch, python_udtf_meta.output_type - ) + with ModuleUDFLoader.use_module_context(module_context): + response_batch = self._process_udtf_with_list_array( + udtf_func, input_batch, python_udtf_meta.output_type + ) # Send the response batch if not started: @@ -2692,49 +3221,35 @@ def _clear_modules_from_location(self, location: str) -> list: """ Clear module cache for the given location. - Acquires per-module import locks to ensure no concurrent import is - in progress for the modules being cleared, preventing race conditions - where sys.modules entries are removed mid-import. + Uses the exclusive side of the module import protocol so cache cleanup + cannot race with sys.path or sys.modules changes during an import. Returns list of cleared module names. """ - cleared = [] + with ModuleUDFLoader._exclusive_module_import(): + cached_module = ModuleUDFLoader._module_cache.pop(location, None) + if cached_module is None: + return [] - with ModuleUDFLoader._module_cache_lock: - keys_to_remove = [ - key for key in ModuleUDFLoader._module_cache - if key[0] == location - ] - - # For each module, acquire its import lock before clearing. - # This ensures no concurrent _get_or_import_module is in progress - # for this (location, module_name) pair. - for key in keys_to_remove: - _, module_name = key - import_lock = ModuleUDFLoader._get_import_lock(module_name) - - with import_lock: - with ModuleUDFLoader._module_cache_lock: - if key in ModuleUDFLoader._module_cache: - del ModuleUDFLoader._module_cache[key] - - modules_to_remove = [ - name for name, mod in sys.modules.items() - if name == module_name or name.startswith(module_name + ".") - or ( - hasattr(mod, "__file__") and mod.__file__ is not None - and mod.__file__.startswith(location) - ) - ] - for mod_name in modules_to_remove: - del sys.modules[mod_name] - if mod_name not in cleared: - cleared.append(mod_name) - - if module_name not in cleared: - cleared.append(module_name) + module_context = getattr( + cached_module, "_doris_module_context", None + ) + udf_modules = module_context[1] if module_context else {} + module_names = set(udf_modules) + module_names.add(cached_module.__name__) + + normalized_location = os.path.realpath(location) + module_names.update( + name + for name, module in list(sys.modules.items()) + if ModuleUDFLoader._is_module_from_location( + module, normalized_location + ) + ) + for module_name in module_names: + sys.modules.pop(module_name, None) - return cleared + return sorted(module_names) class UDAFOperationType(Enum): diff --git a/be/test/runtime/user_function_cache_test.cpp b/be/test/runtime/user_function_cache_test.cpp index 57cd3149547e9a..6c0a39bd59db55 100644 --- a/be/test/runtime/user_function_cache_test.cpp +++ b/be/test/runtime/user_function_cache_test.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -27,9 +28,15 @@ #include #include "common/status.h" +#include "udf/python/python_env.h" +#include "udf/python/python_server.h" +#include "udf/python/python_udf_runtime.h" +#include "util/defer_op.h" namespace doris { +namespace bp = boost::process; + class UserFunctionCacheTest : public ::testing::Test { protected: UserFunctionCache ufc; @@ -512,6 +519,48 @@ TEST_F(UserFunctionCacheTest, LoadEntryFromLibPyZip) { EXPECT_EQ(ufc._entry_map.count(12345), 1); } +TEST_F(UserFunctionCacheTest, DropPythonCacheFailureKeepsEntryAndFilesForRetry) { + constexpr int64_t function_id = 23456; + std::string sub_dir = test_dir_ + "/0"; + std::filesystem::create_directories(sub_dir); + std::string zip_name = "23456.abc123def456abc123def456abc1.retryable_udf.zip"; + std::string zip_path = sub_dir + "/" + zip_name; + std::string extracted_path = zip_path.substr(0, zip_path.size() - 4); + ASSERT_TRUE(create_zip_file(zip_path, {{"main.py", "def evaluate(): return 1"}})); + ASSERT_TRUE(ufc._load_entry_from_lib(sub_dir, zip_name).ok()); + ASSERT_TRUE(std::filesystem::exists(zip_path)); + ASSERT_TRUE(std::filesystem::exists(extracted_path)); + + bp::ipstream output_stream; + std::string sleep_path = + std::filesystem::exists("/bin/sleep") ? "/bin/sleep" : "/usr/bin/sleep"; + bp::child child(sleep_path, "60", bp::std_out > output_stream, bp::std_err > bp::null); + auto process = std::make_shared(std::move(child), std::move(output_stream)); + process->set_uri_for_test("invalid-python-flight-uri"); + + auto& manager = PythonServerManager::instance(); + PythonVersion version("drop-cache-test", test_dir_, sleep_path); + manager.set_process_pool_for_test(version, {process}); + Defer cleanup_process([&] { + manager.set_process_pool_for_test(version, {}); + process->shutdown(); + }); + + auto failed_status = ufc.drop_function_cache(function_id); + EXPECT_FALSE(failed_status.ok()); + EXPECT_EQ(ufc._entry_map.count(function_id), 1); + EXPECT_TRUE(std::filesystem::exists(zip_path)); + EXPECT_TRUE(std::filesystem::exists(extracted_path)); + + manager.set_process_pool_for_test(version, {}); + process->shutdown(); + auto retry_status = ufc.drop_function_cache(function_id); + EXPECT_TRUE(retry_status.ok()) << retry_status.to_string(); + EXPECT_EQ(ufc._entry_map.count(function_id), 0); + EXPECT_FALSE(std::filesystem::exists(zip_path)); + EXPECT_FALSE(std::filesystem::exists(extracted_path)); +} + TEST_F(UserFunctionCacheTest, LoadEntryFromLibZipWithoutPython) { // Create a zip file without Python files - should fail std::string sub_dir = test_dir_ + "/1"; @@ -541,4 +590,4 @@ TEST_F(UserFunctionCacheTest, LoadEntryFromLibInvalidFileName) { EXPECT_FALSE(status.ok()); } -} // namespace doris \ No newline at end of file +} // namespace doris diff --git a/be/test/udf/python/python_server_test.cpp b/be/test/udf/python/python_server_test.cpp index 3ec463bbe1cdfa..639e671ce91cff 100644 --- a/be/test/udf/python/python_server_test.cpp +++ b/be/test/udf/python/python_server_test.cpp @@ -17,6 +17,8 @@ #include "udf/python/python_server.h" +#include +#include #include #include #include @@ -26,20 +28,57 @@ #include #include #include +#include +#include +#include #include #include #include "common/config.h" #include "common/status.h" +#include "core/data_type/data_type_number.h" #include "udf/python/python_env.h" +#include "udf/python/python_udaf_client.h" #include "udf/python/python_udf_client.h" #include "udf/python/python_udf_meta.h" +#include "udf/python/python_udtf_client.h" +#include "util/defer_op.h" namespace doris { namespace fs = std::filesystem; namespace bp = boost::process; +class ActionResultFlightServer final : public arrow::flight::FlightServerBase { +public: + explicit ActionResultFlightServer(std::vector results) + : _results(std::move(results)) {} + + arrow::Status start() { + auto location = arrow::flight::Location::ForGrpcTcp("localhost", 0); + if (!location.ok()) { + return location.status(); + } + return Init(arrow::flight::FlightServerOptions(*location)); + } + + ~ActionResultFlightServer() override { static_cast(Shutdown()); } + + arrow::Status DoAction(const arrow::flight::ServerCallContext&, const arrow::flight::Action&, + std::unique_ptr* result) override { + std::vector flight_results; + flight_results.reserve(_results.size()); + for (const auto& body : _results) { + flight_results.emplace_back(arrow::Buffer::FromString(body)); + } + *result = std::make_unique(std::move(flight_results)); + return arrow::Status::OK(); + } + +private: + std::vector _results; +}; + class PythonServerTest : public ::testing::Test { protected: std::string test_dir_; @@ -189,6 +228,29 @@ class PythonServerTest : public ::testing::Test { ofs.close(); } + Status install_real_python_server() { + fs::path source_root = fs::current_path(); + fs::path server_script; + while (!source_root.empty()) { + server_script = source_root / "be/src/udf/python/python_server.py"; + if (fs::exists(server_script) || source_root == source_root.parent_path()) { + break; + } + source_root = source_root.parent_path(); + } + if (!fs::exists(server_script)) { + return Status::InternalError("Python server script not found: {}", + server_script.string()); + } + + setenv("DORIS_HOME", test_dir_.c_str(), 1); + fs::path plugin_dir = fs::path(test_dir_) / "plugins/python_udf"; + fs::create_directories(plugin_dir); + fs::copy_file(server_script, plugin_dir / "python_server.py", + fs::copy_options::overwrite_existing); + return Status::OK(); + } + ProcessPtr create_sleep_process() { bp::ipstream output_stream; std::string sleep_path = fs::exists("/bin/sleep") ? "/bin/sleep" : "/usr/bin/sleep"; @@ -196,6 +258,125 @@ class PythonServerTest : public ::testing::Test { return std::make_shared(std::move(child), std::move(output_stream)); } + std::optional find_python_udf_interpreter() { + std::vector candidates; + if (const char* configured = std::getenv("DORIS_PYTHON_UDF_TEST_PYTHON")) { + candidates.emplace_back(configured); + } + if (const char* path_env = std::getenv("PATH")) { + std::stringstream paths(path_env); + std::string path; + while (std::getline(paths, path, ':')) { + fs::path python3 = fs::path(path) / "python3"; + if (fs::exists(python3)) { + candidates.emplace_back(python3.string()); + } + } + } + + for (const auto& candidate : candidates) { + if (!fs::exists(candidate)) { + continue; + } + bp::child check(candidate, "-c", "import pandas, pyarrow", bp::std_out > bp::null, + bp::std_err > bp::null); + check.wait(); + if (check.exit_code() == 0) { + return candidate; + } + } + return std::nullopt; + } + + Status start_python_udf_server(const std::string& python, ProcessPtr* process) { + bp::ipstream output_stream; + try { + bp::child child(python, "-u", get_fight_server_path(), get_base_unix_socket_path(), + bp::std_out > output_stream); + auto candidate = + std::make_shared(std::move(child), std::move(output_stream)); + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (std::chrono::steady_clock::now() < deadline) { + if (fs::exists(candidate->get_socket_file_path())) { + *process = std::move(candidate); + return Status::OK(); + } + if (!candidate->is_alive()) { + return Status::InternalError("Python UDF server exited before becoming ready"); + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + return Status::InternalError("Timed out waiting for Python UDF server at {}", + candidate->get_socket_file_path()); + } catch (const std::exception& e) { + return Status::InternalError("Failed to start Python UDF server: {}", e.what()); + } + } + + Status evaluate_int_udf_batch(const PythonUDFClientPtr& client, + const std::vector& inputs, + std::vector* results) { + arrow::Int32Builder builder; + RETURN_DORIS_STATUS_IF_ERROR(builder.AppendValues(inputs)); + std::shared_ptr input_array; + RETURN_DORIS_STATUS_IF_ERROR(builder.Finish(&input_array)); + auto input_batch = arrow::RecordBatch::Make( + arrow::schema({arrow::field("arg0", arrow::int32(), false)}), inputs.size(), + {input_array}); + + std::shared_ptr output_batch; + Status evaluate_status = client->evaluate(*input_batch, &output_batch); + RETURN_IF_ERROR(evaluate_status); + if (!output_batch || output_batch->num_columns() != 1 || + output_batch->num_rows() != inputs.size() || + output_batch->column(0)->type_id() != arrow::Type::INT32) { + return Status::InternalError("Unexpected Python UDF output batch"); + } + + auto result_array = std::static_pointer_cast(output_batch->column(0)); + results->clear(); + results->reserve(inputs.size()); + for (int64_t i = 0; i < result_array->length(); ++i) { + results->push_back(result_array->Value(i)); + } + return Status::OK(); + } + + Status evaluate_int_module_udf_batch(const PythonUDFMeta& meta, const ProcessPtr& process, + const std::vector& inputs, + std::vector* results) { + PythonUDFClientPtr client; + RETURN_IF_ERROR(PythonUDFClient::create(meta, process, &client)); + Status evaluate_status = evaluate_int_udf_batch(client, inputs, results); + static_cast(client->close()); + return evaluate_status; + } + + Status evaluate_int_module_udf(const PythonUDFMeta& meta, const ProcessPtr& process, + int32_t input, int32_t* result) { + std::vector results; + RETURN_IF_ERROR(evaluate_int_module_udf_batch(meta, process, {input}, &results)); + *result = results[0]; + return Status::OK(); + } + + PythonUDFMeta make_int_module_meta(int64_t id, const fs::path& location, + const std::string& symbol, + PythonClientType client_type = PythonClientType::UDF) { + PythonUDFMeta meta; + meta.id = id; + meta.name = symbol; + meta.symbol = symbol; + meta.location = location.string(); + meta.checksum = "test-checksum"; + meta.runtime_version = "test-runtime"; + meta.input_types = {std::make_shared()}; + meta.return_type = std::make_shared(); + meta.type = PythonUDFLoadType::MODULE; + meta.client_type = client_type; + return meta; + } + template Status get_process_with_retry( PythonServerManager& mgr, const PythonVersion& version, @@ -499,6 +680,1060 @@ TEST_F(PythonServerTest, ClearModuleCacheWithoutProcessesIsNoOp) { EXPECT_TRUE(status.ok()) << status.to_string(); } +TEST_F(PythonServerTest, ClearModuleCacheReloadsModuleOnNextUdfExecution) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + fs::path module_dir = fs::path(test_dir_) / "module_cache"; + fs::create_directories(module_dir); + auto write_module = [&module_dir](int offset) { + std::ofstream module(module_dir / "cache_reload_udf.py", std::ios::trunc); + module << "def evaluate(value):\n" + << " return value + " << offset << "\n"; + }; + + write_module(1); + PythonVersion version("test-runtime", fs::path(*python).parent_path().parent_path().string(), + *python); + PythonServerManager mgr; + ProcessPtr process; + Status fork_status = start_python_udf_server(*python, &process); + ASSERT_TRUE(fork_status.ok()) << fork_status.to_string(); + ASSERT_NE(process, nullptr); + mgr.set_process_pool_for_test(version, {process}); + + PythonUDFMeta meta; + meta.id = 1; + meta.name = "cache_reload_udf"; + meta.symbol = "cache_reload_udf.evaluate"; + meta.location = module_dir.string(); + meta.checksum = "test-checksum"; + meta.runtime_version = version.full_version; + meta.input_types = {std::make_shared()}; + meta.return_type = std::make_shared(); + meta.type = PythonUDFLoadType::MODULE; + meta.client_type = PythonClientType::UDF; + + int32_t result = 0; + Status evaluate_status = evaluate_int_module_udf(meta, process, 10, &result); + ASSERT_TRUE(evaluate_status.ok()) << evaluate_status.to_string(); + ASSERT_EQ(result, 11); + + // DROP clears the Python module before UserFunctionCache deletes its extracted directory. + ASSERT_TRUE(mgr.clear_module_cache(meta.location).ok()); + fs::remove_all(module_dir); + fs::create_directories(module_dir); + write_module(100); + + evaluate_status = evaluate_int_module_udf(meta, process, 10, &result); + ASSERT_TRUE(evaluate_status.ok()) << evaluate_status.to_string(); + EXPECT_EQ(result, 110); + mgr.shutdown(); +} + +TEST_F(PythonServerTest, ConcurrentModuleImportsIsolateSameNamedDependencies) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + fs::path module_a_dir = fs::path(test_dir_) / "module_a"; + fs::path module_b_dir = fs::path(test_dir_) / "module_b"; + fs::create_directories(module_a_dir); + fs::create_directories(module_b_dir); + fs::path import_started = fs::path(test_dir_) / "import_started"; + fs::path second_import_started = fs::path(test_dir_) / "second_import_started"; + fs::path second_import_completed = fs::path(test_dir_) / "second_import_completed"; + fs::path allow_import = fs::path(test_dir_) / "allow_import"; + { + std::ofstream dependency(module_a_dir / "pd.py"); + dependency << "import pathlib\n" + << "import time\n" + << "OFFSET = 1\n" + << "pathlib.Path(r'" << import_started.string() << "').touch()\n" + << "allowed = pathlib.Path(r'" << allow_import.string() << "')\n" + << "deadline = time.monotonic() + 10\n" + << "while not allowed.exists() and time.monotonic() < deadline:\n" + << " time.sleep(0.01)\n"; + std::ofstream module(module_a_dir / "first_udf.py"); + module << "def evaluate(value):\n" + << " import pd\n" + << " return value + pd.OFFSET\n"; + } + { + std::ofstream dependency(module_b_dir / "pd.py"); + dependency << "OFFSET = 100\n"; + std::ofstream module(module_b_dir / "second_udf.py"); + module << "import pathlib\n" + << "def evaluate(value):\n" + << " if value == 0:\n" + << " return 0\n" + << " pathlib.Path(r'" << second_import_started.string() << "').touch()\n" + << " import pd\n" + << " pathlib.Path(r'" << second_import_completed.string() << "').touch()\n" + << " return value + pd.OFFSET\n"; + } + + PythonVersion version("test-runtime", fs::path(*python).parent_path().parent_path().string(), + *python); + PythonServerManager mgr; + ProcessPtr process; + Status fork_status = start_python_udf_server(*python, &process); + ASSERT_TRUE(fork_status.ok()) << fork_status.to_string(); + ASSERT_NE(process, nullptr); + mgr.set_process_pool_for_test(version, {process}); + + auto first_meta = make_int_module_meta(1, module_a_dir, "first_udf.evaluate"); + auto second_meta = make_int_module_meta(2, module_b_dir, "second_udf.evaluate"); + + // Keep B's Flight exchange open after loading its entry module. Sending the + // second batch then reaches evaluate() directly instead of loading B again. + PythonUDFClientPtr second_client; + ASSERT_TRUE(PythonUDFClient::create(second_meta, process, &second_client).ok()); + int32_t warmup_result = -1; + std::vector warmup_results; + ASSERT_TRUE(evaluate_int_udf_batch(second_client, {0}, &warmup_results).ok()); + ASSERT_EQ(warmup_results.size(), 1); + warmup_result = warmup_results[0]; + ASSERT_EQ(warmup_result, 0); + + int32_t first_result = 0; + auto first_status = std::async(std::launch::async, [&] { + return evaluate_int_module_udf(first_meta, process, 10, &first_result); + }); + Defer release_first_import {[&] { std::ofstream(allow_import).close(); }}; + auto marker_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!fs::exists(import_started) && std::chrono::steady_clock::now() < marker_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_TRUE(fs::exists(import_started)); + + int32_t second_result = 0; + auto second_status = std::async(std::launch::async, [&] { + std::vector results; + RETURN_IF_ERROR(evaluate_int_udf_batch(second_client, {10}, &results)); + second_result = results[0]; + return Status::OK(); + }); + + // "pd" used to be treated as a server alias. A live sys.modules lookup + // would return A's pd.py here instead of waiting to restore B's context. + marker_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!fs::exists(second_import_started) && + std::chrono::steady_clock::now() < marker_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_TRUE(fs::exists(second_import_started)); + + // A broken lock-free lookup returns A's live pd module and creates this + // marker. The isolated path cannot finish the import until A releases it. + auto completion_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (!fs::exists(second_import_completed) && + std::chrono::steady_clock::now() < completion_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_FALSE(fs::exists(second_import_completed)); + std::ofstream(allow_import).close(); + + ASSERT_TRUE(first_status.get().ok()); + ASSERT_TRUE(second_status.get().ok()); + EXPECT_EQ(first_result, 11); + EXPECT_EQ(second_result, 110); + EXPECT_TRUE(fs::exists(second_import_completed)); + static_cast(second_client->close()); + mgr.shutdown(); +} + +TEST_F(PythonServerTest, UdafAndUdtfModulesAreIsolatedInTheSameProcess) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + auto write_udaf = [](const fs::path& location, int offset) { + fs::create_directories(location); + std::ofstream(location / "shared_udaf_dependency.py") << "OFFSET = " << offset << "\n"; + std::ofstream(location / "udaf_entry.py") + << "from shared_udaf_dependency import OFFSET\n" + << "class SumAgg:\n" + << " def __init__(self): self.total = 0\n" + << " @property\n" + << " def aggregate_state(self):\n" + << " return {'total': complex(self.total), 'tag': bytearray(b'x')}\n" + << " def accumulate(self, value): self.total += value + OFFSET\n" + << " def merge(self, other): self.total += int(other['total'].real)\n" + << " def finish(self): return self.total\n"; + }; + auto write_udtf = [](const fs::path& location, int offset) { + fs::create_directories(location); + std::ofstream(location / "shared_udtf_dependency.py") << "OFFSET = " << offset << "\n"; + std::ofstream(location / "udtf_entry.py") << "from shared_udtf_dependency import OFFSET\n" + << "def evaluate(value):\n" + << " yield value + OFFSET\n"; + }; + + fs::path udaf_a_dir = fs::path(test_dir_) / "same_process_udaf_a"; + fs::path udaf_b_dir = fs::path(test_dir_) / "same_process_udaf_b"; + fs::path udtf_a_dir = fs::path(test_dir_) / "same_process_udtf_a"; + fs::path udtf_b_dir = fs::path(test_dir_) / "same_process_udtf_b"; + fs::path blocking_dir = fs::path(test_dir_) / "blocking_module_import"; + fs::path import_started = fs::path(test_dir_) / "blocking_import_started"; + fs::path allow_import = fs::path(test_dir_) / "allow_blocking_import"; + fs::path object_udaf_dir = fs::path(test_dir_) / "object_state_udaf"; + write_udaf(udaf_a_dir, 1); + write_udaf(udaf_b_dir, 100); + write_udtf(udtf_a_dir, 2); + write_udtf(udtf_b_dir, 200); + fs::create_directories(object_udaf_dir); + std::ofstream(object_udaf_dir / "object_state_dependency.py") + << "class AggregateValue:\n" + << " def __init__(self, total): self.total = total\n"; + std::ofstream(object_udaf_dir / "object_udaf_entry.py") + << "from object_state_dependency import AggregateValue\n" + << "class ObjectStateAgg:\n" + << " def __init__(self): self.total = 0\n" + << " @property\n" + << " def aggregate_state(self): return AggregateValue(self.total)\n" + << " def accumulate(self, value): self.total += value\n" + << " def merge(self, other): self.total += other.total\n" + << " def finish(self): return self.total\n"; + fs::create_directories(blocking_dir); + std::ofstream(blocking_dir / "blocking_udf.py") + << "import pathlib\n" + << "import time\n" + << "started = pathlib.Path(r'" << import_started.string() << "')\n" + << "allowed = pathlib.Path(r'" << allow_import.string() << "')\n" + << "started.touch()\n" + << "deadline = time.monotonic() + 10\n" + << "while not allowed.exists() and time.monotonic() < deadline:\n" + << " time.sleep(0.01)\n" + << "def evaluate(value): return value\n"; + + ProcessPtr process; + ASSERT_TRUE(start_python_udf_server(*python, &process).ok()); + ASSERT_NE(process, nullptr); + + arrow::Int32Builder input_builder; + ASSERT_TRUE(input_builder.Append(10).ok()); + std::shared_ptr input_array; + ASSERT_TRUE(input_builder.Finish(&input_array).ok()); + auto input_schema = arrow::schema({arrow::field("arg0", arrow::int32(), false)}); + auto input_batch = arrow::RecordBatch::Make(input_schema, 1, {input_array}); + + auto udaf_schema = arrow::schema({arrow::field("arg0", arrow::int32(), false), + arrow::field("places", arrow::int64()), + arrow::field("binary_data", arrow::binary())}); + PythonUDAFClientPtr udaf_a; + PythonUDAFClientPtr udaf_b; + ASSERT_TRUE(PythonUDAFClient::create(make_int_module_meta(1, udaf_a_dir, "udaf_entry.SumAgg", + PythonClientType::UDAF), + process, udaf_schema, &udaf_a) + .ok()); + ASSERT_TRUE(PythonUDAFClient::create(make_int_module_meta(2, udaf_b_dir, "udaf_entry.SumAgg", + PythonClientType::UDAF), + process, udaf_schema, &udaf_b) + .ok()); + ASSERT_TRUE(udaf_a->create(101).ok()); + ASSERT_TRUE(udaf_b->create(102).ok()); + ASSERT_TRUE(udaf_a->accumulate(101, true, *input_batch, 0, 1).ok()); + ASSERT_TRUE(udaf_b->accumulate(102, true, *input_batch, 0, 1).ok()); + + auto finalize_int = [](const PythonUDAFClientPtr& client, int64_t place_id) { + std::shared_ptr output; + EXPECT_TRUE(client->finalize(place_id, &output).ok()); + EXPECT_NE(output, nullptr); + if (!output || output->num_columns() != 1 || output->num_rows() != 1) { + return int32_t {0}; + } + return std::static_pointer_cast(output->column(0))->Value(0); + }; + EXPECT_EQ(finalize_int(udaf_a, 101), 11); + EXPECT_EQ(finalize_int(udaf_b, 102), 110); + + // A user-defined aggregate-state object needs its defining module visible + // for both directions. User serialization code must run exactly once. + PythonUDAFClientPtr object_udaf; + ASSERT_TRUE(PythonUDAFClient::create( + make_int_module_meta(6, object_udaf_dir, "object_udaf_entry.ObjectStateAgg", + PythonClientType::UDAF), + process, udaf_schema, &object_udaf) + .ok()); + ASSERT_TRUE(object_udaf->create(103).ok()); + ASSERT_TRUE(object_udaf->create(104).ok()); + ASSERT_TRUE(object_udaf->accumulate(103, true, *input_batch, 0, 1).ok()); + std::shared_ptr object_state; + ASSERT_TRUE(object_udaf->serialize(103, &object_state).ok()); + ASSERT_TRUE(object_udaf->merge(104, object_state).ok()); + EXPECT_EQ(finalize_int(object_udaf, 104), 10); + + // A built-in-only aggregate state, including complex and bytearray values, + // does not need UDF modules restored. SERIALIZE and MERGE must complete + // while another module import owns the process-wide writer lock. + auto blocking_meta = make_int_module_meta(5, blocking_dir, "blocking_udf.evaluate"); + int32_t blocking_result = 0; + auto blocking_status = std::async(std::launch::async, [&] { + return evaluate_int_module_udf(blocking_meta, process, 10, &blocking_result); + }); + auto marker_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!fs::exists(import_started) && std::chrono::steady_clock::now() < marker_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_TRUE(fs::exists(import_started)); + + std::shared_ptr serialized_state; + auto serialize_status = std::async(std::launch::async, + [&] { return udaf_a->serialize(101, &serialized_state); }); + auto serialize_wait = serialize_status.wait_for(std::chrono::seconds(5)); + std::optional serialize_result; + std::optional> merge_status; + std::optional merge_wait; + if (serialize_wait == std::future_status::ready) { + serialize_result = serialize_status.get(); + if (serialize_result->ok()) { + merge_status.emplace(std::async(std::launch::async, + [&] { return udaf_a->merge(101, serialized_state); })); + merge_wait = merge_status->wait_for(std::chrono::seconds(5)); + } + } + std::ofstream(allow_import).close(); + + EXPECT_EQ(serialize_wait, std::future_status::ready); + if (serialize_result) { + EXPECT_TRUE(serialize_result->ok()) << serialize_result->to_string(); + } + if (merge_wait) { + EXPECT_EQ(*merge_wait, std::future_status::ready); + } + if (merge_status) { + auto merge_result = merge_status->get(); + EXPECT_TRUE(merge_result.ok()) << merge_result.to_string(); + } + if (serialize_wait != std::future_status::ready) { + auto delayed_serialize_result = serialize_status.get(); + ADD_FAILURE() << "primitive UDAF serialization waited for the module import: " + << delayed_serialize_result.to_string(); + } + ASSERT_TRUE(blocking_status.get().ok()); + EXPECT_EQ(blocking_result, 10); + if (merge_wait && *merge_wait == std::future_status::ready) { + EXPECT_EQ(finalize_int(udaf_a, 101), 22); + } + + PythonUDTFClientPtr udtf_a; + PythonUDTFClientPtr udtf_b; + ASSERT_TRUE(PythonUDTFClient::create(make_int_module_meta(3, udtf_a_dir, "udtf_entry.evaluate", + PythonClientType::UDTF), + process, &udtf_a) + .ok()); + ASSERT_TRUE(PythonUDTFClient::create(make_int_module_meta(4, udtf_b_dir, "udtf_entry.evaluate", + PythonClientType::UDTF), + process, &udtf_b) + .ok()); + auto evaluate_udtf = [&](const PythonUDTFClientPtr& client) { + std::shared_ptr output; + EXPECT_TRUE(client->evaluate(*input_batch, &output).ok()); + EXPECT_NE(output, nullptr); + if (!output || output->length() != 1 || output->value_length(0) != 1) { + return int32_t {0}; + } + auto values = std::static_pointer_cast(output->values()); + return values->Value(output->value_offset(0)); + }; + EXPECT_EQ(evaluate_udtf(udtf_a), 12); + EXPECT_EQ(evaluate_udtf(udtf_b), 210); + + static_cast(udaf_a->close()); + static_cast(udaf_b->close()); + static_cast(object_udaf->close()); + static_cast(udtf_a->close()); + static_cast(udtf_b->close()); + process->shutdown(); +} + +TEST_F(PythonServerTest, InlineImportWaitsForModuleImportEnvironment) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + fs::path global_module_dir = fs::path(test_dir_) / "global_modules"; + fs::path udf_module_dir = fs::path(test_dir_) / "module_with_shadow"; + fs::create_directories(global_module_dir); + fs::create_directories(udf_module_dir); + fs::path module_import_started = fs::path(test_dir_) / "module_import_started"; + fs::path inline_import_started = fs::path(test_dir_) / "inline_import_started"; + { + std::ofstream dependency(global_module_dir / "shared_inline_dependency.py"); + dependency << "OFFSET = 1\n"; + } + { + std::ofstream dependency(udf_module_dir / "shared_inline_dependency.py"); + dependency << "OFFSET = 100\n"; + } + { + std::ofstream module(udf_module_dir / "slow_module_udf.py"); + module << "import pathlib\n" + << "import time\n" + << "started = pathlib.Path(r'" << module_import_started.string() << "')\n" + << "inline_started = pathlib.Path(r'" << inline_import_started.string() << "')\n" + << "started.touch()\n" + << "deadline = time.monotonic() + 5\n" + << "while not inline_started.exists() and time.monotonic() < deadline:\n" + << " time.sleep(0.01)\n" + << "if not inline_started.exists():\n" + << " raise RuntimeError('inline import did not start')\n" + << "time.sleep(1)\n" + << "def evaluate(value):\n" + << " return value\n"; + } + + const char* original_python_path = std::getenv("PYTHONPATH"); + std::optional saved_python_path; + if (original_python_path) { + saved_python_path = original_python_path; + } + setenv("PYTHONPATH", global_module_dir.c_str(), 1); + + PythonVersion version("test-runtime", fs::path(*python).parent_path().parent_path().string(), + *python); + PythonServerManager mgr; + ProcessPtr process; + Status fork_status = start_python_udf_server(*python, &process); + if (saved_python_path) { + setenv("PYTHONPATH", saved_python_path->c_str(), 1); + } else { + unsetenv("PYTHONPATH"); + } + ASSERT_TRUE(fork_status.ok()) << fork_status.to_string(); + ASSERT_NE(process, nullptr); + mgr.set_process_pool_for_test(version, {process}); + + PythonUDFMeta module_meta; + module_meta.id = 1; + module_meta.name = "slow_module_udf"; + module_meta.symbol = "slow_module_udf.evaluate"; + module_meta.location = udf_module_dir.string(); + module_meta.checksum = "test-checksum"; + module_meta.runtime_version = version.full_version; + module_meta.input_types = {std::make_shared()}; + module_meta.return_type = std::make_shared(); + module_meta.type = PythonUDFLoadType::MODULE; + module_meta.client_type = PythonClientType::UDF; + + PythonUDFMeta inline_meta; + inline_meta.id = 2; + inline_meta.name = "inline_import_udf"; + inline_meta.symbol = "evaluate"; + inline_meta.runtime_version = version.full_version; + inline_meta.inline_code = "open(r'" + inline_import_started.string() + + "', 'w').close()\n" + "import shared_inline_dependency\n" + "def evaluate(value):\n" + " return value + shared_inline_dependency.OFFSET\n"; + inline_meta.input_types = {std::make_shared()}; + inline_meta.return_type = std::make_shared(); + inline_meta.type = PythonUDFLoadType::INLINE; + inline_meta.client_type = PythonClientType::UDF; + + PythonUDFClientPtr inline_client; + ASSERT_TRUE(PythonUDFClient::create(inline_meta, process, &inline_client).ok()); + + int32_t module_result = 0; + auto module_status = std::async(std::launch::async, [&] { + return evaluate_int_module_udf(module_meta, process, 10, &module_result); + }); + auto marker_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!fs::exists(module_import_started) && + std::chrono::steady_clock::now() < marker_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_TRUE(fs::exists(module_import_started)); + + std::vector inline_results; + auto inline_status = std::async(std::launch::async, [&] { + return evaluate_int_udf_batch(inline_client, {10}, &inline_results); + }); + + ASSERT_TRUE(module_status.get().ok()); + ASSERT_TRUE(inline_status.get().ok()); + EXPECT_EQ(module_result, 10); + EXPECT_EQ(inline_results, std::vector({11})); + static_cast(inline_client->close()); + mgr.shutdown(); +} + +TEST_F(PythonServerTest, ModuleImportWaitsForInlineImport) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + fs::path global_module_dir = fs::path(test_dir_) / "slow_global_modules"; + fs::path udf_module_dir = fs::path(test_dir_) / "module_during_inline_import"; + fs::create_directories(global_module_dir); + fs::create_directories(udf_module_dir); + fs::path inline_import_started = fs::path(test_dir_) / "slow_inline_import_started"; + fs::path allow_child_import = fs::path(test_dir_) / "allow_child_import"; + fs::path module_import_started = fs::path(test_dir_) / "waiting_module_import_started"; + { + std::ofstream dependency(global_module_dir / "nested_global_dependency.py"); + dependency << "OFFSET = 2\n"; + } + { + std::ofstream dependency(global_module_dir / "slow_global_dependency.py"); + dependency << "import pathlib\n" + << "import threading\n" + << "import time\n" + << "started = pathlib.Path(r'" << inline_import_started.string() << "')\n" + << "release = pathlib.Path(r'" << allow_child_import.string() << "')\n" + << "started.touch()\n" + << "deadline = time.monotonic() + 5\n" + << "while not release.exists() and time.monotonic() < deadline:\n" + << " time.sleep(0.01)\n" + << "if not release.exists():\n" + << " raise RuntimeError('timed out waiting to start child import')\n" + << "results = []\n" + << "def import_from_child():\n" + << " import nested_global_dependency\n" + << " results.append(nested_global_dependency.OFFSET)\n" + << "thread = threading.Thread(target=import_from_child)\n" + << "thread.start()\n" + << "thread.join(5)\n" + << "if thread.is_alive():\n" + << " raise RuntimeError('child reader waited behind module writer')\n" + << "OFFSET = results[0]\n"; + } + { + std::ofstream module(udf_module_dir / "waiting_module_udf.py"); + module << "open(r'" << module_import_started.string() << "', 'w').close()\n" + << "def evaluate(value):\n" + << " return value\n"; + } + + const char* original_python_path = std::getenv("PYTHONPATH"); + std::optional saved_python_path; + if (original_python_path) { + saved_python_path = original_python_path; + } + setenv("PYTHONPATH", global_module_dir.c_str(), 1); + + PythonVersion version("test-runtime", fs::path(*python).parent_path().parent_path().string(), + *python); + PythonServerManager mgr; + ProcessPtr process; + Status fork_status = start_python_udf_server(*python, &process); + if (saved_python_path) { + setenv("PYTHONPATH", saved_python_path->c_str(), 1); + } else { + unsetenv("PYTHONPATH"); + } + ASSERT_TRUE(fork_status.ok()) << fork_status.to_string(); + ASSERT_NE(process, nullptr); + mgr.set_process_pool_for_test(version, {process}); + + PythonUDFMeta inline_meta; + inline_meta.id = 1; + inline_meta.name = "slow_inline_import_udf"; + inline_meta.symbol = "evaluate"; + inline_meta.runtime_version = version.full_version; + inline_meta.inline_code = + "import slow_global_dependency\n" + "def evaluate(value):\n" + " return value + slow_global_dependency.OFFSET\n"; + inline_meta.input_types = {std::make_shared()}; + inline_meta.return_type = std::make_shared(); + inline_meta.type = PythonUDFLoadType::INLINE; + inline_meta.client_type = PythonClientType::UDF; + + PythonUDFMeta module_meta; + module_meta.id = 2; + module_meta.name = "waiting_module_udf"; + module_meta.symbol = "waiting_module_udf.evaluate"; + module_meta.location = udf_module_dir.string(); + module_meta.checksum = "test-checksum"; + module_meta.runtime_version = version.full_version; + module_meta.input_types = {std::make_shared()}; + module_meta.return_type = std::make_shared(); + module_meta.type = PythonUDFLoadType::MODULE; + module_meta.client_type = PythonClientType::UDF; + + PythonUDFClientPtr inline_client; + ASSERT_TRUE(PythonUDFClient::create(inline_meta, process, &inline_client).ok()); + std::vector inline_results; + auto inline_status = std::async(std::launch::async, [&] { + return evaluate_int_udf_batch(inline_client, {10}, &inline_results); + }); + auto marker_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!fs::exists(inline_import_started) && + std::chrono::steady_clock::now() < marker_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_TRUE(fs::exists(inline_import_started)); + + int32_t module_result = 0; + auto module_status = std::async(std::launch::async, [&] { + return evaluate_int_module_udf(module_meta, process, 10, &module_result); + }); + Defer release_child_import {[&] { std::ofstream(allow_child_import).close(); }}; + + EXPECT_EQ(module_status.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); + EXPECT_FALSE(fs::exists(module_import_started)); + std::ofstream(allow_child_import).close(); + + ASSERT_TRUE(inline_status.get().ok()); + ASSERT_TRUE(module_status.get().ok()); + EXPECT_EQ(inline_results, std::vector({12})); + EXPECT_EQ(module_result, 10); + static_cast(inline_client->close()); + mgr.shutdown(); +} + +TEST_F(PythonServerTest, CachedRuntimeImportsDoNotWaitForAnotherModuleImport) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + fs::path cached_module_dir = fs::path(test_dir_) / "cached_runtime_import"; + fs::path blocking_module_dir = fs::path(test_dir_) / "blocking_module_import"; + fs::create_directories(cached_module_dir / "cached_import_pkg"); + fs::create_directories(blocking_module_dir); + fs::path import_started = fs::path(test_dir_) / "blocking_import_started"; + fs::path cached_batch_completed = fs::path(test_dir_) / "cached_batch_completed"; + fs::path watchdog_expired = fs::path(test_dir_) / "watchdog_expired"; + { + std::ofstream init(cached_module_dir / "cached_import_pkg/__init__.py"); + std::ofstream dependency(cached_module_dir / "cached_import_pkg/cached_dependency.py"); + dependency << "OFFSET = 1\n"; + std::ofstream module(cached_module_dir / "cached_import_pkg/entry.py"); + module << "import importlib\n" + << "import pathlib\n" + << "completed = pathlib.Path(r'" << cached_batch_completed.string() << "')\n" + << "def evaluate(value):\n" + << " dependency = importlib.import_module(\n" + << " '.cached_dependency', __package__)\n" + << " from .cached_dependency import OFFSET\n" + << " builtin_dependency = __import__(\n" + << " 'cached_dependency', globals(), locals(), (), 1)\n" + << " if value == -1:\n" + << " try:\n" + << " __import__(\n" + << " 'cached_import_pkg.cached_dependency',\n" + << " globals(), locals(), (), -1)\n" + << " except ValueError:\n" + << " pass\n" + << " else:\n" + << " raise RuntimeError('negative import level was accepted')\n" + << " if value == 99:\n" + << " completed.touch()\n" + << " return (value + dependency.OFFSET + OFFSET\n" + << " + builtin_dependency.OFFSET)\n"; + } + { + std::ofstream module(blocking_module_dir / "blocking_import_udf.py"); + module << "import pathlib\n" + << "import time\n" + << "started = pathlib.Path(r'" << import_started.string() << "')\n" + << "completed = pathlib.Path(r'" << cached_batch_completed.string() << "')\n" + << "watchdog = pathlib.Path(r'" << watchdog_expired.string() << "')\n" + << "deadline = time.monotonic() + 15\n" + << "started.touch()\n" + << "while not completed.exists() and time.monotonic() < deadline:\n" + << " time.sleep(0.01)\n" + << "if not completed.exists():\n" + << " watchdog.touch()\n" + << "def evaluate(value):\n" + << " return value\n"; + } + + PythonVersion version("test-runtime", fs::path(*python).parent_path().parent_path().string(), + *python); + PythonServerManager mgr; + ProcessPtr process; + Status fork_status = start_python_udf_server(*python, &process); + ASSERT_TRUE(fork_status.ok()) << fork_status.to_string(); + ASSERT_NE(process, nullptr); + mgr.set_process_pool_for_test(version, {process}); + + auto cached_meta = + make_int_module_meta(1, cached_module_dir, "cached_import_pkg.entry.evaluate"); + auto blocking_meta = + make_int_module_meta(2, blocking_module_dir, "blocking_import_udf.evaluate"); + + PythonUDFClientPtr cached_client; + ASSERT_TRUE(PythonUDFClient::create(cached_meta, process, &cached_client).ok()); + std::vector warmup_results; + ASSERT_TRUE(evaluate_int_udf_batch(cached_client, {-1}, &warmup_results).ok()); + ASSERT_EQ(warmup_results, std::vector({2})); + + int32_t blocking_result = 0; + auto blocking_status = std::async(std::launch::async, [&] { + return evaluate_int_module_udf(blocking_meta, process, 10, &blocking_result); + }); + auto marker_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!fs::exists(import_started) && std::chrono::steady_clock::now() < marker_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_TRUE(fs::exists(import_started)); + + std::vector inputs(100); + std::iota(inputs.begin(), inputs.end(), 0); + std::vector results; + auto cached_status = std::async(std::launch::async, [&] { + return evaluate_int_udf_batch(cached_client, inputs, &results); + }); + + ASSERT_TRUE(cached_status.get().ok()); + ASSERT_FALSE(fs::exists(watchdog_expired)) + << "cached function-body imports waited for the process-wide import lock"; + ASSERT_TRUE(fs::exists(cached_batch_completed)); + ASSERT_EQ(results.size(), inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + EXPECT_EQ(results[i], inputs[i] + 3); + } + ASSERT_TRUE(blocking_status.get().ok()); + EXPECT_EQ(blocking_result, 10); + static_cast(cached_client->close()); + mgr.shutdown(); +} + +TEST_F(PythonServerTest, ModuleUdfChildThreadsUseModuleContext) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + fs::path module_dir = fs::path(test_dir_) / "module_thread_context"; + fs::path package_dir = module_dir / "thread_context_pkg"; + fs::path direct_import_completed = fs::path(test_dir_) / "direct_import_completed"; + fs::path runtime_import_completed = fs::path(test_dir_) / "runtime_import_completed"; + fs::path callable_import_completed = fs::path(test_dir_) / "callable_import_completed"; + fs::create_directories(package_dir); + std::ofstream(package_dir / "__init__.py").close(); + { + std::ofstream dependency(package_dir / "builtin_dependency.py"); + dependency << "OFFSET = 2\n"; + } + { + std::ofstream dependency(package_dir / "importlib_dependency.py"); + dependency << "OFFSET = 3\n"; + } + { + std::ofstream dependency(package_dir / "nested_dependency.py"); + dependency << "OFFSET = 5\n"; + } + { + std::ofstream dependency(package_dir / "initial_dependency.py"); + dependency << "OFFSET = 7\n"; + } + { + std::ofstream dependency(package_dir / "direct_dependency.py"); + dependency << "import pathlib\n" + << "pathlib.Path(r'" << direct_import_completed.string() << "').touch()\n" + << "OFFSET = 11\n"; + } + { + std::ofstream dependency(package_dir / "runtime_dependency.py"); + dependency << "import pathlib\n" + << "pathlib.Path(r'" << runtime_import_completed.string() << "').touch()\n" + << "OFFSET = 13\n"; + } + { + std::ofstream dependency(package_dir / "callable_dependency.py"); + dependency << "import pathlib\n" + << "pathlib.Path(r'" << callable_import_completed.string() << "').touch()\n" + << "OFFSET = 17\n"; + } + { + std::ofstream dependency(package_dir / "parent_dependency.py"); + dependency << "import threading\n" + << "results = []\n" + << "def import_from_child_during_parent_import():\n" + << " from .nested_dependency import OFFSET\n" + << " results.append(OFFSET)\n" + << "thread = threading.Thread(target=import_from_child_during_parent_import)\n" + << "thread.start()\n" + << "thread.join(5)\n" + << "if thread.is_alive():\n" + << " raise RuntimeError('nested child import waited for parent import')\n" + << "OFFSET = results[0]\n"; + } + { + std::ofstream module(package_dir / "entry.py"); + module << "import importlib\n" + << "import pathlib\n" + << "import threading\n" + << "initial_results = []\n" + << "def import_during_initialization():\n" + << " from .initial_dependency import OFFSET\n" + << " initial_results.append(OFFSET)\n" + << "initial_thread = threading.Thread(target=import_during_initialization)\n" + << "initial_thread.start()\n" + << "initial_thread.join(5)\n" + << "if initial_thread.is_alive():\n" + << " raise RuntimeError('child import waited for initial module import')\n" + << "INITIAL_OFFSET = initial_results[0]\n" + << "direct_thread = threading.Thread(\n" + << " target=importlib.import_module,\n" + << " args=('.direct_dependency', __package__))\n" + << "direct_thread.start()\n" + << "direct_thread.join(5)\n" + << "if direct_thread.is_alive():\n" + << " raise RuntimeError('direct child import waited for initial module import')\n" + << "if not pathlib.Path(r'" << direct_import_completed.string() << "').exists():\n" + << " raise RuntimeError('direct child import did not run')\n" + << "from .direct_dependency import OFFSET as DIRECT_OFFSET\n" + << "def evaluate(value):\n" + << " runtime_thread = threading.Thread(\n" + << " target=importlib.import_module,\n" + << " args=('.runtime_dependency', __package__))\n" + << " runtime_thread.start()\n" + << " runtime_thread.join(5)\n" + << " if runtime_thread.is_alive():\n" + << " raise RuntimeError('direct runtime import did not finish')\n" + << " if not pathlib.Path(r'" << runtime_import_completed.string() + << "').exists():\n" + << " raise RuntimeError('direct runtime import did not run')\n" + << " from .runtime_dependency import OFFSET as RUNTIME_OFFSET\n" + << " results = []\n" + << " errors = []\n" + << " class ImportTarget:\n" + << " def __eq__(self, other):\n" + << " raise RuntimeError('thread target equality must not run')\n" + << " def __call__(self):\n" + << " try:\n" + << " dependency = importlib.import_module(\n" + << " '.callable_dependency', __package__)\n" + << " results.append(dependency.OFFSET)\n" + << " except Exception as exc:\n" + << " errors.append(str(exc))\n" + << " callable_thread = threading.Thread(target=ImportTarget())\n" + << " callable_thread.start()\n" + << " callable_thread.join(5)\n" + << " def import_with_statement():\n" + << " try:\n" + << " from .builtin_dependency import OFFSET\n" + << " results.append(OFFSET)\n" + << " except Exception as error:\n" + << " errors.append(str(error))\n" + << " def import_with_importlib():\n" + << " try:\n" + << " dependency = importlib.import_module(\n" + << " '.importlib_dependency', __package__)\n" + << " results.append(dependency.OFFSET)\n" + << " except Exception as error:\n" + << " errors.append(str(error))\n" + << " for target in (import_with_statement, import_with_importlib):\n" + << " thread = threading.Thread(target=target)\n" + << " thread.start()\n" + << " thread.join()\n" + << " if errors:\n" + << " raise RuntimeError('; '.join(errors))\n" + << " from .parent_dependency import OFFSET\n" + << " return (value + INITIAL_OFFSET + DIRECT_OFFSET + RUNTIME_OFFSET\n" + << " + sum(results) + OFFSET)\n"; + } + + PythonVersion version("test-runtime", fs::path(*python).parent_path().parent_path().string(), + *python); + PythonServerManager mgr; + ProcessPtr process; + Status fork_status = start_python_udf_server(*python, &process); + ASSERT_TRUE(fork_status.ok()) << fork_status.to_string(); + ASSERT_NE(process, nullptr); + mgr.set_process_pool_for_test(version, {process}); + + PythonUDFMeta meta; + meta.id = 1; + meta.name = "thread_context_udf"; + meta.symbol = "thread_context_pkg.entry.evaluate"; + meta.location = module_dir.string(); + meta.checksum = "test-checksum"; + meta.runtime_version = version.full_version; + meta.input_types = {std::make_shared()}; + meta.return_type = std::make_shared(); + meta.type = PythonUDFLoadType::MODULE; + meta.client_type = PythonClientType::UDF; + + int32_t result = 0; + Status evaluate_status = evaluate_int_module_udf(meta, process, 10, &result); + ASSERT_TRUE(evaluate_status.ok()) << evaluate_status.to_string(); + EXPECT_EQ(result, 68); + mgr.shutdown(); +} + +TEST_F(PythonServerTest, LongLivedWorkerUsesCallingModuleContext) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + fs::path global_module_dir = fs::path(test_dir_) / "global_thread_pool"; + fs::path first_module_dir = fs::path(test_dir_) / "thread_pool_udf_a"; + fs::path second_module_dir = fs::path(test_dir_) / "thread_pool_udf_b"; + fs::create_directories(global_module_dir); + fs::create_directories(first_module_dir / "shared_thread_pkg"); + fs::create_directories(second_module_dir / "shared_thread_pkg"); + { + std::ofstream pool(global_module_dir / "shared_thread_pool.py"); + pool << "from concurrent.futures import ThreadPoolExecutor\n" + << "executor = ThreadPoolExecutor(max_workers=1)\n"; + } + auto write_udf = [](const fs::path& module_dir, int offset) { + fs::path package_dir = module_dir / "shared_thread_pkg"; + std::ofstream(package_dir / "__init__.py").close(); + { + std::ofstream dependency(package_dir / "dependency.py"); + dependency << "OFFSET = " << offset << "\n"; + } + { + std::ofstream entry(package_dir / "entry.py"); + entry << "import shared_thread_pool\n" + << "def import_dependency(value):\n" + << " from .dependency import OFFSET\n" + << " return value + OFFSET\n" + << "def evaluate(value):\n" + << " return shared_thread_pool.executor.submit(\n" + << " import_dependency, value).result(5)\n"; + } + }; + write_udf(first_module_dir, 1); + write_udf(second_module_dir, 100); + + const char* original_python_path = std::getenv("PYTHONPATH"); + std::optional saved_python_path; + if (original_python_path) { + saved_python_path = original_python_path; + } + setenv("PYTHONPATH", global_module_dir.c_str(), 1); + + PythonVersion version("test-runtime", fs::path(*python).parent_path().parent_path().string(), + *python); + PythonServerManager mgr; + ProcessPtr process; + Status fork_status = start_python_udf_server(*python, &process); + if (saved_python_path) { + setenv("PYTHONPATH", saved_python_path->c_str(), 1); + } else { + unsetenv("PYTHONPATH"); + } + ASSERT_TRUE(fork_status.ok()) << fork_status.to_string(); + ASSERT_NE(process, nullptr); + mgr.set_process_pool_for_test(version, {process}); + + int32_t first_result = 0; + int32_t second_result = 0; + ASSERT_TRUE(evaluate_int_module_udf(make_int_module_meta(1, first_module_dir, + "shared_thread_pkg.entry.evaluate"), + process, 10, &first_result) + .ok()); + ASSERT_TRUE(evaluate_int_module_udf(make_int_module_meta(2, second_module_dir, + "shared_thread_pkg.entry.evaluate"), + process, 10, &second_result) + .ok()); + EXPECT_EQ(first_result, 11); + EXPECT_EQ(second_result, 110); + mgr.shutdown(); +} + +TEST_F(PythonServerTest, ClearModuleCacheWaitsForModuleImport) { + auto python = find_python_udf_interpreter(); + if (!python) { + GTEST_SKIP() << "Python with pandas and pyarrow is required"; + } + + ASSERT_TRUE(install_real_python_server().ok()); + + fs::path module_dir = fs::path(test_dir_) / "module_clear_during_import"; + fs::create_directories(module_dir); + fs::path import_started = fs::path(test_dir_) / "clear_import_started"; + fs::path allow_import = fs::path(test_dir_) / "allow_import"; + auto write_module = [&](int offset, bool wait_for_release) { + std::ofstream module(module_dir / "clear_during_import_udf.py", std::ios::trunc); + if (wait_for_release) { + module << "import pathlib\n" + << "import time\n" + << "started = pathlib.Path(r'" << import_started.string() << "')\n" + << "release = pathlib.Path(r'" << allow_import.string() << "')\n" + << "deadline = time.monotonic() + 10\n" + << "started.touch()\n" + << "while not release.exists() and time.monotonic() < deadline:\n" + << " time.sleep(0.01)\n" + << "if not release.exists():\n" + << " raise RuntimeError('timed out waiting to finish import')\n"; + } + module << "def evaluate(value):\n" + << " return value + " << offset << "\n"; + }; + write_module(1, true); + + PythonVersion version("test-runtime", fs::path(*python).parent_path().parent_path().string(), + *python); + PythonServerManager mgr; + ProcessPtr process; + Status fork_status = start_python_udf_server(*python, &process); + ASSERT_TRUE(fork_status.ok()) << fork_status.to_string(); + ASSERT_NE(process, nullptr); + mgr.set_process_pool_for_test(version, {process}); + + PythonUDFMeta meta; + meta.id = 1; + meta.name = "clear_during_import_udf"; + meta.symbol = "clear_during_import_udf.evaluate"; + meta.location = module_dir.string(); + meta.checksum = "test-checksum"; + meta.runtime_version = version.full_version; + meta.input_types = {std::make_shared()}; + meta.return_type = std::make_shared(); + meta.type = PythonUDFLoadType::MODULE; + meta.client_type = PythonClientType::UDF; + + int32_t initial_result = 0; + auto evaluate_status = std::async(std::launch::async, [&] { + return evaluate_int_module_udf(meta, process, 10, &initial_result); + }); + Defer release_import {[&] { std::ofstream(allow_import).close(); }}; + auto marker_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!fs::exists(import_started) && std::chrono::steady_clock::now() < marker_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_TRUE(fs::exists(import_started)); + + auto clear_status = + std::async(std::launch::async, [&] { return mgr.clear_module_cache(meta.location); }); + EXPECT_EQ(clear_status.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); + std::ofstream(allow_import).close(); + + ASSERT_TRUE(evaluate_status.get().ok()); + ASSERT_EQ(initial_result, 11); + auto clear_result = clear_status.get(); + ASSERT_TRUE(clear_result.ok()) << clear_result.to_string(); + + write_module(100, false); + int32_t reloaded_result = 0; + ASSERT_TRUE(evaluate_int_module_udf(meta, process, 10, &reloaded_result).ok()); + EXPECT_EQ(reloaded_result, 110); + mgr.shutdown(); +} + TEST_F(PythonServerTest, BroadcastActionWithInvalidProcessUriReturnsError) { PythonServerManager mgr; PythonVersion version("3.9.16", test_dir_, test_dir_ + "/bin/python3"); @@ -519,6 +1754,34 @@ TEST_F(PythonServerTest, BroadcastActionWithInvalidProcessUriReturnsError) { mgr.shutdown(); } +TEST_F(PythonServerTest, BroadcastActionReportsFailedFlightResults) { + ActionResultFlightServer success_server({R"({"success": true})"}); + ASSERT_TRUE(success_server.start().ok()); + ActionResultFlightServer failed_server( + {R"({"success": false, "error": "cache clear failed"})"}); + ASSERT_TRUE(failed_server.start().ok()); + + PythonServerManager mgr; + PythonVersion version("3.9.16", test_dir_, test_dir_ + "/bin/python3"); + ProcessPtr success_process = create_sleep_process(); + ASSERT_NE(success_process, nullptr); + ASSERT_TRUE(success_process->is_alive()); + success_process->set_uri_for_test(success_server.location().ToString()); + ProcessPtr failed_process = create_sleep_process(); + ASSERT_NE(failed_process, nullptr); + ASSERT_TRUE(failed_process->is_alive()); + failed_process->set_uri_for_test(failed_server.location().ToString()); + + mgr.set_process_pool_for_test(version, {success_process, failed_process}); + auto status = mgr.clear_module_cache("/tmp/test_udf"); + + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("success=1, failed=1"), std::string::npos); + EXPECT_NE(status.to_string().find("cache clear failed"), std::string::npos); + + mgr.shutdown(); +} + // ============================================================================ // PythonServerManager::get_client() - client retrieval test // ============================================================================ diff --git a/regression-test/data/pythonudaf_p0/test_pythonudaf_pkg_isolation.out b/regression-test/data/pythonudaf_p0/test_pythonudaf_pkg_isolation.out index a55082ac7f931a..d2960c05aec1be 100644 --- a/regression-test/data/pythonudaf_p0/test_pythonudaf_pkg_isolation.out +++ b/regression-test/data/pythonudaf_p0/test_pythonudaf_pkg_isolation.out @@ -11,3 +11,6 @@ -- !pkg_isolation_4 -- 6 1006 2006 3006 +-- !pkg_isolation_5 -- +16 106 + diff --git a/regression-test/data/pythonudf_p0/test_pythonudf_pkg_isolation.out b/regression-test/data/pythonudf_p0/test_pythonudf_pkg_isolation.out index a782ca4cc03d2d..7bef78d41832d8 100644 --- a/regression-test/data/pythonudf_p0/test_pythonudf_pkg_isolation.out +++ b/regression-test/data/pythonudf_p0/test_pythonudf_pkg_isolation.out @@ -11,3 +11,6 @@ -- !pkg_isolation_4 -- 20 30 110 210 +-- !pkg_isolation_5 -- +15 105 + diff --git a/regression-test/data/pythonudtf_p0/test_pythonudtf_pkg_isolation.out b/regression-test/data/pythonudtf_p0/test_pythonudtf_pkg_isolation.out index 1b0866050441b1..ff217eafcdf122 100644 --- a/regression-test/data/pythonudtf_p0/test_pythonudtf_pkg_isolation.out +++ b/regression-test/data/pythonudtf_p0/test_pythonudtf_pkg_isolation.out @@ -15,3 +15,7 @@ 1 101 201 301 2 102 202 302 +-- !pkg_isolation_5 -- +11 101 +12 102 + diff --git a/regression-test/suites/pythonudaf_p0/test_pythonudaf_pkg_isolation.groovy b/regression-test/suites/pythonudaf_p0/test_pythonudaf_pkg_isolation.groovy index 8859f7694d6656..8ffda4057888bb 100644 --- a/regression-test/suites/pythonudaf_p0/test_pythonudaf_pkg_isolation.groovy +++ b/regression-test/suites/pythonudaf_p0/test_pythonudaf_pkg_isolation.groovy @@ -86,10 +86,34 @@ suite('test_pythonudaf_pkg_isolation') { // Case 4: All four combinations together qt_pkg_isolation_4 '''SELECT py_pkg_a_sum_x(v), py_pkg_a_sum_y(v), py_pkg_b_sum_x(v), py_pkg_b_sum_y(v) FROM py_udaf_pkg_tbl;''' + // Case 5: Different top-level modules import the same dependency name + sql '''DROP FUNCTION IF EXISTS py_pkg_a_dependency_sum(INT)''' + sql '''DROP FUNCTION IF EXISTS py_pkg_b_dependency_sum(INT)''' + sql """ + CREATE AGGREGATE FUNCTION py_pkg_a_dependency_sum(INT) RETURNS BIGINT PROPERTIES ( + "type" = "PYTHON_UDF", + "file" = "file://${zipA}", + "symbol" = "udaf_a_entry.SumAgg", + "runtime_version" = "${runtime_version}" + ) + """ + sql """ + CREATE AGGREGATE FUNCTION py_pkg_b_dependency_sum(INT) RETURNS BIGINT PROPERTIES ( + "type" = "PYTHON_UDF", + "file" = "file://${zipB}", + "symbol" = "udaf_b_entry.SumAgg", + "runtime_version" = "${runtime_version}" + ) + """ + + qt_pkg_isolation_5 '''SELECT py_pkg_a_dependency_sum(v), py_pkg_b_dependency_sum(v) FROM py_udaf_pkg_tbl;''' + } finally { try_sql('DROP FUNCTION IF EXISTS py_pkg_a_sum_x(INT);') try_sql('DROP FUNCTION IF EXISTS py_pkg_a_sum_y(INT);') try_sql('DROP FUNCTION IF EXISTS py_pkg_b_sum_x(INT);') try_sql('DROP FUNCTION IF EXISTS py_pkg_b_sum_y(INT);') + try_sql('DROP FUNCTION IF EXISTS py_pkg_a_dependency_sum(INT);') + try_sql('DROP FUNCTION IF EXISTS py_pkg_b_dependency_sum(INT);') } } diff --git a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/mypkg/lazy_state.py b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/mypkg/lazy_state.py new file mode 100644 index 00000000000000..ac2bc7ba51c622 --- /dev/null +++ b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/mypkg/lazy_state.py @@ -0,0 +1,29 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +MODULE_ID = "A" + + +def restore_state(value, expected_module_id): + if MODULE_ID != expected_module_id: + raise RuntimeError("UDAF state was restored by the wrong module") + + from shared_udaf_dependency import State + + state = State() + state.value = value + return state diff --git a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/python_udaf_pkg_test.zip b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/python_udaf_pkg_test.zip index 9996f6c72bbeda..2662d9d19a87b5 100644 Binary files a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/python_udaf_pkg_test.zip and b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/python_udaf_pkg_test.zip differ diff --git a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/shared_udaf_dependency.py b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/shared_udaf_dependency.py new file mode 100644 index 00000000000000..19f8c262fcafc1 --- /dev/null +++ b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/shared_udaf_dependency.py @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Load the namespace package so State.__reduce__ can lazily import its child. +import mypkg + +OFFSET = 10 +MODULE_ID = "A" + + +def add_value(total, value): + return total + value + + +def merge_value(total, other_value): + return total + other_value + + +class State: + def __init__(self): + self.value = 0 + + def __reduce__(self): + from mypkg.lazy_state import restore_state + + return restore_state, (self.value, MODULE_ID) diff --git a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/udaf_a_entry.py b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/udaf_a_entry.py new file mode 100644 index 00000000000000..ad6488f10e366d --- /dev/null +++ b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_a/udaf_a_entry.py @@ -0,0 +1,48 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +class SumAgg: + def __init__(self): + from shared_udaf_dependency import State + + self.state = State() + + def init(self): + from shared_udaf_dependency import State + + self.state = State() + + @property + def aggregate_state(self): + return self.state + + def accumulate(self, value): + from shared_udaf_dependency import add_value + + if value is not None: + self.state.value = add_value(self.state.value, value) + + def merge(self, other_state): + from shared_udaf_dependency import merge_value + + if other_state is not None: + self.state.value = merge_value(self.state.value, other_state.value) + + def finish(self): + from shared_udaf_dependency import OFFSET + + return self.state.value + OFFSET diff --git a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/mypkg/lazy_state.py b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/mypkg/lazy_state.py new file mode 100644 index 00000000000000..ae5777e62afdf3 --- /dev/null +++ b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/mypkg/lazy_state.py @@ -0,0 +1,29 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +MODULE_ID = "B" + + +def restore_state(value, expected_module_id): + if MODULE_ID != expected_module_id: + raise RuntimeError("UDAF state was restored by the wrong module") + + from shared_udaf_dependency import State + + state = State() + state.value = value + return state diff --git a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/python_udaf_pkg_test.zip b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/python_udaf_pkg_test.zip index 3a7341fade02cc..d9668ec29e60a5 100644 Binary files a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/python_udaf_pkg_test.zip and b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/python_udaf_pkg_test.zip differ diff --git a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/shared_udaf_dependency.py b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/shared_udaf_dependency.py new file mode 100644 index 00000000000000..45e0ee456b4d67 --- /dev/null +++ b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/shared_udaf_dependency.py @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Load the namespace package so State.__reduce__ can lazily import its child. +import mypkg + +OFFSET = 100 +MODULE_ID = "B" + + +def add_value(total, value): + return total + value + + +def merge_value(total, other_value): + return total + other_value + + +class State: + def __init__(self): + self.value = 0 + + def __reduce__(self): + from mypkg.lazy_state import restore_state + + return restore_state, (self.value, MODULE_ID) diff --git a/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/udaf_b_entry.py b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/udaf_b_entry.py new file mode 100644 index 00000000000000..ad6488f10e366d --- /dev/null +++ b/regression-test/suites/pythonudaf_p0/udaf_scripts/python_udaf_pkg_b/udaf_b_entry.py @@ -0,0 +1,48 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +class SumAgg: + def __init__(self): + from shared_udaf_dependency import State + + self.state = State() + + def init(self): + from shared_udaf_dependency import State + + self.state = State() + + @property + def aggregate_state(self): + return self.state + + def accumulate(self, value): + from shared_udaf_dependency import add_value + + if value is not None: + self.state.value = add_value(self.state.value, value) + + def merge(self, other_state): + from shared_udaf_dependency import merge_value + + if other_state is not None: + self.state.value = merge_value(self.state.value, other_state.value) + + def finish(self): + from shared_udaf_dependency import OFFSET + + return self.state.value + OFFSET diff --git a/regression-test/suites/pythonudf_p0/test_pythonudf_pkg_isolation.groovy b/regression-test/suites/pythonudf_p0/test_pythonudf_pkg_isolation.groovy index 9be912498936b2..6bdbed9147e37e 100644 --- a/regression-test/suites/pythonudf_p0/test_pythonudf_pkg_isolation.groovy +++ b/regression-test/suites/pythonudf_p0/test_pythonudf_pkg_isolation.groovy @@ -75,10 +75,34 @@ suite("test_pythonudf_pkg_isolation") { // Case 4: All four combinations together qt_pkg_isolation_4 """SELECT py_pkg_a_mod_x(10), py_pkg_a_mod_y(10), py_pkg_b_mod_x(10), py_pkg_b_mod_y(10);""" + // Case 5: Different top-level modules import the same dependency name + sql """DROP FUNCTION IF EXISTS py_pkg_a_dependency(INT)""" + sql """DROP FUNCTION IF EXISTS py_pkg_b_dependency(INT)""" + sql """ + CREATE FUNCTION py_pkg_a_dependency(INT) RETURNS INT PROPERTIES ( + "type" = "PYTHON_UDF", + "file" = "file://${zipA}", + "symbol" = "udf_a_entry.evaluate", + "runtime_version" = "${runtime_version}" + ) + """ + sql """ + CREATE FUNCTION py_pkg_b_dependency(INT) RETURNS INT PROPERTIES ( + "type" = "PYTHON_UDF", + "file" = "file://${zipB}", + "symbol" = "udf_b_entry.evaluate", + "runtime_version" = "${runtime_version}" + ) + """ + + qt_pkg_isolation_5 """SELECT py_pkg_a_dependency(5), py_pkg_b_dependency(5);""" + } finally { try_sql("DROP FUNCTION IF EXISTS py_pkg_a_mod_x(INT);") try_sql("DROP FUNCTION IF EXISTS py_pkg_a_mod_y(INT);") try_sql("DROP FUNCTION IF EXISTS py_pkg_b_mod_x(INT);") try_sql("DROP FUNCTION IF EXISTS py_pkg_b_mod_y(INT);") + try_sql("DROP FUNCTION IF EXISTS py_pkg_a_dependency(INT);") + try_sql("DROP FUNCTION IF EXISTS py_pkg_b_dependency(INT);") } } diff --git a/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/python_udf_pkg_test.zip b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/python_udf_pkg_test.zip index e76530f21e1d4a..a300a319e87f55 100644 Binary files a/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/python_udf_pkg_test.zip and b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/python_udf_pkg_test.zip differ diff --git a/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/shared_udf_dependency.py b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/shared_udf_dependency.py new file mode 100644 index 00000000000000..ed714a69bef599 --- /dev/null +++ b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/shared_udf_dependency.py @@ -0,0 +1,18 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +OFFSET = 10 diff --git a/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/udf_a_entry.py b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/udf_a_entry.py new file mode 100644 index 00000000000000..bd4552aff4c1b5 --- /dev/null +++ b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_a/udf_a_entry.py @@ -0,0 +1,23 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +def evaluate(value): + from shared_udf_dependency import OFFSET + + if value is None: + return None + return value + OFFSET diff --git a/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/python_udf_pkg_test.zip b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/python_udf_pkg_test.zip index c1ff0ab123ebe7..f2b89f04b18b68 100644 Binary files a/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/python_udf_pkg_test.zip and b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/python_udf_pkg_test.zip differ diff --git a/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/shared_udf_dependency.py b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/shared_udf_dependency.py new file mode 100644 index 00000000000000..1e15ccd9fbe06c --- /dev/null +++ b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/shared_udf_dependency.py @@ -0,0 +1,18 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +OFFSET = 100 diff --git a/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/udf_b_entry.py b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/udf_b_entry.py new file mode 100644 index 00000000000000..bd4552aff4c1b5 --- /dev/null +++ b/regression-test/suites/pythonudf_p0/udf_scripts/python_udf_pkg_b/udf_b_entry.py @@ -0,0 +1,23 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +def evaluate(value): + from shared_udf_dependency import OFFSET + + if value is None: + return None + return value + OFFSET diff --git a/regression-test/suites/pythonudtf_p0/test_pythonudtf_pkg_isolation.groovy b/regression-test/suites/pythonudtf_p0/test_pythonudtf_pkg_isolation.groovy index d877d69a096d70..a7e06edb61e483 100644 --- a/regression-test/suites/pythonudtf_p0/test_pythonudtf_pkg_isolation.groovy +++ b/regression-test/suites/pythonudtf_p0/test_pythonudtf_pkg_isolation.groovy @@ -120,10 +120,44 @@ suite("test_pythonudtf_pkg_isolation") { ORDER BY ax.c, ay.c, bx.c, b_y.c; """ + // Case 5: Different top-level modules import the same dependency name + sql """DROP FUNCTION IF EXISTS py_pkg_a_dependency_t(INT)""" + sql """DROP FUNCTION IF EXISTS py_pkg_b_dependency_t(INT)""" + sql """ + CREATE TABLES FUNCTION py_pkg_a_dependency_t(INT) + RETURNS ARRAY + PROPERTIES ( + "type" = "PYTHON_UDF", + "file" = "file://${zipA}", + "symbol" = "udtf_a_entry.process", + "runtime_version" = "${runtime_version}" + ) + """ + sql """ + CREATE TABLES FUNCTION py_pkg_b_dependency_t(INT) + RETURNS ARRAY + PROPERTIES ( + "type" = "PYTHON_UDF", + "file" = "file://${zipB}", + "symbol" = "udtf_b_entry.process", + "runtime_version" = "${runtime_version}" + ) + """ + + qt_pkg_isolation_5 """ + SELECT a.c, b.c + FROM py_udtf_pkg_tbl + LATERAL VIEW py_pkg_a_dependency_t(v) a AS c + LATERAL VIEW py_pkg_b_dependency_t(v) b AS c + ORDER BY a.c, b.c; + """ + } finally { try_sql("DROP FUNCTION IF EXISTS py_pkg_a_t_x(INT);") try_sql("DROP FUNCTION IF EXISTS py_pkg_a_t_y(INT);") try_sql("DROP FUNCTION IF EXISTS py_pkg_b_t_x(INT);") try_sql("DROP FUNCTION IF EXISTS py_pkg_b_t_y(INT);") + try_sql("DROP FUNCTION IF EXISTS py_pkg_a_dependency_t(INT);") + try_sql("DROP FUNCTION IF EXISTS py_pkg_b_dependency_t(INT);") } } diff --git a/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/python_udtf_pkg_test.zip b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/python_udtf_pkg_test.zip index 0828cd94cb2512..78eac5a3695ffe 100644 Binary files a/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/python_udtf_pkg_test.zip and b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/python_udtf_pkg_test.zip differ diff --git a/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/shared_udtf_dependency.py b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/shared_udtf_dependency.py new file mode 100644 index 00000000000000..ed714a69bef599 --- /dev/null +++ b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/shared_udtf_dependency.py @@ -0,0 +1,18 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +OFFSET = 10 diff --git a/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/udtf_a_entry.py b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/udtf_a_entry.py new file mode 100644 index 00000000000000..e7bcee8b77e47c --- /dev/null +++ b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_a/udtf_a_entry.py @@ -0,0 +1,22 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +def process(value): + from shared_udtf_dependency import OFFSET + + if value is not None: + yield value + OFFSET diff --git a/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/python_udtf_pkg_test.zip b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/python_udtf_pkg_test.zip index 5f937b546f2315..d015cb60621938 100644 Binary files a/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/python_udtf_pkg_test.zip and b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/python_udtf_pkg_test.zip differ diff --git a/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/shared_udtf_dependency.py b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/shared_udtf_dependency.py new file mode 100644 index 00000000000000..1e15ccd9fbe06c --- /dev/null +++ b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/shared_udtf_dependency.py @@ -0,0 +1,18 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +OFFSET = 100 diff --git a/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/udtf_b_entry.py b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/udtf_b_entry.py new file mode 100644 index 00000000000000..e7bcee8b77e47c --- /dev/null +++ b/regression-test/suites/pythonudtf_p0/udtf_scripts/python_udtf_pkg_b/udtf_b_entry.py @@ -0,0 +1,22 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +def process(value): + from shared_udtf_dependency import OFFSET + + if value is not None: + yield value + OFFSET