be/src/exprs/function/ai/embed.h

Source
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#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/defer_op.h"
#include "util/jsonb_utils.h"
#include "util/s3_uri.h"
#include "util/s3_util.h"

namespace doris {
class FunctionEmbed : public AIFunction<FunctionEmbed> {
public:
    static constexpr auto name = "embed";

    static constexpr size_t number_of_arguments = 2;

    static constexpr auto system_prompt = "";

    DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
        return std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>()));
    }

    using PreparedFunctionImpl::execute;

    AIResource select_ai_resource(const TAIResource& resource, PrimitiveType input_type) const {
        bool has_complete_multimodal_embed_properties = _has_complete_resource_properties(
                resource.embed_mm_endpoint, resource.embed_mm_provider_type,
                resource.embed_mm_model_name, resource.embed_mm_api_key);
        if (input_type == PrimitiveType::TYPE_JSONB && has_complete_multimodal_embed_properties) {
            return AIResource::from_multimodal_embed(resource);
        }
        bool has_complete_embed_properties = _has_complete_resource_properties(
                resource.embed_endpoint, resource.embed_provider_type, resource.embed_model_name,
                resource.embed_api_key);
        return has_complete_embed_properties ? AIResource::from_embed(resource)
                                             : AIResource(resource);
    }

    Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                   uint32_t result, size_t input_rows_count, const AIResource& config,
                   std::shared_ptr<AIAdapter>& adapter) const {
        if (arguments.size() != 2) {
            return Status::InvalidArgument("Function EMBED expects 2 arguments, but got {}",
                                           arguments.size());
        }

        const auto& input = block.get_by_position(arguments[1]);
        ColumnUInt8::MutablePtr result_null_map;
        if (input.type->is_nullable()) {
            const auto& [column, is_const] = unpack_if_const(input.column);
            const auto& nullable =
                    assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*column);
            result_null_map = ColumnUInt8::create(input_rows_count, 0);
            VectorizedUtils::update_null_map(result_null_map->get_data(),
                                             nullable.get_null_map_data(), is_const);
        }

        record_input(context, input_rows_count, result_null_map);

        if (result_null_map &&
            !simd::contain_zero(result_null_map->get_data().data(), input_rows_count)) {
            block.get_by_position(result).column =
                    block.get_by_position(result).type->create_column_const(input_rows_count,
                                                                            Field());
            return Status::OK();
        }

        ColumnPtr input_column =
                input.unnest_nullable(input.type->is_nullable() ? input.get_nullable_column_info()
                                                                : NullableColumnInfo {},
                                      false)
                        .column;
        PrimitiveType input_type = remove_nullable(input.type)->get_primitive_type();
        if (input_type == PrimitiveType::TYPE_JSONB) {
            return _execute_multimodal_embed(context, block, result, input_rows_count, config,
                                             adapter, input_column, std::move(result_null_map));
        }
        if (input_type == PrimitiveType::TYPE_STRING || input_type == PrimitiveType::TYPE_VARCHAR ||
            input_type == PrimitiveType::TYPE_CHAR) {
            return _execute_text_embed(context, block, result, input_rows_count, config, adapter,
                                       input_column, std::move(result_null_map));
        }
        return Status::InvalidArgument(
                "Function EMBED expects the second argument to be STRING or JSON, but got type {}",
                input.type->get_name());
    }

    static FunctionPtr create() { return std::make_shared<FunctionEmbed>(); }

private:
    static bool _has_complete_resource_properties(std::string_view endpoint,
                                                  std::string_view provider_type,
                                                  std::string_view model_name,
                                                  std::string_view api_key) {
        return !endpoint.empty() && !provider_type.empty() && !model_name.empty() &&
               (provider_type == "LOCAL" || !api_key.empty());
    }

    static int32_t _get_embed_max_batch_size(FunctionContext* context) {
        QueryContext* query_ctx = context->state()->get_query_ctx();
        DORIS_CHECK(query_ctx != nullptr);

        return query_ctx->query_options().embed_max_batch_size;
    }

    Status _execute_text_embed(FunctionContext* context, Block& block, uint32_t result,
                               size_t input_rows_count, const AIResource& config,
                               std::shared_ptr<AIAdapter>& adapter, const ColumnPtr& input_column,
                               ColumnUInt8::MutablePtr result_null_map) const {
        auto col_result = ColumnArray::create(
                ColumnNullable::create(ColumnFloat32::create(), ColumnUInt8::create()));
        std::vector<std::string> batch_prompts;
        size_t current_batch_size = 0;
        const int32_t max_batch_size = _get_embed_max_batch_size(context);
        const size_t max_context_window_size =
                static_cast<size_t>(get_ai_context_window_size(context));
        const NullMap* null_map = result_null_map ? &result_null_map->get_data() : nullptr;
        const Columns prompt_columns {input_column};

        for (size_t i = 0; i < input_rows_count; ++i) {
            if (null_map && (*null_map)[i]) {
                continue;
            }

            std::string prompt;
            RETURN_IF_ERROR(build_prompt(prompt_columns, i, prompt));

            const size_t prompt_size = prompt.size();

            if (prompt_size > max_context_window_size) {
                // flush history batch
                RETURN_IF_ERROR(_flush_text_embedding_batch(batch_prompts, *col_result, config,
                                                            adapter, context));
                current_batch_size = 0;

                batch_prompts.emplace_back(std::move(prompt));
                RETURN_IF_ERROR(_flush_text_embedding_batch(batch_prompts, *col_result, config,
                                                            adapter, context));
                continue;
            }

            if (!batch_prompts.empty() &&
                (current_batch_size + prompt_size > max_context_window_size ||
                 batch_prompts.size() >= static_cast<size_t>(max_batch_size))) {
                RETURN_IF_ERROR(_flush_text_embedding_batch(batch_prompts, *col_result, config,
                                                            adapter, context));
                current_batch_size = 0;
            }

            batch_prompts.emplace_back(std::move(prompt));
            current_batch_size += prompt_size;
        }

        RETURN_IF_ERROR(
                _flush_text_embedding_batch(batch_prompts, *col_result, config, adapter, context));

        block.replace_by_position(result, _expand_and_wrap_nullable_result(
                                                  std::move(col_result), std::move(result_null_map),
                                                  input_rows_count));
        return Status::OK();
    }

    Status _execute_multimodal_embed(FunctionContext* context, Block& block, uint32_t result,
                                     size_t input_rows_count, const AIResource& config,
                                     std::shared_ptr<AIAdapter>& adapter,
                                     const ColumnPtr& input_column,
                                     ColumnUInt8::MutablePtr result_null_map) const {
        auto col_result = ColumnArray::create(
                ColumnNullable::create(ColumnFloat32::create(), ColumnUInt8::create()));
        std::vector<MultimodalType> batch_media_types;
        std::vector<std::string> batch_media_content_types;
        std::vector<std::string> batch_media_urls;
        const NullMap* null_map = result_null_map ? &result_null_map->get_data() : nullptr;

        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 int32_t max_batch_size = _get_embed_max_batch_size(context);

        for (size_t i = 0; i < input_rows_count; ++i) {
            if (null_map && (*null_map)[i]) {
                continue;
            }

            rapidjson::Document file_input;
            RETURN_IF_ERROR(_parse_file_input(*input_column, i, file_input));

            std::string content_type;
            MultimodalType media_type;
            RETURN_IF_ERROR(_infer_media_type(file_input, content_type, media_type));

            std::string media_url;
            RETURN_IF_ERROR(_resolve_media_url(file_input, ttl_seconds, media_url));

            if (!batch_media_urls.empty() &&
                batch_media_urls.size() >= static_cast<size_t>(max_batch_size)) {
                RETURN_IF_ERROR(_flush_multimodal_embedding_batch(
                        batch_media_types, batch_media_content_types, batch_media_urls, *col_result,
                        config, adapter, context));
            }

            batch_media_types.emplace_back(media_type);
            batch_media_content_types.emplace_back(std::move(content_type));
            batch_media_urls.emplace_back(std::move(media_url));
        }

        RETURN_IF_ERROR(_flush_multimodal_embedding_batch(
                batch_media_types, batch_media_content_types, batch_media_urls, *col_result, config,
                adapter, context));

        block.replace_by_position(result, _expand_and_wrap_nullable_result(
                                                  std::move(col_result), std::move(result_null_map),
                                                  input_rows_count));
        return Status::OK();
    }

    // EMBED-private helper.
    // Sends one embedding request with a prebuilt request body and validates returned row count.
    Status _execute_prebuilt_embedding_request(const std::string& request_body,
                                               std::vector<std::vector<float>>& results,
                                               size_t expected_size, const AIResource& config,
                                               std::shared_ptr<AIAdapter>& adapter,
                                               FunctionContext* context) const {
        std::string response;
#ifdef BE_TEST
        if (config.provider_type == "MOCK") {
            results.clear();
            results.reserve(expected_size);
            for (size_t i = 0; i < expected_size; ++i) {
                results.emplace_back(std::initializer_list<float> {0, 1, 2, 3, 4});
            }
            return Status::OK();
        }
#endif

        RETURN_IF_ERROR(
                this->send_request_to_llm(request_body, response, config, adapter, context));

        const auto& statistics = get_executor(context)->statistics();
        const auto parse_start = statistics ? MonotonicNanos() : 0;
        bool parse_failed = true;
        Defer record_parse([&] {
            if (statistics) {
                statistics->record_parse(MonotonicNanos() - parse_start, parse_failed);
            }
        });
        RETURN_IF_ERROR(adapter->parse_embedding_response(response, results));
        if (results.empty()) {
            return Status::InternalError("AI returned empty result");
        }
        if (results.size() != expected_size) [[unlikely]] {
            return Status::InternalError(
                    "AI embedding returned {} results, but {} inputs were sent", results.size(),
                    expected_size);
        }
        parse_failed = false;
        return Status::OK();
    }

    // EMBED-private helper.
    // Flushes one accumulated text embedding batch into the output array column.
    Status _flush_text_embedding_batch(std::vector<std::string>& batch_prompts,
                                       ColumnArray& col_result, const AIResource& config,
                                       std::shared_ptr<AIAdapter>& adapter,
                                       FunctionContext* context) const {
        if (batch_prompts.empty()) {
            return Status::OK();
        }

        std::string request_body;
        RETURN_IF_ERROR(adapter->build_embedding_request(batch_prompts, request_body));
        std::vector<std::vector<float>> batch_results;
        RETURN_IF_ERROR(_execute_prebuilt_embedding_request(
                request_body, batch_results, batch_prompts.size(), config, adapter, context));
        for (const auto& batch_result : batch_results) {
            _insert_embedding_result(col_result, batch_result);
        }
        batch_prompts.clear();
        return Status::OK();
    }

    // EMBED-private helper.
    // Flushes one accumulated multimodal embedding batch into the output array column.
    Status _flush_multimodal_embedding_batch(std::vector<MultimodalType>& batch_media_types,
                                             std::vector<std::string>& batch_media_content_types,
                                             std::vector<std::string>& batch_media_urls,
                                             ColumnArray& col_result, const AIResource& config,
                                             std::shared_ptr<AIAdapter>& adapter,
                                             FunctionContext* context) const {
        if (batch_media_urls.empty()) {
            return Status::OK();
        }

        std::string request_body;
        RETURN_IF_ERROR(adapter->build_multimodal_embedding_request(
                batch_media_types, batch_media_urls, batch_media_content_types, request_body));

        std::vector<std::vector<float>> batch_results;
        RETURN_IF_ERROR(_execute_prebuilt_embedding_request(
                request_body, batch_results, batch_media_urls.size(), config, adapter, context));
        for (const auto& batch_result : batch_results) {
            _insert_embedding_result(col_result, batch_result);
        }
        batch_media_types.clear();
        batch_media_content_types.clear();
        batch_media_urls.clear();
        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 ColumnPtr _expand_and_wrap_nullable_result(ColumnArray::MutablePtr result,
                                                      ColumnUInt8::MutablePtr result_null_map,
                                                      size_t input_rows_count) {
        if (!result_null_map) {
            return result;
        }

        auto& offsets = result->get_offsets();
        size_t compact_row = offsets.size();
        offsets.resize(input_rows_count);
        // For example, embedding rows 1 and 3 produces compact offsets [5, 10]. Given
        // result_null_map [1, 0, 1, 0, 1], expand them to [0, 5, 5, 10, 10], where NULL rows
        // reuse the previous offset. Fill backwards to avoid overwriting unread compact offsets.
        for (size_t row = input_rows_count; row-- > 0;) {
            if (result_null_map->get_data()[row]) {
                offsets[row] = compact_row == 0 ? 0 : offsets[compact_row - 1];
            } else {
                offsets[row] = offsets[--compact_row];
            }
        }
        return ColumnNullable::create(std::move(result), std::move(result_null_map));
    }

    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, std::string& content_type,
                                    MultimodalType& media_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 IColumn& file_column, size_t row_num,
                                    rapidjson::Document& file_input) {
        StringRef file_ref = file_column.get_data_at(row_num);
        std::string file_json = JsonbToJson::jsonb_to_json_string(file_ref.data, file_ref.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 = DORIS_TRY(S3ClientFactory::instance().create(s3_client_conf));

        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, .prefix = ""},
                                                      ttl_seconds);
        return Status::OK();
    }
};

}; // namespace doris

Coverage Report

Created: 2026-10-14 19:55

Line	Count	Source
1		// Licensed to the Apache Software Foundation (ASF) under one
2		// or more contributor license agreements. See the NOTICE file
3		// distributed with this work for additional information
4		// regarding copyright ownership. The ASF licenses this file
5		// to you under the Apache License, Version 2.0 (the
6		// "License"); you may not use this file except in compliance
7		// with the License. You may obtain a copy of the License at
8		//
9		// http://www.apache.org/licenses/LICENSE-2.0
10		//
11		// Unless required by applicable law or agreed to in writing,
12		// software distributed under the License is distributed on an
13		// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14		// KIND, either express or implied. See the License for the
15		// specific language governing permissions and limitations
16		// under the License.
17
18		#pragma once
19
20		#include <glog/logging.h>
21		#include <rapidjson/document.h>
22
23		#include <string_view>
24
25		#include "core/data_type/data_type_nullable.h"
26		#include "core/data_type/primitive_type.h"
27		#include "exprs/function/ai/ai_functions.h"
28		#include "util/defer_op.h"
29		#include "util/jsonb_utils.h"
30		#include "util/s3_uri.h"
31		#include "util/s3_util.h"
32
33		namespace doris {
34		class FunctionEmbed : public AIFunction<FunctionEmbed> {
35		public:
36		static constexpr auto name = "embed";
37
38		static constexpr size_t number_of_arguments = 2;
39
40		static constexpr auto system_prompt = "";
41
42	5	DataTypePtr get_nested_return_type_impl(const DataTypes& /arguments/) const {
43	5	return std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>()));
44	5	}
45
46		using PreparedFunctionImpl::execute;
47
48	23	AIResource select_ai_resource(const TAIResource& resource, PrimitiveType input_type) const {
49	23	bool has_complete_multimodal_embed_properties = _has_complete_resource_properties(
50	23	resource.embed_mm_endpoint, resource.embed_mm_provider_type,
51	23	resource.embed_mm_model_name, resource.embed_mm_api_key);
52	23	if (input_type == PrimitiveType::TYPE_JSONB && has_complete_multimodal_embed_properties) {
53	1	return AIResource::from_multimodal_embed(resource);
54	1	}
55	22	bool has_complete_embed_properties = _has_complete_resource_properties(
56	22	resource.embed_endpoint, resource.embed_provider_type, resource.embed_model_name,
57	22	resource.embed_api_key);
58	22	return has_complete_embed_properties ? AIResource::from_embed(resource)
59	22	: AIResource(resource);
60	23	}
61
62		Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
63		uint32_t result, size_t input_rows_count, const AIResource& config,
64	29	std::shared_ptr<AIAdapter>& adapter) const {
65	29	if (arguments.size() != 2) {
66	1	return Status::InvalidArgument("Function EMBED expects 2 arguments, but got {}",
67	1	arguments.size());
68	1	}
69
70	28	const auto& input = block.get_by_position(arguments[1]);
71	28	ColumnUInt8::MutablePtr result_null_map;
72	28	if (input.type->is_nullable()) {
73	10	const auto& [column, is_const] = unpack_if_const(input.column);
74	10	const auto& nullable =
75	10	assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*column);
76	10	result_null_map = ColumnUInt8::create(input_rows_count, 0);
77	10	VectorizedUtils::update_null_map(result_null_map->get_data(),
78	10	nullable.get_null_map_data(), is_const);
79	10	}
80
81	28	record_input(context, input_rows_count, result_null_map);
82
83	28	if (result_null_map &&
84	28	!simd::contain_zero(result_null_map->get_data().data(), input_rows_count)) {
85	2	block.get_by_position(result).column =
86	2	block.get_by_position(result).type->create_column_const(input_rows_count,
87	2	Field());
88	2	return Status::OK();
89	2	}
90
91	26	ColumnPtr input_column =
92	26	input.unnest_nullable(input.type->is_nullable() ? input.get_nullable_column_info()
93	26	: NullableColumnInfo {},
94	26	false)
95	26	.column;
96	26	PrimitiveType input_type = remove_nullable(input.type)->get_primitive_type();
97	26	if (input_type == PrimitiveType::TYPE_JSONB) {
98	14	return _execute_multimodal_embed(context, block, result, input_rows_count, config,
99	14	adapter, input_column, std::move(result_null_map));
100	14	}
101	12	if (input_type == PrimitiveType::TYPE_STRING \|\| input_type == PrimitiveType::TYPE_VARCHAR \|\|
102	12	input_type == PrimitiveType::TYPE_CHAR) {
103	11	return _execute_text_embed(context, block, result, input_rows_count, config, adapter,
104	11	input_column, std::move(result_null_map));
105	11	}
106	1	return Status::InvalidArgument(
107	1	"Function EMBED expects the second argument to be STRING or JSON, but got type {}",
108	1	input.type->get_name());
109	12	}
110
111	24	static FunctionPtr create() { return std::make_shared<FunctionEmbed>(); }
112
113		private:
114		static bool _has_complete_resource_properties(std::string_view endpoint,
115		std::string_view provider_type,
116		std::string_view model_name,
117	45	std::string_view api_key) {
118	45	return !endpoint.empty() && !provider_type.empty() && !model_name.empty() &&
119	45	(provider_type == "LOCAL" \|\| !api_key.empty());
120	45	}
121
122	25	static int32_t _get_embed_max_batch_size(FunctionContext* context) {
123	25	QueryContext* query_ctx = context->state()->get_query_ctx();
124	25	DORIS_CHECK(query_ctx != nullptr);
125
126	25	return query_ctx->query_options().embed_max_batch_size;
127	25	}
128
129		Status _execute_text_embed(FunctionContext* context, Block& block, uint32_t result,
130		size_t input_rows_count, const AIResource& config,
131		std::shared_ptr<AIAdapter>& adapter, const ColumnPtr& input_column,
132	11	ColumnUInt8::MutablePtr result_null_map) const {
133	11	auto col_result = ColumnArray::create(
134	11	ColumnNullable::create(ColumnFloat32::create(), ColumnUInt8::create()));
135	11	std::vector<std::string> batch_prompts;
136	11	size_t current_batch_size = 0;
137	11	const int32_t max_batch_size = _get_embed_max_batch_size(context);
138	11	const size_t max_context_window_size =
139	11	static_cast<size_t>(get_ai_context_window_size(context));
140	11	const NullMap* null_map = result_null_map ? &result_null_map->get_data() : nullptr;
141	11	const Columns prompt_columns {input_column};
142
143	49	for (size_t i = 0; i < input_rows_count; ++i) {
144	38	if (null_map && (*null_map)[i]) {
145	14	continue;
146	14	}
147
148	24	std::string prompt;
149	24	RETURN_IF_ERROR(build_prompt(prompt_columns, i, prompt));
150
151	24	const size_t prompt_size = prompt.size();
152
153	24	if (prompt_size > max_context_window_size) {
154		// flush history batch
155	10	RETURN_IF_ERROR(_flush_text_embedding_batch(batch_prompts, *col_result, config,
156	10	adapter, context));
157	10	current_batch_size = 0;
158
159	10	batch_prompts.emplace_back(std::move(prompt));
160	10	RETURN_IF_ERROR(_flush_text_embedding_batch(batch_prompts, *col_result, config,
161	10	adapter, context));
162	10	continue;
163	10	}
164
165	14	if (!batch_prompts.empty() &&
166	14	(current_batch_size + prompt_size > max_context_window_size \|\|
167	8	batch_prompts.size() >= static_cast<size_t>(max_batch_size))) {
168	3	RETURN_IF_ERROR(_flush_text_embedding_batch(batch_prompts, *col_result, config,
169	3	adapter, context));
170	3	current_batch_size = 0;
171	3	}
172
173	14	batch_prompts.emplace_back(std::move(prompt));
174	14	current_batch_size += prompt_size;
175	14	}
176
177	11	RETURN_IF_ERROR(
178	11	_flush_text_embedding_batch(batch_prompts, *col_result, config, adapter, context));
179
180	11	block.replace_by_position(result, _expand_and_wrap_nullable_result(
181	11	std::move(col_result), std::move(result_null_map),
182	11	input_rows_count));
183	11	return Status::OK();
184	11	}
185
186		Status _execute_multimodal_embed(FunctionContext* context, Block& block, uint32_t result,
187		size_t input_rows_count, const AIResource& config,
188		std::shared_ptr<AIAdapter>& adapter,
189		const ColumnPtr& input_column,
190	14	ColumnUInt8::MutablePtr result_null_map) const {
191	14	auto col_result = ColumnArray::create(
192	14	ColumnNullable::create(ColumnFloat32::create(), ColumnUInt8::create()));
193	14	std::vector<MultimodalType> batch_media_types;
194	14	std::vector<std::string> batch_media_content_types;
195	14	std::vector<std::string> batch_media_urls;
196	14	const NullMap* null_map = result_null_map ? &result_null_map->get_data() : nullptr;
197
198	14	int64_t ttl_seconds = 3600;
199	14	QueryContext* query_ctx = context->state()->get_query_ctx();
200	14	if (query_ctx && query_ctx->query_options().__isset.file_presigned_url_ttl_seconds) {
201	14	ttl_seconds = query_ctx->query_options().file_presigned_url_ttl_seconds;
202	14	if (ttl_seconds <= 0) {
203	1	ttl_seconds = 3600;
204	1	}
205	14	}
206
207	14	const int32_t max_batch_size = _get_embed_max_batch_size(context);
208
209	39	for (size_t i = 0; i < input_rows_count; ++i) {
210	29	if (null_map && (*null_map)[i]) {
211	3	continue;
212	3	}
213
214	26	rapidjson::Document file_input;
215	26	RETURN_IF_ERROR(_parse_file_input(*input_column, i, file_input));
216
217	26	std::string content_type;
218	26	MultimodalType media_type;
219	26	RETURN_IF_ERROR(_infer_media_type(file_input, content_type, media_type));
220
221	24	std::string media_url;
222	24	RETURN_IF_ERROR(_resolve_media_url(file_input, ttl_seconds, media_url));
223
224	22	if (!batch_media_urls.empty() &&
225	22	batch_media_urls.size() >= static_cast<size_t>(max_batch_size)) {
226	6	RETURN_IF_ERROR(_flush_multimodal_embedding_batch(
227	6	batch_media_types, batch_media_content_types, batch_media_urls, *col_result,
228	6	config, adapter, context));
229	6	}
230
231	22	batch_media_types.emplace_back(media_type);
232	22	batch_media_content_types.emplace_back(std::move(content_type));
233	22	batch_media_urls.emplace_back(std::move(media_url));
234	22	}
235
236	10	RETURN_IF_ERROR(_flush_multimodal_embedding_batch(
237	10	batch_media_types, batch_media_content_types, batch_media_urls, *col_result, config,
238	10	adapter, context));
239
240	10	block.replace_by_position(result, _expand_and_wrap_nullable_result(
241	10	std::move(col_result), std::move(result_null_map),
242	10	input_rows_count));
243	10	return Status::OK();
244	10	}
245
246		// EMBED-private helper.
247		// Sends one embedding request with a prebuilt request body and validates returned row count.
248		Status _execute_prebuilt_embedding_request(const std::string& request_body,
249		std::vector<std::vector<float>>& results,
250		size_t expected_size, const AIResource& config,
251		std::shared_ptr<AIAdapter>& adapter,
252	0	FunctionContext* context) const {
253	0	std::string response;
254		#ifdef BE_TEST
255		if (config.provider_type == "MOCK") {
256		results.clear();
257		results.reserve(expected_size);
258		for (size_t i = 0; i < expected_size; ++i) {
259		results.emplace_back(std::initializer_list<float> {0, 1, 2, 3, 4});
260		}
261		return Status::OK();
262		}
263		#endif
264
265	0	RETURN_IF_ERROR(
266	0	this->send_request_to_llm(request_body, response, config, adapter, context));
267
268	0	const auto& statistics = get_executor(context)->statistics();
269	0	const auto parse_start = statistics ? MonotonicNanos() : 0;
270	0	bool parse_failed = true;
271	32	Defer record_parse([&] {
272	32	if (statistics) {
273	25	statistics->record_parse(MonotonicNanos() - parse_start, parse_failed);
274	25	}
275	32	});
276	0	RETURN_IF_ERROR(adapter->parse_embedding_response(response, results));
277	0	if (results.empty()) {
278	0	return Status::InternalError("AI returned empty result");
279	0	}
280	0	if (results.size() != expected_size) [[unlikely]] {
281	0	return Status::InternalError(
282	0	"AI embedding returned {} results, but {} inputs were sent", results.size(),
283	0	expected_size);
284	0	}
285	0	parse_failed = false;
286	0	return Status::OK();
287	0	}
288
289		// EMBED-private helper.
290		// Flushes one accumulated text embedding batch into the output array column.
291		Status _flush_text_embedding_batch(std::vector<std::string>& batch_prompts,
292		ColumnArray& col_result, const AIResource& config,
293		std::shared_ptr<AIAdapter>& adapter,
294	34	FunctionContext* context) const {
295	34	if (batch_prompts.empty()) {
296	15	return Status::OK();
297	15	}
298
299	19	std::string request_body;
300	19	RETURN_IF_ERROR(adapter->build_embedding_request(batch_prompts, request_body));
301	19	std::vector<std::vector<float>> batch_results;
302	19	RETURN_IF_ERROR(_execute_prebuilt_embedding_request(
303	19	request_body, batch_results, batch_prompts.size(), config, adapter, context));
304	24	for (const auto& batch_result : batch_results) {
305	24	_insert_embedding_result(col_result, batch_result);
306	24	}
307	19	batch_prompts.clear();
308	19	return Status::OK();
309	19	}
310
311		// EMBED-private helper.
312		// Flushes one accumulated multimodal embedding batch into the output array column.
313		Status _flush_multimodal_embedding_batch(std::vector<MultimodalType>& batch_media_types,
314		std::vector<std::string>& batch_media_content_types,
315		std::vector<std::string>& batch_media_urls,
316		ColumnArray& col_result, const AIResource& config,
317		std::shared_ptr<AIAdapter>& adapter,
318	16	FunctionContext* context) const {
319	16	if (batch_media_urls.empty()) {
320	0	return Status::OK();
321	0	}
322
323	16	std::string request_body;
324	16	RETURN_IF_ERROR(adapter->build_multimodal_embedding_request(
325	16	batch_media_types, batch_media_urls, batch_media_content_types, request_body));
326
327	16	std::vector<std::vector<float>> batch_results;
328	16	RETURN_IF_ERROR(_execute_prebuilt_embedding_request(
329	16	request_body, batch_results, batch_media_urls.size(), config, adapter, context));
330	22	for (const auto& batch_result : batch_results) {
331	22	_insert_embedding_result(col_result, batch_result);
332	22	}
333	16	batch_media_types.clear();
334	16	batch_media_content_types.clear();
335	16	batch_media_urls.clear();
336	16	return Status::OK();
337	16	}
338
339		static void _insert_embedding_result(ColumnArray& col_array,
340	46	const std::vector<float>& float_result) {
341	46	auto& offsets = col_array.get_offsets();
342	46	auto& nested_nullable_col = assert_cast<ColumnNullable&>(col_array.get_data());
343	46	auto& nested_col =
344	46	assert_cast<ColumnFloat32&>(*(nested_nullable_col.get_nested_column_ptr()));
345	46	nested_col.reserve(nested_col.size() + float_result.size());
346
347	46	size_t current_offset = nested_col.size();
348	46	nested_col.insert_many_raw_data(reinterpret_cast<const char*>(float_result.data()),
349	46	float_result.size());
350	46	offsets.push_back(current_offset + float_result.size());
351	46	auto& null_map = nested_nullable_col.get_null_map_column();
352	46	null_map.insert_many_vals(0, float_result.size());
353	46	}
354
355		static ColumnPtr _expand_and_wrap_nullable_result(ColumnArray::MutablePtr result,
356		ColumnUInt8::MutablePtr result_null_map,
357	21	size_t input_rows_count) {
358	21	if (!result_null_map) {
359	13	return result;
360	13	}
361
362	8	auto& offsets = result->get_offsets();
363	8	size_t compact_row = offsets.size();
364	8	offsets.resize(input_rows_count);
365		// For example, embedding rows 1 and 3 produces compact offsets [5, 10]. Given
366		// result_null_map [1, 0, 1, 0, 1], expand them to [0, 5, 5, 10, 10], where NULL rows
367		// reuse the previous offset. Fill backwards to avoid overwriting unread compact offsets.
368	44	for (size_t row = input_rows_count; row-- > 0;) {
369	36	if (result_null_map->get_data()[row]) {
370	17	offsets[row] = compact_row == 0 ? 0 : offsets[compact_row - 1];
371	19	} else {
372	19	offsets[row] = offsets[--compact_row];
373	19	}
374	36	}
375	8	return ColumnNullable::create(std::move(result), std::move(result_null_map));
376	21	}
377
378	87	static bool _starts_with_ignore_case(std::string_view s, std::string_view prefix) {
379	87	if (s.size() < prefix.size()) {
380	0	return false;
381	0	}
382	438	return std::equal(prefix.begin(), prefix.end(), s.begin(), [](char a, char b) {
383	438	return std::tolower(static_cast<unsigned char>(a)) ==
384	438	std::tolower(static_cast<unsigned char>(b));
385	438	});
386	87	}
387
388		static Status _infer_media_type(const rapidjson::Value& file_input, std::string& content_type,
389	26	MultimodalType& media_type) {
390	26	RETURN_IF_ERROR(_get_required_string_field(file_input, "content_type", content_type));
391
392	25	if (_starts_with_ignore_case(content_type, "image/")) {
393	14	media_type = MultimodalType::IMAGE;
394	14	return Status::OK();
395	14	} else if (_starts_with_ignore_case(content_type, "video/")) {
396	8	media_type = MultimodalType::VIDEO;
397	8	return Status::OK();
398	8	} else if (_starts_with_ignore_case(content_type, "audio/")) {
399	2	media_type = MultimodalType::AUDIO;
400	2	return Status::OK();
401	2	}
402
403	1	return Status::InvalidArgument("Unsupported content_type for EMBED: {}", content_type);
404	25	}
405
406		// Parse the FILE-like JSONB argument into a JSON object for downstream field reads.
407		static Status _parse_file_input(const IColumn& file_column, size_t row_num,
408	26	rapidjson::Document& file_input) {
409	26	StringRef file_ref = file_column.get_data_at(row_num);
410	26	std::string file_json = JsonbToJson::jsonb_to_json_string(file_ref.data, file_ref.size);
411	26	file_input.Parse(file_json.c_str());
412	26	DORIS_CHECK(!file_input.HasParseError() && file_input.IsObject());
413	26	return Status::OK();
414	26	}
415
416		// TODO(lzq): After support FILE type, We should use the interface provided by FILE to get the fields
417		// replacing this function
418		static Status _get_required_string_field(const rapidjson::Value& obj, const char* field_name,
419	55	std::string& value) {
420	55	auto iter = obj.FindMember(field_name);
421	55	if (iter == obj.MemberEnd() \|\| !iter->value.IsString()) {
422	3	return Status::InvalidArgument(
423	3	"EMBED file json field '{}' is required and must be a string", field_name);
424	3	}
425	52	value = iter->value.GetString();
426	52	if (value.empty()) {
427	0	return Status::InvalidArgument("EMBED file json field '{}' can not be empty",
428	0	field_name);
429	0	}
430	52	return Status::OK();
431	52	}
432
433		static Status init_s3_client_conf_from_json(const rapidjson::Value& file_input,
434	3	S3ClientConf& s3_client_conf) {
435	3	std::string endpoint;
436	3	RETURN_IF_ERROR(_get_required_string_field(file_input, "endpoint", endpoint));
437	2	std::string region;
438	2	RETURN_IF_ERROR(_get_required_string_field(file_input, "region", region));
439
440	4	auto get_optional_string_field = [&](const char* field_name, std::string& value) {
441	4	auto iter = file_input.FindMember(field_name);
442	4	if (iter == file_input.MemberEnd() \|\| iter->value.IsNull()) {
443	0	return;
444	0	}
445	4	DORIS_CHECK(iter->value.IsString());
446	4	value = iter->value.GetString();
447	4	};
448
449	1	get_optional_string_field("ak", s3_client_conf.ak);
450	1	get_optional_string_field("sk", s3_client_conf.sk);
451	1	get_optional_string_field("role_arn", s3_client_conf.role_arn);
452	1	get_optional_string_field("external_id", s3_client_conf.external_id);
453	1	s3_client_conf.endpoint = endpoint;
454	1	s3_client_conf.region = region;
455
456	1	return Status::OK();
457	2	}
458
459		Status _resolve_media_url(const rapidjson::Value& file_input, int64_t ttl_seconds,
460	24	std::string& media_url) const {
461	24	std::string uri;
462	24	RETURN_IF_ERROR(_get_required_string_field(file_input, "uri", uri));
463
464		// If it's a direct http/https URL, use it as-is
465	24	if (_starts_with_ignore_case(uri, "http://") \|\| _starts_with_ignore_case(uri, "https://")) {
466	21	media_url = uri;
467	21	return Status::OK();
468	21	}
469
470	3	S3ClientConf s3_client_conf;
471	3	RETURN_IF_ERROR(init_s3_client_conf_from_json(file_input, s3_client_conf));
472	1	auto s3_client = DORIS_TRY(S3ClientFactory::instance().create(s3_client_conf));
473
474	1	S3URI s3_uri(uri);
475	1	RETURN_IF_ERROR(s3_uri.parse());
476	1	std::string bucket = s3_uri.get_bucket();
477	1	std::string key = s3_uri.get_key();
478	1	DORIS_CHECK(!bucket.empty() && !key.empty());
479	1	media_url = s3_client->generate_presigned_url({.bucket = bucket, .key = key, .prefix = ""},
480	1	ttl_seconds);
481	1	return Status::OK();
482	1	}
483		};
484
485		}; // namespace doris