[Enhancement](udf) clear cache when droping function

linrrzqqq · linrrzqqq · commit 06484c7c12fd · 2026-02-11T02:32:36.000+08:00
diff --git a/be/src/agent/task_worker_pool.cpp b/be/src/agent/task_worker_pool.cpp
@@ -88,6 +88,7 @@
 #include "runtime/index_policy/index_policy_mgr.h"
 #include "runtime/memory/global_memory_arbitrator.h"
 #include "runtime/snapshot_loader.h"
+#include "runtime/user_function_cache.h"
 #include "service/backend_options.h"
 #include "util/brpc_client_cache.h"
 #include "util/debug_points.h"
@@ -2396,12 +2397,17 @@ void clean_trash_callback(StorageEngine& engine, const TAgentTaskRequest& req) {
 }
 
 void clean_udf_cache_callback(const TAgentTaskRequest& req) {
+    const auto& clean_req = req.clean_udf_cache_req;
+
     if (doris::config::enable_java_support) {
-        LOG(INFO) << "clean udf cache start: " << req.clean_udf_cache_req.function_signature;
-        static_cast<void>(
-                Jni::Util::clean_udf_class_load_cache(req.clean_udf_cache_req.function_signature));
-        LOG(INFO) << "clean udf cache  finish: " << req.clean_udf_cache_req.function_signature;
+        static_cast<void>(Jni::Util::clean_udf_class_load_cache(clean_req.function_signature));
+    }
+
+    if (clean_req.__isset.function_id && clean_req.function_id > 0) {
+        UserFunctionCache::instance()->drop_function_cache(clean_req.function_id);
     }
+
+    LOG(INFO) << "clean udf cache finish: function_signature=" << clean_req.function_signature;
 }
 
 void report_index_policy_callback(const ClusterInfo* cluster_info) {
diff --git a/be/src/runtime/user_function_cache.cpp b/be/src/runtime/user_function_cache.cpp
@@ -46,6 +46,7 @@
 #include "util/dynamic_util.h"
 #include "util/md5.h"
 #include "util/string_util.h"
+#include "udf/python/python_server.h"
 
 namespace doris {
 
@@ -111,7 +112,15 @@ UserFunctionCacheEntry::~UserFunctionCacheEntry() {
 
     // delete library file if should_delete_library is set
     if (should_delete_library.load()) {
-        unlink(lib_file.c_str());
+        if (type == LibType::PY_ZIP) {
+            // For Python UDF, we need to delete both the unzipped directory and the original zip file.
+            auto st = io::global_local_filesystem()->delete_directory_or_file(lib_file);
+
+            st = io::global_local_filesystem()->delete_file(lib_file + ".zip");
+
+        } else {
+            unlink(lib_file.c_str());
+        }
     }
 }
 
@@ -174,10 +183,20 @@ Status UserFunctionCache::_load_entry_from_lib(const std::string& dir, const std
                      << ", other_checksum info: = " << it->second->debug_string();
         return Status::InternalError("duplicate function id");
     }
+
+    std::string full_path = dir + "/" + file;
     // create a cache entry and put it into entry map
-    std::shared_ptr<UserFunctionCacheEntry> entry = UserFunctionCacheEntry::create_shared(
-            function_id, checksum, dir + "/" + file, lib_type);
+    std::shared_ptr<UserFunctionCacheEntry> entry =
+            UserFunctionCacheEntry::create_shared(function_id, checksum, full_path, lib_type);
     entry->is_downloaded = true;
+
+    // For Python UDF, _check_cache_is_python_udf has already unzipped the file.
+    // Set lib_file to the unzipped directory.
+    if (lib_type == LibType::PY_ZIP) {
+        entry->lib_file = full_path.substr(0, full_path.size() - 4);
+        entry->is_unziped = true;
+    }
+
     _entry_map[function_id] = entry;
 
     return Status::OK();
@@ -547,4 +566,29 @@ Status UserFunctionCache::_check_and_return_default_java_udf_url(const std::stri
     return Status::OK();
 }
 
+void UserFunctionCache::drop_function_cache(int64_t fid) {
+    std::shared_ptr<UserFunctionCacheEntry> entry = nullptr;
+    {
+        std::lock_guard<std::mutex> l(_cache_lock);
+        auto it = _entry_map.find(fid);
+        if (it == _entry_map.end()) {
+            return;
+        }
+        entry = it->second;
+        _entry_map.erase(it);
+    }
+
+    // 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();
+        }
+    }
+
+    // Mark for deletion, destructor will delete the files
+    entry->should_delete_library.store(true);
+}
+
 } // namespace doris
diff --git a/be/src/runtime/user_function_cache.h b/be/src/runtime/user_function_cache.h
@@ -62,6 +62,9 @@ class UserFunctionCache {
     Status get_pypath(int64_t fid, const std::string& url, const std::string& checksum,
                       std::string* libpath);
 
+    // Drop the cached function library by function id.
+    void drop_function_cache(int64_t fid);
+
 #ifndef BE_TEST
 private:
 #endif
diff --git a/be/src/udf/python/python_server.cpp b/be/src/udf/python/python_server.cpp
@@ -27,6 +27,7 @@
 #include <boost/asio.hpp>
 #include <boost/process.hpp>
 
+#include "arrow/flight/client.h"
 #include "common/config.h"
 #include "udf/python/python_udaf_client.h"
 #include "udf/python/python_udf_client.h"
@@ -273,6 +274,78 @@ void PythonServerManager::shutdown() {
     _process_pools.clear();
 }
 
+Status PythonServerManager::clear_module_cache(const std::string& location) {
+    if (location.empty()) {
+        return Status::InvalidArgument("Empty location for clear_module_cache");
+    }
+
+    std::lock_guard<std::mutex> lock(_pools_mutex);
+
+    std::string body = fmt::format(R"({{"location": "{}"}})", location);
+
+    int success_count = 0;
+    int fail_count = 0;
+    bool has_active_process = false;
+
+    for (auto& [version, pool] : _process_pools) {
+        for (auto& process : pool) {
+            if (!process || !process->is_alive()) {
+                continue;
+            }
+            has_active_process = true;
+            try {
+                auto loc_result = arrow::flight::Location::Parse(process->get_uri());
+                if (!loc_result.ok()) [[unlikely]] {
+                    fail_count++;
+                    continue;
+                }
+
+                auto client_result = arrow::flight::FlightClient::Connect(*loc_result);
+                if (!client_result.ok()) [[unlikely]] {
+                    fail_count++;
+                    continue;
+                }
+                auto client = std::move(*client_result);
+
+                arrow::flight::Action action;
+                action.type = "clear_module_cache";
+                action.body = arrow::Buffer::FromString(body);
+
+                auto result_stream = client->DoAction(action);
+                if (!result_stream.ok()) {
+                    fail_count++;
+                    continue;
+                }
+
+                auto result = (*result_stream)->Next();
+                if (result.ok() && *result) {
+                    success_count++;
+                } else {
+                    fail_count++;
+                }
+
+            } catch (...) {
+                fail_count++;
+            }
+        }
+    }
+
+    if (!has_active_process) {
+        return Status::OK();
+    }
+
+    LOG(INFO) << "clear_module_cache completed for location=" << location
+              << ", success=" << success_count << ", failed=" << fail_count;
+
+    if (fail_count > 0) {
+        return Status::InternalError(
+                "clear_module_cache failed for location={}, success={}, failed={}",
+                location, success_count, fail_count);
+    }
+
+    return Status::OK();
+}
+
 // Explicit template instantiation for UDF, UDAF and UDTF clients
 template Status PythonServerManager::get_client<PythonUDFClient>(
         const PythonUDFMeta& func_meta, const PythonVersion& version,
diff --git a/be/src/udf/python/python_server.h b/be/src/udf/python/python_server.h
@@ -48,6 +48,9 @@ class PythonServerManager {
 
     Status get_process(const PythonVersion& version, ProcessPtr* process);
 
+    // Clear Python module cache for a specific UDF location across all processes
+    Status clear_module_cache(const std::string& location);
+
     Status ensure_pool_initialized(const PythonVersion& version);
 
     void shutdown();
diff --git a/be/src/udf/python/python_server.py b/be/src/udf/python/python_server.py
@@ -2442,6 +2442,97 @@ def do_exchange(
         else:
             raise ValueError(f"Unsupported client type: {python_udf_meta.client_type}")
 
+    def do_action(
+        self,
+        context: flight.ServerCallContext,
+        action: flight.Action,
+    ):
+        """
+        Handle Flight actions for cache management.
+
+        Supported actions:
+        - "clear_module_cache": Clear Python module cache for a specific location
+          Body: JSON with "location" field (the UDF cache directory path)
+        """
+        action_type = action.type
+
+        if action_type == "clear_module_cache":
+            yield from self._handle_clear_module_cache(action.body.to_pybytes())
+        else:
+            raise flight.FlightUnavailableError(f"Unknown action: {action_type}")
+
+    def _handle_clear_module_cache(self, body: bytes):
+        """
+        Clear Python module cache for a specific UDF location.
+
+        This removes modules from sys.modules that were loaded from the specified
+        location, allowing fresh imports when a new UDF with the same module name
+        is created.
+        """
+        try:
+            params = json.loads(body.decode("utf-8"))
+            location = params.get("location", "")
+
+            if not location:
+                yield flight.Result(b'{"success": false, "error": "empty location"}')
+                return
+
+            cleared_modules = self._clear_modules_from_location(location)
+
+            result = {
+                "success": True,
+                "cleared_modules": cleared_modules,
+                "location": location,
+            }
+            yield flight.Result(json.dumps(result).encode("utf-8"))
+
+        except Exception as e:
+            logging.error("clear_module_cache failed: %s", e)
+            yield flight.Result(json.dumps({
+                "success": False,
+                "error": str(e)
+            }).encode("utf-8"))
+
+    def _clear_modules_from_location(self, location: str) -> list:
+        """
+        Remove all modules from sys.modules that were loaded from the given location.
+
+        Returns list of cleared module names.
+        """
+        cleared = []
+        modules_to_remove = []
+
+        # Find all modules loaded from this location
+        for name, module in list(sys.modules.items()):
+            if module is None:
+                continue
+
+            # Check if module was loaded from this location
+            module_file = getattr(module, "__file__", None)
+            module_path = getattr(module, "__path__", None)
+
+            if module_file and module_file.startswith(location):
+                modules_to_remove.append(name)
+            elif module_path:
+                # For packages, check __path__
+                for p in module_path:
+                    if p.startswith(location):
+                        modules_to_remove.append(name)
+                        break
+
+        # Remove modules (do this separately to avoid modifying dict during iteration)
+        for name in modules_to_remove:
+            del sys.modules[name]
+            cleared.append(name)
+
+        # Also clear any import locks for these modules
+        with ModuleUDFLoader._import_locks_lock:
+            for name in cleared:
+                if name in ModuleUDFLoader._import_locks:
+                    del ModuleUDFLoader._import_locks[name]
+
+        return cleared
+
 
 class UDAFOperationType(Enum):
     """Enum representing UDAF operation types."""
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropFunctionCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropFunctionCommand.java
@@ -22,7 +22,9 @@
 import org.apache.doris.analysis.StmtType;
 import org.apache.doris.catalog.Database;
 import org.apache.doris.catalog.Env;
+import org.apache.doris.catalog.Function;
 import org.apache.doris.catalog.FunctionSearchDesc;
+import org.apache.doris.common.AnalysisException;
 import org.apache.doris.common.ErrorCode;
 import org.apache.doris.common.ErrorReport;
 import org.apache.doris.mysql.privilege.PrivPredicate;
@@ -71,6 +73,34 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
         }
         argsDef.analyze();
         FunctionSearchDesc function = new FunctionSearchDesc(functionName, argsDef.getArgTypes(), argsDef.isVariadic());
+
+        // Get function id before dropping, for cleaning cached library files in BE
+        long functionId = -1;
+        try {
+            Function fn = null;
+            if (SetType.GLOBAL.equals(setType)) {
+                fn = Env.getCurrentEnv().getGlobalFunctionMgr().getFunction(function);
+            } else {
+                String dbName = functionName.getDb();
+                if (dbName == null) {
+                    dbName = ctx.getDatabase();
+                    functionName.setDb(dbName);
+                }
+                Database db = Env.getCurrentInternalCatalog().getDbNullable(dbName);
+                if (db != null) {
+                    fn = db.getFunction(function);
+                }
+            }
+            if (fn != null) {
+                functionId = fn.getId();
+            } else {
+                LOG.warn("Function not found: {}, setType: {}", function.getName(), setType);
+            }
+        } catch (AnalysisException e) {
+            LOG.warn("Function not found when getting function id: {}, error: {}",
+                    function.getName(), e.getMessage());
+        }
+
         if (SetType.GLOBAL.equals(setType)) {
             Env.getCurrentEnv().getGlobalFunctionMgr().dropFunction(function, ifExists);
         } else {
@@ -90,9 +120,10 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
         String functionSignature = getSignatureString();
         AgentBatchTask batchTask = new AgentBatchTask();
         for (Backend backend : backendsInfo.values()) {
-            CleanUDFCacheTask cleanUDFCacheTask = new CleanUDFCacheTask(backend.getId(), functionSignature);
+            CleanUDFCacheTask cleanUDFCacheTask = new CleanUDFCacheTask(backend.getId(), functionSignature, functionId);
             batchTask.addTask(cleanUDFCacheTask);
-            LOG.info("clean udf cache in be {}, beId {}", backend.getHost(), backend.getId());
+            LOG.info("clean udf cache in be {}, beId {}, functionId {}",
+                    backend.getHost(), backend.getId(), functionId);
         }
         AgentTaskExecutor.submit(batchTask);
     }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/task/CleanUDFCacheTask.java b/fe/fe-core/src/main/java/org/apache/doris/task/CleanUDFCacheTask.java
@@ -27,15 +27,18 @@
 public class CleanUDFCacheTask extends AgentTask {
     private static final Logger LOG = LogManager.getLogger(CleanUDFCacheTask.class);
     private String functionSignature;
+    private long functionId;
 
-    public CleanUDFCacheTask(long backendId, String signature) {
+    public CleanUDFCacheTask(long backendId, String signature, long functionId) {
         super(null, backendId, TTaskType.CLEAN_UDF_CACHE, -1, -1, -1, -1, -1, -1, -1);
         this.functionSignature = signature;
+        this.functionId = functionId;
     }
 
     public TCleanUDFCacheReq toThrift() {
         TCleanUDFCacheReq req = new TCleanUDFCacheReq();
         req.setFunctionSignature(this.functionSignature);
+        req.setFunctionId(this.functionId);
         return req;
     }
 }
diff --git a/gensrc/thrift/AgentService.thrift b/gensrc/thrift/AgentService.thrift
@@ -162,6 +162,7 @@ struct TCleanTrashReq {}
 
 struct TCleanUDFCacheReq {
     1: optional string function_signature //function_name(arg_type)
+    2: optional i64 function_id // function id for cleaning cached library files
 }
 
 enum TCompressionType {

Original file line number	Diff line number	Diff line change
`@@ -27,15 +27,18 @@`
`27`	`27`	`public class CleanUDFCacheTask extends AgentTask {`
`28`	`28`	`private static final Logger LOG = LogManager.getLogger(CleanUDFCacheTask.class);`
`29`	`29`	`private String functionSignature;`
	`30`	`+ private long functionId;`
`30`	`31`
`31`		`- public CleanUDFCacheTask(long backendId, String signature) {`
	`32`	`+ public CleanUDFCacheTask(long backendId, String signature, long functionId) {`
`32`	`33`	`super(null, backendId, TTaskType.CLEAN_UDF_CACHE, -1, -1, -1, -1, -1, -1, -1);`
`33`	`34`	`this.functionSignature = signature;`
	`35`	`+ this.functionId = functionId;`
`34`	`36`	`}`
`35`	`37`
`36`	`38`	`public TCleanUDFCacheReq toThrift() {`
`37`	`39`	`TCleanUDFCacheReq req = new TCleanUDFCacheReq();`
`38`	`40`	`req.setFunctionSignature(this.functionSignature);`
	`41`	`+ req.setFunctionId(this.functionId);`
`39`	`42`	`return req;`
`40`	`43`	`}`
`41`	`44`	`}`
Original file line number	Diff line number	Diff line change
`@@ -162,6 +162,7 @@ struct TCleanTrashReq {}`
`162`	`162`
`163`	`163`	`struct TCleanUDFCacheReq {`
`164`	`164`	`1: optional string function_signature //function_name(arg_type)`
	`165`	`+ 2: optional i64 function_id // function id for cleaning cached library files`
`165`	`166`	`}`
`166`	`167`
`167`	`168`	`enum TCompressionType {`