Skip to content

Commit b149584

Browse files
authored
chore(kg): s/Neo4j/Graph/ for spec types for future reuse by other KGs. (#337)
1 parent c0c1295 commit b149584

File tree

3 files changed

+54
-49
lines changed

3 files changed

+54
-49
lines changed

examples/docs_to_kg/main.py

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -96,35 +96,35 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
9696
"document_node",
9797
cocoindex.storages.Neo4j(
9898
connection=conn_spec,
99-
mapping=cocoindex.storages.Neo4jNode(label="Document")),
99+
mapping=cocoindex.storages.GraphNode(label="Document")),
100100
primary_key_fields=["filename"],
101101
)
102102
entity_relationship.export(
103103
"entity_relationship",
104104
cocoindex.storages.Neo4j(
105105
connection=conn_spec,
106-
mapping=cocoindex.storages.Neo4jRelationship(
106+
mapping=cocoindex.storages.GraphRelationship(
107107
rel_type="RELATIONSHIP",
108-
source=cocoindex.storages.Neo4jRelationshipEnd(
108+
source=cocoindex.storages.GraphRelationshipEnd(
109109
label="Entity",
110110
fields=[
111-
cocoindex.storages.Neo4jFieldMapping(
111+
cocoindex.storages.GraphFieldMapping(
112112
field_name="subject", node_field_name="value"),
113-
cocoindex.storages.Neo4jFieldMapping(
113+
cocoindex.storages.GraphFieldMapping(
114114
field_name="subject_embedding", node_field_name="embedding"),
115115
]
116116
),
117-
target=cocoindex.storages.Neo4jRelationshipEnd(
117+
target=cocoindex.storages.GraphRelationshipEnd(
118118
label="Entity",
119119
fields=[
120-
cocoindex.storages.Neo4jFieldMapping(
120+
cocoindex.storages.GraphFieldMapping(
121121
field_name="object", node_field_name="value"),
122-
cocoindex.storages.Neo4jFieldMapping(
122+
cocoindex.storages.GraphFieldMapping(
123123
field_name="object_embedding", node_field_name="embedding"),
124124
]
125125
),
126126
nodes={
127-
"Entity": cocoindex.storages.Neo4jRelationshipNode(
127+
"Entity": cocoindex.storages.GraphRelationshipNode(
128128
primary_key_fields=["value"],
129129
vector_indexes=[
130130
cocoindex.VectorIndexDef(
@@ -142,15 +142,15 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
142142
"entity_mention",
143143
cocoindex.storages.Neo4j(
144144
connection=conn_spec,
145-
mapping=cocoindex.storages.Neo4jRelationship(
145+
mapping=cocoindex.storages.GraphRelationship(
146146
rel_type="MENTION",
147-
source=cocoindex.storages.Neo4jRelationshipEnd(
147+
source=cocoindex.storages.GraphRelationshipEnd(
148148
label="Document",
149-
fields=[cocoindex.storages.Neo4jFieldMapping("filename")],
149+
fields=[cocoindex.storages.GraphFieldMapping("filename")],
150150
),
151-
target=cocoindex.storages.Neo4jRelationshipEnd(
151+
target=cocoindex.storages.GraphRelationshipEnd(
152152
label="Entity",
153-
fields=[cocoindex.storages.Neo4jFieldMapping(
153+
fields=[cocoindex.storages.GraphFieldMapping(
154154
field_name="entity", node_field_name="value")],
155155
),
156156
),

python/cocoindex/storages.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,44 +29,44 @@ class Neo4jConnection:
2929
db: str | None = None
3030

3131
@dataclass
32-
class Neo4jFieldMapping:
32+
class GraphFieldMapping:
3333
"""Mapping for a Neo4j field."""
3434
field_name: str
3535
# Field name for the node in the Knowledge Graph.
3636
# If unspecified, it's the same as `field_name`.
3737
node_field_name: str | None = None
3838

3939
@dataclass
40-
class Neo4jRelationshipEnd:
40+
class GraphRelationshipEnd:
4141
"""Spec for a Neo4j node type."""
4242
label: str
43-
fields: list[Neo4jFieldMapping]
43+
fields: list[GraphFieldMapping]
4444

4545
@dataclass
46-
class Neo4jRelationshipNode:
46+
class GraphRelationshipNode:
4747
"""Spec for a Neo4j node type."""
4848
primary_key_fields: Sequence[str]
4949
vector_indexes: Sequence[index.VectorIndexDef] = ()
5050

5151
@dataclass
52-
class Neo4jNode:
52+
class GraphNode:
5353
"""Spec for a Neo4j node type."""
5454
kind = "Node"
5555

5656
label: str
5757

5858
@dataclass
59-
class Neo4jRelationship:
59+
class GraphRelationship:
6060
"""Spec for a Neo4j relationship."""
6161
kind = "Relationship"
6262

6363
rel_type: str
64-
source: Neo4jRelationshipEnd
65-
target: Neo4jRelationshipEnd
66-
nodes: dict[str, Neo4jRelationshipNode] | None = None
64+
source: GraphRelationshipEnd
65+
target: GraphRelationshipEnd
66+
nodes: dict[str, GraphRelationshipNode] | None = None
6767

6868
class Neo4j(op.StorageSpec):
6969
"""Graph storage powered by Neo4j."""
7070

7171
connection: AuthEntryReference
72-
mapping: Neo4jNode | Neo4jRelationship
72+
mapping: GraphNode | GraphRelationship

src/ops/storages/neo4j.rs

Lines changed: 30 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ pub struct ConnectionSpec {
1919
}
2020

2121
#[derive(Debug, Deserialize)]
22-
pub struct FieldMapping {
22+
pub struct GraphFieldMappingSpec {
2323
field_name: FieldName,
2424

2525
/// Field name for the node in the Knowledge Graph.
@@ -28,48 +28,48 @@ pub struct FieldMapping {
2828
node_field_name: Option<FieldName>,
2929
}
3030

31-
impl FieldMapping {
31+
impl GraphFieldMappingSpec {
3232
fn get_node_field_name(&self) -> &FieldName {
3333
self.node_field_name.as_ref().unwrap_or(&self.field_name)
3434
}
3535
}
3636

3737
#[derive(Debug, Deserialize)]
38-
pub struct RelationshipEndSpec {
38+
pub struct GraphRelationshipEndSpec {
3939
label: String,
40-
fields: Vec<FieldMapping>,
40+
fields: Vec<GraphFieldMappingSpec>,
4141
}
4242

4343
#[derive(Debug, Deserialize)]
44-
pub struct RelationshipNodeSpec {
44+
pub struct GraphRelationshipNodeSpec {
4545
#[serde(flatten)]
4646
index_options: spec::IndexOptions,
4747
}
4848

4949
#[derive(Debug, Deserialize)]
50-
pub struct NodeSpec {
50+
pub struct GraphNodeSpec {
5151
label: String,
5252
}
5353

5454
#[derive(Debug, Deserialize)]
55-
pub struct RelationshipSpec {
55+
pub struct GraphRelationshipSpec {
5656
rel_type: String,
57-
source: RelationshipEndSpec,
58-
target: RelationshipEndSpec,
59-
nodes: Option<BTreeMap<String, RelationshipNodeSpec>>,
57+
source: GraphRelationshipEndSpec,
58+
target: GraphRelationshipEndSpec,
59+
nodes: Option<BTreeMap<String, GraphRelationshipNodeSpec>>,
6060
}
6161

6262
#[derive(Debug, Deserialize)]
6363
#[serde(tag = "kind")]
64-
pub enum RowMappingSpec {
65-
Relationship(RelationshipSpec),
66-
Node(NodeSpec),
64+
pub enum GraphMappingSpec {
65+
Relationship(GraphRelationshipSpec),
66+
Node(GraphNodeSpec),
6767
}
6868

6969
#[derive(Debug, Deserialize)]
7070
pub struct Spec {
7171
connection: AuthEntryReference,
72-
mapping: RowMappingSpec,
72+
mapping: GraphMappingSpec,
7373
}
7474

7575
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
@@ -101,10 +101,12 @@ impl ElementType {
101101
}
102102
}
103103

104-
fn from_mapping_spec(spec: &RowMappingSpec) -> Self {
104+
fn from_mapping_spec(spec: &GraphMappingSpec) -> Self {
105105
match spec {
106-
RowMappingSpec::Relationship(spec) => ElementType::Relationship(spec.rel_type.clone()),
107-
RowMappingSpec::Node(spec) => ElementType::Node(spec.label.clone()),
106+
GraphMappingSpec::Relationship(spec) => {
107+
ElementType::Relationship(spec.rel_type.clone())
108+
}
109+
GraphMappingSpec::Node(spec) => ElementType::Node(spec.label.clone()),
108110
}
109111
}
110112

@@ -393,7 +395,7 @@ impl ExportContext {
393395
key_fields.iter().map(|f| &f.name),
394396
);
395397
let result = match spec.mapping {
396-
RowMappingSpec::Node(node_spec) => {
398+
GraphMappingSpec::Node(node_spec) => {
397399
let delete_cypher = formatdoc! {"
398400
OPTIONAL MATCH (old_node:{label} {key_fields_literal})
399401
WITH old_node
@@ -433,7 +435,7 @@ impl ExportContext {
433435
tgt_fields: None,
434436
}
435437
}
436-
RowMappingSpec::Relationship(rel_spec) => {
438+
GraphMappingSpec::Relationship(rel_spec) => {
437439
let delete_cypher = formatdoc! {"
438440
OPTIONAL MATCH (old_src)-[old_rel:{rel_type} {key_fields_literal}]->(old_tgt)
439441
@@ -687,8 +689,8 @@ impl RelationshipSetupState {
687689
}
688690
let mut dependent_node_labels = vec![];
689691
match &spec.mapping {
690-
RowMappingSpec::Node(_) => {}
691-
RowMappingSpec::Relationship(rel_spec) => {
692+
GraphMappingSpec::Node(_) => {}
693+
GraphMappingSpec::Relationship(rel_spec) => {
692694
let (src_label_info, tgt_label_info) = end_nodes_label_info.ok_or_else(|| {
693695
anyhow!(
694696
"Expect `end_nodes_label_info` existing for relationship `{}`",
@@ -1079,12 +1081,15 @@ impl Factory {
10791081
struct DependentNodeLabelAnalyzer<'a> {
10801082
label_name: &'a str,
10811083
fields: IndexMap<&'a str, AnalyzedGraphFieldMapping>,
1082-
remaining_fields: HashMap<&'a str, &'a FieldMapping>,
1084+
remaining_fields: HashMap<&'a str, &'a GraphFieldMappingSpec>,
10831085
index_options: Option<&'a IndexOptions>,
10841086
}
10851087

10861088
impl<'a> DependentNodeLabelAnalyzer<'a> {
1087-
fn new(rel_spec: &'a RelationshipSpec, rel_end_spec: &'a RelationshipEndSpec) -> Result<Self> {
1089+
fn new(
1090+
rel_spec: &'a GraphRelationshipSpec,
1091+
rel_end_spec: &'a GraphRelationshipEndSpec,
1092+
) -> Result<Self> {
10881093
Ok(Self {
10891094
label_name: rel_end_spec.label.as_str(),
10901095
fields: IndexMap::new(),
@@ -1181,7 +1186,7 @@ impl StorageFactoryBase for Factory {
11811186
let setup_key = GraphElement::from_spec(&spec);
11821187

11831188
let (value_fields_info, rel_end_label_info) = match &spec.mapping {
1184-
RowMappingSpec::Node(_) => (
1189+
GraphMappingSpec::Node(_) => (
11851190
value_fields_schema
11861191
.into_iter()
11871192
.enumerate()
@@ -1193,7 +1198,7 @@ impl StorageFactoryBase for Factory {
11931198
.collect(),
11941199
None,
11951200
),
1196-
RowMappingSpec::Relationship(rel_spec) => {
1201+
GraphMappingSpec::Relationship(rel_spec) => {
11971202
let mut src_label_analyzer =
11981203
DependentNodeLabelAnalyzer::new(&rel_spec, &rel_spec.source)?;
11991204
let mut tgt_label_analyzer =

0 commit comments

Comments
 (0)