[Enhancement](udf) clear cache when droping function

Author: linrrzqqq

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) {

be/src/runtime/user_function_cache.cpp

@@ -42,6 +42,7 @@
 #include "io/fs/local_file_system.h"
 #include "runtime/exec_env.h"
 #include "runtime/plugin/cloud_plugin_downloader.h"
+#include "udf/python/python_server.h"
 #include "util/defer_op.h"
 #include "util/dynamic_util.h"
 #include "util/md5.h"
@@ -111,7 +112,20 @@ 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");
+
+            if (!st.ok()) [[unlikely]] {
+                LOG(WARNING) << "failed to delete python udf files, lib_file=" << lib_file << ": "
+                             << st.to_string();
+            }
+
+        } else {
+            unlink(lib_file.c_str());
+        }
     }
 }
 
@@ -174,10 +188,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 +571,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

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

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,

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();