Skip to content

Commit c6184c8

Browse files
zhanglei1949zhanglei1949
authored
fix(interactive): Fix the property getter for primary key (#4319)
- Unify the getter for primary key properties with the common property getters. - Refactor the code, let `MutableFragment`, `GraphDB`, `GraphDBSession` and `ReadTransaction` all have methods to retrieve getters for vertex_id and vertex properties. - DO we need to keep `VertexIdVertexAccessor` and `VertexIdPathAccessor`?
1 parent f41ad88 commit c6184c8

File tree

11 files changed

+129
-72
lines changed

11 files changed

+129
-72
lines changed

‎flex/engines/graph_db/database/graph_db.cc‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -305,7 +305,12 @@ const Schema& GraphDB::schema() const { return graph_.schema(); }
305305

306306
std::shared_ptr<ColumnBase> GraphDB::get_vertex_property_column(
307307
uint8_t label, const std::string& col_name) const {
308-
return graph_.get_vertex_table(label).get_column(col_name);
308+
return graph_.get_vertex_property_column(label, col_name);
309+
}
310+
311+
std::shared_ptr<RefColumnBase> GraphDB::get_vertex_id_column(
312+
uint8_t label) const {
313+
return graph_.get_vertex_id_column(label);
309314
}
310315

311316
AppWrapper GraphDB::CreateApp(uint8_t app_type, int thread_id) {

‎flex/engines/graph_db/database/graph_db.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,8 @@ class GraphDB {
137137
std::shared_ptr<ColumnBase> get_vertex_property_column(
138138
uint8_t label, const std::string& col_name) const;
139139

140+
std::shared_ptr<RefColumnBase> get_vertex_id_column(uint8_t label) const;
141+
140142
AppWrapper CreateApp(uint8_t app_type, int thread_id);
141143

142144
void GetAppInfo(Encoder& result);

‎flex/engines/graph_db/database/graph_db_session.cc‎

Lines changed: 1 addition & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -79,33 +79,7 @@ std::shared_ptr<ColumnBase> GraphDBSession::get_vertex_property_column(
7979

8080
std::shared_ptr<RefColumnBase> GraphDBSession::get_vertex_id_column(
8181
uint8_t label) const {
82-
if (db_.graph().lf_indexers_[label].get_type() == PropertyType::kInt64) {
83-
return std::make_shared<TypedRefColumn<int64_t>>(
84-
dynamic_cast<const TypedColumn<int64_t>&>(
85-
db_.graph().lf_indexers_[label].get_keys()));
86-
} else if (db_.graph().lf_indexers_[label].get_type() ==
87-
PropertyType::kInt32) {
88-
return std::make_shared<TypedRefColumn<int32_t>>(
89-
dynamic_cast<const TypedColumn<int32_t>&>(
90-
db_.graph().lf_indexers_[label].get_keys()));
91-
} else if (db_.graph().lf_indexers_[label].get_type() ==
92-
PropertyType::kUInt64) {
93-
return std::make_shared<TypedRefColumn<uint64_t>>(
94-
dynamic_cast<const TypedColumn<uint64_t>&>(
95-
db_.graph().lf_indexers_[label].get_keys()));
96-
} else if (db_.graph().lf_indexers_[label].get_type() ==
97-
PropertyType::kUInt32) {
98-
return std::make_shared<TypedRefColumn<uint32_t>>(
99-
dynamic_cast<const TypedColumn<uint32_t>&>(
100-
db_.graph().lf_indexers_[label].get_keys()));
101-
} else if (db_.graph().lf_indexers_[label].get_type() ==
102-
PropertyType::kStringView) {
103-
return std::make_shared<TypedRefColumn<std::string_view>>(
104-
dynamic_cast<const TypedColumn<std::string_view>&>(
105-
db_.graph().lf_indexers_[label].get_keys()));
106-
} else {
107-
return nullptr;
108-
}
82+
return db_.get_vertex_id_column(label);
10983
}
11084

11185
Result<std::vector<char>> GraphDBSession::Eval(const std::string& input) {

‎flex/engines/graph_db/database/read_transaction.h‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -290,11 +290,41 @@ class ReadTransaction {
290290

291291
const MutablePropertyFragment& graph() const;
292292

293+
/*
294+
* @brief Get the handle of the vertex property column, only for non-primary
295+
* key columns.
296+
*/
293297
const std::shared_ptr<ColumnBase> get_vertex_property_column(
294298
uint8_t label, const std::string& col_name) const {
295299
return graph_.get_vertex_table(label).get_column(col_name);
296300
}
297301

302+
/**
303+
* @brief Get the handle of the vertex property column, including the primary
304+
* key.
305+
* @tparam T The type of the column.
306+
* @param label The label of the vertex.
307+
* @param col_name The name of the column.
308+
*/
309+
template <typename T>
310+
const std::shared_ptr<TypedRefColumn<T>> get_vertex_ref_property_column(
311+
uint8_t label, const std::string& col_name) const {
312+
auto pk = graph_.schema().get_vertex_primary_key(label);
313+
CHECK(pk.size() == 1) << "Only support single primary key";
314+
if (col_name == std::get<1>(pk[0])) {
315+
return std::dynamic_pointer_cast<TypedRefColumn<T>>(
316+
graph_.get_vertex_id_column(label));
317+
} else {
318+
auto ptr = graph_.get_vertex_table(label).get_column(col_name);
319+
if (ptr) {
320+
return std::dynamic_pointer_cast<TypedRefColumn<T>>(
321+
CreateRefColumn(ptr));
322+
} else {
323+
return nullptr;
324+
}
325+
}
326+
}
327+
298328
class vertex_iterator {
299329
public:
300330
vertex_iterator(label_t label, vid_t cur, vid_t num,

‎flex/engines/graph_db/runtime/adhoc/var.cc‎

Lines changed: 5 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -56,25 +56,9 @@ Var::Var(const ReadTransaction& txn, const Context& ctx,
5656
if (pt.has_id()) {
5757
getter_ = std::make_shared<VertexGIdPathAccessor>(ctx, tag);
5858
} else if (pt.has_key()) {
59-
if (pt.key().name() == "id") {
60-
if (type_ == RTAnyType::kStringValue) {
61-
getter_ =
62-
std::make_shared<VertexIdPathAccessor<std::string_view>>(
63-
txn, ctx, tag);
64-
} else if (type_ == RTAnyType::kI32Value) {
65-
getter_ = std::make_shared<VertexIdPathAccessor<int32_t>>(
66-
txn, ctx, tag);
67-
} else if (type_ == RTAnyType::kI64Value) {
68-
getter_ = std::make_shared<VertexIdPathAccessor<int64_t>>(
69-
txn, ctx, tag);
70-
} else {
71-
LOG(FATAL) << "not support for "
72-
<< static_cast<int>(type_.type_enum_);
73-
}
74-
} else {
75-
getter_ = create_vertex_property_path_accessor(txn, ctx, tag, type_,
76-
pt.key().name());
77-
}
59+
getter_ = create_vertex_property_path_accessor(txn, ctx, tag, type_,
60+
pt.key().name());
61+
7862
} else if (pt.has_label()) {
7963
getter_ = create_vertex_label_path_accessor(ctx, tag);
8064
} else {
@@ -126,23 +110,8 @@ Var::Var(const ReadTransaction& txn, const Context& ctx,
126110
if (pt.has_id()) {
127111
getter_ = std::make_shared<VertexGIdVertexAccessor>();
128112
} else if (pt.has_key()) {
129-
if (pt.key().name() == "id") {
130-
if (type_ == RTAnyType::kStringValue) {
131-
getter_ =
132-
std::make_shared<VertexIdVertexAccessor<std::string_view>>(
133-
txn);
134-
} else if (type_ == RTAnyType::kI32Value) {
135-
getter_ = std::make_shared<VertexIdVertexAccessor<int32_t>>(txn);
136-
} else if (type_ == RTAnyType::kI64Value) {
137-
getter_ = std::make_shared<VertexIdVertexAccessor<int64_t>>(txn);
138-
} else {
139-
LOG(FATAL) << "not support for "
140-
<< static_cast<int>(type_.type_enum_);
141-
}
142-
} else {
143-
getter_ = create_vertex_property_vertex_accessor(txn, type_,
144-
pt.key().name());
145-
}
113+
getter_ = create_vertex_property_vertex_accessor(txn, type_,
114+
pt.key().name());
146115
} else if (pt.has_label()) {
147116
getter_ = std::make_shared<VertexLabelVertexAccessor>();
148117
} else {

‎flex/engines/graph_db/runtime/common/accessors.h‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -156,9 +156,8 @@ class VertexPropertyPathAccessor : public IAccessor {
156156
int label_num = txn.schema().vertex_label_num();
157157
property_columns_.resize(label_num, nullptr);
158158
for (int i = 0; i < label_num; ++i) {
159-
property_columns_[i] = dynamic_cast<const TypedColumn<elem_t>*>(
160-
txn.get_vertex_property_column(static_cast<label_t>(i), prop_name)
161-
.get());
159+
property_columns_[i] = txn.template get_vertex_ref_property_column<T>(
160+
static_cast<label_t>(i), prop_name);
162161
}
163162
}
164163

@@ -205,7 +204,7 @@ class VertexPropertyPathAccessor : public IAccessor {
205204

206205
private:
207206
const IVertexColumn& vertex_col_;
208-
std::vector<const TypedColumn<elem_t>*> property_columns_;
207+
std::vector<std::shared_ptr<TypedRefColumn<elem_t>>> property_columns_;
209208
};
210209

211210
class VertexLabelPathAccessor : public IAccessor {
@@ -323,9 +322,8 @@ class VertexPropertyVertexAccessor : public IAccessor {
323322
int label_num = txn.schema().vertex_label_num();
324323
property_columns_.resize(label_num, nullptr);
325324
for (int i = 0; i < label_num; ++i) {
326-
property_columns_[i] = dynamic_cast<const TypedColumn<elem_t>*>(
327-
txn.get_vertex_property_column(static_cast<label_t>(i), prop_name)
328-
.get());
325+
property_columns_[i] = txn.template get_vertex_ref_property_column<T>(
326+
static_cast<label_t>(i), prop_name);
329327
}
330328
}
331329

@@ -366,7 +364,7 @@ class VertexPropertyVertexAccessor : public IAccessor {
366364
}
367365

368366
private:
369-
std::vector<const TypedColumn<elem_t>*> property_columns_;
367+
std::vector<std::shared_ptr<TypedRefColumn<elem_t>>> property_columns_;
370368
};
371369

372370
class EdgeIdPathAccessor : public IAccessor {

‎flex/engines/graph_db/runtime/common/columns/vertex_columns.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,8 @@ class OptionalSLVertexColumn : public IVertexColumn {
215215

216216
ISigColumn* generate_signature() const override;
217217

218+
label_t label() const { return label_; }
219+
218220
private:
219221
friend class OptionalSLVertexColumnBuilder;
220222
label_t label_;

‎flex/interactive/sdk/python/gs_interactive/tests/conftest.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,20 @@ def create_partial_modern_graph(interactive_session):
334334
delete_running_graph(interactive_session, graph_id)
335335

336336

337+
@pytest.fixture(scope="function")
338+
def create_graph_with_custom_pk_name(interactive_session):
339+
modern_graph_custom_pk_name = modern_graph_full.copy()
340+
for vertex_type in modern_graph_custom_pk_name["schema"]["vertex_types"]:
341+
vertex_type["properties"][0]["property_name"] = "custom_id"
342+
vertex_type["primary_keys"] = ["custom_id"]
343+
create_graph_request = CreateGraphRequest.from_dict(modern_graph_custom_pk_name)
344+
resp = interactive_session.create_graph(create_graph_request)
345+
assert resp.is_ok()
346+
graph_id = resp.get_value().graph_id
347+
yield graph_id
348+
delete_running_graph(interactive_session, graph_id)
349+
350+
337351
def wait_job_finish(sess: Session, job_id: str):
338352
assert job_id is not None
339353
while True:

‎flex/interactive/sdk/python/gs_interactive/tests/test_robustness.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,3 +265,27 @@ def test_call_proc_in_cypher(interactive_session, neo4j_session, create_modern_g
265265
for record in result:
266266
cnt += 1
267267
assert cnt == 8
268+
269+
270+
def test_custom_pk_name(
271+
interactive_session, neo4j_session, create_graph_with_custom_pk_name
272+
):
273+
print("[Test custom pk name]")
274+
import_data_to_full_modern_graph(
275+
interactive_session, create_graph_with_custom_pk_name
276+
)
277+
start_service_on_graph(interactive_session, create_graph_with_custom_pk_name)
278+
result = neo4j_session.run(
279+
"MATCH (n: person) where n.custom_id = 4 return n.custom_id;"
280+
)
281+
records = result.fetch(10)
282+
for record in records:
283+
print(record)
284+
assert len(records) == 1
285+
286+
result = neo4j_session.run(
287+
"MATCH (n:person)-[e]-(v:person) where v.custom_id = 1 return count(e);"
288+
)
289+
records = result.fetch(1)
290+
assert len(records) == 1 and records[0]["$f0"] == 2
291+
start_service_on_graph(interactive_session, "1")

‎flex/storages/rt_mutable_graph/mutable_property_fragment.cc‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,4 +496,38 @@ const CsrBase* MutablePropertyFragment::get_ie_csr(label_t label,
496496
return ie_[index];
497497
}
498498

499+
std::shared_ptr<ColumnBase> MutablePropertyFragment::get_vertex_property_column(
500+
uint8_t label, const std::string& prop) const {
501+
return vertex_data_[label].get_column(prop);
502+
}
503+
504+
std::shared_ptr<RefColumnBase> MutablePropertyFragment::get_vertex_id_column(
505+
uint8_t label) const {
506+
if (lf_indexers_[label].get_type() == PropertyType::kInt64) {
507+
return std::make_shared<TypedRefColumn<int64_t>>(
508+
dynamic_cast<const TypedColumn<int64_t>&>(
509+
lf_indexers_[label].get_keys()));
510+
} else if (lf_indexers_[label].get_type() == PropertyType::kInt32) {
511+
return std::make_shared<TypedRefColumn<int32_t>>(
512+
dynamic_cast<const TypedColumn<int32_t>&>(
513+
lf_indexers_[label].get_keys()));
514+
} else if (lf_indexers_[label].get_type() == PropertyType::kUInt64) {
515+
return std::make_shared<TypedRefColumn<uint64_t>>(
516+
dynamic_cast<const TypedColumn<uint64_t>&>(
517+
lf_indexers_[label].get_keys()));
518+
} else if (lf_indexers_[label].get_type() == PropertyType::kUInt32) {
519+
return std::make_shared<TypedRefColumn<uint32_t>>(
520+
dynamic_cast<const TypedColumn<uint32_t>&>(
521+
lf_indexers_[label].get_keys()));
522+
} else if (lf_indexers_[label].get_type() == PropertyType::kStringView) {
523+
return std::make_shared<TypedRefColumn<std::string_view>>(
524+
dynamic_cast<const TypedColumn<std::string_view>&>(
525+
lf_indexers_[label].get_keys()));
526+
} else {
527+
LOG(ERROR) << "Unsupported vertex id type: "
528+
<< lf_indexers_[label].get_type();
529+
return nullptr;
530+
}
531+
}
532+
499533
} // namespace gs

0 commit comments

Comments
 (0)