xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: tidy query APIs and execution (#431)

This commit is contained in:
Qinren Zhou 2026-06-01 21:07:35 +08:00 committed by GitHub
parent 02bfb31cf5
commit 8e8bb81db0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 260 additions and 152 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
src/binding/c/c_api.cc
View File

@ -5908,6 +5908,29 @@ zvec_error_code_t zvec_sub_query_set_query_vector(
return ZVEC_OK;
}
zvec_error_code_t zvec_sub_query_set_sparse_vector(
zvec_sub_query_t *query, const uint32_t *indices, const float *values,
size_t count) {
if (!query || (!indices && count > 0) || (!values && count > 0)) {
SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT,
"Sub-vector query, indices or values pointer is null");
return ZVEC_ERROR_INVALID_ARGUMENT;
}
auto *ptr = reinterpret_cast<zvec::SubQuery *>(query);
auto &payload = std::get<zvec::VectorClause>(ptr->target_.clause_);
if (count == 0) {
payload.sparse_indices_.clear();
payload.sparse_values_.clear();
return ZVEC_OK;
}
payload.sparse_indices_.assign(
reinterpret_cast<const char *>(indices), count * sizeof(uint32_t));
payload.sparse_values_.assign(
reinterpret_cast<const char *>(values), count * sizeof(float));
return ZVEC_OK;
}
zvec_error_code_t zvec_sub_query_set_sparse_indices(
zvec_sub_query_t *query, const uint32_t *indices, size_t count) {
if (!query || (!indices && count > 0)) {

72
src/db/collection.cc
View File

@ -14,8 +14,6 @@
#include <atomic>
#include <cstdint>
#include <filesystem>
#include <future>
#include <memory>
#include <mutex>
#include <set>
@ -36,6 +34,7 @@
#include <zvec/db/status.h>
#include "db/common/constants.h"
#include "db/common/file_helper.h"
#include "db/common/global_resource.h"
#include "db/common/profiler.h"
#include "db/common/typedef.h"
#include "db/index/common/delete_store.h"
@ -1581,7 +1580,8 @@ Status CollectionImpl::DeleteByFilter(const std::string &filter) {
query.output_fields_ = std::vector<std::string>{};
query.include_doc_id_ = true;
auto ret = sql_engine_->execute(schema_, query, get_all_segments());
auto ret =
sql_engine_->execute(schema_, std::move(query), get_all_segments());
if (!ret.has_value()) {
return ret.error();
}
@ -1619,7 +1619,7 @@ Result<DocPtrList> CollectionImpl::Query(const SearchQuery &query) const {
return DocPtrList();
}
return sql_engine_->execute(schema_, sanitized, segments);
return sql_engine_->execute(schema_, std::move(sanitized), segments);
}
Result<DocPtrList> CollectionImpl::Query(const MultiQuery &query) const {
@ -1628,13 +1628,14 @@ Result<DocPtrList> CollectionImpl::Query(const MultiQuery &query) const {
CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false);
if (query.queries.size() < 2) {
return tl::make_unexpected(
Status::InvalidArgument("Query requires at least 2 sub-queries"));
return tl::make_unexpected(Status::InvalidArgument(
"Invalid query: MultiQuery requires at least 2 sub-queries, got ",
query.queries.size()));
}
if (!query.reranker) {
return tl::make_unexpected(
Status::InvalidArgument("Reranker is required for multi-vector query"));
return tl::make_unexpected(Status::InvalidArgument(
"Invalid query: MultiQuery requires a reranker"));
}
auto segments = get_all_segments();
@ -1642,10 +1643,15 @@ Result<DocPtrList> CollectionImpl::Query(const MultiQuery &query) const {
return DocPtrList();
}
struct PendingQuery {
std::string field_name;
SearchQuery query;
};
// Convert each SubQuery to a SearchQuery and validate.
std::set<std::string> seen_fields;
std::vector<SearchQuery> converted_queries;
converted_queries.reserve(query.queries.size());
std::vector<PendingQuery> pending_queries;
pending_queries.reserve(query.queries.size());
for (const auto &sub : query.queries) {
const auto &target = sub.target_;
@ -1670,27 +1676,39 @@ Result<DocPtrList> CollectionImpl::Query(const MultiQuery &query) const {
auto s = sq.validate_and_sanitize(field_schema);
CHECK_RETURN_STATUS_EXPECTED(s);
converted_queries.push_back(std::move(sq));
}
// Execute each sub-query concurrently and collect results per field.
std::vector<std::future<Result<DocPtrList>>> futures;
futures.reserve(converted_queries.size());
for (const auto &sq : converted_queries) {
futures.push_back(std::async(std::launch::async, [&]() {
auto engine = sqlengine::SQLEngine::create(std::make_shared<Profiler>());
return engine->execute(schema_, sq, segments);
}));
pending_queries.push_back({target.field_name_, std::move(sq)});
}
std::map<std::string, DocPtrList> query_results;
for (size_t i = 0; i < converted_queries.size(); ++i) {
auto result = futures[i].get();
if (!result.has_value()) {
return tl::make_unexpected(result.error());
auto execute_query = [&](PendingQuery &pending) -> Result<DocPtrList> {
auto engine = sqlengine::SQLEngine::create(std::make_shared<Profiler>());
return engine->execute(schema_, std::move(pending.query), segments);
};
std::vector<Result<DocPtrList>> results(pending_queries.size());
// Single-segment queries have no segment-level fanout; multi-segment queries
// already use the query pool per sub-query.
if (segments.size() == 1) {
auto group = GlobalResource::Instance().query_thread_pool()->make_group();
for (size_t i = 0; i < pending_queries.size(); ++i) {
group->execute(
[&, i]() { results[i] = execute_query(pending_queries[i]); });
}
query_results[converted_queries[i].target_.field_name_] =
std::move(result.value());
group->wait_finish();
} else {
for (size_t i = 0; i < pending_queries.size(); ++i) {
results[i] = execute_query(pending_queries[i]);
}
}
for (size_t i = 0; i < pending_queries.size(); ++i) {
if (!results[i]) {
return tl::make_unexpected(results[i].error());
}
query_results[pending_queries[i].field_name] =
std::move(results[i].value());
}
// Merge and rerank results

29
src/db/common/profiler.h
View File

@ -14,6 +14,8 @@
#pragma once
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <zvec/ailego/encoding/json.h>
#include <zvec/ailego/logger/logger.h>
@ -198,4 +200,31 @@ class ScopedLatency {
Profiler::Ptr profiler_;
};
//! RAII helper that closes a profiler stage on every exit path.
class ScopedProfilerStage {
public:
ScopedProfilerStage(Profiler::Ptr profiler, const std::string &name)
: profiler_(std::move(profiler)) {
active_ = profiler_ && profiler_->open_stage(name) == 0;
}
~ScopedProfilerStage() {
close();
}
void close() {
if (active_) {
profiler_->close_stage();
active_ = false;
}
}
ScopedProfilerStage(const ScopedProfilerStage &) = delete;
ScopedProfilerStage &operator=(const ScopedProfilerStage &) = delete;
private:
Profiler::Ptr profiler_;
bool active_{false};
};
} // namespace zvec

41
src/db/sqlengine/analyzer/query_analyzer.cc
View File

@ -21,8 +21,6 @@
#include <zvec/db/config.h>
#include <zvec/db/status.h>
#include <zvec/db/type.h>
#include "db/common/constants.h"
#include "db/common/error_code.h"
#include "db/index/common/type_helper.h"
#include "db/sqlengine/analyzer/query_node.h"
#include "db/sqlengine/common/util.h"
@ -515,42 +513,35 @@ Status QueryAnalyzer::check_and_convert_vector(
vector_field_name);
}
std::string vector_term;
uint32_t dimension = vector_meta->dimension();
std::string vector_sparse_indices;
std::string vector_sparse_values;
QueryParams::Ptr query_params;
const QueryNode::Ptr &vector_value_node = query_rel_node->right();
// for pb request
if (vector_value_node->op() == QueryNodeOp::Q_VECTOR_MATRIX_VALUE) {
// for format vector = [,,,]
const QueryVectorMatrixNode::Ptr &vector_node =
QueryVectorMatrixNode::Ptr vector_node =
std::dynamic_pointer_cast<QueryVectorMatrixNode>(vector_value_node);
// we only have vector matrix, other info is not available
vector_term = vector_node->matrix();
vector_sparse_indices = vector_node->sparse_indices();
vector_sparse_values = vector_node->sparse_values();
query_params = vector_node->query_params();
// Consume the vector payload; this node is detached from the search
// condition after conversion.
auto vector_data = vector_node->take_node();
auto core_data_type =
DataTypeCodeBook::to_data_type(vector_meta->data_type());
if (core_data_type == core::IndexMeta::DataType::DT_UNDEFINED) {
return Status::InvalidArgument("invalid data type:",
(int)vector_meta->data_type());
}
*vector_cond = std::make_shared<QueryInfo::QueryVectorCondInfo>(
vector_meta, vector_data->take_matrix(), core_data_type, dimension,
vector_data->take_sparse_indices(), vector_data->take_sparse_values(),
vector_data->take_query_params());
return Status::OK();
} else {
return Status::InvalidArgument("invalid vector value node. op[",
vector_value_node->op_name(), "], text[",
vector_value_node->text(), "]");
}
auto core_data_type =
DataTypeCodeBook::to_data_type(vector_meta->data_type());
if (core_data_type == core::IndexMeta::DataType::DT_UNDEFINED) {
return Status::InvalidArgument("invalid data type:",
(int)vector_meta->data_type());
}
*vector_cond = std::make_shared<QueryInfo::QueryVectorCondInfo>(
vector_meta, vector_term, core_data_type, dimension,
std::move(vector_sparse_indices), std::move(vector_sparse_values),
std::move(query_params));
return Status::OK();
}

4
src/db/sqlengine/analyzer/query_info.h
View File

@ -47,13 +47,13 @@ class QueryInfo {
using Ptr = std::shared_ptr<QueryVectorCondInfo>;
QueryVectorCondInfo(const FieldSchema *vector_schema,
const std::string &vector_term,
std::string vector_term,
core::IndexMeta::DataType core_data_type, int dimension,
std::string vector_sparse_indices,
std::string vector_sparse_values,
QueryParams::Ptr query_params)
: vector_schema_(vector_schema),
vector_term_(vector_term),
vector_term_(std::move(vector_term)),
data_type_(core_data_type),
dimension_(dimension),
vector_sparse_indices_(std::move(vector_sparse_indices)),

10
src/db/sqlengine/analyzer/query_node.h
View File

@ -166,7 +166,7 @@ class QueryNode : public Generic_Node<QueryNodeOp, QueryNode> {
QueryNode::Ptr detach_from_invert_cond(QueryInfo *query_info_ptr);
virtual std::string text() const override;
std::string text() const override;
virtual void set_text(std::string /*new_val*/) {
/* for QueryConstantNode only */
@ -218,8 +218,12 @@ class QueryVectorMatrixNode : public QueryNode {
return node_->query_params();
}
std::shared_ptr<VectorMatrixNode> take_node() {
return std::move(node_);
}
private:
std::shared_ptr<const VectorMatrixNode> node_{nullptr};
std::shared_ptr<VectorMatrixNode> node_{nullptr};
};
class QueryConstantNode : public QueryNode {
@ -261,7 +265,7 @@ class QueryFuncNode : public QueryNode {
using Ptr = std::shared_ptr<QueryFuncNode>;
QueryFuncNode();
virtual ~QueryFuncNode() = default;
~QueryFuncNode() override = default;
void set_func_name_node(QueryNode::Ptr func_name_node);
const QueryNode::Ptr &get_func_name_node() const;

22
src/db/sqlengine/parser/node.h
View File

@ -118,7 +118,7 @@ class Node : public Generic_Node<NodeOp, Node> {
NodeType type();
virtual std::string text() const override;
std::string text() const override;
std::string to_string();
private:
@ -137,7 +137,7 @@ class RangeNode : public Node {
RangeNode();
RangeNode(bool m_min_equal, bool m_max_equal);
virtual ~RangeNode() = default;
~RangeNode() override = default;
void set_min_equal(bool value);
void set_max_equal(bool value);
@ -171,18 +171,34 @@ class VectorMatrixNode : public Node {
return matrix_;
}
std::string take_matrix() {
return std::move(matrix_);
}
const std::string &sparse_indices() const {
return sparse_indices_;
}
std::string take_sparse_indices() {
return std::move(sparse_indices_);
}
const std::string &sparse_values() const {
return sparse_values_;
}
std::string take_sparse_values() {
return std::move(sparse_values_);
}
const QueryParams::Ptr &query_params() const {
return query_params_;
}
QueryParams::Ptr take_query_params() {
return std::move(query_params_);
}
std::string text() const override {
// do not distinguish between matrix and vector
static std::string txt = "[...]";
@ -231,7 +247,7 @@ class FuncNode : public Node {
using Ptr = std::shared_ptr<FuncNode>;
FuncNode();
virtual ~FuncNode() = default;
~FuncNode() override = default;
void set_func_name_node(Node::Ptr func_name_node);
const Node::Ptr &get_func_name_node();

38
src/db/sqlengine/parser/sql_info_helper.cc
View File

@ -26,16 +26,17 @@ namespace zvec::sqlengine {
using namespace zvec;
Node::Ptr handle_vector(const SearchQuery &request, std::string * /*err_msg*/) {
const auto *vc = request.target_.get_vector_clause();
Node::Ptr handle_vector(SearchQuery *request) {
auto *vc = request->target_.get_vector_clause();
if (vc == nullptr) {
return nullptr;
}
Node::Ptr rel_exp = std::make_shared<Node>(NodeOp::T_EQ);
rel_exp->set_left(std::make_shared<IDNode>(request.target_.field_name_));
rel_exp->set_left(std::make_shared<IDNode>(request->target_.field_name_));
rel_exp->set_right(std::make_shared<VectorMatrixNode>(
vc->query_vector_, vc->sparse_indices_, vc->sparse_values_,
request.target_.query_params_));
std::move(vc->query_vector_), std::move(vc->sparse_indices_),
std::move(vc->sparse_values_),
std::move(request->target_.query_params_)));
return rel_exp;
}
@ -65,18 +66,18 @@ void handle_query_field(const SearchQuery *query, SelectInfo *selected_info) {
}
}
bool SQLInfoHelper::MessageToSQLInfo(const SearchQuery *query,
Node::Ptr filter_node,
std::shared_ptr<GroupBy> group_by,
sqlengine::SQLInfo::Ptr *sql_info,
std::string *err_msg) {
Result<sqlengine::SQLInfo::Ptr> SQLInfoHelper::BuildSQLInfoFromSearchQuery(
SearchQuery query, Node::Ptr filter_node,
std::shared_ptr<GroupBy> group_by) {
Node::Ptr index_params_node_ptr = nullptr;
if (const auto *vc = query->target_.get_vector_clause();
if (const auto *vc = query.target_.get_vector_clause();
vc != nullptr &&
(!vc->query_vector_.empty() || !vc->sparse_indices_.empty())) {
index_params_node_ptr = handle_vector(*query, err_msg);
index_params_node_ptr = handle_vector(&query);
if (index_params_node_ptr == nullptr) {
return false;
return tl::make_unexpected(Status::InvalidArgument(
"Failed to build vector condition for field: ",
query.target_.field_name_));
}
}
@ -92,13 +93,13 @@ bool SQLInfoHelper::MessageToSQLInfo(const SearchQuery *query,
}
SelectInfo::Ptr select_info = std::make_shared<SelectInfo>("");
handle_query_field(query, select_info.get());
handle_query_field(&query, select_info.get());
select_info->set_search_cond(cond_expr);
uint32_t topk = query->topk_;
uint32_t topk = query.topk_;
select_info->set_limit(topk);
select_info->set_include_vector(query->include_vector_);
select_info->set_include_doc_id(query->include_doc_id_);
select_info->set_include_vector(query.include_vector_);
select_info->set_include_doc_id(query.include_doc_id_);
select_info->set_group_by(std::move(group_by));
//
@ -111,8 +112,7 @@ bool SQLInfoHelper::MessageToSQLInfo(const SearchQuery *query,
// select_info->add_order_by_elem(std::move(orderby_elem_info));
// }
*sql_info = std::make_shared<SQLInfo>(SQLInfo::SQLType::SELECT, select_info);
return true;
return std::make_shared<SQLInfo>(SQLInfo::SQLType::SELECT, select_info);
}
} // namespace zvec::sqlengine

11
src/db/sqlengine/parser/sql_info_helper.h
View File

@ -15,6 +15,7 @@
#pragma once
#include <zvec/db/query.h>
#include <zvec/db/status.h>
#include "db/sqlengine/common/group_by.h"
#include "db/sqlengine/parser/node.h"
#include "db/sqlengine/parser/sql_info.h"
@ -23,11 +24,11 @@ namespace zvec::sqlengine {
class SQLInfoHelper {
public:
//! Perform QueryRequest to sql info conversion:
static bool MessageToSQLInfo(const SearchQuery *query, Node::Ptr filter_node,
std::shared_ptr<GroupBy> group_by,
sqlengine::SQLInfo::Ptr *sql_info,
std::string *err_msg);
//! Build SQLInfo from SearchQuery. Takes query by value so callers may copy
//! or move it; vector payloads can be moved while building SQLInfo.
static Result<sqlengine::SQLInfo::Ptr> BuildSQLInfoFromSearchQuery(
SearchQuery query, Node::Ptr filter_node,
std::shared_ptr<GroupBy> group_by);
};
} // namespace zvec::sqlengine

2
src/db/sqlengine/sqlengine.h
View File

@ -27,7 +27,7 @@ class SQLEngine {
virtual ~SQLEngine();
virtual Result<DocPtrList> execute(
CollectionSchema::Ptr collection, const SearchQuery &query,
CollectionSchema::Ptr collection, SearchQuery query,
const std::vector<Segment::Ptr> &segments) = 0;
virtual Result<GroupResults> execute_group_by(

61
src/db/sqlengine/sqlengine_impl.cc
View File

@ -56,13 +56,13 @@ SQLEngineImpl::SQLEngineImpl(zvec::Profiler::Ptr profiler)
: profiler_(std::move(profiler)) {}
Result<DocPtrList> SQLEngineImpl::execute(
CollectionSchema::Ptr collection, const SearchQuery &query,
CollectionSchema::Ptr collection, SearchQuery query,
const std::vector<Segment::Ptr> &segments) {
if (segments.empty()) {
return DocPtrList{};
}
auto query_info = parse_request(collection, query, nullptr);
auto query_info = build_query_info(collection, std::move(query), nullptr);
if (!query_info) {
return tl::make_unexpected(query_info.error());
}
@ -100,7 +100,7 @@ Result<GroupResults> SQLEngineImpl::execute_group_by(
}
SearchQuery query = from_group_by(group_by_query);
auto query_info = parse_request(
auto query_info = build_query_info(
collection, query,
std::make_shared<GroupBy>(group_by_query.group_by_field_name_,
group_by_query.group_topk_,
@ -231,24 +231,21 @@ Result<FtsCondInfo::Ptr> SQLEngineImpl::parse_fts_query(
Result<QueryInfo::Ptr> SQLEngineImpl::parse_sql_info(
const CollectionSchema &schema, const SQLInfo::Ptr &sql_info) {
profiler_->open_stage("analyze stage");
ScopedProfilerStage stage_guard(profiler_, "analyze_sql_info");
QueryAnalyzer analyzer;
auto query_info = analyzer.analyze(schema, sql_info);
if (!query_info) {
return tl::make_unexpected(Status::InvalidArgument(
"Analyze SQL info failed: ", query_info.error().c_str()));
}
profiler_->close_stage();
LOG_DEBUG("query_info: [%s]", query_info.value()->to_string().c_str());
return query_info.value();
}
Result<QueryInfo::Ptr> SQLEngineImpl::parse_request(
CollectionSchema::Ptr collection, const SearchQuery &request,
Result<QueryInfo::Ptr> SQLEngineImpl::build_query_info(
CollectionSchema::Ptr collection, SearchQuery request,
std::shared_ptr<GroupBy> group_by) {
profiler_->open_stage("message_to_sqlinfo");
sqlengine::SQLInfo::Ptr sql_info;
std::string err_msg;
ScopedProfilerStage stage_guard(profiler_, "build_sql_info");
Node::Ptr filter_node;
if (!request.filter_.empty()) {
ZVecParser::Ptr parser = ZVecParser::create();
@ -271,19 +268,7 @@ Result<QueryInfo::Ptr> SQLEngineImpl::parse_request(
}
}
sqlengine::SQLInfoHelper::MessageToSQLInfo(&request, std::move(filter_node),
std::move(group_by), &sql_info,
&err_msg);
profiler_->close_stage();
if (!err_msg.empty()) {
LOG_ERROR("QueryAgent, message to sql info failed, err_msg: %s",
err_msg.c_str());
return tl::make_unexpected(Status::InvalidArgument(
"Convert message to SQL info failed: ", err_msg));
}
// If the request carries an FTS query, parse it and attach to SelectInfo
// so that query_analyzer can propagate it to QueryInfo.
FtsCondInfo::Ptr fts_cond_info;
if (const auto *fts_clause = request.target_.get_fts_clause()) {
auto fts_result =
parse_fts_query(collection, request.target_.field_name_, *fts_clause,
@ -291,13 +276,30 @@ Result<QueryInfo::Ptr> SQLEngineImpl::parse_request(
if (!fts_result) {
return tl::make_unexpected(fts_result.error());
}
auto select_info =
std::dynamic_pointer_cast<SelectInfo>(sql_info->base_info());
select_info->set_fts_cond_info(std::move(fts_result.value()));
fts_cond_info = std::move(fts_result.value());
}
LOG_DEBUG("Sql info is %s", sql_info->to_string().c_str());
return parse_sql_info(*collection, std::move(sql_info));
auto sql_info = sqlengine::SQLInfoHelper::BuildSQLInfoFromSearchQuery(
std::move(request), std::move(filter_node), std::move(group_by));
if (!sql_info) {
return tl::make_unexpected(sql_info.error());
}
// Attach FTS info to SelectInfo so query_analyzer can propagate it to
// QueryInfo.
if (fts_cond_info) {
auto select_info =
std::dynamic_pointer_cast<SelectInfo>(sql_info.value()->base_info());
if (!select_info) {
return tl::make_unexpected(Status::InternalError(
"BuildSQLInfoFromSearchQuery did not produce SelectInfo"));
}
select_info->set_fts_cond_info(std::move(fts_cond_info));
}
LOG_DEBUG("Sql info is %s", sql_info.value()->to_string().c_str());
stage_guard.close();
return parse_sql_info(*collection, std::move(sql_info.value()));
}
Result<std::unique_ptr<arrow::RecordBatchReader>>
@ -306,7 +308,7 @@ SQLEngineImpl::search_by_query_info(
std::vector<sqlengine::QueryInfo::Ptr> *query_infos) {
global_init();
profiler_->open_stage("plan stage");
ScopedProfilerStage stage_guard(profiler_, "build_query_plan");
QueryPlanner planner(collection.get());
auto plan_info =
planner.make_plan(segments, profiler_->trace_id(), query_infos);
@ -314,7 +316,6 @@ SQLEngineImpl::search_by_query_info(
LOG_ERROR("plan query_info failed: [%s]", plan_info.error().c_str());
return tl::make_unexpected(plan_info.error());
}
profiler_->close_stage();
// LOG_DEBUG("plan_info: [%s]", plan_info->to_string().c_str());
return plan_info.value()->execute_to_reader();
}

12
src/db/sqlengine/sqlengine_impl.h
View File

@ -34,10 +34,10 @@ class SQLEngineImpl : public SQLEngine {
public:
SQLEngineImpl(zvec::Profiler::Ptr profiler);
//! Parse pb request
Result<QueryInfo::Ptr> parse_request(CollectionSchema::Ptr collection,
const SearchQuery &request,
std::shared_ptr<GroupBy> group_by);
//! Build analyzed query info from a structured search query.
Result<QueryInfo::Ptr> build_query_info(CollectionSchema::Ptr collection,
SearchQuery request,
std::shared_ptr<GroupBy> group_by);
//! Perform search with given query_info, segments and index filter
Result<std::unique_ptr<arrow::RecordBatchReader>> search_by_query_info(
@ -46,7 +46,7 @@ class SQLEngineImpl : public SQLEngine {
std::vector<sqlengine::QueryInfo::Ptr> *query_infos);
Result<DocPtrList> execute(
CollectionSchema::Ptr collection, const SearchQuery &query,
CollectionSchema::Ptr collection, SearchQuery query,
const std::vector<Segment::Ptr> &segments) override;
Result<GroupResults> execute_group_by(
@ -79,4 +79,4 @@ class SQLEngineImpl : public SQLEngine {
std::string execution_time_info_{};
};
} // namespace zvec::sqlengine
} // namespace zvec::sqlengine

12
src/include/zvec/c_api.h
View File

@ -2167,6 +2167,18 @@ zvec_sub_query_get_field_name(const zvec_sub_query_t *query);
ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_sub_query_set_query_vector(
zvec_sub_query_t *query, const void *data, size_t size);
/**
* @brief Set sparse vector indices and values
* @param query Sub-query pointer
* @param indices Array of uint32_t indices
* @param values Array of float values
* @param count Number of sparse vector entries
* @return zvec_error_code_t Error code
*/
ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_sub_query_set_sparse_vector(
zvec_sub_query_t *query, const uint32_t *indices, const float *values,
size_t count);
/**
* @brief Set sparse vector indices
* @param query Sub-query pointer

13
tests/c/c_api_test.c
View File

@ -4562,6 +4562,19 @@ void test_multi_vector_query_with_rrf_reranker(void) {
zvec_free((char *)got_fields);
}
zvec_sub_query_t *sparse_query = zvec_sub_query_create();
TEST_ASSERT(sparse_query != NULL);
uint32_t sparse_indices[] = {1, 3};
float sparse_values[] = {0.25f, 0.75f};
err = zvec_sub_query_set_sparse_vector(sparse_query, sparse_indices,
sparse_values, 2);
TEST_ASSERT(err == ZVEC_OK);
err = zvec_sub_query_set_sparse_vector(sparse_query, NULL, sparse_values, 2);
TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT);
err = zvec_sub_query_set_sparse_vector(sparse_query, sparse_indices, NULL, 2);
TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT);
zvec_sub_query_destroy(sparse_query);
zvec_multi_query_destroy(mvq2);
teardown_multi_query_fixture(&f);

22
tests/db/sqlengine/optimizer_test.cc
View File

@ -100,7 +100,7 @@ TEST_F(OptimizerTest, Basic) {
query.filter_ = "age > 200";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -123,7 +123,7 @@ TEST_F(OptimizerTest, Case1) {
query.filter_ = "age > 12";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -146,7 +146,7 @@ TEST_F(OptimizerTest, Case2_1) {
query.filter_ = "age > 100 and age > 101 or age > 102";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -170,7 +170,7 @@ TEST_F(OptimizerTest, Case2_2) {
query.filter_ = "age > 100 or age > 90";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -194,7 +194,7 @@ TEST_F(OptimizerTest, Case3_1) {
query.filter_ = "age > 100 and age > 101 and age > 10";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -218,7 +218,7 @@ TEST_F(OptimizerTest, Case3_2) {
query.filter_ = "age > 10 and age > 11 and age > 100";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -241,7 +241,7 @@ TEST_F(OptimizerTest, Case3_3) {
query.filter_ = "(age > 10 or age > 11) and age > 100";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -265,7 +265,7 @@ TEST_F(OptimizerTest, Case3_4) {
query.filter_ = "age > 10 and (age > 101 and (age > 10 and age > 10))";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -289,7 +289,7 @@ TEST_F(OptimizerTest, Case4) {
query.filter_ = "age in (10, 20)";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
@ -304,7 +304,7 @@ TEST_F(OptimizerTest, Case4) {
// in and optimizable, optimize optimizable
query.filter_ = "age in (10, 20) and age > 100";
ret = engine->parse_request(schema, query, nullptr);
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
query_info = ret.value();
optimized = optimizer->optimize(segment.get(), query_info.get());
@ -312,7 +312,7 @@ TEST_F(OptimizerTest, Case4) {
// in or optimizable, not optimized
query.filter_ = "age in (10, 20) or age > 100";
ret = engine->parse_request(schema, query, nullptr);
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
query_info = ret.value();
optimized = optimizer->optimize(segment.get(), query_info.get());

38
tests/db/sqlengine/query_info_test.cc
View File

@ -96,7 +96,7 @@ TEST_F(QueryInfoTest, BasicQueryRequest) {
query.target_.query_params_->set_radius(0.8F);
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
@ -137,7 +137,7 @@ TEST_F(QueryInfoTest, QueryRequestWithFilter) {
query.filter_ = "name<3 or name=4 or 1-dash_score_field='test'";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
@ -221,7 +221,7 @@ TEST_F(QueryInfoTest, QueryRequestWithIncludeVector) {
query.include_vector_ = true;
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
@ -263,7 +263,7 @@ TEST_F(QueryInfoTest, OR_ANCESTOR) {
query.filter_ = "name=1 and (name=2 or name=3)";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
}
@ -281,7 +281,7 @@ TEST_F(QueryInfoTest, QueryRequestWithInFilter) {
"name=3 or name in (1, 2, 3) or category not in (\"a\", \"b\", \"c\")";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
@ -370,23 +370,23 @@ TEST_F(QueryInfoTest, QueryRequestWithInFilterWrong) {
query.target_.query_params_->set_radius(0.8F);
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
query.filter_ = ("name in ()");
ret = engine->parse_request(schema, query, nullptr);
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
query.filter_ = ("name in (\"a\", 2, 3)");
ret = engine->parse_request(schema, query, nullptr);
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
query.filter_ = ("name in (1.1, 2, 3)");
ret = engine->parse_request(schema, query, nullptr);
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
query.filter_ = ("category in (1.1, \"b\")");
ret = engine->parse_request(schema, query, nullptr);
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
}
@ -410,7 +410,7 @@ TEST_F(QueryInfoTest, QueryRequestWithInFilterNum1024) {
query.filter_ = filter_str;
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
@ -451,7 +451,7 @@ TEST_F(QueryInfoTest, QueryRequestWithFilter_contain) {
)";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
@ -598,7 +598,7 @@ TEST_F(QueryInfoTest, SelectNonExistField) {
query.include_vector_ = false;
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr("not defined in schema"));
@ -616,7 +616,7 @@ TEST_F(QueryInfoTest, ContainAllExceedLimit) {
}
query.filter_ += ")";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr(
@ -635,7 +635,7 @@ TEST_F(QueryInfoTest, ContainAnyExceedLimit) {
}
query.filter_ += ")";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr(
@ -647,7 +647,7 @@ TEST_F(QueryInfoTest, ArrayLengthNonExistField) {
query.topk_ = 200;
query.filter_ = "array_length(not_exist_field) > 1";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr("array_length argument not found in schema"));
@ -658,7 +658,7 @@ TEST_F(QueryInfoTest, ArrayLengthOnNonArrayField) {
query.topk_ = 200;
query.filter_ = "array_length(name) > 1";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr("array_length only support array"));
@ -669,7 +669,7 @@ TEST_F(QueryInfoTest, ArrayLengthInvalidArgument) {
query.topk_ = 200;
query.filter_ = "array_length(name_array) > '1'";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(
ret.error().message(),
@ -681,7 +681,7 @@ TEST_F(QueryInfoTest, ArrayLengthInvalidOp) {
query.topk_ = 200;
query.filter_ = "array_length(name_array) like '%'";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(), testing::HasSubstr("syntax error"));
}

2
tests/db/sqlengine/simple_rewriter_test.cc
View File

@ -195,7 +195,7 @@ class SimpleRewriterTest : public testing::Test {
query.filter_ = filter;
auto engine = std::make_shared<SQLEngineImpl>(profiler_);
auto ret = engine->parse_request(schema, query, nullptr);
auto ret = engine->build_query_info(schema, query, nullptr);
// ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();