You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@doris.apache.org by zc...@apache.org on 2026/04/14 09:22:33 UTC
(doris) branch master updated: [Feature](multimodal_embed) Support multimodal(file) embed (#62147)
This is an automated email from the ASF dual-hosted git repository.
zclll pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 443e1acf438 [Feature](multimodal_embed) Support multimodal(file) embed (#62147)
443e1acf438 is described below
commit 443e1acf438fd95400314b8d2cad8edb95a8373a
Author: linrrarity <li...@selectdb.com>
AuthorDate: Tue Apr 14 17:22:26 2026 +0800
[Feature](multimodal_embed) Support multimodal(file) embed (#62147)
Support image/video/audio embedding from json.
Supported multimodal_embed AI provider : Gemini(image, video, audio),
Qwen(image, video), Voyage(image, video), Jina(image)
Supported authentication methods include IAM + EXTERNAL_ID / AK + SK
```text
-- ima role, video
mysql> SELECT ARRAY_SIZE(
-> EMBED(
-> 'qwen_mul_embed',
-> CAST('{
'> "uri": "s3://selectdb-qa-test-3/lzq-multimodal-test/video/45944deac7d96c872f559f2ef94ea0a9.mp4",
'> "content_type": "video/mp4",
'> "provider": "AWS",
'> "endpoint": "s3.us-east-1.amazonaws.com",
'> "region": "us-east-1",
'> "role_arn": "arn:aws:iam::447051187841:role/lzq-role-test",
'> "external_id": "1001"
'> }' AS JSON)
-> )
-> ) AS video_embed_size;
+------------------+
| video_embed_size |
+------------------+
| 2560 |
+------------------+
-- ima role, video
mysql> SELECT ARRAY_SIZE(
-> EMBED(
-> 'qwen_mul_embed',
-> CAST('{
'> "uri": "s3://selectdb-qa-test-3/lzq-multimodal-test/image/ai-agg.png",
'> "content_type": "image/png",
'> "provider": "AWS",
'> "endpoint": "s3.us-east-1.amazonaws.com",
'> "region": "us-east-1",
'> "role_arn": "arn:aws:iam::447051187841:role/lzq-role-test",
'> "external_id": "1001"
'> }' AS JSON)
-> )
-> ) AS img_embed_size;
+----------------+
| img_embed_size |
+----------------+
| 2560 |
+----------------+
-- gemini multimodal embed
mysql> SELECT ARRAY_SIZE(
-> EMBED(
-> 'gemini_mul_embed',
-> CAST('{
'> "uri": "s3://selectdb-qa-test-3/lzq-multimodal-test/image/ai-agg.png",
'> "content_type": "image/png",
'> "provider": "AWS",
'> "endpoint": "s3.us-east-1.amazonaws.com",
'> "region": "us-east-1",
'> "role_arn": "arn:aws:iam::447051187841:role/lzq-role-test",
'> "external_id": "1001"
'> }' AS JSON)
-> )
-> ) AS img_embed_size;
+----------------+
| img_embed_size |
+----------------+
| 3072 |
+----------------+
-- ak, sk
Doris> SELECT ARRAY_SIZE(
-> EMBED(
-> 'qwen_mul_embed',
-> CAST('{
'> "uri": "s3://<bucket>/<key>/test_img.png",
'> "content_type": "image/png",
'> "endpoint": "cos.ap-hongkong.myqcloud.com",
'> "region": "ap-hongkong",
'> "ak": "<ak>",
'> "sk": "<sk>"
'> }' AS JSON)
-> )
-> ) AS img_embed_size_ak_sk;
+----------------------+
| img_embed_size_ak_sk |
+----------------------+
| 2560 |
+----------------------+
```
---
be/src/exprs/function/ai/ai_adapter.h | 292 ++++++++++++++-
be/src/exprs/function/ai/ai_filter.h | 39 ++
be/src/exprs/function/ai/ai_functions.h | 155 ++++----
be/src/exprs/function/ai/ai_similarity.h | 41 ++
be/src/exprs/function/ai/embed.h | 223 +++++++++++
be/test/ai/ai_adapter_test.cpp | 202 +++++++++-
be/test/ai/ai_function_test.cpp | 323 +++++++++++++++-
be/test/ai/embed_test.cpp | 416 ++++++++++++++++++++-
.../java/org/apache/doris/catalog/Resource.java | 8 +
.../property/constants/AIProperties.java | 2 +-
.../org/apache/doris/nereids/StatementContext.java | 15 +
.../trees/expressions/functions/agg/AIAgg.java | 1 +
.../trees/expressions/functions/ai/AIFunction.java | 1 +
.../trees/expressions/functions/ai/Embed.java | 48 ++-
.../main/java/org/apache/doris/qe/Coordinator.java | 26 +-
.../java/org/apache/doris/qe/SessionVariable.java | 9 +
.../doris/qe/runtime/ThriftPlansBuilder.java | 30 +-
gensrc/thrift/PaloInternalService.thrift | 4 +
.../suites/ai_p0/test_ai_functions.groovy | 5 +-
19 files changed, 1721 insertions(+), 119 deletions(-)
diff --git a/be/src/exprs/function/ai/ai_adapter.h b/be/src/exprs/function/ai/ai_adapter.h
index b2de31377fc..0244261a3ed 100644
--- a/be/src/exprs/function/ai/ai_adapter.h
+++ b/be/src/exprs/function/ai/ai_adapter.h
@@ -88,6 +88,20 @@ struct AIResource {
}
};
+enum class MultimodalType { IMAGE, VIDEO, AUDIO };
+
+inline const char* multimodal_type_to_string(MultimodalType type) {
+ switch (type) {
+ case MultimodalType::IMAGE:
+ return "image";
+ case MultimodalType::VIDEO:
+ return "video";
+ case MultimodalType::AUDIO:
+ return "audio";
+ }
+ return "unknown";
+}
+
class AIAdapter {
public:
virtual ~AIAdapter() = default;
@@ -127,6 +141,13 @@ public:
_config.provider_type);
}
+ virtual Status build_multimodal_embedding_request(MultimodalType /*media_type*/,
+ const std::string& /*media_url*/,
+ std::string& /*request_body*/) const {
+ return Status::NotSupported("{} does not support multimodal Embed feature.",
+ _config.provider_type);
+ }
+
virtual Status parse_embedding_response(const std::string& response_body,
std::vector<std::vector<float>>& results) const {
return Status::NotSupported("{} does not support the Embed feature.",
@@ -195,6 +216,63 @@ public:
return Status::OK();
}
+ Status build_multimodal_embedding_request(MultimodalType media_type,
+ const std::string& media_url,
+ std::string& request_body) const override {
+ if (media_type != MultimodalType::IMAGE && media_type != MultimodalType::VIDEO)
+ [[unlikely]] {
+ return Status::InvalidArgument(
+ "VoyageAI only supports image/video multimodal embed, got {}",
+ multimodal_type_to_string(media_type));
+ }
+ if (_config.dimensions != -1) {
+ LOG(WARNING) << "VoyageAI multimodal embedding currently ignores dimensions parameter, "
+ << "model=" << _config.model_name << ", dimensions=" << _config.dimensions;
+ }
+
+ rapidjson::Document doc;
+ doc.SetObject();
+ auto& allocator = doc.GetAllocator();
+
+ /*{
+ "inputs": [
+ {
+ "content": [
+ {"type": "image_url", "image_url": "<url>"}
+ ]
+ }
+ ],
+ "model": "voyage-multimodal-3.5"
+ }*/
+ doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator);
+
+ rapidjson::Value inputs(rapidjson::kArrayType);
+ rapidjson::Value input(rapidjson::kObjectType);
+ rapidjson::Value content(rapidjson::kArrayType);
+
+ rapidjson::Value media_item(rapidjson::kObjectType);
+ if (media_type == MultimodalType::IMAGE) {
+ media_item.AddMember("type", "image_url", allocator);
+ media_item.AddMember("image_url", rapidjson::Value(media_url.c_str(), allocator),
+ allocator);
+ } else {
+ media_item.AddMember("type", "video_url", allocator);
+ media_item.AddMember("video_url", rapidjson::Value(media_url.c_str(), allocator),
+ allocator);
+ }
+ content.PushBack(media_item, allocator);
+
+ input.AddMember("content", content, allocator);
+ inputs.PushBack(input, allocator);
+ doc.AddMember("inputs", inputs, allocator);
+
+ rapidjson::StringBuffer buffer;
+ rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
+ doc.Accept(writer);
+ request_body = buffer.GetString();
+ return Status::OK();
+ }
+
Status parse_embedding_response(const std::string& response_body,
std::vector<std::vector<float>>& results) const override {
rapidjson::Document doc;
@@ -355,6 +433,13 @@ public:
return Status::OK();
}
+ Status build_multimodal_embedding_request(MultimodalType /*media_type*/,
+ const std::string& /*media_url*/,
+ std::string& /*request_body*/) const override {
+ return Status::NotSupported("{} does not support multimodal Embed feature.",
+ _config.provider_type);
+ }
+
Status parse_embedding_response(const std::string& response_body,
std::vector<std::vector<float>>& results) const override {
rapidjson::Document doc;
@@ -703,6 +788,13 @@ public:
return Status::OK();
}
+ Status build_multimodal_embedding_request(MultimodalType /*media_type*/,
+ const std::string& /*media_url*/,
+ std::string& /*request_body*/) const override {
+ return Status::NotSupported("{} does not support multimodal Embed feature.",
+ _config.provider_type);
+ }
+
protected:
bool supports_dimension_param(const std::string& model_name) const override {
return !(model_name == "text-embedding-ada-002");
@@ -779,6 +871,103 @@ protected:
};
class QwenAdapter : public OpenAIAdapter {
+public:
+ Status build_multimodal_embedding_request(MultimodalType media_type,
+ const std::string& media_url,
+ std::string& request_body) const override {
+ if (media_type != MultimodalType::IMAGE && media_type != MultimodalType::VIDEO) {
+ return Status::InvalidArgument(
+ "QWEN only supports image/video multimodal embed, got {}",
+ multimodal_type_to_string(media_type));
+ }
+
+ rapidjson::Document doc;
+ doc.SetObject();
+ auto& allocator = doc.GetAllocator();
+
+ /*{
+ "model": "tongyi-embedding-vision-plus",
+ "input": {
+ "contents": [
+ {"image": "<url>"}
+ ]
+ }
+ "parameters": {
+ "dimension": 512
+ }
+ }*/
+ doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator);
+ rapidjson::Value input(rapidjson::kObjectType);
+ rapidjson::Value contents(rapidjson::kArrayType);
+
+ rapidjson::Value media_item(rapidjson::kObjectType);
+ if (media_type == MultimodalType::IMAGE) {
+ media_item.AddMember("image", rapidjson::Value(media_url.c_str(), allocator),
+ allocator);
+ } else {
+ media_item.AddMember("video", rapidjson::Value(media_url.c_str(), allocator),
+ allocator);
+ }
+ contents.PushBack(media_item, allocator);
+
+ input.AddMember("contents", contents, allocator);
+ doc.AddMember("input", input, allocator);
+ if (_config.dimensions != -1 && supports_dimension_param(_config.model_name)) {
+ rapidjson::Value parameters(rapidjson::kObjectType);
+ std::string param_name = get_dimension_param_name();
+ rapidjson::Value dimension_name(param_name.c_str(), allocator);
+ parameters.AddMember(dimension_name, _config.dimensions, allocator);
+ doc.AddMember("parameters", parameters, allocator);
+ }
+
+ rapidjson::StringBuffer buffer;
+ rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
+ doc.Accept(writer);
+ request_body = buffer.GetString();
+ return Status::OK();
+ }
+
+ Status parse_embedding_response(const std::string& response_body,
+ std::vector<std::vector<float>>& results) const override {
+ rapidjson::Document doc;
+ doc.Parse(response_body.c_str());
+
+ if (doc.HasParseError() || !doc.IsObject()) [[unlikely]] {
+ return Status::InternalError("Failed to parse {} response: {}", _config.provider_type,
+ response_body);
+ }
+ // Qwen multimodal embedding usually returns:
+ // {
+ // "output": {
+ // "embeddings": [
+ // {"index":0, "embedding":[...], "type":"image|video|text"},
+ // ...
+ // ]
+ // }
+ // }
+ //
+ // In text-only or compatibility endpoints, Qwen may also return OpenAI-style
+ // "data":[{"embedding":[...]}]. For compatibility we first parse native
+ // output.embeddings and then fallback to OpenAIAdapter parser.
+ if (doc.HasMember("output") && doc["output"].IsObject() &&
+ doc["output"].HasMember("embeddings") && doc["output"]["embeddings"].IsArray()) {
+ const auto& embeddings = doc["output"]["embeddings"];
+ results.reserve(embeddings.Size());
+ for (rapidjson::SizeType i = 0; i < embeddings.Size(); i++) {
+ if (!embeddings[i].HasMember("embedding") ||
+ !embeddings[i]["embedding"].IsArray()) {
+ return Status::InternalError("Invalid {} response format: {}",
+ _config.provider_type, response_body);
+ }
+ std::transform(embeddings[i]["embedding"].Begin(), embeddings[i]["embedding"].End(),
+ std::back_inserter(results.emplace_back()),
+ [](const auto& val) { return val.GetFloat(); });
+ }
+ return Status::OK();
+ }
+ return OpenAIAdapter::parse_embedding_response(response_body, results);
+ }
+
protected:
bool supports_dimension_param(const std::string& model_name) const override {
static const std::unordered_set<std::string> no_dimension_models = {
@@ -789,6 +978,55 @@ protected:
std::string get_dimension_param_name() const override { return "dimension"; }
};
+class JinaAdapter : public VoyageAIAdapter {
+public:
+ Status build_multimodal_embedding_request(MultimodalType media_type,
+ const std::string& media_url,
+ std::string& request_body) const override {
+ if (media_type != MultimodalType::IMAGE && media_type != MultimodalType::VIDEO)
+ [[unlikely]] {
+ return Status::InvalidArgument(
+ "JINA only supports image/video multimodal embed, got {}",
+ multimodal_type_to_string(media_type));
+ }
+
+ rapidjson::Document doc;
+ doc.SetObject();
+ auto& allocator = doc.GetAllocator();
+
+ /*{
+ "model": "jina-embeddings-v4",
+ "task": "text-matching",
+ "input": [
+ {"image": "<url>"}
+ ]
+ }*/
+ doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator);
+ doc.AddMember("task", "text-matching", allocator);
+
+ rapidjson::Value input(rapidjson::kArrayType);
+ rapidjson::Value media_item(rapidjson::kObjectType);
+ if (media_type == MultimodalType::IMAGE) {
+ media_item.AddMember("image", rapidjson::Value(media_url.c_str(), allocator),
+ allocator);
+ } else {
+ media_item.AddMember("video", rapidjson::Value(media_url.c_str(), allocator),
+ allocator);
+ }
+ input.PushBack(media_item, allocator);
+ if (_config.dimensions != -1 && supports_dimension_param(_config.model_name)) {
+ doc.AddMember("dimensions", _config.dimensions, allocator);
+ }
+ doc.AddMember("input", input, allocator);
+
+ rapidjson::StringBuffer buffer;
+ rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
+ doc.Accept(writer);
+ request_body = buffer.GetString();
+ return Status::OK();
+ }
+};
+
class BaichuanAdapter : public OpenAIAdapter {
protected:
bool supports_dimension_param(const std::string& model_name) const override { return false; }
@@ -969,6 +1207,51 @@ public:
return Status::OK();
}
+ Status build_multimodal_embedding_request(MultimodalType media_type,
+ const std::string& media_url,
+ std::string& request_body) const override {
+ const char* mime_type = nullptr;
+ switch (media_type) {
+ case MultimodalType::IMAGE:
+ mime_type = "image/png";
+ break;
+ case MultimodalType::AUDIO:
+ mime_type = "audio/mpeg";
+ break;
+ case MultimodalType::VIDEO:
+ mime_type = "video/mp4";
+ break;
+ }
+
+ rapidjson::Document doc;
+ doc.SetObject();
+ auto& allocator = doc.GetAllocator();
+
+ std::string model_name = _config.model_name;
+ if (!model_name.starts_with("models/")) {
+ model_name = "models/" + model_name;
+ }
+ doc.AddMember("model", rapidjson::Value(model_name.c_str(), allocator), allocator);
+ add_dimension_params(doc, allocator);
+
+ rapidjson::Value content(rapidjson::kObjectType);
+ rapidjson::Value parts(rapidjson::kArrayType);
+ rapidjson::Value part(rapidjson::kObjectType);
+ rapidjson::Value file_data(rapidjson::kObjectType);
+ file_data.AddMember("mime_type", rapidjson::Value(mime_type, allocator), allocator);
+ file_data.AddMember("file_uri", rapidjson::Value(media_url.c_str(), allocator), allocator);
+ part.AddMember("file_data", file_data, allocator);
+ parts.PushBack(part, allocator);
+ content.AddMember("parts", parts, allocator);
+ doc.AddMember("content", content, allocator);
+
+ rapidjson::StringBuffer buffer;
+ rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
+ doc.Accept(writer);
+ request_body = buffer.GetString();
+ return Status::OK();
+ }
+
Status parse_embedding_response(const std::string& response_body,
std::vector<std::vector<float>>& results) const override {
rapidjson::Document doc;
@@ -1135,6 +1418,12 @@ public:
return Status::OK();
}
+ Status build_multimodal_embedding_request(MultimodalType /*media_type*/,
+ const std::string& /*media_url*/,
+ std::string& /*request_body*/) const override {
+ return Status::OK();
+ }
+
Status parse_embedding_response(const std::string& response_body,
std::vector<std::vector<float>>& results) const override {
rapidjson::Document doc;
@@ -1166,6 +1455,7 @@ public:
{"MINIMAX", []() { return std::make_shared<MinimaxAdapter>(); }},
{"ZHIPU", []() { return std::make_shared<ZhipuAdapter>(); }},
{"QWEN", []() { return std::make_shared<QwenAdapter>(); }},
+ {"JINA", []() { return std::make_shared<JinaAdapter>(); }},
{"BAICHUAN", []() { return std::make_shared<BaichuanAdapter>(); }},
{"ANTHROPIC", []() { return std::make_shared<AnthropicAdapter>(); }},
{"GEMINI", []() { return std::make_shared<GeminiAdapter>(); }},
@@ -1177,4 +1467,4 @@ public:
}
};
-} // namespace doris
\ No newline at end of file
+} // namespace doris
diff --git a/be/src/exprs/function/ai/ai_filter.h b/be/src/exprs/function/ai/ai_filter.h
index ac34dba3bbc..f00c42337a7 100644
--- a/be/src/exprs/function/ai/ai_filter.h
+++ b/be/src/exprs/function/ai/ai_filter.h
@@ -17,6 +17,10 @@
#pragma once
+#include <algorithm>
+#include <cctype>
+#include <cstdlib>
+
#include "exprs/function/ai/ai_functions.h"
namespace doris {
@@ -38,6 +42,41 @@ public:
return std::make_shared<DataTypeBool>();
}
+ Status execute_with_adapter(FunctionContext* context, Block& block,
+ const ColumnNumbers& arguments, uint32_t result,
+ size_t input_rows_count, const TAIResource& config,
+ std::shared_ptr<AIAdapter>& adapter) const {
+ auto col_result = ColumnUInt8::create();
+
+ for (size_t i = 0; i < input_rows_count; ++i) {
+ std::string prompt;
+ RETURN_IF_ERROR(build_prompt(block, arguments, i, prompt));
+
+ std::string string_result;
+ RETURN_IF_ERROR(
+ execute_single_request(prompt, string_result, config, adapter, context));
+
+#ifdef BE_TEST
+ const char* test_result = std::getenv("AI_TEST_RESULT");
+ if (test_result != nullptr) {
+ string_result = test_result;
+ } else {
+ string_result = "0";
+ }
+#endif
+
+ std::string_view trimmed = doris::trim(string_result);
+ if (trimmed != "1" && trimmed != "0") {
+ return Status::RuntimeError("Failed to parse boolean value: " + string_result);
+ }
+
+ col_result->insert_value(static_cast<UInt8>(trimmed == "1"));
+ }
+
+ block.replace_by_position(result, std::move(col_result));
+ return Status::OK();
+ }
+
static FunctionPtr create() { return std::make_shared<FunctionAIFilter>(); }
};
} // namespace doris
\ No newline at end of file
diff --git a/be/src/exprs/function/ai/ai_functions.h b/be/src/exprs/function/ai/ai_functions.h
index d8c3659dcc9..0386689091c 100644
--- a/be/src/exprs/function/ai/ai_functions.h
+++ b/be/src/exprs/function/ai/ai_functions.h
@@ -19,6 +19,7 @@
#include <gen_cpp/FrontendService.h>
#include <gen_cpp/PaloInternalService_types.h>
+#include <glog/logging.h>
#include <algorithm>
#include <cctype>
@@ -43,6 +44,7 @@
#include "runtime/query_context.h"
#include "runtime/runtime_state.h"
#include "service/http/http_client.h"
+#include "util/security.h"
#include "util/threadpool.h"
namespace doris {
@@ -73,10 +75,6 @@ public:
Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
uint32_t result, size_t input_rows_count) const override {
- DataTypePtr return_type_impl =
- assert_cast<const Derived&>(*this).get_return_type_impl(DataTypes());
- MutableColumnPtr col_result = return_type_impl->create_column();
-
TAIResource config;
std::shared_ptr<AIAdapter> adapter;
if (Status status = assert_cast<const Derived*>(this)->_init_from_resource(
@@ -85,91 +83,42 @@ public:
return status;
}
+ return assert_cast<const Derived&>(*this).execute_with_adapter(
+ context, block, arguments, result, input_rows_count, config, adapter);
+ }
+
+protected:
+ // Derived classes can override this method for non-text/default behavior.
+ // The base implementation keeps previous text-oriented processing unchanged.
+ Status execute_with_adapter(FunctionContext* context, Block& block,
+ const ColumnNumbers& arguments, uint32_t result,
+ size_t input_rows_count, const TAIResource& config,
+ std::shared_ptr<AIAdapter>& adapter) const {
+ DataTypePtr return_type_impl =
+ assert_cast<const Derived&>(*this).get_return_type_impl(DataTypes());
+ if (return_type_impl->get_primitive_type() != PrimitiveType::TYPE_STRING) {
+ return Status::InternalError("{} must override execute for non-string return type",
+ get_name());
+ }
+ MutableColumnPtr col_result = ColumnString::create();
+
for (size_t i = 0; i < input_rows_count; ++i) {
// Build AI prompt text
std::string prompt;
RETURN_IF_ERROR(
assert_cast<const Derived&>(*this).build_prompt(block, arguments, i, prompt));
- // Execute a single AI request and get the result
- if (return_type_impl->get_primitive_type() == PrimitiveType::TYPE_ARRAY) {
- // Array(Float) for AI_EMBED
- std::vector<float> float_result;
- RETURN_IF_ERROR(
- execute_single_request(prompt, float_result, config, adapter, context));
-
- auto& col_array = assert_cast<ColumnArray&>(*col_result);
- auto& offsets = col_array.get_offsets();
- auto& nested_nullable_col = assert_cast<ColumnNullable&>(col_array.get_data());
- auto& nested_col =
- assert_cast<ColumnFloat32&>(*(nested_nullable_col.get_nested_column_ptr()));
- nested_col.reserve(nested_col.size() + float_result.size());
-
- size_t current_offset = nested_col.size();
- nested_col.insert_many_raw_data(reinterpret_cast<const char*>(float_result.data()),
- float_result.size());
- offsets.push_back(current_offset + float_result.size());
- auto& null_map = nested_nullable_col.get_null_map_column();
- null_map.insert_many_vals(0, float_result.size());
- } else {
- std::string string_result;
- RETURN_IF_ERROR(
- execute_single_request(prompt, string_result, config, adapter, context));
-
- switch (return_type_impl->get_primitive_type()) {
- case PrimitiveType::TYPE_STRING: { // string
- assert_cast<ColumnString&>(*col_result)
- .insert_data(string_result.data(), string_result.size());
- break;
- }
- case PrimitiveType::TYPE_BOOLEAN: { // boolean for AI_FILTER
-#ifdef BE_TEST
- const char* test_result = std::getenv("AI_TEST_RESULT");
- if (test_result != nullptr) {
- string_result = test_result;
- } else {
- string_result = "0";
- }
-#endif
- trim_string(string_result);
- if (string_result != "1" && string_result != "0") {
- return Status::RuntimeError("Failed to parse boolean value: " +
- string_result);
- }
- assert_cast<ColumnUInt8&>(*col_result)
- .insert_value(static_cast<UInt8>(string_result == "1"));
- break;
- }
- case PrimitiveType::TYPE_FLOAT: { // float for AI_SIMILARITY
-#ifdef BE_TEST
- const char* test_result = std::getenv("AI_TEST_RESULT");
- if (test_result != nullptr) {
- string_result = test_result;
- } else {
- string_result = "0.0";
- }
-#endif
- trim_string(string_result);
- try {
- float float_value = std::stof(string_result);
- assert_cast<ColumnFloat32&>(*col_result).insert_value(float_value);
- } catch (...) {
- return Status::RuntimeError("Failed to parse float value: " +
- string_result);
- }
- break;
- }
- default:
- return Status::InternalError("Unsupported ReturnType for AIFunction");
- }
- }
+ std::string string_result;
+ RETURN_IF_ERROR(
+ execute_single_request(prompt, string_result, config, adapter, context));
+ assert_cast<ColumnString&>(*col_result)
+ .insert_data(string_result.data(), string_result.size());
}
block.replace_by_position(result, std::move(col_result));
return Status::OK();
}
-protected:
// The endpoint `v1/completions` does not support `system_prompt`.
// To ensure a clear structure and stable AI results.
// Convert from `v1/completions` to `v1/chat/completions`
@@ -181,17 +130,6 @@ protected:
}
}
-private:
- // Trim whitespace and newlines from string
- static void trim_string(std::string& str) {
- str.erase(str.begin(), std::find_if(str.begin(), str.end(),
- [](unsigned char ch) { return !std::isspace(ch); }));
- str.erase(std::find_if(str.rbegin(), str.rend(),
- [](unsigned char ch) { return !std::isspace(ch); })
- .base(),
- str.end());
- }
-
// The ai resource must be literal
Status _init_from_resource(FunctionContext* context, const Block& block,
const ColumnNumbers& arguments, TAIResource& config,
@@ -228,7 +166,7 @@ private:
Status do_send_request(HttpClient* client, const std::string& request_body,
std::string& response, const TAIResource& config,
std::shared_ptr<AIAdapter>& adapter, FunctionContext* context) const {
- RETURN_IF_ERROR(client->init(config.endpoint));
+ RETURN_IF_ERROR(client->init(config.endpoint, false));
QueryContext* query_ctx = context->state()->get_query_ctx();
int64_t remaining_query_time = query_ctx->get_remaining_query_time_seconds();
@@ -242,7 +180,22 @@ private:
RETURN_IF_ERROR(adapter->set_authentication(client));
}
- return client->execute_post_request(request_body, &response);
+ Status st = client->execute_post_request(request_body, &response);
+ long http_status = client->get_http_status();
+
+ if (!st.ok()) {
+ LOG(INFO) << "AI HTTP request failed before status validation, provider="
+ << config.provider_type << ", model=" << config.model_name
+ << ", endpoint=" << mask_token(config.endpoint)
+ << ", exec_status=" << st.to_string() << ", response_body=" << response;
+ return st;
+ }
+ if (http_status != 200) {
+ return Status::HttpError(
+ "http status code is not 200, code={}, url={}, response_body={}", http_status,
+ mask_token(config.endpoint), response);
+ }
+ return Status::OK();
}
// Sends the request with retry mechanism for handling transient failures
@@ -311,6 +264,26 @@ private:
result = std::move(results[0]);
return Status::OK();
}
+
+ // Sends a pre-built embedding request body and parses the float vector result.
+ // Used when the request body has already been constructed (e.g., multimodal embedding).
+ Status execute_embedding_request(const std::string& request_body, std::vector<float>& result,
+ const TAIResource& config, std::shared_ptr<AIAdapter>& adapter,
+ FunctionContext* context) const {
+ std::vector<std::vector<float>> results;
+ std::string response;
+ if (config.provider_type == "MOCK") {
+ response = "{\"embedding\": [0, 1, 2, 3, 4]}";
+ } else {
+ RETURN_IF_ERROR(send_request_to_llm(request_body, response, config, adapter, context));
+ }
+ RETURN_IF_ERROR(adapter->parse_embedding_response(response, results));
+ if (results.empty()) {
+ return Status::InternalError("AI returned empty result");
+ }
+ result = std::move(results[0]);
+ return Status::OK();
+ }
};
-} // namespace doris
\ No newline at end of file
+} // namespace doris
diff --git a/be/src/exprs/function/ai/ai_similarity.h b/be/src/exprs/function/ai/ai_similarity.h
index bb6aed79a4b..21c96781341 100644
--- a/be/src/exprs/function/ai/ai_similarity.h
+++ b/be/src/exprs/function/ai/ai_similarity.h
@@ -17,6 +17,10 @@
#pragma once
+#include <algorithm>
+#include <cctype>
+#include <cstdlib>
+
#include "exprs/function/ai/ai_functions.h"
namespace doris {
@@ -41,6 +45,43 @@ public:
return std::make_shared<DataTypeFloat32>();
}
+ Status execute_with_adapter(FunctionContext* context, Block& block,
+ const ColumnNumbers& arguments, uint32_t result,
+ size_t input_rows_count, const TAIResource& config,
+ std::shared_ptr<AIAdapter>& adapter) const {
+ auto col_result = ColumnFloat32::create();
+
+ for (size_t i = 0; i < input_rows_count; ++i) {
+ std::string prompt;
+ RETURN_IF_ERROR(build_prompt(block, arguments, i, prompt));
+
+ std::string string_result;
+ RETURN_IF_ERROR(
+ execute_single_request(prompt, string_result, config, adapter, context));
+
+#ifdef BE_TEST
+ const char* test_result = std::getenv("AI_TEST_RESULT");
+ if (test_result != nullptr) {
+ string_result = test_result;
+ } else {
+ string_result = "0.0";
+ }
+#endif
+
+ std::string_view trimmed = doris::trim(string_result);
+ float float_value = 0;
+ auto [ptr, ec] = fast_float::from_chars(trimmed.data(), trimmed.data() + trimmed.size(),
+ float_value);
+ if (ec != std::errc() || ptr != trimmed.data() + trimmed.size()) [[unlikely]] {
+ return Status::RuntimeError("Failed to parse float value: " + string_result);
+ }
+ assert_cast<ColumnFloat32&>(*col_result).insert_value(float_value);
+ }
+
+ block.replace_by_position(result, std::move(col_result));
+ return Status::OK();
+ }
+
static FunctionPtr create() { return std::make_shared<FunctionAISimilarity>(); }
Status build_prompt(const Block& block, const ColumnNumbers& arguments, size_t row_num,
diff --git a/be/src/exprs/function/ai/embed.h b/be/src/exprs/function/ai/embed.h
index 8c5988dc4cd..f5f2176381b 100644
--- a/be/src/exprs/function/ai/embed.h
+++ b/be/src/exprs/function/ai/embed.h
@@ -17,8 +17,17 @@
#pragma once
+#include <glog/logging.h>
+#include <rapidjson/document.h>
+
+#include <string_view>
+
+#include "core/data_type/data_type_nullable.h"
#include "core/data_type/primitive_type.h"
#include "exprs/function/ai/ai_functions.h"
+#include "util/jsonb_utils.h"
+#include "util/s3_uri.h"
+#include "util/s3_util.h"
namespace doris {
class FunctionEmbed : public AIFunction<FunctionEmbed> {
@@ -33,7 +42,221 @@ public:
return std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>()));
}
+ Status execute_with_adapter(FunctionContext* context, Block& block,
+ const ColumnNumbers& arguments, uint32_t result,
+ size_t input_rows_count, const TAIResource& config,
+ std::shared_ptr<AIAdapter>& adapter) const {
+ if (arguments.size() != 2) {
+ return Status::InvalidArgument("Function EMBED expects 2 arguments, but got {}",
+ arguments.size());
+ }
+
+ PrimitiveType input_type =
+ remove_nullable(block.get_by_position(arguments[1]).type)->get_primitive_type();
+ if (input_type == PrimitiveType::TYPE_JSONB) {
+ return _execute_multimodal_embed(context, block, arguments, result, input_rows_count,
+ config, adapter);
+ }
+ if (input_type == PrimitiveType::TYPE_STRING || input_type == PrimitiveType::TYPE_VARCHAR ||
+ input_type == PrimitiveType::TYPE_CHAR) {
+ return _execute_text_embed(context, block, arguments, result, input_rows_count, config,
+ adapter);
+ }
+ return Status::InvalidArgument(
+ "Function EMBED expects the second argument to be STRING or JSON, but got type {}",
+ block.get_by_position(arguments[1]).type->get_name());
+ }
+
static FunctionPtr create() { return std::make_shared<FunctionEmbed>(); }
+
+private:
+ Status _execute_text_embed(FunctionContext* context, Block& block,
+ const ColumnNumbers& arguments, uint32_t result,
+ size_t input_rows_count, const TAIResource& config,
+ std::shared_ptr<AIAdapter>& adapter) const {
+ auto col_result = ColumnArray::create(
+ ColumnNullable::create(ColumnFloat32::create(), ColumnUInt8::create()));
+
+ for (size_t i = 0; i < input_rows_count; ++i) {
+ std::string prompt;
+ RETURN_IF_ERROR(build_prompt(block, arguments, i, prompt));
+
+ std::vector<float> float_result;
+ RETURN_IF_ERROR(execute_single_request(prompt, float_result, config, adapter, context));
+ _insert_embedding_result(*col_result, float_result);
+ }
+
+ block.replace_by_position(result, std::move(col_result));
+ return Status::OK();
+ }
+
+ Status _execute_multimodal_embed(FunctionContext* context, Block& block,
+ const ColumnNumbers& arguments, uint32_t result,
+ size_t input_rows_count, const TAIResource& config,
+ std::shared_ptr<AIAdapter>& adapter) const {
+ auto col_result = ColumnArray::create(
+ ColumnNullable::create(ColumnFloat32::create(), ColumnUInt8::create()));
+
+ int64_t ttl_seconds = 3600;
+ QueryContext* query_ctx = context->state()->get_query_ctx();
+ if (query_ctx && query_ctx->query_options().__isset.file_presigned_url_ttl_seconds) {
+ ttl_seconds = query_ctx->query_options().file_presigned_url_ttl_seconds;
+ if (ttl_seconds <= 0) {
+ ttl_seconds = 3600;
+ }
+ }
+
+ const ColumnWithTypeAndName& file_column = block.get_by_position(arguments[1]);
+ for (size_t i = 0; i < input_rows_count; ++i) {
+ rapidjson::Document file_input;
+ RETURN_IF_ERROR(_parse_file_input(file_column, i, file_input));
+
+ MultimodalType media_type;
+ RETURN_IF_ERROR(_infer_media_type(file_input, media_type));
+
+ std::string media_url;
+ RETURN_IF_ERROR(_resolve_media_url(file_input, ttl_seconds, media_url));
+
+ std::string request_body;
+ RETURN_IF_ERROR(adapter->build_multimodal_embedding_request(media_type, media_url,
+ request_body));
+
+ std::vector<float> float_result;
+ RETURN_IF_ERROR(execute_embedding_request(request_body, float_result, config, adapter,
+ context));
+ _insert_embedding_result(*col_result, float_result);
+ }
+
+ block.replace_by_position(result, std::move(col_result));
+ return Status::OK();
+ }
+
+ static void _insert_embedding_result(ColumnArray& col_array,
+ const std::vector<float>& float_result) {
+ auto& offsets = col_array.get_offsets();
+ auto& nested_nullable_col = assert_cast<ColumnNullable&>(col_array.get_data());
+ auto& nested_col =
+ assert_cast<ColumnFloat32&>(*(nested_nullable_col.get_nested_column_ptr()));
+ nested_col.reserve(nested_col.size() + float_result.size());
+
+ size_t current_offset = nested_col.size();
+ nested_col.insert_many_raw_data(reinterpret_cast<const char*>(float_result.data()),
+ float_result.size());
+ offsets.push_back(current_offset + float_result.size());
+ auto& null_map = nested_nullable_col.get_null_map_column();
+ null_map.insert_many_vals(0, float_result.size());
+ }
+
+ static bool _starts_with_ignore_case(std::string_view s, std::string_view prefix) {
+ if (s.size() < prefix.size()) {
+ return false;
+ }
+ return std::equal(prefix.begin(), prefix.end(), s.begin(), [](char a, char b) {
+ return std::tolower(static_cast<unsigned char>(a)) ==
+ std::tolower(static_cast<unsigned char>(b));
+ });
+ }
+
+ static Status _infer_media_type(const rapidjson::Value& file_input,
+ MultimodalType& media_type) {
+ std::string content_type;
+ RETURN_IF_ERROR(_get_required_string_field(file_input, "content_type", content_type));
+
+ if (_starts_with_ignore_case(content_type, "image/")) {
+ media_type = MultimodalType::IMAGE;
+ return Status::OK();
+ } else if (_starts_with_ignore_case(content_type, "video/")) {
+ media_type = MultimodalType::VIDEO;
+ return Status::OK();
+ } else if (_starts_with_ignore_case(content_type, "audio/")) {
+ media_type = MultimodalType::AUDIO;
+ return Status::OK();
+ }
+
+ return Status::InvalidArgument("Unsupported content_type for EMBED: {}", content_type);
+ }
+
+ // Parse the FILE-like JSONB argument into a JSON object for downstream field reads.
+ static Status _parse_file_input(const ColumnWithTypeAndName& file_column, size_t row_num,
+ rapidjson::Document& file_input) {
+ std::string file_json =
+ JsonbToJson::jsonb_to_json_string(file_column.column->get_data_at(row_num).data,
+ file_column.column->get_data_at(row_num).size);
+ file_input.Parse(file_json.c_str());
+ DORIS_CHECK(!file_input.HasParseError() && file_input.IsObject());
+ return Status::OK();
+ }
+
+ // TODO(lzq): After support FILE type, We should use the interface provided by FILE to get the fields
+ // replacing this function
+ static Status _get_required_string_field(const rapidjson::Value& obj, const char* field_name,
+ std::string& value) {
+ auto iter = obj.FindMember(field_name);
+ if (iter == obj.MemberEnd() || !iter->value.IsString()) {
+ return Status::InvalidArgument(
+ "EMBED file json field '{}' is required and must be a string", field_name);
+ }
+ value = iter->value.GetString();
+ if (value.empty()) {
+ return Status::InvalidArgument("EMBED file json field '{}' can not be empty",
+ field_name);
+ }
+ return Status::OK();
+ }
+
+ static Status init_s3_client_conf_from_json(const rapidjson::Value& file_input,
+ S3ClientConf& s3_client_conf) {
+ std::string endpoint;
+ RETURN_IF_ERROR(_get_required_string_field(file_input, "endpoint", endpoint));
+ std::string region;
+ RETURN_IF_ERROR(_get_required_string_field(file_input, "region", region));
+
+ auto get_optional_string_field = [&](const char* field_name, std::string& value) {
+ auto iter = file_input.FindMember(field_name);
+ if (iter == file_input.MemberEnd() || iter->value.IsNull()) {
+ return;
+ }
+ DORIS_CHECK(iter->value.IsString());
+ value = iter->value.GetString();
+ };
+
+ get_optional_string_field("ak", s3_client_conf.ak);
+ get_optional_string_field("sk", s3_client_conf.sk);
+ get_optional_string_field("role_arn", s3_client_conf.role_arn);
+ get_optional_string_field("external_id", s3_client_conf.external_id);
+ s3_client_conf.endpoint = endpoint;
+ s3_client_conf.region = region;
+
+ return Status::OK();
+ }
+
+ Status _resolve_media_url(const rapidjson::Value& file_input, int64_t ttl_seconds,
+ std::string& media_url) const {
+ std::string uri;
+ RETURN_IF_ERROR(_get_required_string_field(file_input, "uri", uri));
+
+ // If it's a direct http/https URL, use it as-is
+ if (_starts_with_ignore_case(uri, "http://") || _starts_with_ignore_case(uri, "https://")) {
+ media_url = uri;
+ return Status::OK();
+ }
+
+ S3ClientConf s3_client_conf;
+ RETURN_IF_ERROR(init_s3_client_conf_from_json(file_input, s3_client_conf));
+ auto s3_client = S3ClientFactory::instance().create(s3_client_conf);
+ if (s3_client == nullptr) {
+ return Status::InternalError("Failed to create S3 client for EMBED file input");
+ }
+
+ S3URI s3_uri(uri);
+ RETURN_IF_ERROR(s3_uri.parse());
+ std::string bucket = s3_uri.get_bucket();
+ std::string key = s3_uri.get_key();
+ DORIS_CHECK(!bucket.empty() && !key.empty());
+ media_url = s3_client->generate_presigned_url({.bucket = bucket, .key = key}, ttl_seconds,
+ s3_client_conf);
+ return Status::OK();
+ }
};
}; // namespace doris
diff --git a/be/test/ai/ai_adapter_test.cpp b/be/test/ai/ai_adapter_test.cpp
index 6956af5e3fd..aaaf64e7eaa 100644
--- a/be/test/ai/ai_adapter_test.cpp
+++ b/be/test/ai/ai_adapter_test.cpp
@@ -561,9 +561,9 @@ TEST(AI_ADAPTER_TEST, unsupported_provider_type) {
}
TEST(AI_ADAPTER_TEST, adapter_factory_all_types) {
- std::vector<std::string> types = {"LOCAL", "OPENAI", "MOONSHOT", "DEEPSEEK",
- "MINIMAX", "ZHIPU", "QWEN", "BAICHUAN",
- "ANTHROPIC", "GEMINI", "VOYAGEAI", "MOCK"};
+ std::vector<std::string> types = {"LOCAL", "OPENAI", "MOONSHOT", "DEEPSEEK", "MINIMAX",
+ "ZHIPU", "QWEN", "JINA", "BAICHUAN", "ANTHROPIC",
+ "GEMINI", "VOYAGEAI", "MOCK"};
for (const auto& type : types) {
auto adapter = doris::AIAdapterFactory::create_adapter(type);
ASSERT_TRUE(adapter != nullptr) << "Adapter not found for type: " << type;
@@ -688,4 +688,200 @@ TEST(AI_ADAPTER_TEST, voyage_adapter_chat_test) {
"[NOT_IMPLEMENTED_ERROR]VoyageAI don't support text generation");
}
+TEST(AI_ADAPTER_TEST, qwen_multimodal_embedding_request_image) {
+ QwenAdapter adapter;
+ TAIResource config;
+ config.model_name = "tongyi-embedding-vision-plus";
+ config.dimensions = 1024;
+ adapter.init(config);
+
+ std::string request_body;
+ Status st = adapter.build_multimodal_embedding_request(MultimodalType::IMAGE,
+ "https://a/b/c.png", request_body);
+ ASSERT_TRUE(st.ok()) << st.to_string();
+
+ rapidjson::Document doc;
+ doc.Parse(request_body.c_str());
+ ASSERT_FALSE(doc.HasParseError()) << request_body;
+ ASSERT_TRUE(doc.IsObject());
+ ASSERT_TRUE(doc.HasMember("model"));
+ ASSERT_STREQ(doc["model"].GetString(), "tongyi-embedding-vision-plus");
+ ASSERT_TRUE(doc.HasMember("input"));
+ ASSERT_TRUE(doc["input"].HasMember("contents"));
+ const auto& contents = doc["input"]["contents"];
+ ASSERT_TRUE(contents.IsArray());
+ ASSERT_EQ(contents.Size(), 1);
+ ASSERT_TRUE(contents[0].HasMember("image"));
+ ASSERT_STREQ(contents[0]["image"].GetString(), "https://a/b/c.png");
+ ASSERT_TRUE(doc.HasMember("parameters"));
+ ASSERT_TRUE(doc["parameters"].IsObject());
+ ASSERT_TRUE(doc["parameters"].HasMember("dimension"));
+ ASSERT_EQ(doc["parameters"]["dimension"].GetInt(), 1024);
+}
+
+TEST(AI_ADAPTER_TEST, qwen_multimodal_embedding_request_video) {
+ QwenAdapter adapter;
+ TAIResource config;
+ config.model_name = "tongyi-embedding-vision-plus";
+ config.dimensions = 1024;
+ adapter.init(config);
+
+ std::string request_body;
+ Status st = adapter.build_multimodal_embedding_request(MultimodalType::VIDEO,
+ "https://a/b/c.mp4", request_body);
+ ASSERT_TRUE(st.ok()) << st.to_string();
+
+ rapidjson::Document doc;
+ doc.Parse(request_body.c_str());
+ ASSERT_FALSE(doc.HasParseError()) << request_body;
+ ASSERT_TRUE(doc.IsObject());
+ ASSERT_TRUE(doc.HasMember("model"));
+ ASSERT_STREQ(doc["model"].GetString(), "tongyi-embedding-vision-plus");
+ ASSERT_TRUE(doc.HasMember("input"));
+ ASSERT_TRUE(doc["input"].HasMember("contents"));
+ const auto& contents = doc["input"]["contents"];
+ ASSERT_TRUE(contents.IsArray());
+ ASSERT_EQ(contents.Size(), 1);
+ ASSERT_TRUE(contents[0].HasMember("video"));
+ ASSERT_STREQ(contents[0]["video"].GetString(), "https://a/b/c.mp4");
+ ASSERT_TRUE(doc.HasMember("parameters"));
+ ASSERT_TRUE(doc["parameters"].IsObject());
+ ASSERT_TRUE(doc["parameters"].HasMember("dimension"));
+ ASSERT_EQ(doc["parameters"]["dimension"].GetInt(), 1024);
+}
+
+TEST(AI_ADAPTER_TEST, qwen_multimodal_embedding_request_audio_not_supported) {
+ QwenAdapter adapter;
+ TAIResource config;
+ config.model_name = "tongyi-embedding-vision-plus";
+ adapter.init(config);
+
+ std::string request_body;
+ Status st = adapter.build_multimodal_embedding_request(MultimodalType::AUDIO,
+ "https://a/b/c.mp3", request_body);
+ ASSERT_FALSE(st.ok());
+ ASSERT_THAT(st.to_string(),
+ ::testing::HasSubstr("QWEN only supports image/video multimodal embed"));
+ ASSERT_THAT(st.to_string(), ::testing::HasSubstr("audio"));
+}
+
+TEST(AI_ADAPTER_TEST, voyage_multimodal_embedding_request) {
+ VoyageAIAdapter adapter;
+ TAIResource config;
+ config.model_name = "voyage-multimodal-3.5";
+ config.dimensions = 2048;
+ adapter.init(config);
+
+ std::string request_body;
+ Status st = adapter.build_multimodal_embedding_request(MultimodalType::VIDEO,
+ "https://a/b/c.mp4", request_body);
+ ASSERT_TRUE(st.ok()) << st.to_string();
+
+ rapidjson::Document doc;
+ doc.Parse(request_body.c_str());
+ ASSERT_FALSE(doc.HasParseError()) << request_body;
+ ASSERT_TRUE(doc.IsObject());
+ ASSERT_TRUE(doc.HasMember("inputs"));
+ const auto& inputs = doc["inputs"];
+ ASSERT_TRUE(inputs.IsArray());
+ ASSERT_EQ(inputs.Size(), 1);
+ ASSERT_TRUE(inputs[0].HasMember("content"));
+ const auto& content = inputs[0]["content"];
+ ASSERT_TRUE(content.IsArray());
+ ASSERT_EQ(content.Size(), 1);
+ ASSERT_TRUE(content[0].HasMember("type"));
+ ASSERT_STREQ(content[0]["type"].GetString(), "video_url");
+ ASSERT_TRUE(content[0].HasMember("video_url"));
+ ASSERT_STREQ(content[0]["video_url"].GetString(), "https://a/b/c.mp4");
+ ASSERT_FALSE(doc.HasMember("dimensions"));
+ ASSERT_FALSE(doc.HasMember("output_dimension"));
+}
+
+TEST(AI_ADAPTER_TEST, jina_multimodal_embedding_request) {
+ JinaAdapter adapter;
+ TAIResource config;
+ config.model_name = "jina-embeddings-v4";
+ config.dimensions = 512;
+ adapter.init(config);
+
+ std::string request_body;
+ Status st = adapter.build_multimodal_embedding_request(MultimodalType::IMAGE,
+ "https://a/b/c.jpg", request_body);
+ ASSERT_TRUE(st.ok()) << st.to_string();
+
+ rapidjson::Document doc;
+ doc.Parse(request_body.c_str());
+ ASSERT_FALSE(doc.HasParseError()) << request_body;
+ ASSERT_TRUE(doc.IsObject());
+ ASSERT_TRUE(doc.HasMember("task"));
+ ASSERT_STREQ(doc["task"].GetString(), "text-matching");
+ ASSERT_TRUE(doc.HasMember("input"));
+ const auto& input = doc["input"];
+ ASSERT_TRUE(input.IsArray());
+ ASSERT_EQ(input.Size(), 1);
+ ASSERT_TRUE(input[0].HasMember("image"));
+ ASSERT_STREQ(input[0]["image"].GetString(), "https://a/b/c.jpg");
+ ASSERT_TRUE(doc.HasMember("dimensions"));
+ ASSERT_EQ(doc["dimensions"].GetInt(), 512);
+}
+
+TEST(AI_ADAPTER_TEST, multimodal_provider_support) {
+ OpenAIAdapter openai_adapter;
+ TAIResource openai_config;
+ openai_config.provider_type = "OPENAI";
+ openai_adapter.init(openai_config);
+
+ std::string request_body;
+ Status st = openai_adapter.build_multimodal_embedding_request(
+ MultimodalType::IMAGE, "https://a/b/c.png", request_body);
+ ASSERT_FALSE(st.ok());
+ ASSERT_THAT(st.to_string(), ::testing::HasSubstr("does not support multimodal Embed"));
+}
+
+TEST(AI_ADAPTER_TEST, gemini_multimodal_embedding_request) {
+ GeminiAdapter gemini_adapter;
+ TAIResource gemini_config;
+ gemini_config.provider_type = "GEMINI";
+ gemini_config.model_name = "gemini-embedding-2-preview";
+ gemini_config.dimensions = 768;
+ gemini_adapter.init(gemini_config);
+
+ struct GeminiMultimodalCase {
+ MultimodalType media_type;
+ const char* media_url;
+ const char* mime_type;
+ };
+ const std::vector<GeminiMultimodalCase> test_cases = {
+ {MultimodalType::IMAGE, "https://a/b/c.png", "image/png"},
+ {MultimodalType::AUDIO, "https://a/b/c.mp3", "audio/mpeg"},
+ {MultimodalType::VIDEO, "https://a/b/c.mp4", "video/mp4"},
+ };
+
+ for (const auto& test_case : test_cases) {
+ std::string request_body;
+ Status st = gemini_adapter.build_multimodal_embedding_request(
+ test_case.media_type, test_case.media_url, request_body);
+ ASSERT_TRUE(st.ok()) << st.to_string();
+
+ rapidjson::Document doc;
+ doc.Parse(request_body.c_str());
+ ASSERT_FALSE(doc.HasParseError()) << request_body;
+ ASSERT_TRUE(doc.IsObject());
+ ASSERT_TRUE(doc.HasMember("model"));
+ ASSERT_STREQ(doc["model"].GetString(), "models/gemini-embedding-2-preview");
+ ASSERT_TRUE(doc.HasMember("outputDimensionality"));
+ ASSERT_EQ(doc["outputDimensionality"].GetInt(), 768);
+ ASSERT_TRUE(doc.HasMember("content"));
+ ASSERT_TRUE(doc["content"].HasMember("parts"));
+ ASSERT_TRUE(doc["content"]["parts"].IsArray());
+ ASSERT_EQ(doc["content"]["parts"].Size(), 1);
+ ASSERT_TRUE(doc["content"]["parts"][0].HasMember("file_data"));
+ ASSERT_TRUE(doc["content"]["parts"][0]["file_data"].IsObject());
+ ASSERT_STREQ(doc["content"]["parts"][0]["file_data"]["mime_type"].GetString(),
+ test_case.mime_type);
+ ASSERT_STREQ(doc["content"]["parts"][0]["file_data"]["file_uri"].GetString(),
+ test_case.media_url);
+ }
+}
+
} // namespace doris
diff --git a/be/test/ai/ai_function_test.cpp b/be/test/ai/ai_function_test.cpp
index 415e9d2d222..8643c3d5c54 100644
--- a/be/test/ai/ai_function_test.cpp
+++ b/be/test/ai/ai_function_test.cpp
@@ -15,16 +15,22 @@
// specific language governing permissions and limitations
// under the License.
+#include <arpa/inet.h>
#include <gen_cpp/PaloInternalService_types.h>
#include <gtest/gtest.h>
+#include <netinet/in.h>
+#include <sys/socket.h>
+#include <unistd.h>
#include <string>
+#include <thread>
#include "core/block/block.h"
#include "core/column/column_array.h"
#include "core/column/column_string.h"
#include "core/column/column_vector.h"
#include "core/data_type/data_type_array.h"
+#include "exprs/function/ai/ai_adapter.h"
#include "exprs/function/ai/ai_classify.h"
#include "exprs/function/ai/ai_extract.h"
#include "exprs/function/ai/ai_filter.h"
@@ -41,6 +47,118 @@
namespace doris {
+class FunctionAITransportTestHelper : public FunctionAISentiment {
+public:
+ using FunctionAISentiment::do_send_request;
+ using FunctionAISentiment::execute_embedding_request;
+};
+
+class EmptyEmbeddingResultAdapter : public MockAdapter {
+public:
+ Status parse_embedding_response(const std::string& /*response_body*/,
+ std::vector<std::vector<float>>& /*results*/) const override {
+ return Status::OK();
+ }
+};
+
+class OneShotHttpServer {
+public:
+ OneShotHttpServer(int status_code, std::string response_body)
+ : _status_code(status_code), _response_body(std::move(response_body)) {
+ _listen_fd = socket(AF_INET, SOCK_STREAM, 0);
+ DCHECK_GE(_listen_fd, 0);
+
+ int reuse_addr = 1;
+ int ret = setsockopt(_listen_fd, SOL_SOCKET, SO_REUSEADDR, &reuse_addr, sizeof(reuse_addr));
+ DCHECK_EQ(ret, 0);
+
+ sockaddr_in addr {};
+ addr.sin_family = AF_INET;
+ addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
+ addr.sin_port = 0;
+ ret = bind(_listen_fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr));
+ DCHECK_EQ(ret, 0);
+ ret = listen(_listen_fd, 1);
+ DCHECK_EQ(ret, 0);
+
+ socklen_t addr_len = sizeof(addr);
+ ret = getsockname(_listen_fd, reinterpret_cast<sockaddr*>(&addr), &addr_len);
+ DCHECK_EQ(ret, 0);
+ _port = ntohs(addr.sin_port);
+
+ _thread = std::thread([this] { _serve_once(); });
+ }
+
+ ~OneShotHttpServer() {
+ if (_thread.joinable()) {
+ _thread.join();
+ }
+ if (_listen_fd >= 0) {
+ close(_listen_fd);
+ }
+ }
+
+ std::string endpoint() const { return "http://127.0.0.1:" + std::to_string(_port); }
+
+ std::string join_and_get_request() {
+ if (_thread.joinable()) {
+ _thread.join();
+ }
+ return _request;
+ }
+
+private:
+ void _serve_once() {
+ int client_fd = accept(_listen_fd, nullptr, nullptr);
+ DCHECK_GE(client_fd, 0);
+
+ char buffer[4096];
+ while (true) {
+ ssize_t read_bytes = recv(client_fd, buffer, sizeof(buffer), 0);
+ if (read_bytes <= 0) {
+ break;
+ }
+ _request.append(buffer, read_bytes);
+ if (_request.find("\r\n\r\n") != std::string::npos) {
+ auto header_end = _request.find("\r\n\r\n");
+ size_t body_length = 0;
+ size_t length_pos = _request.find("Content-Length:");
+ if (length_pos != std::string::npos) {
+ size_t value_begin = length_pos + sizeof("Content-Length:") - 1;
+ size_t value_end = _request.find("\r\n", value_begin);
+ std::string length_str = _request.substr(value_begin, value_end - value_begin);
+ body_length = std::stoul(length_str);
+ }
+ if (_request.size() >= header_end + 4 + body_length) {
+ break;
+ }
+ }
+ }
+
+ std::string status_text = _status_code == 200 ? "OK" : "Internal Server Error";
+ std::string response = "HTTP/1.1 " + std::to_string(_status_code) + " " + status_text +
+ "\r\nContent-Type: application/json\r\nContent-Length: " +
+ std::to_string(_response_body.size()) +
+ "\r\nConnection: close\r\n\r\n" + _response_body;
+ size_t sent_bytes = 0;
+ while (sent_bytes < response.size()) {
+ ssize_t sent =
+ send(client_fd, response.data() + sent_bytes, response.size() - sent_bytes, 0);
+ DCHECK_GT(sent, 0);
+ sent_bytes += sent;
+ }
+ shutdown(client_fd, SHUT_RDWR);
+ close(client_fd);
+ }
+
+ int _listen_fd = -1;
+ uint16_t _port = 0;
+ int _status_code;
+ std::string _response_body;
+ std::string _request;
+ std::thread _thread;
+};
+
TEST(AIFunctionTest, AISummarizeTest) {
FunctionAISummarize function;
@@ -369,6 +487,43 @@ TEST(AIFunctionTest, AISimilarityTrimWhitespace) {
unsetenv("AI_TEST_RESULT");
}
+TEST(AIFunctionTest, AISimilarityInvalidValue) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ std::vector<std::string> invalid_cases = {"abc", "1.2x", "", " ", "\n\t", "1,2"};
+
+ for (const auto& invalid_value : invalid_cases) {
+ setenv("AI_TEST_RESULT", invalid_value.c_str(), 1);
+
+ std::vector<std::string> resources = {"mock_resource"};
+ std::vector<std::string> text1 = {"Test text 1"};
+ std::vector<std::string> text2 = {"Test text 2"};
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+ auto col_text1 = ColumnHelper::create_column<DataTypeString>(text1);
+ auto col_text2 = ColumnHelper::create_column<DataTypeString>(text2);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(col_text1), std::make_shared<DataTypeString>(), "text1"});
+ block.insert({std::move(col_text2), std::make_shared<DataTypeString>(), "text2"});
+ block.insert({nullptr, std::make_shared<DataTypeFloat32>(), "result"});
+
+ ColumnNumbers arguments = {0, 1, 2};
+ size_t result_idx = 3;
+
+ auto similarity_func = FunctionAISimilarity::create();
+ Status exec_status = similarity_func->execute_impl(ctx.get(), block, arguments, result_idx,
+ text1.size());
+
+ ASSERT_FALSE(exec_status.ok())
+ << "Should have failed for invalid value: '" << invalid_value << "'";
+ ASSERT_NE(exec_status.to_string().find("Failed to parse float value"), std::string::npos);
+ }
+
+ unsetenv("AI_TEST_RESULT");
+}
+
TEST(AIFunctionTest, AIFilterTest) {
FunctionAIFilter function;
@@ -417,6 +572,39 @@ TEST(AIFunctionTest, AIFilterExecuteTest) {
ASSERT_TRUE(val == 0);
}
+TEST(AIFunctionTest, AIFilterExecuteMultipleRows) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ setenv("AI_TEST_RESULT", " 1 ", 1);
+
+ std::vector<std::string> resources = {"mock_resource", "mock_resource"};
+ std::vector<std::string> texts = {"This is valid.", "This is also valid."};
+ 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<DataTypeBool>(), "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto filter_func = FunctionAIFilter::create();
+ Status exec_status =
+ filter_func->execute_impl(ctx.get(), block, arguments, result_idx, texts.size());
+
+ unsetenv("AI_TEST_RESULT");
+
+ ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
+ const auto& res_col =
+ assert_cast<const ColumnUInt8&>(*block.get_by_position(result_idx).column);
+ ASSERT_EQ(res_col.size(), 2);
+ ASSERT_EQ(res_col.get_data()[0], 1);
+ ASSERT_EQ(res_col.get_data()[1], 1);
+}
+
TEST(AIFunctionTest, AIFilterTrimWhitespace) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
@@ -662,4 +850,137 @@ TEST(AIFunctionTest, NormalizeEndpointNoopForOtherPaths) {
ASSERT_EQ(resource.endpoint, "https://localhost/v1/responses");
}
-} // namespace doris
\ No newline at end of file
+TEST(AIFunctionTest, DoSendRequestTransportError) {
+ TQueryOptions query_options = create_fake_query_options();
+ query_options.__set_query_timeout(5);
+ auto query_ctx = MockQueryContext::create(TUniqueId(), ExecEnv::GetInstance(), 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, {}, {});
+
+ TAIResource config;
+ config.endpoint = "http://127.0.0.1:1";
+ config.provider_type = "OPENAI";
+ config.model_name = "test-model";
+ config.api_key = "secret";
+
+ std::shared_ptr<AIAdapter> adapter = std::make_shared<OpenAIAdapter>();
+ adapter->init(config);
+
+ HttpClient client;
+ std::string response;
+ FunctionAITransportTestHelper helper;
+ Status st = helper.do_send_request(&client, "{}", response, config, adapter, ctx.get());
+
+ ASSERT_FALSE(st.ok());
+ ASSERT_EQ(response, "");
+}
+
+TEST(AIFunctionTest, DoSendRequestNon200) {
+ TQueryOptions query_options = create_fake_query_options();
+ query_options.__set_query_timeout(5);
+ auto query_ctx = MockQueryContext::create(TUniqueId(), ExecEnv::GetInstance(), 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, {}, {});
+
+ OneShotHttpServer server(500, R"({"error":"bad request"})");
+
+ TAIResource config;
+ config.endpoint = server.endpoint();
+ config.provider_type = "OPENAI";
+ config.model_name = "test-model";
+ config.api_key = "secret";
+
+ std::shared_ptr<AIAdapter> adapter = std::make_shared<OpenAIAdapter>();
+ adapter->init(config);
+
+ HttpClient client;
+ std::string response;
+ FunctionAITransportTestHelper helper;
+ Status st = helper.do_send_request(&client, R"({"message":"hello"})", response, config, adapter,
+ ctx.get());
+
+ ASSERT_FALSE(st.ok());
+ ASSERT_NE(st.to_string().find("http status code is not 200"), std::string::npos);
+ ASSERT_EQ(response, R"({"error":"bad request"})");
+
+ std::string request = server.join_and_get_request();
+ ASSERT_NE(request.find("Authorization: Bearer secret"), std::string::npos);
+ ASSERT_NE(request.find("Content-Type: application/json"), std::string::npos);
+}
+
+TEST(AIFunctionTest, DoSendRequestSuccess) {
+ TQueryOptions query_options = create_fake_query_options();
+ query_options.__set_query_timeout(5);
+ auto query_ctx = MockQueryContext::create(TUniqueId(), ExecEnv::GetInstance(), 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, {}, {});
+
+ OneShotHttpServer server(200, R"({"ok":true})");
+
+ TAIResource config;
+ config.endpoint = server.endpoint();
+ config.provider_type = "OPENAI";
+ config.model_name = "test-model";
+ config.api_key = "secret";
+
+ std::shared_ptr<AIAdapter> adapter = std::make_shared<OpenAIAdapter>();
+ adapter->init(config);
+
+ HttpClient client;
+ std::string response;
+ FunctionAITransportTestHelper helper;
+ Status st = helper.do_send_request(&client, R"({"message":"hello"})", response, config, adapter,
+ ctx.get());
+
+ ASSERT_TRUE(st.ok()) << st.to_string();
+ ASSERT_EQ(response, R"({"ok":true})");
+
+ std::string request = server.join_and_get_request();
+ ASSERT_NE(request.find("POST "), std::string::npos);
+ ASSERT_NE(request.find(R"({"message":"hello"})"), std::string::npos);
+}
+
+TEST(AIFunctionTest, ExecuteEmbeddingRequestMockSuccess) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ TAIResource config;
+ config.provider_type = "MOCK";
+ std::shared_ptr<AIAdapter> adapter = std::make_shared<MockAdapter>();
+ adapter->init(config);
+
+ FunctionAITransportTestHelper helper;
+ std::vector<float> result;
+ Status st = helper.execute_embedding_request("{}", result, config, adapter, ctx.get());
+
+ ASSERT_TRUE(st.ok()) << st.to_string();
+ ASSERT_EQ(result.size(), 5);
+ for (size_t i = 0; i < result.size(); ++i) {
+ ASSERT_FLOAT_EQ(result[i], static_cast<float>(i));
+ }
+}
+
+TEST(AIFunctionTest, ExecuteEmbeddingRequestEmptyResult) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ TAIResource config;
+ config.provider_type = "MOCK";
+ std::shared_ptr<AIAdapter> adapter = std::make_shared<EmptyEmbeddingResultAdapter>();
+ adapter->init(config);
+
+ FunctionAITransportTestHelper helper;
+ std::vector<float> result;
+ Status st = helper.execute_embedding_request("{}", result, config, adapter, ctx.get());
+
+ ASSERT_FALSE(st.ok());
+ ASSERT_NE(st.to_string().find("AI returned empty result"), std::string::npos);
+}
+
+} // namespace doris
diff --git a/be/test/ai/embed_test.cpp b/be/test/ai/embed_test.cpp
index 0f3ef7fcd92..a296d21daea 100644
--- a/be/test/ai/embed_test.cpp
+++ b/be/test/ai/embed_test.cpp
@@ -25,7 +25,11 @@
#include <string>
#include <vector>
+#include "core/data_type/data_type_jsonb.h"
+#include "core/data_type/data_type_number.h"
+#include "core/value/jsonb_value.h"
#include "exprs/function/ai/ai_adapter.h"
+#include "io/fs/obj_storage_client.h"
#include "testutil/column_helper.h"
#include "testutil/mock/mock_runtime_state.h"
@@ -40,6 +44,101 @@ private:
std::string _content_type;
};
+class MockEmbedObjStorageClient : public io::ObjStorageClient {
+public:
+ io::ObjectStorageUploadResponse create_multipart_upload(
+ const io::ObjectStoragePathOptions& /*opts*/) override {
+ return {};
+ }
+
+ io::ObjectStorageResponse put_object(const io::ObjectStoragePathOptions& /*opts*/,
+ std::string_view /*stream*/) override {
+ return io::ObjectStorageResponse::OK();
+ }
+
+ io::ObjectStorageUploadResponse upload_part(const io::ObjectStoragePathOptions& /*opts*/,
+ std::string_view /*stream*/,
+ int /*part_num*/) override {
+ return {};
+ }
+
+ io::ObjectStorageResponse complete_multipart_upload(
+ const io::ObjectStoragePathOptions& /*opts*/,
+ const std::vector<io::ObjectCompleteMultiPart>& /*completed_parts*/) override {
+ return io::ObjectStorageResponse::OK();
+ }
+
+ io::ObjectStorageHeadResponse head_object(
+ const io::ObjectStoragePathOptions& /*opts*/) override {
+ return {};
+ }
+
+ io::ObjectStorageResponse get_object(const io::ObjectStoragePathOptions& /*opts*/,
+ void* /*buffer*/, size_t /*offset*/, size_t /*bytes_read*/,
+ size_t* /*size_return*/) override {
+ return io::ObjectStorageResponse::OK();
+ }
+
+ io::ObjectStorageResponse list_objects(const io::ObjectStoragePathOptions& /*opts*/,
+ std::vector<io::FileInfo>* /*files*/) override {
+ return io::ObjectStorageResponse::OK();
+ }
+
+ io::ObjectStorageResponse delete_objects(const io::ObjectStoragePathOptions& /*opts*/,
+ std::vector<std::string> /*objs*/) override {
+ return io::ObjectStorageResponse::OK();
+ }
+
+ io::ObjectStorageResponse delete_object(const io::ObjectStoragePathOptions& /*opts*/) override {
+ return io::ObjectStorageResponse::OK();
+ }
+
+ io::ObjectStorageResponse delete_objects_recursively(
+ const io::ObjectStoragePathOptions& /*opts*/) override {
+ return io::ObjectStorageResponse::OK();
+ }
+
+ std::string generate_presigned_url(const io::ObjectStoragePathOptions& opts,
+ int64_t expiration_secs, const S3ClientConf& conf) override {
+ last_opts = opts;
+ last_expiration_secs = expiration_secs;
+ last_conf = conf;
+ return fmt::format("mock-s3://{}/{}?ttl={}", opts.bucket, opts.key, expiration_secs);
+ }
+
+ io::ObjectStoragePathOptions last_opts;
+ int64_t last_expiration_secs = 0;
+ S3ClientConf last_conf;
+};
+
+static ColumnString::MutablePtr create_jsonb_column(const std::vector<std::string>& json_rows) {
+ auto column = ColumnString::create();
+ for (const auto& json_row : json_rows) {
+ JsonBinaryValue jsonb_value;
+ Status st = jsonb_value.from_json_string(json_row);
+ EXPECT_TRUE(st.ok()) << st.to_string();
+ column->insert_data(jsonb_value.value(), jsonb_value.size());
+ }
+ return column;
+}
+
+static void assert_mock_embedding_column(const ColumnArray& col_array, size_t row_count) {
+ const auto& offsets = col_array.get_offsets();
+ ASSERT_EQ(offsets.size(), row_count);
+
+ const auto& nested_nullable_col = assert_cast<const ColumnNullable&>(col_array.get_data());
+ const auto& nested_col =
+ assert_cast<const ColumnFloat32&>(*nested_nullable_col.get_nested_column_ptr());
+ ASSERT_EQ(nested_col.size(), row_count * 5);
+
+ for (size_t row = 0; row < row_count; ++row) {
+ ASSERT_EQ(offsets[row], (row + 1) * 5);
+ for (size_t i = 0; i < 5; ++i) {
+ ASSERT_FLOAT_EQ(nested_col.get_element(row * 5 + i), static_cast<float>(i));
+ }
+ }
+}
+
TEST(EMBED_TEST, embed_function_build_test) {
FunctionEmbed function;
@@ -73,7 +172,10 @@ TEST(EMBED_TEST, embed_function_test) {
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"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
ColumnNumbers arguments = {0, 1};
size_t result_idx = 2;
@@ -96,6 +198,316 @@ TEST(EMBED_TEST, embed_function_test) {
}
}
+TEST(EMBED_TEST, embed_function_text_multi_rows) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ std::vector<std::string> resources = {"mock_resource", "mock_resource"};
+ std::vector<std::string> texts = {"test input 1", "test input 2"};
+ 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<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status =
+ embed_func->execute_impl(ctx.get(), block, arguments, result_idx, texts.size());
+
+ ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
+ const auto& col_array =
+ assert_cast<const ColumnArray&>(*block.get_by_position(result_idx).column);
+ assert_mock_embedding_column(col_array, texts.size());
+}
+
+TEST(EMBED_TEST, embed_function_multimodal_direct_url) {
+ TQueryOptions query_options = create_fake_query_options();
+ query_options.__set_file_presigned_url_ttl_seconds(0);
+ auto query_ctx = MockQueryContext::create(TUniqueId(), ExecEnv::GetInstance(), query_options);
+ query_ctx->set_mock_ai_resource();
+ 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 = {"mock_resource", "mock_resource", "mock_resource"};
+ std::vector<std::string> file_json_rows = {
+ R"({"content_type":"image/png","uri":"https://example.com/a.png"})",
+ R"({"content_type":"video/mp4","uri":"https://example.com/b.mp4"})",
+ R"({"content_type":"audio/mpeg","uri":"https://example.com/c.mp3"})"};
+
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+ auto col_file = create_jsonb_column(file_json_rows);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(col_file), std::make_shared<DataTypeJsonb>(), "file"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status = embed_func->execute_impl(ctx.get(), block, arguments, result_idx,
+ file_json_rows.size());
+
+ ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
+ const auto& col_array =
+ assert_cast<const ColumnArray&>(*block.get_by_position(result_idx).column);
+ assert_mock_embedding_column(col_array, file_json_rows.size());
+}
+
+TEST(EMBED_TEST, embed_function_multimodal_s3_presigned_url) {
+ TQueryOptions query_options = create_fake_query_options();
+ query_options.__set_file_presigned_url_ttl_seconds(123);
+ auto query_ctx = MockQueryContext::create(TUniqueId(), ExecEnv::GetInstance(), query_options);
+ query_ctx->set_mock_ai_resource();
+ TQueryGlobals query_globals;
+ RuntimeState runtime_state(TUniqueId(), 0, query_options, query_globals, nullptr,
+ query_ctx.get());
+ auto ctx = FunctionContext::create_context(&runtime_state, {}, {});
+
+ auto mock_client = std::make_shared<MockEmbedObjStorageClient>();
+ S3ClientFactory::instance().set_client_creator_for_test(
+ [mock_client](const S3ClientConf&) { return mock_client; });
+
+ std::vector<std::string> resources = {"mock_resource"};
+ std::vector<std::string> file_json_rows = {R"({
+ "content_type":"image/png",
+ "uri":"s3://test-bucket/path/to/image.png",
+ "endpoint":"cos.ap-beijing.myqcloud.com",
+ "region":"ap-beijing",
+ "ak":"test-ak",
+ "sk":"test-sk",
+ "role_arn":"test-role",
+ "external_id":"test-external-id"
+ })"};
+
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+ auto col_file = create_jsonb_column(file_json_rows);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(col_file), std::make_shared<DataTypeJsonb>(), "file"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status = embed_func->execute_impl(ctx.get(), block, arguments, result_idx,
+ file_json_rows.size());
+
+ S3ClientFactory::instance().clear_client_creator_for_test();
+
+ ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
+ const auto& col_array =
+ assert_cast<const ColumnArray&>(*block.get_by_position(result_idx).column);
+ assert_mock_embedding_column(col_array, file_json_rows.size());
+
+ ASSERT_EQ(mock_client->last_opts.bucket, "test-bucket");
+ ASSERT_EQ(mock_client->last_opts.key, "path/to/image.png");
+ ASSERT_EQ(mock_client->last_expiration_secs, 123);
+ ASSERT_EQ(mock_client->last_conf.endpoint, "cos.ap-beijing.myqcloud.com");
+ ASSERT_EQ(mock_client->last_conf.region, "ap-beijing");
+ ASSERT_EQ(mock_client->last_conf.ak, "test-ak");
+ ASSERT_EQ(mock_client->last_conf.sk, "test-sk");
+ ASSERT_EQ(mock_client->last_conf.role_arn, "test-role");
+ ASSERT_EQ(mock_client->last_conf.external_id, "test-external-id");
+}
+
+TEST(EMBED_TEST, embed_function_multimodal_s3_missing_endpoint) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ std::vector<std::string> resources = {"mock_resource"};
+ std::vector<std::string> file_json_rows = {R"({
+ "content_type":"image/png",
+ "uri":"s3://test-bucket/path/to/image.png",
+ "region":"ap-beijing"
+ })"};
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+ auto col_file = create_jsonb_column(file_json_rows);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(col_file), std::make_shared<DataTypeJsonb>(), "file"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status = embed_func->execute_impl(ctx.get(), block, arguments, result_idx,
+ file_json_rows.size());
+
+ ASSERT_FALSE(exec_status.ok());
+ ASSERT_NE(exec_status.to_string().find("field 'endpoint' is required"), std::string::npos);
+}
+
+TEST(EMBED_TEST, embed_function_multimodal_s3_missing_region) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ std::vector<std::string> resources = {"mock_resource"};
+ std::vector<std::string> file_json_rows = {R"({
+ "content_type":"image/png",
+ "uri":"s3://test-bucket/path/to/image.png",
+ "endpoint":"cos.ap-beijing.myqcloud.com"
+ })"};
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+ auto col_file = create_jsonb_column(file_json_rows);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(col_file), std::make_shared<DataTypeJsonb>(), "file"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status = embed_func->execute_impl(ctx.get(), block, arguments, result_idx,
+ file_json_rows.size());
+
+ ASSERT_FALSE(exec_status.ok());
+ ASSERT_NE(exec_status.to_string().find("field 'region' is required"), std::string::npos);
+}
+
+TEST(EMBED_TEST, embed_function_wrong_argument_count) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ std::vector<std::string> resources = {"mock_resource"};
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0};
+ size_t result_idx = 1;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status =
+ embed_func->execute_impl(ctx.get(), block, arguments, result_idx, resources.size());
+
+ ASSERT_FALSE(exec_status.ok());
+ ASSERT_NE(exec_status.to_string().find("Function EMBED expects 2 arguments"),
+ std::string::npos);
+}
+
+TEST(EMBED_TEST, embed_function_invalid_input_type) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ std::vector<std::string> resources = {"mock_resource"};
+ std::vector<Int32> ids = {1};
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+ auto col_ids = ColumnHelper::create_column<DataTypeInt32>(ids);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(col_ids), std::make_shared<DataTypeInt32>(), "id"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status =
+ embed_func->execute_impl(ctx.get(), block, arguments, result_idx, resources.size());
+
+ ASSERT_FALSE(exec_status.ok());
+ ASSERT_NE(exec_status.to_string().find(
+ "Function EMBED expects the second argument to be STRING or JSON"),
+ std::string::npos);
+}
+
+TEST(EMBED_TEST, embed_function_missing_required_json_field) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ std::vector<std::string> resources = {"mock_resource"};
+ std::vector<std::string> file_json_rows = {R"({"uri":"https://example.com/a.png"})"};
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+ auto col_file = create_jsonb_column(file_json_rows);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(col_file), std::make_shared<DataTypeJsonb>(), "file"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status = embed_func->execute_impl(ctx.get(), block, arguments, result_idx,
+ file_json_rows.size());
+
+ ASSERT_FALSE(exec_status.ok());
+ ASSERT_NE(exec_status.to_string().find("field 'content_type' is required"), std::string::npos);
+}
+
+TEST(EMBED_TEST, embed_function_unsupported_content_type) {
+ auto runtime_state = std::make_unique<MockRuntimeState>();
+ auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
+
+ std::vector<std::string> resources = {"mock_resource"};
+ std::vector<std::string> file_json_rows = {
+ R"({"content_type":"text/plain","uri":"https://example.com/a.txt"})"};
+ auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
+ auto col_file = create_jsonb_column(file_json_rows);
+
+ Block block;
+ block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
+ block.insert({std::move(col_file), std::make_shared<DataTypeJsonb>(), "file"});
+ block.insert(
+ {nullptr,
+ std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())),
+ "result"});
+
+ ColumnNumbers arguments = {0, 1};
+ size_t result_idx = 2;
+
+ auto embed_func = FunctionEmbed::create();
+ Status exec_status = embed_func->execute_impl(ctx.get(), block, arguments, result_idx,
+ file_json_rows.size());
+
+ ASSERT_FALSE(exec_status.ok());
+ ASSERT_NE(exec_status.to_string().find("Unsupported content_type for EMBED"),
+ std::string::npos);
+}
+
TEST(EMBED_TEST, local_adapter_embedding_request) {
LocalAdapter adapter;
TAIResource config;
@@ -727,4 +1139,4 @@ TEST(EMBED_TEST, minimax_adapter_embedding_request) {
ASSERT_STREQ(doc["texts"][0].GetString(), "embed with minimax");
}
-} // namespace doris
\ No newline at end of file
+} // namespace doris
diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/Resource.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/Resource.java
index 274d9af23fa..96e510c1a83 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/catalog/Resource.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/Resource.java
@@ -28,6 +28,7 @@ import org.apache.doris.nereids.trees.plans.commands.CreateResourceCommand;
import org.apache.doris.nereids.trees.plans.commands.info.CreateResourceInfo;
import org.apache.doris.persist.gson.GsonPostProcessable;
import org.apache.doris.persist.gson.GsonUtils;
+import org.apache.doris.qe.ConnectContext;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableMap;
@@ -295,4 +296,11 @@ public abstract class Resource implements Writable, GsonPostProcessable {
}
public void applyDefaultProperties() {}
+
+ public static void registerUsedAIResourceName(String resourceName) {
+ ConnectContext ctx = ConnectContext.get();
+ if (ctx != null && ctx.getStatementContext() != null) {
+ ctx.getStatementContext().registerUsedAIResourceName(resourceName);
+ }
+ }
}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java
index 8bddd094954..6019aaef303 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java
@@ -55,7 +55,7 @@ public class AIProperties extends BaseProperties {
public static final List<String> REQUIRED_FIELDS = Arrays.asList(ENDPOINT, PROVIDER_TYPE, MODEL_NAME);
public static final List<String> PROVIDERS
= Arrays.asList("OPENAI", "LOCAL", "GEMINI", "DEEPSEEK", "ANTHROPIC",
- "MOONSHOT", "QWEN", "MINIMAX", "ZHIPU", "BAICHUAN", "VOYAGEAI");
+ "MOONSHOT", "QWEN", "MINIMAX", "ZHIPU", "BAICHUAN", "VOYAGEAI", "JINA");
public static void requiredAIProperties(Map<String, String> properties) throws DdlException {
// Check required field
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java
index a65da6ef8c3..520ef17fb0b 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java
@@ -69,6 +69,7 @@ import org.apache.doris.system.Backend;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Stopwatch;
+import com.google.common.base.Strings;
import com.google.common.base.Supplier;
import com.google.common.base.Suppliers;
import com.google.common.base.Throwables;
@@ -84,10 +85,12 @@ import java.io.Closeable;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.Collection;
+import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@@ -316,6 +319,7 @@ public class StatementContext implements Closeable {
private boolean hasNestedColumns;
private final Set<CTEId> mustInlineCTE = new HashSet<>();
+ private final Set<String> usedAIResourceNames = new LinkedHashSet<>();
private final Map<String, Integer> lowerCaseTableNamesCache = Maps.newHashMap();
private final Map<String, Integer> lowerCaseDatabaseNamesCache = Maps.newHashMap();
@@ -465,6 +469,17 @@ public class StatementContext implements Closeable {
return connectContext;
}
+ public Set<String> getUsedAIResourceNames() {
+ return Collections.unmodifiableSet(usedAIResourceNames);
+ }
+
+ public void registerUsedAIResourceName(String resourceName) {
+ if (Strings.isNullOrEmpty(resourceName)) {
+ throw new AnalysisException("AI resource name can not be empty");
+ }
+ usedAIResourceNames.add(resourceName);
+ }
+
public void setOriginStatement(OriginStatement originStatement) {
this.originStatement = originStatement;
if (originStatement != null && sqlCacheContext != null) {
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java
index a609019ab00..833f1b7db4d 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java
@@ -90,6 +90,7 @@ public class AIAgg extends NullableAggregateFunction
if (!(resource instanceof AIResource)) {
throw new AnalysisException("AI resource '" + resourceName + "' does not exist");
}
+ Resource.registerUsedAIResourceName(resourceName);
}
}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java
index 2826991208c..7399f2a348f 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java
@@ -61,6 +61,7 @@ public abstract class AIFunction extends ScalarFunction
if (!(resource instanceof AIResource)) {
throw new AnalysisException("AI resource '" + resourceName + "' does not exist");
}
+ Resource.registerUsedAIResourceName(resourceName);
}
}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java
index 3ce12e72043..e3e8725856f 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java
@@ -17,12 +17,18 @@
package org.apache.doris.nereids.trees.expressions.functions.ai;
+import org.apache.doris.catalog.AIResource;
+import org.apache.doris.catalog.Env;
import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.catalog.Resource;
+import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.literal.StringLikeLiteral;
import org.apache.doris.nereids.trees.expressions.literal.StringLiteral;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.types.ArrayType;
import org.apache.doris.nereids.types.FloatType;
+import org.apache.doris.nereids.types.JsonType;
import org.apache.doris.nereids.types.StringType;
import org.apache.doris.nereids.types.VarcharType;
@@ -41,7 +47,11 @@ public class Embed extends AIFunction {
FunctionSignature.ret(ArrayType.of(FloatType.INSTANCE))
.args(StringType.INSTANCE, StringType.INSTANCE),
FunctionSignature.ret(ArrayType.of(FloatType.INSTANCE))
- .args(VarcharType.SYSTEM_DEFAULT, VarcharType.SYSTEM_DEFAULT)
+ .args(VarcharType.SYSTEM_DEFAULT, VarcharType.SYSTEM_DEFAULT),
+ FunctionSignature.ret(ArrayType.of(FloatType.INSTANCE))
+ .args(StringType.INSTANCE, JsonType.INSTANCE),
+ FunctionSignature.ret(ArrayType.of(FloatType.INSTANCE))
+ .args(VarcharType.SYSTEM_DEFAULT, JsonType.INSTANCE)
);
/**
@@ -60,7 +70,7 @@ public class Embed extends AIFunction {
@Override
public Embed withChildren(List<Expression> children) {
- Preconditions.checkArgument(children.size() == 1 || children.size() == 2,
+ Preconditions.checkArgument(children.size() >= 1 && children.size() <= 2,
"Function EMBED only accepts 1 or 2 arguments");
if (children.size() == 1) {
return new Embed(new StringLiteral(getResourceName()),
@@ -83,4 +93,38 @@ public class Embed extends AIFunction {
public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
return visitor.visitEmbed(this, context);
}
+
+ @Override
+ public void checkLegalityBeforeTypeCoercion() {
+ if (arity() == 1) {
+ return;
+ }
+ if (arity() == 2) {
+ String aiResourceName = requireStringLiteral(child(0), "resource name",
+ "AI Function must accept literal for the resource name.");
+ validateAIResource(aiResourceName);
+ return;
+ }
+ throw new AnalysisException("Function EMBED only accepts 1 or 2 arguments");
+ }
+
+ private static String requireStringLiteral(Expression arg, String argName, String errorMsg) {
+ if (!(arg instanceof StringLikeLiteral)) {
+ throw new AnalysisException(errorMsg);
+ }
+ String value = ((StringLikeLiteral) arg).getStringValue();
+ if (value == null || value.isEmpty()) {
+ throw new AnalysisException("EMBED " + argName + " can not be empty.");
+ }
+ return value;
+ }
+
+ private static void validateAIResource(String resourceName) {
+ Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName);
+ if (!(resource instanceof AIResource)) {
+ throw new AnalysisException("AI resource '" + resourceName + "' does not exist");
+ }
+ Resource.registerUsedAIResourceName(resourceName);
+ }
+
}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java
index d73dad2fd6b..5181a77b042 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java
@@ -1667,6 +1667,21 @@ public class Coordinator implements CoordInterface {
return backend.getArrowFlightAddress();
}
+ private Map<String, TAIResource> getNeededAiResources() {
+ Map<String, TAIResource> aiResourceMap = Maps.newLinkedHashMap();
+ if (context == null || context.getStatementContext() == null) {
+ return aiResourceMap;
+ }
+ for (String resourceName : context.getStatementContext().getUsedAIResourceNames()) {
+ Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName);
+ if (!(resource instanceof AIResource)) {
+ throw new IllegalStateException("AI resource '" + resourceName + "' does not exist");
+ }
+ aiResourceMap.put(resourceName, ((AIResource) resource).toThrift());
+ }
+ return aiResourceMap;
+ }
+
// estimate if this fragment contains UnionNode
private boolean containsUnionNode(PlanNode node) {
if (node instanceof UnionNode) {
@@ -3280,15 +3295,7 @@ public class Coordinator implements CoordInterface {
}
// 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);
+ params.setAiResources(getNeededAiResources());
res.put(instanceExecParam.host, params);
res.get(instanceExecParam.host).setBucketSeqToInstanceIdx(new HashMap<Integer, Integer>());
res.get(instanceExecParam.host).setShuffleIdxToInstanceIdx(new HashMap<Integer, Integer>());
@@ -3536,4 +3543,3 @@ public class Coordinator implements CoordInterface {
this.queryOptions.setEnableProfile(isSafe && queryOptions.isEnableProfile());
}
}
-
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
index 78fc4367b66..6f86e2bc71a 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
@@ -958,6 +958,7 @@ public class SessionVariable implements Serializable, Writable {
public static final String ENABLE_STRICT_CAST = "enable_strict_cast";
public static final String DEFAULT_AI_RESOURCE = "default_ai_resource";
+ public static final String FILE_PRESIGNED_URL_TTL_SECONDS = "file_presigned_url_ttl_seconds";
public static final String HNSW_EF_SEARCH = "hnsw_ef_search";
public static final String HNSW_CHECK_RELATIVE_DISTANCE = "hnsw_check_relative_distance";
public static final String HNSW_BOUNDED_QUEUE = "hnsw_bounded_queue";
@@ -3455,6 +3456,13 @@ public class SessionVariable implements Serializable, Writable {
})
public String defaultAIResource = "";
+ @VarAttrDef.VarAttr(name = FILE_PRESIGNED_URL_TTL_SECONDS, needForward = true,
+ description = {
+ "EMBED 多模态场景中，S3 预签名 URL 的有效期（秒）。",
+ "Expiration time in seconds for S3 presigned URL used by multimodal EMBED."
+ })
+ public long filePresignedUrlTtlSeconds = 3600;
+
public void setEnableEsParallelScroll(boolean enableESParallelScroll) {
this.enableESParallelScroll = enableESParallelScroll;
}
@@ -5427,6 +5435,7 @@ public class SessionVariable implements Serializable, Writable {
tResult.setEnableParquetFilePageCache(enableParquetFilePageCache);
tResult.setEnableOrcFilterByMinMax(enableOrcFilterByMinMax);
tResult.setEnablePaimonCppReader(enablePaimonCppReader);
+ tResult.setFilePresignedUrlTtlSeconds(filePresignedUrlTtlSeconds);
tResult.setCheckOrcInitSargsSuccess(checkOrcInitSargsSuccess);
tResult.setTruncateCharOrVarcharColumns(truncateCharOrVarcharColumns);
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java
index 41b4ac9e765..4f8499970c9 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java
@@ -24,6 +24,7 @@ import org.apache.doris.catalog.Env;
import org.apache.doris.catalog.Resource;
import org.apache.doris.common.Config;
import org.apache.doris.datasource.FileQueryScanNode;
+import org.apache.doris.nereids.StatementContext;
import org.apache.doris.nereids.trees.plans.distribute.DistributedPlan;
import org.apache.doris.nereids.trees.plans.distribute.PipelineDistributedPlan;
import org.apache.doris.nereids.trees.plans.distribute.worker.DistributedPlanWorker;
@@ -213,6 +214,27 @@ public class ThriftPlansBuilder {
});
}
+ static Map<String, TAIResource> collectAiResources(ConnectContext connectContext) {
+ Map<String, TAIResource> aiResourceMap = Maps.newLinkedHashMap();
+ if (connectContext == null) {
+ return aiResourceMap;
+ }
+
+ StatementContext statementContext = connectContext.getStatementContext();
+ if (statementContext == null) {
+ return aiResourceMap;
+ }
+
+ for (String resourceName : statementContext.getUsedAIResourceNames()) {
+ Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName);
+ if (!(resource instanceof AIResource)) {
+ throw new IllegalStateException("AI resource '" + resourceName + "' does not exist");
+ }
+ aiResourceMap.put(resourceName, ((AIResource) resource).toThrift());
+ }
+ return aiResourceMap;
+ }
+
private static void setParamsForOlapTableSink(List<PipelineDistributedPlan> distributedPlans,
Map<DistributedPlanWorker, TPipelineFragmentParamsList> fragmentsGroupByWorker,
CoordinatorContext coordinatorContext) {
@@ -426,13 +448,7 @@ public class ThriftPlansBuilder {
params.setShuffleIdxToInstanceIdx(computeDestIdToInstanceId(fragmentPlan, w, instanceToIndex));
// Only 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);
+ params.setAiResources(collectAiResources(connectContext));
return params;
});
diff --git a/gensrc/thrift/PaloInternalService.thrift b/gensrc/thrift/PaloInternalService.thrift
index 65c567365b6..968971c395e 100644
--- a/gensrc/thrift/PaloInternalService.thrift
+++ b/gensrc/thrift/PaloInternalService.thrift
@@ -439,6 +439,7 @@ struct TQueryOptions {
200: optional bool enable_adjust_conjunct_order_by_cost;
// Use paimon-cpp to read Paimon splits on BE
201: optional bool enable_paimon_cpp_reader = false;
+
// Whether all fragments of this query are assigned to a single backend.
// When true, the streaming aggregation operator can use more aggressive
// hash table expansion thresholds since all data is local.
@@ -474,6 +475,9 @@ struct TQueryOptions {
210: optional double max_scan_mem_ratio = 0.3;
211: optional bool enable_adaptive_scan = false;
212: optional bool enable_local_exchange_before_agg = true;
+
+ 213: optional i64 file_presigned_url_ttl_seconds = 3600;
+
// For cloud, to control if the content would be written into file cache
// In write path, to control if the content would be written into file cache.
// In read path, read from file cache or remote storage when execute query.
diff --git a/regression-test/suites/ai_p0/test_ai_functions.groovy b/regression-test/suites/ai_p0/test_ai_functions.groovy
index 7ff9ff7777a..57cc2040e14 100644
--- a/regression-test/suites/ai_p0/test_ai_functions.groovy
+++ b/regression-test/suites/ai_p0/test_ai_functions.groovy
@@ -100,7 +100,10 @@ suite("test_ai_functions") {
sql """${sql_text}"""
assertTrue(false)
} catch (Exception e) {
- assertTrue(e.getMessage().contains("timeout") || e.getMessage().contains("requested URL returned error"))
+ assertTrue(e.getMessage().contains("timeout")
+ || e.getMessage().contains("requested URL returned error")
+ || e.getMessage().contains("http status code is not 200"),
+ "Unexpected exception message: " + e.getMessage())
} finally {
sql """UNSET VARIABLE query_timeout;"""
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@doris.apache.org
For additional commands, e-mail: commits-help@doris.apache.org