diff --git a/src/query/v2/plan/operator.cpp b/src/query/v2/plan/operator.cpp index 068ca9192..196f405aa 100644 --- a/src/query/v2/plan/operator.cpp +++ b/src/query/v2/plan/operator.cpp @@ -92,6 +92,7 @@ extern const Event ScanAllByLabelPropertyRangeOperator; extern const Event ScanAllByLabelPropertyValueOperator; extern const Event ScanAllByLabelPropertyOperator; extern const Event ScanAllByIdOperator; +extern const Event ScanAllByPrimaryKeyOperator; extern const Event ExpandOperator; extern const Event ExpandVariableOperator; extern const Event ConstructNamedPathOperator; @@ -545,6 +546,20 @@ UniqueCursorPtr ScanAllByLabelProperty::MakeCursor(utils::MemoryResource *mem) c throw QueryRuntimeException("ScanAllByLabelProperty is not supported"); } +ScanAllByPrimaryKey::ScanAllByPrimaryKey(const std::shared_ptr &input, Symbol output_symbol, + storage::v3::LabelId label, std::vector primary_key, + storage::v3::View view) + : ScanAll(input, output_symbol, view), label_(label), primary_key_(primary_key) { + MG_ASSERT(primary_key.front()); +} + +ACCEPT_WITH_INPUT(ScanAllByPrimaryKey) + +UniqueCursorPtr ScanAllByPrimaryKey::MakeCursor(utils::MemoryResource *mem) const { + // EventCounter::IncrementCounter(EventCounter::ScanAllByPrimaryKeyOperator); + throw QueryRuntimeException("ScanAllByPrimaryKey cursur is yet to be implemented."); +} + ScanAllById::ScanAllById(const std::shared_ptr &input, Symbol output_symbol, Expression *expression, storage::v3::View view) : ScanAll(input, output_symbol, view), expression_(expression) { diff --git a/src/query/v2/plan/operator.lcp b/src/query/v2/plan/operator.lcp index d6277809c..680bb4de3 100644 --- a/src/query/v2/plan/operator.lcp +++ b/src/query/v2/plan/operator.lcp @@ -111,6 +111,7 @@ class ScanAllByLabelPropertyRange; class ScanAllByLabelPropertyValue; class ScanAllByLabelProperty; class ScanAllById; +class ScanAllByPrimaryKey; class Expand; class ExpandVariable; class ConstructNamedPath; @@ -141,7 +142,7 @@ class Foreach; using LogicalOperatorCompositeVisitor = utils::CompositeVisitor< Once, CreateNode, CreateExpand, ScanAll, ScanAllByLabel, ScanAllByLabelPropertyRange, ScanAllByLabelPropertyValue, - ScanAllByLabelProperty, ScanAllById, + ScanAllByLabelProperty, ScanAllById, ScanAllByPrimaryKey, Expand, ExpandVariable, ConstructNamedPath, Filter, Produce, Delete, SetProperty, SetProperties, SetLabels, RemoveProperty, RemoveLabels, EdgeUniquenessFilter, Accumulate, Aggregate, Skip, Limit, OrderBy, Merge, @@ -841,7 +842,28 @@ given label and property. (:serialize (:slk)) (:clone)) +(lcp:define-class scan-all-by-primary-key (scan-all) + ((label "::storage::v3::LabelId" :scope :public) + (primary-key "std::vector" :scope :public) + (expression "Expression *" :scope :public + :slk-save #'slk-save-ast-pointer + :slk-load (slk-load-ast-pointer "Expression"))) + (:documentation + "ScanAll producing a single node with specified by the label and primary key") + (:public + #>cpp + ScanAllByPrimaryKey() {} + ScanAllByPrimaryKey(const std::shared_ptr &input, + Symbol output_symbol, + storage::v3::LabelId label, + std::vector primary_key, + storage::v3::View view = storage::v3::View::OLD); + bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override; + UniqueCursorPtr MakeCursor(utils::MemoryResource *) const override; + cpp<#) + (:serialize (:slk)) + (:clone)) (lcp:define-class scan-all-by-id (scan-all) ((expression "Expression *" :scope :public diff --git a/src/query/v2/plan/rewrite/index_lookup.hpp b/src/query/v2/plan/rewrite/index_lookup.hpp index 57ddba54e..a861153e9 100644 --- a/src/query/v2/plan/rewrite/index_lookup.hpp +++ b/src/query/v2/plan/rewrite/index_lookup.hpp @@ -25,8 +25,10 @@ #include +#include "query/v2/frontend/ast/ast.hpp" #include "query/v2/plan/operator.hpp" #include "query/v2/plan/preprocess.hpp" +#include "storage/v3/id_types.hpp" DECLARE_int64(query_vertex_count_to_expand_existing); @@ -584,6 +586,26 @@ class IndexLookupRewriter final : public HierarchicalLogicalOperatorVisitor { // Without labels, we cannot generate any indexed ScanAll. return nullptr; } + + // First, try to see if we can find a vertex based on the possibly + // supplied primary key. + auto property_filters = filters_.PropertyFilters(node_symbol); + storage::v3::LabelId prim_label; + std::vector primary_key; + + if (!property_filters.empty()) { + for (const auto &label : labels) { + if (db_->LabelIndexExists(GetLabel(label))) { + prim_label = GetLabel(label); + primary_key = db_->ExtractPrimaryKey(prim_label, property_filters); + break; + } + } + if (!primary_key.empty()) { + return std::make_unique(input, node_symbol, prim_label, primary_key); + } + } + auto found_index = FindBestLabelPropertyIndex(node_symbol, bound_symbols); if (found_index && // Use label+property index if we satisfy max_vertex_count. diff --git a/src/query/v2/plan/vertex_count_cache.hpp b/src/query/v2/plan/vertex_count_cache.hpp index a7bfbdf85..70b8d324b 100644 --- a/src/query/v2/plan/vertex_count_cache.hpp +++ b/src/query/v2/plan/vertex_count_cache.hpp @@ -12,9 +12,11 @@ /// @file #pragma once +#include #include #include "query/v2/bindings/typed_value.hpp" +#include "query/v2/plan/preprocess.hpp" #include "query/v2/shard_request_manager.hpp" #include "storage/v3/conversions.hpp" #include "storage/v3/id_types.hpp" @@ -52,11 +54,31 @@ class VertexCountCache { return 1; } - // For now return true if label is primary label bool LabelIndexExists(storage::v3::LabelId label) { return shard_request_manager_->IsPrimaryLabel(label); } bool LabelPropertyIndexExists(storage::v3::LabelId /*label*/, storage::v3::PropertyId /*property*/) { return false; } + std::vector ExtractPrimaryKey(storage::v3::LabelId label, + std::vector property_filters) { + std::vector pk; + const auto schema = shard_request_manager_->GetSchemaForLabel(label); + + std::vector schema_properties; + schema_properties.reserve(schema.size()); + + std::transform(schema.begin(), schema.end(), std::back_inserter(schema_properties), + [](const auto &schema_elem) { return schema_elem.property_id; }); + + for (const auto &property_filter : property_filters) { + const auto &property_id = NameToProperty(property_filter.property_filter->property_.name); + if (std::find(schema_properties.begin(), schema_properties.end(), property_id) != schema_properties.end()) { + pk.push_back(property_filter.expression); + } + } + + return pk.size() == schema_properties.size() ? pk : std::vector{}; + } + msgs::ShardRequestManagerInterface *shard_request_manager_; }; diff --git a/src/query/v2/shard_request_manager.hpp b/src/query/v2/shard_request_manager.hpp index a73201046..44ea02bba 100644 --- a/src/query/v2/shard_request_manager.hpp +++ b/src/query/v2/shard_request_manager.hpp @@ -131,6 +131,7 @@ class ShardRequestManagerInterface { virtual const std::string &EdgeTypeToName(memgraph::storage::v3::EdgeTypeId type) const = 0; virtual bool IsPrimaryLabel(LabelId label) const = 0; virtual bool IsPrimaryKey(LabelId primary_label, PropertyId property) const = 0; + virtual std::vector GetSchemaForLabel(LabelId label) const = 0; }; // TODO(kostasrim)rename this class template @@ -244,6 +245,10 @@ class ShardRequestManager : public ShardRequestManagerInterface { }) != schema_it->second.end(); } + std::vector GetSchemaForLabel(LabelId label) const override { + return shards_map_.schemas.at(label); + } + bool IsPrimaryLabel(LabelId label) const override { return shards_map_.label_spaces.contains(label); } // TODO(kostasrim) Simplify return result