nprasad2077/doris
1
0 Fork 0
mirror of https://github.com/apache/doris.git synced 2026-11-04 07:31:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Enhancement](ai_func) Skip Null inputs in AI functions (#66242)

Browse Source 
Problem Summary:

The framework's default NULL implementation unwraps Nullable arguments
and executes AI functions for every input row. For partially NULL
inputs, the nested placeholder values of NULL rows are still included in
prompts and sent to external AI providers.

This causes unnecessary remote requests and token consumption. It also
requires special handling for embedding results to preserve the original
row order without duplicating large embedding vectors.

This PR:

- Disables the framework's default NULL implementation for all
`AIFunction` subclasses.
- Determines the Nullable return type in the `AIFunction` base class.
- Extracts nested prompt columns and merges argument null maps in the
common AI execution path.
- Skips NULL rows before building prompts or sending requests.
- Restores NULL rows in the final result while preserving the original
row order.
- Handles both text and multimodal Nullable inputs for `EMBED`.
- Expands embedding array offsets in place, avoiding a copy of the
nested Float32 embedding data.
- Returns a constant NULL column when all input rows are NULL.


### Release note

Fix AI scalar functions to skip NULL input rows instead of sending their
placeholder values to external AI providers.
This commit is contained in:
linrrarity authored and GitHub committed 2026-08-05 11:05:43 +08:00
1 parent 5a24e633b1
commit a7953ebc3a
16 files changed
+777 -203
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

be/src/exprs/function/ai/ai_adapter.h
+20
View File
@@ -1568,6 +1568,14 @@ public:
// Mock adapter used only for UT to bypass real HTTP calls and return deterministic data.
class MockAdapter : public AIAdapter {
public:
#ifdef BE_TEST
static void clear_embedding_inputs_for_test() { _embedding_inputs_for_test().clear(); }
static const std::vector<std::string>& get_embedding_inputs_for_test() {
return _embedding_inputs_for_test();
}
#endif
Status set_authentication(HttpClient* client) const override { return Status::OK(); }
Status build_request_payload(const std::vector<std::string>& inputs,
@@ -1583,6 +1591,10 @@ public:
Status build_embedding_request(const std::vector<std::string>& inputs,
std::string& request_body) const override {
#ifdef BE_TEST
auto& embedding_inputs = _embedding_inputs_for_test();
embedding_inputs.insert(embedding_inputs.end(), inputs.begin(), inputs.end());
#endif
return Status::OK();
}
@@ -1612,6 +1624,14 @@ public:
[](const auto& val) { return val.GetFloat(); });
return Status::OK();
}
private:
#ifdef BE_TEST
static std::vector<std::string>& _embedding_inputs_for_test() {
static thread_local std::vector<std::string> embedding_inputs;
return embedding_inputs;
}
#endif
};
class AIAdapterFactory {
be/src/exprs/function/ai/ai_classify.h
+2 -2
View File
@@ -37,13 +37,13 @@ public:
static constexpr size_t number_of_arguments = 3;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeString>();
}
static FunctionPtr create() { return std::make_shared<FunctionAIClassify>(); }
Status build_prompt(const Block& block, const ColumnNumbers& arguments, size_t row_num,
Status build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const override;
};
} // namespace doris
be/src/exprs/function/ai/ai_extract.h
+2 -2
View File
@@ -38,13 +38,13 @@ public:
static constexpr size_t number_of_arguments = 3;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeString>();
}
static FunctionPtr create() { return std::make_shared<FunctionAIExtract>(); }
Status build_prompt(const Block& block, const ColumnNumbers& arguments, size_t row_num,
Status build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const override;
};
be/src/exprs/function/ai/ai_filter.h
+1 -1
View File
@@ -38,7 +38,7 @@ public:
static constexpr size_t number_of_arguments = 2;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeBool>();
}
be/src/exprs/function/ai/ai_fix_grammar.h
+1 -1
View File
@@ -38,7 +38,7 @@ public:
static constexpr size_t number_of_arguments = 2;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeString>();
}
be/src/exprs/function/ai/ai_functions.cpp
+52 -112
View File
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
#include "core/column/column_array.h"
#include "core/column/column_array_view.h"
#include "exprs/function/ai/ai_classify.h"
#include "exprs/function/ai/ai_extract.h"
#include "exprs/function/ai/ai_filter.h"
@@ -30,151 +30,93 @@
#include "exprs/function/simple_function_factory.h"
namespace doris {
Status FunctionAIClassify::build_prompt(const Block& block, const ColumnNumbers& arguments,
size_t row_num, std::string& prompt) const {
// Get the text column
const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);
StringRef text = text_column.column->get_data_at(row_num);
std::string text_str = std::string(text.data, text.size);
// Get the labels array column
const ColumnWithTypeAndName& labels_column = block.get_by_position(arguments[2]);
const auto& [array_column, array_row_num] =
check_column_const_set_readability(*labels_column.column, row_num);
const auto* col_array = check_and_get_column<ColumnArray>(*array_column);
if (col_array == nullptr) {
static Status format_labels(const ColumnPtr& labels_column, size_t row_num,
std::string_view function_name, std::string& labels_str) {
auto readable_column = check_column_const_set_readability(*labels_column, row_num);
if (!is_column<ColumnArray>(*readable_column.first)) {
return Status::InternalError(
"labels argument for {} must be Array(String) or Array(Varchar)", name);
"labels argument for {} must be Array(String) or Array(Varchar)", function_name);
}
std::vector<std::string> label_values;
const auto& data = col_array->get_data();
const auto& offsets = col_array->get_offsets();
size_t start = array_row_num > 0 ? offsets[array_row_num - 1] : 0;
size_t end = offsets[array_row_num];
for (size_t i = start; i < end; ++i) {
Field field;
data.get(i, field);
label_values.emplace_back(field.template get<TYPE_STRING>());
}
std::string labels_str = "[";
for (size_t i = 0; i < label_values.size(); ++i) {
if (i > 0) {
auto labels_view = ColumnArrayView<TYPE_STRING>::create(labels_column);
auto labels = labels_view[row_num];
labels_str = "[";
bool is_first_label = true;
for (size_t i = 0; i < labels.size(); ++i) {
if (labels.is_null_at(i)) {
continue;
}
if (!is_first_label) {
labels_str += ", ";
}
labels_str += "\"" + label_values[i] + "\"";
StringRef label = labels.value_at(i);
labels_str += "\"";
labels_str.append(label.data, label.size);
labels_str += "\"";
is_first_label = false;
}
labels_str += "]";
return Status::OK();
}
Status FunctionAIClassify::build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const {
// Get the text column
StringRef text = prompt_columns[0]->get_data_at(row_num);
std::string text_str = std::string(text.data, text.size);
std::string labels_str;
RETURN_IF_ERROR(format_labels(prompt_columns[1], row_num, name, labels_str));
prompt = "Labels: " + labels_str + "\nText: " + text_str;
return Status::OK();
}
Status FunctionAIExtract::build_prompt(const Block& block, const ColumnNumbers& arguments,
size_t row_num, std::string& prompt) const {
Status FunctionAIExtract::build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const {
// Get the text column
const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);
StringRef text = text_column.column->get_data_at(row_num);
StringRef text = prompt_columns[0]->get_data_at(row_num);
std::string text_str = std::string(text.data, text.size);
// Get the labels array column
const ColumnWithTypeAndName& labels_column = block.get_by_position(arguments[2]);
const auto& [array_column, array_row_num] =
check_column_const_set_readability(*labels_column.column, row_num);
const auto* col_array = check_and_get_column<ColumnArray>(*array_column);
if (col_array == nullptr) {
return Status::InternalError(
"labels argument for {} must be Array(String) or Array(Varchar)", name);
}
std::vector<std::string> label_values;
const auto& offsets = col_array->get_offsets();
const auto& data = col_array->get_data();
size_t start = array_row_num > 0 ? offsets[array_row_num - 1] : 0;
size_t end = offsets[array_row_num];
for (size_t i = start; i < end; ++i) {
Field field;
data.get(i, field);
label_values.emplace_back(field.template get<TYPE_STRING>());
}
std::string labels_str = "[";
for (size_t i = 0; i < label_values.size(); ++i) {
if (i > 0) {
labels_str += ", ";
}
labels_str += "\"" + label_values[i] + "\"";
}
labels_str += "]";
std::string labels_str;
RETURN_IF_ERROR(format_labels(prompt_columns[1], row_num, name, labels_str));
prompt = "Labels: " + labels_str + "\nText: " + text_str;
return Status::OK();
}
Status FunctionAIGenerate::build_prompt(const Block& block, const ColumnNumbers& arguments,
size_t row_num, std::string& prompt) const {
const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);
StringRef text_ref = text_column.column->get_data_at(row_num);
Status FunctionAIGenerate::build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const {
StringRef text_ref = prompt_columns[0]->get_data_at(row_num);
prompt = std::string(text_ref.data, text_ref.size);
return Status::OK();
}
Status FunctionAIMask::build_prompt(const Block& block, const ColumnNumbers& arguments,
size_t row_num, std::string& prompt) const {
Status FunctionAIMask::build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const {
// Get the text column
const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);
StringRef text = text_column.column->get_data_at(row_num);
StringRef text = prompt_columns[0]->get_data_at(row_num);
std::string text_str = std::string(text.data, text.size);
// Get the labels array column
const ColumnWithTypeAndName& labels_column = block.get_by_position(arguments[2]);
const auto& [array_column, array_row_num] =
check_column_const_set_readability(*labels_column.column, row_num);
const auto* col_array = check_and_get_column<ColumnArray>(*array_column);
if (col_array == nullptr) {
return Status::InternalError(
"labels argument for {} must be Array(String) or Array(Varchar)", name);
}
std::vector<std::string> label_values;
const auto& offsets = col_array->get_offsets();
const auto& data = col_array->get_data();
size_t start = array_row_num > 0 ? offsets[array_row_num - 1] : 0;
size_t end = offsets[array_row_num];
for (size_t i = start; i < end; ++i) {
Field field;
data.get(i, field);
label_values.emplace_back(field.template get<TYPE_STRING>());
}
std::string labels_str = "[";
for (size_t i = 0; i < label_values.size(); ++i) {
if (i > 0) {
labels_str += ", ";
}
labels_str += "\"" + label_values[i] + "\"";
}
labels_str += "]";
std::string labels_str;
RETURN_IF_ERROR(format_labels(prompt_columns[1], row_num, name, labels_str));
prompt = "Labels: " + labels_str + "\nText: " + text_str;
return Status::OK();
}
Status FunctionAISimilarity::build_prompt(const Block& block, const ColumnNumbers& arguments,
size_t row_num, std::string& prompt) const {
Status FunctionAISimilarity::build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const {
// text1
const ColumnWithTypeAndName& text_column_1 = block.get_by_position(arguments[1]);
StringRef text_1 = text_column_1.column.get()->get_data_at(row_num);
StringRef text_1 = prompt_columns[0]->get_data_at(row_num);
std::string text_str_1 = std::string(text_1.data, text_1.size);
// text2
const ColumnWithTypeAndName& text_column_2 = block.get_by_position(arguments[2]);
StringRef text_2 = text_column_2.column.get()->get_data_at(row_num);
StringRef text_2 = prompt_columns[1]->get_data_at(row_num);
std::string text_str_2 = std::string(text_2.data, text_2.size);
prompt = "Text 1: " + text_str_1 + "\nText 2: " + text_str_2;
@@ -182,16 +124,14 @@ Status FunctionAISimilarity::build_prompt(const Block& block, const ColumnNumber
return Status::OK();
}
Status FunctionAITranslate::build_prompt(const Block& block, const ColumnNumbers& arguments,
size_t row_num, std::string& prompt) const {
Status FunctionAITranslate::build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const {
// text
const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);
StringRef text = text_column.column.get()->get_data_at(row_num);
StringRef text = prompt_columns[0]->get_data_at(row_num);
std::string text_str = std::string(text.data, text.size);
// target language
const ColumnWithTypeAndName& lang_column = block.get_by_position(arguments[2]);
StringRef lang = lang_column.column.get()->get_data_at(row_num);
StringRef lang = prompt_columns[1]->get_data_at(row_num);
std::string target_lang = std::string(lang.data, lang.size);
prompt = "Translate the following text to " + target_lang + ".\nText: " + text_str;
be/src/exprs/function/ai/ai_functions.h
+89 -28
View File
@@ -36,9 +36,11 @@
#include "core/column/column_nullable.h"
#include "core/cow.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_number.h"
#include "core/data_type/define_primitive_type.h"
#include "core/data_type/primitive_type.h"
#include "exec/common/util.hpp"
#include "exprs/function/ai/ai_adapter.h"
#include "exprs/function/function.h"
#include "runtime/query_context.h"
@@ -65,10 +67,21 @@ public:
bool is_blockable() const override { return true; }
virtual Status build_prompt(const Block& block, const ColumnNumbers& arguments, size_t row_num,
bool use_default_implementation_for_nulls() const final { return false; }
DataTypePtr get_return_type_impl(const DataTypes& arguments) const final {
bool has_nullable_argument = std::ranges::any_of(
arguments, [](const auto& argument) { return argument->is_nullable(); });
DataTypePtr return_type =
assert_cast<const Derived&>(*this).get_nested_return_type_impl(arguments);
return has_nullable_argument ? make_nullable(return_type) : return_type;
}
using PreparedFunctionImpl::execute;
virtual Status build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const {
const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);
StringRef text_ref = text_column.column->get_data_at(row_num);
StringRef text_ref = prompt_columns[0]->get_data_at(row_num);
prompt = std::string(text_ref.data, text_ref.size);
return Status::OK();
@@ -76,6 +89,13 @@ public:
Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
uint32_t result, size_t input_rows_count) const override {
if (block.get_by_position(arguments[0]).column->only_null()) {
block.get_by_position(result).column =
block.get_by_position(result).type->create_column_const(input_rows_count,
Field());
return Status::OK();
}
TAIResource config;
std::shared_ptr<AIAdapter> adapter;
if (Status status = this->_init_from_resource(context, block, arguments, config, adapter);
@@ -83,8 +103,8 @@ public:
return status;
}
return assert_cast<const Derived&>(*this).execute_with_adapter(
context, block, arguments, result, input_rows_count, config, adapter);
return assert_cast<const Derived&>(*this).execute(context, block, arguments, result,
input_rows_count, config, adapter);
}
protected:
@@ -98,20 +118,6 @@ protected:
return query_ctx->query_options().ai_context_window_size;
}
// Derived classes can override this method for non-text/default behavior.
// The base implementation handles all string-input/string-output batchable functions.
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 = assert_cast<const Derived&>(*this).create_result_column();
RETURN_IF_ERROR(execute_batched_prompts(context, block, arguments, input_rows_count, config,
adapter, *col_result));
block.replace_by_position(result, std::move(col_result));
return Status::OK();
}
MutableColumnPtr create_result_column() const { return ColumnString::create(); }
// Provider-reusable hook for AI functions(string) -> string.
@@ -285,19 +291,50 @@ protected:
// Provider-reusable helper for string-returning functions.
// Runs the common batch execution flow; derived classes only need to define how one batch of
// string results is inserted into the final output column.
Status execute_batched_prompts(FunctionContext* context, Block& block,
const ColumnNumbers& arguments, size_t input_rows_count,
const TAIResource& config, std::shared_ptr<AIAdapter>& adapter,
IColumn& col_result) const {
Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
uint32_t result, size_t input_rows_count, const TAIResource& config,
std::shared_ptr<AIAdapter>& adapter) const {
Columns prompt_columns;
prompt_columns.reserve(arguments.size() - 1);
ColumnUInt8::MutablePtr result_null_map;
for (size_t i = 1; i < arguments.size(); ++i) {
const auto& argument = block.get_by_position(arguments[i]);
if (argument.type->is_nullable()) {
const auto& [column, is_const] = unpack_if_const(argument.column);
const auto& nullable =
assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*column);
if (!result_null_map) {
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);
}
prompt_columns.emplace_back(argument.unnest_nullable().column);
}
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();
}
auto col_result = assert_cast<const Derived&>(*this).create_result_column();
std::vector<std::string> batch_prompts;
size_t current_batch_size = 2; // []
const size_t max_batch_prompt_size =
static_cast<size_t>(get_ai_context_window_size(context));
const NullMap* null_map = result_null_map ? &result_null_map->get_data() : nullptr;
for (size_t i = 0; i < input_rows_count; ++i) {
if (null_map && (*null_map)[i]) {
continue;
}
std::string prompt;
RETURN_IF_ERROR(
assert_cast<const Derived&>(*this).build_prompt(block, arguments, i, prompt));
assert_cast<const Derived&>(*this).build_prompt(prompt_columns, i, prompt));
size_t entry_size = estimate_batch_entry_size(batch_prompts.size(), prompt);
if (entry_size > max_batch_prompt_size) {
@@ -306,7 +343,7 @@ protected:
RETURN_IF_ERROR(this->execute_batch_request(batch_prompts, batch_results,
config, adapter, context));
RETURN_IF_ERROR(assert_cast<const Derived&>(*this).append_batch_results(
batch_results, col_result));
batch_results, *col_result));
batch_prompts.clear();
current_batch_size = 2;
}
@@ -317,7 +354,7 @@ protected:
RETURN_IF_ERROR(this->execute_batch_request(single_prompts, single_results, config,
adapter, context));
RETURN_IF_ERROR(assert_cast<const Derived&>(*this).append_batch_results(
single_results, col_result));
single_results, *col_result));
continue;
}
@@ -328,7 +365,7 @@ protected:
RETURN_IF_ERROR(this->execute_batch_request(batch_prompts, batch_results, config,
adapter, context));
RETURN_IF_ERROR(assert_cast<const Derived&>(*this).append_batch_results(
batch_results, col_result));
batch_results, *col_result));
batch_prompts.clear();
current_batch_size = 2;
additional_size = entry_size;
@@ -343,8 +380,32 @@ protected:
RETURN_IF_ERROR(this->execute_batch_request(batch_prompts, batch_results, config,
adapter, context));
RETURN_IF_ERROR(assert_cast<const Derived&>(*this).append_batch_results(batch_results,
col_result));
*col_result));
}
if (!result_null_map) {
block.replace_by_position(result, std::move(col_result));
return Status::OK();
}
if (!simd::contain_one(result_null_map->get_data().data(), input_rows_count)) {
block.replace_by_position(result, ColumnNullable::create(std::move(col_result),
std::move(result_null_map)));
return Status::OK();
}
auto nested_result = col_result->clone_empty();
size_t result_row = 0;
for (UInt8 is_null : result_null_map->get_data()) {
if (is_null) {
nested_result->insert_default();
} else {
nested_result->insert_from(*col_result, result_row++);
}
}
block.replace_by_position(result, ColumnNullable::create(std::move(nested_result),
std::move(result_null_map)));
return Status::OK();
}
be/src/exprs/function/ai/ai_generate.h
+2 -2
View File
@@ -36,13 +36,13 @@ public:
static constexpr size_t number_of_arguments = 2;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeString>();
}
static FunctionPtr create() { return std::make_shared<FunctionAIGenerate>(); }
Status build_prompt(const Block& block, const ColumnNumbers& arguments, size_t row_num,
Status build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const override;
};
be/src/exprs/function/ai/ai_mask.h
+2 -2
View File
@@ -37,13 +37,13 @@ public:
static constexpr size_t number_of_arguments = 3;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeString>();
}
static FunctionPtr create() { return std::make_shared<FunctionAIMask>(); }
Status build_prompt(const Block& block, const ColumnNumbers& arguments, size_t row_num,
Status build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const override;
};
be/src/exprs/function/ai/ai_sentiment.h
+1 -1
View File
@@ -36,7 +36,7 @@ public:
static constexpr size_t number_of_arguments = 2;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeString>();
}
be/src/exprs/function/ai/ai_similarity.h
+2 -2
View File
@@ -41,13 +41,13 @@ public:
static constexpr size_t number_of_arguments = 3;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeFloat32>();
}
static FunctionPtr create() { return std::make_shared<FunctionAISimilarity>(); }
Status build_prompt(const Block& block, const ColumnNumbers& arguments, size_t row_num,
Status build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const override;
private:
be/src/exprs/function/ai/ai_summarize.h
+1 -1
View File
@@ -37,7 +37,7 @@ public:
static constexpr size_t number_of_arguments = 2;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeString>();
}
be/src/exprs/function/ai/ai_translate.h
+2 -2
View File
@@ -35,13 +35,13 @@ public:
"corresponding item, with no explanation, markdown, or extra text.";
static constexpr size_t number_of_arguments = 3;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeString>();
}
static FunctionPtr create() { return std::make_shared<FunctionAITranslate>(); }
Status build_prompt(const Block& block, const ColumnNumbers& arguments, size_t row_num,
Status build_prompt(const Columns& prompt_columns, size_t row_num,
std::string& prompt) const override;
};
be/src/exprs/function/ai/embed.h
+84 -27
View File
@@ -38,33 +38,53 @@ public:
static constexpr auto system_prompt = "";
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
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 {
using PreparedFunctionImpl::execute;
Status execute(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();
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);
}
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().column;
PrimitiveType input_type = remove_nullable(input.type)->get_primitive_type();
if (input_type == PrimitiveType::TYPE_JSONB) {
return _execute_multimodal_embed(context, block, arguments, result, input_rows_count,
config, adapter);
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, arguments, result, input_rows_count, config,
adapter);
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 {}",
block.get_by_position(arguments[1]).type->get_name());
input.type->get_name());
}
static FunctionPtr create() { return std::make_shared<FunctionEmbed>(); }
@@ -77,10 +97,10 @@ private:
return query_ctx->query_options().embed_max_batch_size;
}
Status _execute_text_embed(FunctionContext* context, Block& block,
const ColumnNumbers& arguments, uint32_t result,
Status _execute_text_embed(FunctionContext* context, Block& block, uint32_t result,
size_t input_rows_count, const TAIResource& config,
std::shared_ptr<AIAdapter>& adapter) const {
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;
@@ -88,10 +108,16 @@ private:
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(block, arguments, i, prompt));
RETURN_IF_ERROR(build_prompt(prompt_columns, i, prompt));
const size_t prompt_size = prompt.size();
@@ -122,19 +148,23 @@ private:
RETURN_IF_ERROR(
_flush_text_embedding_batch(batch_prompts, *col_result, config, adapter, context));
block.replace_by_position(result, std::move(col_result));
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,
const ColumnNumbers& arguments, uint32_t result,
Status _execute_multimodal_embed(FunctionContext* context, Block& block, uint32_t result,
size_t input_rows_count, const TAIResource& config,
std::shared_ptr<AIAdapter>& adapter) const {
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();
@@ -147,10 +177,13 @@ private:
const int32_t max_batch_size = _get_embed_max_batch_size(context);
const ColumnWithTypeAndName& file_column = block.get_by_position(arguments[1]);
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(file_column, i, file_input));
RETURN_IF_ERROR(_parse_file_input(*input_column, i, file_input));
std::string content_type;
MultimodalType media_type;
@@ -175,7 +208,9 @@ private:
batch_media_types, batch_media_content_types, batch_media_urls, *col_result, config,
adapter, context));
block.replace_by_position(result, std::move(col_result));
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();
}
@@ -279,6 +314,29 @@ private:
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;
@@ -308,11 +366,10 @@ private:
}
// 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,
static Status _parse_file_input(const IColumn& 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);
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();
be/test/ai/ai_function_test.cpp
+328 -13
View File
@@ -27,6 +27,7 @@
#include "core/block/block.h"
#include "core/column/column_array.h"
#include "core/column/column_const.h"
#include "core/column/column_nullable.h"
#include "core/column/column_string.h"
#include "core/column/column_vector.h"
@@ -43,6 +44,7 @@
#include "exprs/function/ai/ai_summarize.h"
#include "exprs/function/ai/ai_translate.h"
#include "exprs/function/ai/embed.h"
#include "exprs/function/simple_function_factory.h"
#include "testutil/column_helper.h"
#include "testutil/mock/mock_runtime_state.h"
@@ -63,7 +65,7 @@ public:
using AIFunction<FunctionAIFilterBatchTestHelper>::execute_batch_request;
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr get_nested_return_type_impl(const DataTypes& /*arguments*/) const {
return std::make_shared<DataTypeBool>();
}
@@ -183,16 +185,19 @@ private:
};
namespace {
MutableColumnPtr create_string_array_column(const std::vector<std::vector<std::string>>& rows) {
MutableColumnPtr create_string_array_column(const std::vector<std::vector<std::string>>& rows,
const std::vector<UInt8>& null_map = {}) {
auto nested_column = ColumnString::create();
auto null_map_column = ColumnUInt8::create();
auto offsets_column = ColumnOffset64::create();
IColumn::Offset offset = 0;
size_t element = 0;
for (const auto& row : rows) {
for (const auto& value : row) {
nested_column->insert_data(value.data(), value.size());
null_map_column->insert_value(0);
null_map_column->insert_value(null_map.empty() ? 0 : null_map[element]);
++element;
}
offset += row.size();
offsets_column->insert_value(offset);
@@ -202,6 +207,21 @@ MutableColumnPtr create_string_array_column(const std::vector<std::vector<std::s
ColumnNullable::create(std::move(nested_column), std::move(null_map_column)),
std::move(offsets_column));
}
FunctionBasePtr get_ai_function(const std::string& name, const Block& arguments,
const DataTypePtr& return_type) {
return SimpleFunctionFactory::instance().get_function(
name, arguments.get_columns_with_type_and_name(), return_type);
}
Columns get_prompt_columns(const Block& block, const ColumnNumbers& arguments) {
Columns prompt_columns;
prompt_columns.reserve(arguments.size() - 1);
for (size_t i = 1; i < arguments.size(); ++i) {
prompt_columns.emplace_back(block.get_by_position(arguments[i]).unnest_nullable().column);
}
return prompt_columns;
}
} // namespace
TEST(AIFunctionTest, AISummarizeTest) {
@@ -219,7 +239,7 @@ TEST(AIFunctionTest, AISummarizeTest) {
ColumnNumbers arguments = {0, 1};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt, "This is a test document that needs to be summarized.");
@@ -240,7 +260,7 @@ TEST(AIFunctionTest, AISentimentTest) {
ColumnNumbers arguments = {0, 1};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt, "I really enjoyed the doris community!");
@@ -266,7 +286,7 @@ TEST(AIFunctionTest, AIMaskTest) {
ColumnNumbers arguments = {0, 1, 2};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt,
@@ -289,7 +309,7 @@ TEST(AIFunctionTest, AIGenerateTest) {
ColumnNumbers arguments = {0, 1};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt, "Write a poem about spring");
@@ -310,7 +330,7 @@ TEST(AIFunctionTest, AIFixGrammarTest) {
ColumnNumbers arguments = {0, 1};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt, "She don't like apples");
@@ -335,7 +355,7 @@ TEST(AIFunctionTest, AIExtractTest) {
ColumnNumbers arguments = {0, 1, 2};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt,
@@ -362,7 +382,7 @@ TEST(AIFunctionTest, AIClassifyTest) {
ColumnNumbers arguments = {0, 1, 2};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt,
@@ -370,6 +390,38 @@ TEST(AIFunctionTest, AIClassifyTest) {
"Text: This product exceeded my expectations");
}
TEST(AIFunctionTest, NullableLabelElementsAreSkipped) {
std::vector<std::string> texts = {"good product"};
auto labels = create_string_array_column({{"positive", "unused-null", "negative"}},
std::vector<UInt8> {0, 1, 0});
Block block;
block.insert({ColumnHelper::create_column<DataTypeString>({"resource_name"}),
std::make_shared<DataTypeString>(), "resource"});
block.insert({ColumnHelper::create_column<DataTypeString>(texts),
std::make_shared<DataTypeString>(), "text"});
block.insert({std::move(labels),
std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>()), "labels"});
Columns prompt_columns = get_prompt_columns(block, {0, 1, 2});
std::string prompt;
const std::string expected =
"Labels: [\"positive\", \"negative\"]\n"
"Text: good product";
FunctionAIClassify classify;
ASSERT_TRUE(classify.build_prompt(prompt_columns, 0, prompt).ok());
EXPECT_EQ(prompt, expected);
FunctionAIExtract extract;
ASSERT_TRUE(extract.build_prompt(prompt_columns, 0, prompt).ok());
EXPECT_EQ(prompt, expected);
FunctionAIMask mask;
ASSERT_TRUE(mask.build_prompt(prompt_columns, 0, prompt).ok());
EXPECT_EQ(prompt, expected);
}
TEST(AIFunctionTest, AITranslateTest) {
FunctionAITranslate function;
@@ -388,7 +440,7 @@ TEST(AIFunctionTest, AITranslateTest) {
ColumnNumbers arguments = {0, 1, 2};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt,
@@ -414,7 +466,7 @@ TEST(AIFunctionTest, AISimilarityTest) {
ColumnNumbers arguments = {0, 1, 2};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt, "Text 1: I like this dish\nText 2: This dish is very good");
@@ -593,7 +645,7 @@ TEST(AIFunctionTest, AIFilterTest) {
ColumnNumbers arguments = {0, 1};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt(get_prompt_columns(block, arguments), 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt, "This is a valid sentence.");
@@ -1143,6 +1195,269 @@ TEST(AIFunctionTest, AIStringFunctionBatchExecuteTest) {
unsetenv("AI_TEST_RESULT");
}
TEST(AIFunctionTest, NullableStringResultThroughPreparedFunction) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
setenv("AI_TEST_RESULT", R"(["answer-a","answer-c"])", 1);
std::vector<std::string> texts = {"unused-null", "text-a", "unused-null", "text-c",
"unused-null"};
std::vector<UInt8> null_map = {1, 0, 1, 0, 1};
Block block;
block.insert({ColumnHelper::create_column<DataTypeString>(
std::vector<std::string>(texts.size(), "mock_resource")),
std::make_shared<DataTypeString>(), "resource"});
block.insert({ColumnHelper::create_nullable_column<DataTypeString>(texts, null_map),
make_nullable(std::make_shared<DataTypeString>()), "text"});
auto return_type = make_nullable(std::make_shared<DataTypeString>());
auto function = get_ai_function("ai_generate", block, return_type);
ASSERT_NE(function, nullptr);
EXPECT_TRUE(function->get_return_type()->equals(*return_type));
block.insert({nullptr, return_type, "result"});
Status status = function->execute(ctx.get(), block, {0, 1}, 2, texts.size());
unsetenv("AI_TEST_RESULT");
ASSERT_TRUE(status.ok()) << status.to_string();
const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(2).column);
const auto& nested = assert_cast<const ColumnString&>(result.get_nested_column());
ASSERT_EQ(result.size(), texts.size());
for (size_t row = 0; row < null_map.size(); ++row) {
EXPECT_EQ(result.is_null_at(row), null_map[row] != 0);
}
EXPECT_EQ(nested.get_data_at(1).to_string(), "answer-a");
EXPECT_EQ(nested.get_data_at(3).to_string(), "answer-c");
}
TEST(AIFunctionTest, NullableInputWithoutNullsThroughPreparedFunction) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
setenv("AI_TEST_RESULT", R"(["answer-a","answer-b","answer-c"])", 1);
std::vector<std::string> texts = {"text-a", "text-b", "text-c"};
Block block;
block.insert({ColumnHelper::create_column<DataTypeString>(
std::vector<std::string>(texts.size(), "mock_resource")),
std::make_shared<DataTypeString>(), "resource"});
block.insert({ColumnHelper::create_nullable_column<DataTypeString>(
texts, std::vector<UInt8>(texts.size(), 0)),
make_nullable(std::make_shared<DataTypeString>()), "text"});
auto return_type = make_nullable(std::make_shared<DataTypeString>());
auto function = get_ai_function("ai_generate", block, return_type);
ASSERT_NE(function, nullptr);
block.insert({nullptr, return_type, "result"});
Status status = function->execute(ctx.get(), block, {0, 1}, 2, texts.size());
unsetenv("AI_TEST_RESULT");
ASSERT_TRUE(status.ok()) << status.to_string();
const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(2).column);
const auto& nested = assert_cast<const ColumnString&>(result.get_nested_column());
ASSERT_EQ(result.size(), texts.size());
for (size_t row = 0; row < texts.size(); ++row) {
EXPECT_FALSE(result.is_null_at(row));
}
EXPECT_EQ(nested.get_data_at(0).to_string(), "answer-a");
EXPECT_EQ(nested.get_data_at(1).to_string(), "answer-b");
EXPECT_EQ(nested.get_data_at(2).to_string(), "answer-c");
}
TEST(AIFunctionTest, NullableBoolResultThroughPreparedFunction) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
setenv("AI_TEST_RESULT", R"(["1","0"])", 1);
std::vector<std::string> texts = {"unused-null", "valid", "unused-null", "invalid"};
std::vector<UInt8> null_map = {1, 0, 1, 0};
Block block;
block.insert({ColumnHelper::create_column<DataTypeString>(
std::vector<std::string>(texts.size(), "mock_resource")),
std::make_shared<DataTypeString>(), "resource"});
block.insert({ColumnHelper::create_nullable_column<DataTypeString>(texts, null_map),
make_nullable(std::make_shared<DataTypeString>()), "text"});
auto return_type = make_nullable(std::make_shared<DataTypeBool>());
auto function = get_ai_function("ai_filter", block, return_type);
ASSERT_NE(function, nullptr);
EXPECT_TRUE(function->get_return_type()->equals(*return_type));
block.insert({nullptr, return_type, "result"});
Status status = function->execute(ctx.get(), block, {0, 1}, 2, texts.size());
unsetenv("AI_TEST_RESULT");
ASSERT_TRUE(status.ok()) << status.to_string();
const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(2).column);
const auto& nested = assert_cast<const ColumnUInt8&>(result.get_nested_column());
EXPECT_TRUE(result.is_null_at(0));
EXPECT_EQ(nested.get_element(1), 1);
EXPECT_TRUE(result.is_null_at(2));
EXPECT_EQ(nested.get_element(3), 0);
}
TEST(AIFunctionTest, NullableFloatResultMergesArgumentNullMaps) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
setenv("AI_TEST_RESULT", R"(["0.5","1.5"])", 1);
std::vector<std::string> text1 = {"left-a", "unused-null", "left-c", "left-d"};
std::vector<std::string> text2 = {"right-a", "right-b", "unused-null", "right-d"};
std::vector<UInt8> null_map1 = {0, 1, 0, 0};
std::vector<UInt8> null_map2 = {0, 0, 1, 0};
Block block;
block.insert({ColumnHelper::create_column<DataTypeString>(
std::vector<std::string>(text1.size(), "mock_resource")),
std::make_shared<DataTypeString>(), "resource"});
block.insert({ColumnHelper::create_nullable_column<DataTypeString>(text1, null_map1),
make_nullable(std::make_shared<DataTypeString>()), "text1"});
block.insert({ColumnHelper::create_nullable_column<DataTypeString>(text2, null_map2),
make_nullable(std::make_shared<DataTypeString>()), "text2"});
auto return_type = make_nullable(std::make_shared<DataTypeFloat32>());
auto function = get_ai_function("ai_similarity", block, return_type);
ASSERT_NE(function, nullptr);
EXPECT_TRUE(function->get_return_type()->equals(*return_type));
block.insert({nullptr, return_type, "result"});
Status status = function->execute(ctx.get(), block, {0, 1, 2}, 3, text1.size());
unsetenv("AI_TEST_RESULT");
ASSERT_TRUE(status.ok()) << status.to_string();
const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(3).column);
const auto& nested = assert_cast<const ColumnFloat32&>(result.get_nested_column());
EXPECT_FALSE(result.is_null_at(0));
EXPECT_FLOAT_EQ(nested.get_element(0), 0.5f);
EXPECT_TRUE(result.is_null_at(1));
EXPECT_TRUE(result.is_null_at(2));
EXPECT_FALSE(result.is_null_at(3));
EXPECT_FLOAT_EQ(nested.get_element(3), 1.5f);
}
TEST(AIFunctionTest, NullableArrayArgumentThroughPreparedFunction) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
setenv("AI_TEST_RESULT", R"(["positive"])", 1);
std::vector<std::string> texts = {"unused-null", "good product", "unused-null"};
std::vector<UInt8> labels_null_map = {1, 0, 1};
auto labels = create_string_array_column({{}, {"positive", "negative"}, {}});
Block block;
block.insert({ColumnHelper::create_column<DataTypeString>(
std::vector<std::string>(texts.size(), "mock_resource")),
std::make_shared<DataTypeString>(), "resource"});
block.insert({ColumnHelper::create_column<DataTypeString>(texts),
std::make_shared<DataTypeString>(), "text"});
block.insert(
{ColumnNullable::create(std::move(labels),
ColumnHelper::create_column<DataTypeUInt8>(labels_null_map)),
make_nullable(std::make_shared<DataTypeArray>(
make_nullable(std::make_shared<DataTypeString>()))),
"labels"});
auto return_type = make_nullable(std::make_shared<DataTypeString>());
auto function = get_ai_function("ai_classify", block, return_type);
ASSERT_NE(function, nullptr);
block.insert({nullptr, return_type, "result"});
Status status = function->execute(ctx.get(), block, {0, 1, 2}, 3, texts.size());
unsetenv("AI_TEST_RESULT");
ASSERT_TRUE(status.ok()) << status.to_string();
const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(3).column);
const auto& nested = assert_cast<const ColumnString&>(result.get_nested_column());
EXPECT_TRUE(result.is_null_at(0));
EXPECT_EQ(nested.get_data_at(1).to_string(), "positive");
EXPECT_TRUE(result.is_null_at(2));
}
TEST(AIFunctionTest, NullableLabelElementsThroughPreparedFunction) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
setenv("AI_TEST_RESULT", R"(["positive"])", 1);
auto labels = create_string_array_column({{"positive", "unused-null", "negative"}},
std::vector<UInt8> {0, 1, 0});
Block block;
block.insert({ColumnHelper::create_column<DataTypeString>({"mock_resource"}),
std::make_shared<DataTypeString>(), "resource"});
block.insert({ColumnHelper::create_column<DataTypeString>({"good product"}),
std::make_shared<DataTypeString>(), "text"});
block.insert({std::move(labels),
std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>()), "labels"});
auto return_type = std::make_shared<DataTypeString>();
auto function = get_ai_function("ai_classify", block, return_type);
ASSERT_NE(function, nullptr);
block.insert({nullptr, return_type, "result"});
Status status = function->execute(ctx.get(), block, {0, 1, 2}, 3, 1);
unsetenv("AI_TEST_RESULT");
ASSERT_TRUE(status.ok()) << status.to_string();
const auto& result = assert_cast<const ColumnString&>(*block.get_by_position(3).column);
ASSERT_EQ(result.size(), 1);
EXPECT_EQ(result.get_data_at(0).to_string(), "positive");
}
TEST(AIFunctionTest, AllNullConstArgumentReturnsConstNull) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
constexpr size_t row_count = 5;
auto resource = ColumnConst::create(
ColumnHelper::create_column<DataTypeString>({"mock_resource"}), row_count);
auto nullable_text =
ColumnHelper::create_nullable_column<DataTypeString>({""}, std::vector<UInt8> {1});
auto text = ColumnConst::create(std::move(nullable_text), row_count);
Block block;
block.insert({std::move(resource), std::make_shared<DataTypeString>(), "resource"});
block.insert({std::move(text), make_nullable(std::make_shared<DataTypeString>()), "text"});
auto return_type = make_nullable(std::make_shared<DataTypeString>());
auto function = get_ai_function("ai_generate", block, return_type);
ASSERT_NE(function, nullptr);
block.insert({nullptr, return_type, "result"});
Status status = function->execute(ctx.get(), block, {0, 1}, 2, row_count);
ASSERT_TRUE(status.ok()) << status.to_string();
ASSERT_TRUE(is_column_const(*block.get_by_position(2).column));
ColumnPtr full_result = block.get_by_position(2).column->convert_to_full_column_if_const();
const auto& nullable_result = assert_cast<const ColumnNullable&>(*full_result);
ASSERT_EQ(nullable_result.size(), row_count);
for (size_t row = 0; row < row_count; ++row) {
EXPECT_TRUE(nullable_result.is_null_at(row));
}
}
TEST(AIFunctionTest, NullResourceReturnsConstNullBeforeLookup) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
constexpr size_t row_count = 3;
auto nullable_resource =
ColumnHelper::create_nullable_column<DataTypeString>({""}, std::vector<UInt8> {1});
auto resource = ColumnConst::create(std::move(nullable_resource), row_count);
auto text = ColumnHelper::create_column<DataTypeString>(
std::vector<std::string>(row_count, "prompt"));
Block block;
block.insert(
{std::move(resource), make_nullable(std::make_shared<DataTypeString>()), "resource"});
block.insert({std::move(text), std::make_shared<DataTypeString>(), "text"});
auto return_type = make_nullable(std::make_shared<DataTypeString>());
auto function = get_ai_function("ai_generate", block, return_type);
ASSERT_NE(function, nullptr);
block.insert({nullptr, return_type, "result"});
Status status = function->execute(ctx.get(), block, {0, 1}, 2, row_count);
ASSERT_TRUE(status.ok()) << status.to_string();
ASSERT_TRUE(is_column_const(*block.get_by_position(2).column));
EXPECT_TRUE(block.get_by_position(2).column->only_null());
}
TEST(AIFunctionTest, MissingAIResourcesMetadataTest) {
auto query_ctx = MockQueryContext::create();
TQueryOptions query_options;
be/test/ai/embed_test.cpp
+188 -7
View File
@@ -26,10 +26,12 @@
#include <string>
#include <vector>
#include "core/column/column_const.h"
#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 "exprs/function/simple_function_factory.h"
#include "io/fs/obj_storage_client.h"
#include "testutil/column_helper.h"
#include "testutil/mock/mock_runtime_state.h"
@@ -151,6 +153,25 @@ static ColumnString::MutablePtr create_jsonb_column(const std::vector<std::strin
return column;
}
static ColumnPtr create_nullable_jsonb_column(const std::vector<std::string>& json_rows,
const std::vector<UInt8>& null_map) {
EXPECT_EQ(json_rows.size(), null_map.size());
auto column = ColumnString::create();
auto null_column = ColumnUInt8::create();
for (size_t i = 0; i < json_rows.size(); ++i) {
if (null_map[i]) {
column->insert_default();
} else {
JsonBinaryValue jsonb_value;
Status st = jsonb_value.from_json_string(json_rows[i]);
EXPECT_TRUE(st.ok()) << st.to_string();
column->insert_data(jsonb_value.value(), jsonb_value.size());
}
null_column->insert_value(null_map[i]);
}
return ColumnNullable::create(std::move(column), std::move(null_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);
@@ -168,6 +189,39 @@ static void assert_mock_embedding_column(const ColumnArray& col_array, size_t ro
}
}
static void assert_mock_nullable_embedding_column(const IColumn& column,
const std::vector<UInt8>& expected_null_map) {
const auto& nullable_column = assert_cast<const ColumnNullable&>(column);
ASSERT_EQ(nullable_column.size(), expected_null_map.size());
const auto& col_array = assert_cast<const ColumnArray&>(nullable_column.get_nested_column());
const auto& offsets = col_array.get_offsets();
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());
size_t expected_offset = 0;
for (size_t row = 0; row < expected_null_map.size(); ++row) {
ASSERT_EQ(nullable_column.is_null_at(row), expected_null_map[row] != 0);
if (expected_null_map[row]) {
ASSERT_EQ(offsets[row], expected_offset);
continue;
}
expected_offset += 5;
ASSERT_EQ(offsets[row], expected_offset);
for (size_t i = 0; i < 5; ++i) {
ASSERT_FLOAT_EQ(nested_col.get_element(expected_offset - 5 + i), static_cast<float>(i));
}
}
ASSERT_EQ(nested_col.size(), expected_offset);
}
static FunctionBasePtr get_embed_function(const Block& block, const DataTypePtr& return_type) {
return SimpleFunctionFactory::instance().get_function(
"embed", block.get_columns_with_type_and_name(), return_type);
}
TEST(EMBED_TEST, embed_function_build_test) {
FunctionEmbed function;
@@ -183,7 +237,7 @@ TEST(EMBED_TEST, embed_function_build_test) {
ColumnNumbers arguments = {0, 1};
std::string prompt;
Status status = function.build_prompt(block, arguments, 0, prompt);
Status status = function.build_prompt({block.get_by_position(arguments[1]).column}, 0, prompt);
ASSERT_TRUE(status.ok());
ASSERT_EQ(prompt, "this is a test prompt");
@@ -297,6 +351,133 @@ TEST(EMBED_TEST, embed_function_multimodal_direct_url) {
assert_mock_embedding_column(col_array, file_json_rows.size());
}
TEST(EMBED_TEST, embed_function_partial_null_through_framework) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
std::vector<std::string> resources(5, "mock_resource");
std::vector<std::string> texts = {"", "text-a", "", "text-c", ""};
std::vector<UInt8> null_map = {1, 0, 1, 0, 1};
auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
auto col_text = ColumnHelper::create_nullable_column<DataTypeString>(texts, null_map);
auto return_type = make_nullable(
std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())));
Block block;
block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
block.insert({std::move(col_text), make_nullable(std::make_shared<DataTypeString>()), "text"});
auto function = get_embed_function(block, return_type);
ASSERT_NE(function, nullptr);
EXPECT_TRUE(function->get_return_type()->equals(*return_type));
block.insert({nullptr, return_type, "result"});
const size_t result_idx = 2;
MockAdapter::clear_embedding_inputs_for_test();
Status exec_status = function->execute(ctx.get(), block, {0, 1}, result_idx, texts.size());
ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
EXPECT_THAT(MockAdapter::get_embedding_inputs_for_test(),
::testing::ElementsAre("text-a", "text-c"));
assert_mock_nullable_embedding_column(*block.get_by_position(result_idx).column, null_map);
}
TEST(EMBED_TEST, embed_function_multimodal_partial_null_through_framework) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
std::vector<std::string> resources(5, "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"})", ""};
std::vector<UInt8> null_map = {1, 0, 1, 0, 1};
auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
auto col_file = create_nullable_jsonb_column(file_json_rows, null_map);
auto return_type = make_nullable(
std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())));
Block block;
block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
block.insert({std::move(col_file), make_nullable(std::make_shared<DataTypeJsonb>()), "file"});
auto function = get_embed_function(block, return_type);
ASSERT_NE(function, nullptr);
block.insert({nullptr, return_type, "result"});
const size_t result_idx = 2;
Status exec_status =
function->execute(ctx.get(), block, {0, 1}, result_idx, file_json_rows.size());
ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
assert_mock_nullable_embedding_column(*block.get_by_position(result_idx).column, null_map);
}
TEST(EMBED_TEST, embed_function_all_null_const_nullable_through_framework) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
constexpr size_t row_count = 5;
std::vector<std::string> resources(row_count, "mock_resource");
auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
auto nullable_text =
ColumnHelper::create_nullable_column<DataTypeString>({""}, std::vector<UInt8> {1});
auto col_text = ColumnConst::create(std::move(nullable_text), row_count);
auto return_type = make_nullable(
std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())));
Block block;
block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
block.insert({std::move(col_text), make_nullable(std::make_shared<DataTypeString>()), "text"});
auto function = get_embed_function(block, return_type);
ASSERT_NE(function, nullptr);
block.insert({nullptr, return_type, "result"});
const size_t result_idx = 2;
Status exec_status = function->execute(ctx.get(), block, {0, 1}, result_idx, row_count);
ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
const auto& result_column = block.get_by_position(result_idx).column;
ASSERT_TRUE(is_column_const(*result_column));
EXPECT_TRUE(result_column->only_null());
ColumnPtr full_result = result_column->convert_to_full_column_if_const();
assert_mock_nullable_embedding_column(*full_result, std::vector<UInt8>(row_count, 1));
}
TEST(EMBED_TEST, embed_function_null_rows_across_batches_through_framework) {
TQueryOptions query_options = create_fake_query_options();
query_options.__set_embed_max_batch_size(2);
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> texts = {"", "text-a", "", "text-b", "", "text-c",
"", "text-d", "", "text-e", ""};
std::vector<UInt8> null_map = {1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1};
std::vector<std::string> resources(texts.size(), "mock_resource");
auto col_resource = ColumnHelper::create_column<DataTypeString>(resources);
auto col_text = ColumnHelper::create_nullable_column<DataTypeString>(texts, null_map);
auto return_type = make_nullable(
std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())));
Block block;
block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"});
block.insert({std::move(col_text), make_nullable(std::make_shared<DataTypeString>()), "text"});
auto function = get_embed_function(block, return_type);
ASSERT_NE(function, nullptr);
block.insert({nullptr, return_type, "result"});
const size_t result_idx = 2;
Status exec_status = function->execute(ctx.get(), block, {0, 1}, result_idx, texts.size());
ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
assert_mock_nullable_embedding_column(*block.get_by_position(result_idx).column, null_map);
}
TEST(EMBED_TEST, embed_function_multimodal_batch_request) {
auto runtime_state = std::make_unique<MockRuntimeState>();
auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {});
@@ -327,8 +508,8 @@ TEST(EMBED_TEST, embed_function_multimodal_batch_request) {
ColumnNumbers arguments = {0, 1};
size_t result_idx = 2;
FunctionEmbed embed_func;
Status exec_status = embed_func.execute_with_adapter(ctx.get(), block, arguments, result_idx,
file_json_rows.size(), config, adapter);
Status exec_status = embed_func.execute(ctx.get(), block, arguments, result_idx,
file_json_rows.size(), config, adapter);
ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
EXPECT_THAT(counting_adapter->batch_sizes, ::testing::ElementsAre(3));
@@ -373,8 +554,8 @@ TEST(EMBED_TEST, embed_function_multimodal_batch_split_by_session_variable) {
ColumnNumbers arguments = {0, 1};
size_t result_idx = 2;
FunctionEmbed embed_func;
Status exec_status = embed_func.execute_with_adapter(ctx.get(), block, arguments, result_idx,
file_json_rows.size(), config, adapter);
Status exec_status = embed_func.execute(ctx.get(), block, arguments, result_idx,
file_json_rows.size(), config, adapter);
ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
EXPECT_THAT(counting_adapter->batch_sizes, ::testing::ElementsAre(2, 1));
@@ -416,8 +597,8 @@ TEST(EMBED_TEST, embed_function_text_batch_split_by_session_variable) {
ColumnNumbers arguments = {0, 1};
size_t result_idx = 2;
FunctionEmbed embed_func;
Status exec_status = embed_func.execute_with_adapter(ctx.get(), block, arguments, result_idx,
texts.size(), config, adapter);
Status exec_status = embed_func.execute(ctx.get(), block, arguments, result_idx, texts.size(),
config, adapter);
ASSERT_TRUE(exec_status.ok()) << exec_status.to_string();
EXPECT_THAT(counting_adapter->batch_sizes, ::testing::ElementsAre(2, 1));