[Fix](ai) Fix _exec_plan_fragment_impl meet unknown error when call AI_Functions

Author: linrrzqqq

be/src/runtime/query_context.h

@@ -269,11 +269,12 @@ class QueryContext : public std::enable_shared_from_this<QueryContext> {
                 std::make_unique<std::map<std::string, TAIResource>>(std::move(ai_resources));
     }
 
-    const std::map<std::string, TAIResource>& get_ai_resources() const {
+    Status get_ai_resources(const std::map<std::string, TAIResource>** ai_resources) const {
         if (_ai_resources == nullptr) {
-            throw Status::InternalError("AI resources not found");
+            return Status::InternalError("AI resources not found");
         }
-        return *_ai_resources;
+        *ai_resources = _ai_resources.get();
+        return Status::OK();
     }
 
     std::unordered_map<TNetworkAddress, std::shared_ptr<PBackendService_Stub>>

be/src/vec/aggregate_functions/aggregate_function_ai_agg.h

@@ -141,9 +141,10 @@ class AggregateFunctionAIAggData {
             _task = task_ref.to_string();
 
             std::string resource_name = resource_name_ref.to_string();
-            const std::map<std::string, TAIResource>& ai_resources = _ctx->get_ai_resources();
-            auto it = ai_resources.find(resource_name);
-            if (it == ai_resources.end()) {
+            const std::map<std::string, TAIResource>* ai_resources = nullptr;
+            THROW_IF_ERROR(_ctx->get_ai_resources(&ai_resources));
+            auto it = ai_resources->find(resource_name);
+            if (it == ai_resources->end()) {
                 throw Exception(ErrorCode::NOT_FOUND, "AI resource not found: " + resource_name);
             }
             _ai_config = it->second;

be/src/vec/functions/ai/ai_functions.h

@@ -190,10 +190,10 @@ class AIFunction : public IFunction {
         StringRef resource_name_ref = resource_column.column->get_data_at(0);
         std::string resource_name = std::string(resource_name_ref.data, resource_name_ref.size);
 
-        const std::map<std::string, TAIResource>& ai_resources =
-                context->state()->get_query_ctx()->get_ai_resources();
-        auto it = ai_resources.find(resource_name);
-        if (it == ai_resources.end()) {
+        const std::map<std::string, TAIResource>* ai_resources = nullptr;
+        RETURN_IF_ERROR(context->state()->get_query_ctx()->get_ai_resources(&ai_resources));
+        auto it = ai_resources->find(resource_name);
+        if (it == ai_resources->end()) {
             return Status::InvalidArgument("AI resource not found: " + resource_name);
         }
         config = it->second;

be/test/ai/aggregate_function_ai_agg_test.cpp

@@ -413,4 +413,32 @@ TEST_F(AggregateFunctionAIAggTest, mock_resource_send_request_test) {
     _agg_function->destroy(place);
 }
 
+TEST_F(AggregateFunctionAIAggTest, missing_ai_resources_metadata_test) {
+    auto empty_query_ctx = MockQueryContext::create();
+    _agg_function->set_query_context(empty_query_ctx.get());
+
+    std::vector<std::string> resources = {"resource_name"};
+    std::vector<std::string> texts = {"test input"};
+    std::vector<std::string> task = {"summarize"};
+    auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+    auto col_text = ColumnHelper::create_column<DataTypeString>(texts);
+    auto col_task = ColumnHelper::create_column<DataTypeString>(task);
+
+    std::unique_ptr<char[]> memory(new char[_agg_function->size_of_data()]);
+    AggregateDataPtr place = memory.get();
+    _agg_function->create(place);
+
+    const IColumn* columns[3] = {col_resource.get(), col_text.get(), col_task.get()};
+
+    try {
+        _agg_function->add(place, columns, 0, _arena);
+        FAIL() << "Expected exception for missing AI resources";
+    } catch (const Exception& e) {
+        EXPECT_EQ(e.code(), ErrorCode::INTERNAL_ERROR);
+        EXPECT_NE(e.to_string().find("AI resources metadata missing"), std::string::npos);
+    }
+
+    _agg_function->destroy(place);
+}
+
 } // namespace doris::vectorized

be/test/ai/ai_function_test.cpp

@@ -551,6 +551,35 @@ TEST(AIFunctionTest, MockResourceSendRequest) {
     ASSERT_EQ(val, "this is a mock response. test input");
 }
 
+TEST(AIFunctionTest, MissingAIResourcesMetadataTest) {
+    auto query_ctx = MockQueryContext::create();
+    TQueryOptions query_options;
+    TQueryGlobals query_globals;
+    RuntimeState runtime_state(TUniqueId(), 0, query_options, query_globals, nullptr,
+                               query_ctx.get());
+    auto ctx = FunctionContext::create_context(&runtime_state, {}, {});
+
+    std::vector<std::string> resources = {"resource_name"};
+    std::vector<std::string> texts = {"test"};
+    auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+    auto col_text = ColumnHelper::create_column<DataTypeString>(texts);
+
+    Block block;
+    block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+    block.insert({std::move(col_text), std::make_shared<DataTypeString>(), "text"});
+    block.insert({nullptr, std::make_shared<DataTypeString>(), "result"});
+
+    ColumnNumbers arguments = {0, 1};
+    size_t result_idx = 2;
+
+    auto sentiment_func = FunctionAISentiment::create();
+    Status exec_status =
+            sentiment_func->execute_impl(ctx.get(), block, arguments, result_idx, texts.size());
+
+    ASSERT_FALSE(exec_status.ok());
+    ASSERT_NE(exec_status.to_string().find("AI resources metadata missing"), std::string::npos);
+}
+
 TEST(AIFunctionTest, ReturnTypeTest) {
     FunctionAIClassify func_classify;
     DataTypes args;

fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java

@@ -19,8 +19,10 @@
 
 import org.apache.doris.analysis.DescriptorTable;
 import org.apache.doris.analysis.StorageBackend;
+import org.apache.doris.catalog.AIResource;
 import org.apache.doris.catalog.Env;
 import org.apache.doris.catalog.FsBroker;
+import org.apache.doris.catalog.Resource;
 import org.apache.doris.common.Config;
 import org.apache.doris.common.MarkedCountDownLatch;
 import org.apache.doris.common.Pair;
@@ -89,6 +91,7 @@
 import org.apache.doris.system.SystemInfoService;
 import org.apache.doris.task.LoadEtlTask;
 import org.apache.doris.thrift.PaloInternalServiceVersion;
+import org.apache.doris.thrift.TAIResource;
 import org.apache.doris.thrift.TBrokerScanRange;
 import org.apache.doris.thrift.TDataSinkType;
 import org.apache.doris.thrift.TDescriptorTable;
@@ -3235,6 +3238,17 @@ Map<TNetworkAddress, TPipelineFragmentParams> toThrift(int backendNum) {
                     if (ignoreDataDistribution) {
                         params.setParallelInstances(parallelTasksNum);
                     }
+
+                    // Used for AI Functions
+                    Map<String, TAIResource> aiResourceMap = Maps.newLinkedHashMap();
+                    for (Resource resource : Env.getCurrentEnv().getResourceMgr()
+                                                    .getResource(Resource.ResourceType.AI)) {
+                        if (resource instanceof AIResource) {
+                            aiResourceMap.put(resource.getName(), ((AIResource) resource).toThrift());
+                        }
+                    }
+
+                    params.setAiResources(aiResourceMap);
                     res.put(instanceExecParam.host, params);
                     res.get(instanceExecParam.host).setBucketSeqToInstanceIdx(new HashMap<Integer, Integer>());
                     res.get(instanceExecParam.host).setShuffleIdxToInstanceIdx(new HashMap<Integer, Integer>());