query_modules/vector_search_module.cpp
Rama master · réplica de lectura del repositorio original.
// Copyright 2026 Memgraph Ltd.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
// License, and you may not use this file except in compliance with the Business Source License.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.
#include <cmath>
#include <iostream>
#include <string_view>
#include <mgp.hpp>
namespace VectorSearch {
static constexpr std::string_view kProcedureSearch = "search";
static constexpr std::string_view kProcedureSearchEdges = "search_edges";
static constexpr std::string_view kParameterIndexName = "index_name";
static constexpr std::string_view kParameterResultSetSize = "result_set_size";
static constexpr std::string_view kParameterQueryVector = "query_vector";
static constexpr std::string_view kReturnNode = "node";
static constexpr std::string_view kReturnEdge = "edge";
static constexpr std::string_view kReturnDistance = "distance";
static constexpr std::string_view kReturnSimilarity = "similarity";
static constexpr std::string_view kProcedureShowIndexInfo = "show_index_info";
static constexpr std::string_view kReturnIndexName = "index_name";
static constexpr std::string_view kReturnLabel = "label";
static constexpr std::string_view kReturnProperty = "property";
static constexpr std::string_view kMetric = "metric";
static constexpr std::string_view kReturnDimension = "dimension";
static constexpr std::string_view kReturnCapacity = "capacity";
static constexpr std::string_view kReturnSize = "size";
static constexpr std::string_view kReturnScalarKind = "scalar_kind";
static constexpr std::string_view kReturnIndexType = "index_type";
static constexpr std::string_view kProcedureCosineSimilarity = "cosine_similarity";
static constexpr std::string_view kParameterVector1 = "vector1";
static constexpr std::string_view kParameterVector2 = "vector2";
void Search(mgp_list *args, mgp_graph *memgraph_graph, mgp_result *result, mgp_memory *memory);
void SearchEdges(mgp_list *args, mgp_graph *memgraph_graph, mgp_result *result, mgp_memory *memory);
void ShowIndexInfo(mgp_list *args, mgp_graph *memgraph_graph, mgp_result *result, mgp_memory *memory);
void CosineSimilarityFunction(mgp_list *args, mgp_func_context *ctx, mgp_func_result *res, mgp_memory *memory);
} // namespace VectorSearch
void VectorSearch::Search(mgp_list *args, mgp_graph *memgraph_graph, mgp_result *result, mgp_memory *memory) {
mgp::MemoryDispatcherGuard guard{memory};
const auto record_factory = mgp::RecordFactory(result);
auto arguments = mgp::List(args);
try {
const auto index_name = arguments[0].ValueString();
const auto result_set_size = arguments[1].ValueInt();
auto query_vector = arguments[2].ValueList();
auto results = mgp::SearchVectorIndex(memgraph_graph, index_name, query_vector, result_set_size);
for (const auto &result : results) {
auto record = record_factory.NewRecord();
auto result_list = result.ValueList();
record.Insert(VectorSearch::kReturnNode.data(), result_list[0].ValueNode());
record.Insert(VectorSearch::kReturnDistance.data(), result_list[1].ValueDouble());
record.Insert(VectorSearch::kReturnSimilarity.data(), result_list[2].ValueDouble());
}
} catch (const std::exception &e) {
record_factory.SetErrorMessage(e.what());
}
}
void VectorSearch::SearchEdges(mgp_list *args, mgp_graph *memgraph_graph, mgp_result *result, mgp_memory *memory) {
mgp::MemoryDispatcherGuard guard{memory};
const auto record_factory = mgp::RecordFactory(result);
auto arguments = mgp::List(args);
try {
const auto index_name = arguments[0].ValueString();
const auto result_set_size = arguments[1].ValueInt();
auto query_vector = arguments[2].ValueList();
auto results = mgp::SearchVectorIndexOnEdges(memgraph_graph, index_name, query_vector, result_set_size);
for (const auto &result : results) {
auto record = record_factory.NewRecord();
auto result_list = result.ValueList();
record.Insert(VectorSearch::kReturnEdge.data(), result_list[0].ValueRelationship());
record.Insert(VectorSearch::kReturnDistance.data(), result_list[1].ValueDouble());
record.Insert(VectorSearch::kReturnSimilarity.data(), result_list[2].ValueDouble());
}
} catch (const std::exception &e) {
record_factory.SetErrorMessage(e.what());
}
}
void VectorSearch::ShowIndexInfo(mgp_list *args, mgp_graph *memgraph_graph, mgp_result *result, mgp_memory *memory) {
mgp::MemoryDispatcherGuard guard{memory};
const auto record_factory = mgp::RecordFactory(result);
auto arguments = mgp::List(args);
try {
const auto index_info = mgp::GetVectorIndexInfo(memgraph_graph);
for (const auto &info : index_info) {
auto record = record_factory.NewRecord();
auto info_list = info.ValueList();
record.Insert(VectorSearch::kReturnIndexName.data(), info_list[0].ValueString());
record.Insert(VectorSearch::kReturnLabel.data(), info_list[1].ValueString());
record.Insert(VectorSearch::kReturnProperty.data(), info_list[2].ValueString());
record.Insert(VectorSearch::kMetric.data(), info_list[3].ValueString());
record.Insert(VectorSearch::kReturnDimension.data(), info_list[4].ValueInt());
record.Insert(VectorSearch::kReturnCapacity.data(), info_list[5].ValueInt());
record.Insert(VectorSearch::kReturnSize.data(), info_list[6].ValueInt());
record.Insert(VectorSearch::kReturnScalarKind.data(), info_list[7].ValueString());
record.Insert(VectorSearch::kReturnIndexType.data(), info_list[8].ValueString());
}
} catch (const std::exception &e) {
record_factory.SetErrorMessage(e.what());
}
}
void VectorSearch::CosineSimilarityFunction(mgp_list *args, mgp_func_context * /*ctx*/, mgp_func_result *res,
mgp_memory *memory) {
try {
mgp::MemoryDispatcherGuard guard{memory};
const auto vector1 = mgp::Value(mgp::ref_type, mgp::list_at(args, 0)).ValueList();
const auto vector2 = mgp::Value(mgp::ref_type, mgp::list_at(args, 1)).ValueList();
if (vector1.Size() == 0 || vector1.Size() != vector2.Size()) {
throw std::invalid_argument("Vectors must be non-empty and have the same dimension");
}
auto dot_product = 0.0;
auto magnitude1 = 0.0;
auto magnitude2 = 0.0;
for (auto i = 0; i < vector1.Size(); ++i) {
auto get_numeric_value = [&](const mgp::Value &val) {
if (val.IsDouble()) {
return val.ValueDouble();
}
if (val.IsInt()) {
return static_cast<double>(val.ValueInt());
}
throw std::invalid_argument("Vector elements must be numeric (int or double)");
};
const auto val1 = get_numeric_value(vector1[i]);
const auto val2 = get_numeric_value(vector2[i]);
dot_product += val1 * val2;
magnitude1 += val1 * val1;
magnitude2 += val2 * val2;
}
magnitude1 = std::sqrt(magnitude1);
magnitude2 = std::sqrt(magnitude2);
if (magnitude1 == 0.0 || magnitude2 == 0.0) [[unlikely]] {
throw std::invalid_argument("Cannot calculate cosine similarity for zero vectors");
}
const auto cosine_similarity = dot_product / (magnitude1 * magnitude2);
auto result = mgp::Result(res);
result.SetValue(cosine_similarity);
} catch (const std::exception &e) {
mgp::func_result_set_error_msg(res, e.what(), memory);
}
}
extern "C" int mgp_init_module(struct mgp_module *module, struct mgp_memory *memory) {
try {
mgp::MemoryDispatcherGuard guard{memory};
AddProcedure(VectorSearch::Search,
VectorSearch::kProcedureSearch,
mgp::ProcedureType::Read,
{
mgp::Parameter(VectorSearch::kParameterIndexName, mgp::Type::String),
mgp::Parameter(VectorSearch::kParameterResultSetSize, mgp::Type::Int),
mgp::Parameter(VectorSearch::kParameterQueryVector, {mgp::Type::List, mgp::Type::Any}),
},
{
mgp::Return(VectorSearch::kReturnNode, mgp::Type::Node),
mgp::Return(VectorSearch::kReturnDistance, mgp::Type::Double),
mgp::Return(VectorSearch::kReturnSimilarity, mgp::Type::Double),
},
module,
memory);
AddProcedure(VectorSearch::ShowIndexInfo,
VectorSearch::kProcedureShowIndexInfo,
mgp::ProcedureType::Read,
{},
{
mgp::Return(VectorSearch::kReturnIndexName, mgp::Type::String),
mgp::Return(VectorSearch::kReturnLabel, mgp::Type::String),
mgp::Return(VectorSearch::kReturnProperty, mgp::Type::String),
mgp::Return(VectorSearch::kMetric, mgp::Type::String),
mgp::Return(VectorSearch::kReturnDimension, mgp::Type::Int),
mgp::Return(VectorSearch::kReturnCapacity, mgp::Type::Int),
mgp::Return(VectorSearch::kReturnSize, mgp::Type::Int),
mgp::Return(VectorSearch::kReturnScalarKind, mgp::Type::String),
mgp::Return(VectorSearch::kReturnIndexType, mgp::Type::String),
},
module,
memory);
AddProcedure(VectorSearch::SearchEdges,
VectorSearch::kProcedureSearchEdges,
mgp::ProcedureType::Read,
{
mgp::Parameter(VectorSearch::kParameterIndexName, mgp::Type::String),
mgp::Parameter(VectorSearch::kParameterResultSetSize, mgp::Type::Int),
mgp::Parameter(VectorSearch::kParameterQueryVector, {mgp::Type::List, mgp::Type::Any}),
},
{
mgp::Return(VectorSearch::kReturnEdge, mgp::Type::Relationship),
mgp::Return(VectorSearch::kReturnDistance, mgp::Type::Double),
mgp::Return(VectorSearch::kReturnSimilarity, mgp::Type::Double),
},
module,
memory);
mgp::AddFunction(VectorSearch::CosineSimilarityFunction,
VectorSearch::kProcedureCosineSimilarity,
{
mgp::Parameter(VectorSearch::kParameterVector1, {mgp::Type::List, mgp::Type::Any}),
mgp::Parameter(VectorSearch::kParameterVector2, {mgp::Type::List, mgp::Type::Any}),
},
module,
memory);
} catch (const std::exception &e) {
std::cerr << "Error while initializing query module: " << e.what() << '\n';
return 1;
}
return 0;
}
extern "C" int mgp_shutdown_module() { return 0; }