Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions examples/docs_to_kg/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,35 +96,35 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
"document_node",
cocoindex.storages.Neo4j(
connection=conn_spec,
mapping=cocoindex.storages.Neo4jNode(label="Document")),
mapping=cocoindex.storages.GraphNode(label="Document")),
primary_key_fields=["filename"],
)
entity_relationship.export(
"entity_relationship",
cocoindex.storages.Neo4j(
connection=conn_spec,
mapping=cocoindex.storages.Neo4jRelationship(
mapping=cocoindex.storages.GraphRelationship(
rel_type="RELATIONSHIP",
source=cocoindex.storages.Neo4jRelationshipEnd(
source=cocoindex.storages.GraphRelationshipEnd(
label="Entity",
fields=[
cocoindex.storages.Neo4jFieldMapping(
cocoindex.storages.GraphFieldMapping(
field_name="subject", node_field_name="value"),
cocoindex.storages.Neo4jFieldMapping(
cocoindex.storages.GraphFieldMapping(
field_name="subject_embedding", node_field_name="embedding"),
]
),
target=cocoindex.storages.Neo4jRelationshipEnd(
target=cocoindex.storages.GraphRelationshipEnd(
label="Entity",
fields=[
cocoindex.storages.Neo4jFieldMapping(
cocoindex.storages.GraphFieldMapping(
field_name="object", node_field_name="value"),
cocoindex.storages.Neo4jFieldMapping(
cocoindex.storages.GraphFieldMapping(
field_name="object_embedding", node_field_name="embedding"),
]
),
nodes={
"Entity": cocoindex.storages.Neo4jRelationshipNode(
"Entity": cocoindex.storages.GraphRelationshipNode(
primary_key_fields=["value"],
vector_indexes=[
cocoindex.VectorIndexDef(
Expand All @@ -142,15 +142,15 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
"entity_mention",
cocoindex.storages.Neo4j(
connection=conn_spec,
mapping=cocoindex.storages.Neo4jRelationship(
mapping=cocoindex.storages.GraphRelationship(
rel_type="MENTION",
source=cocoindex.storages.Neo4jRelationshipEnd(
source=cocoindex.storages.GraphRelationshipEnd(
label="Document",
fields=[cocoindex.storages.Neo4jFieldMapping("filename")],
fields=[cocoindex.storages.GraphFieldMapping("filename")],
),
target=cocoindex.storages.Neo4jRelationshipEnd(
target=cocoindex.storages.GraphRelationshipEnd(
label="Entity",
fields=[cocoindex.storages.Neo4jFieldMapping(
fields=[cocoindex.storages.GraphFieldMapping(
field_name="entity", node_field_name="value")],
),
),
Expand Down
20 changes: 10 additions & 10 deletions python/cocoindex/storages.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,44 +29,44 @@ class Neo4jConnection:
db: str | None = None

@dataclass
class Neo4jFieldMapping:
class GraphFieldMapping:
"""Mapping for a Neo4j field."""
field_name: str
# Field name for the node in the Knowledge Graph.
# If unspecified, it's the same as `field_name`.
node_field_name: str | None = None

@dataclass
class Neo4jRelationshipEnd:
class GraphRelationshipEnd:
"""Spec for a Neo4j node type."""
label: str
fields: list[Neo4jFieldMapping]
fields: list[GraphFieldMapping]

@dataclass
class Neo4jRelationshipNode:
class GraphRelationshipNode:
"""Spec for a Neo4j node type."""
primary_key_fields: Sequence[str]
vector_indexes: Sequence[index.VectorIndexDef] = ()

@dataclass
class Neo4jNode:
class GraphNode:
"""Spec for a Neo4j node type."""
kind = "Node"

label: str

@dataclass
class Neo4jRelationship:
class GraphRelationship:
"""Spec for a Neo4j relationship."""
kind = "Relationship"

rel_type: str
source: Neo4jRelationshipEnd
target: Neo4jRelationshipEnd
nodes: dict[str, Neo4jRelationshipNode] | None = None
source: GraphRelationshipEnd
target: GraphRelationshipEnd
nodes: dict[str, GraphRelationshipNode] | None = None

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

connection: AuthEntryReference
mapping: Neo4jNode | Neo4jRelationship
mapping: GraphNode | GraphRelationship
55 changes: 30 additions & 25 deletions src/ops/storages/neo4j.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ pub struct ConnectionSpec {
}

#[derive(Debug, Deserialize)]
pub struct FieldMapping {
pub struct GraphFieldMappingSpec {
field_name: FieldName,

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

impl FieldMapping {
impl GraphFieldMappingSpec {
fn get_node_field_name(&self) -> &FieldName {
self.node_field_name.as_ref().unwrap_or(&self.field_name)
}
}

#[derive(Debug, Deserialize)]
pub struct RelationshipEndSpec {
pub struct GraphRelationshipEndSpec {
label: String,
fields: Vec<FieldMapping>,
fields: Vec<GraphFieldMappingSpec>,
}

#[derive(Debug, Deserialize)]
pub struct RelationshipNodeSpec {
pub struct GraphRelationshipNodeSpec {
#[serde(flatten)]
index_options: spec::IndexOptions,
}

#[derive(Debug, Deserialize)]
pub struct NodeSpec {
pub struct GraphNodeSpec {
label: String,
}

#[derive(Debug, Deserialize)]
pub struct RelationshipSpec {
pub struct GraphRelationshipSpec {
rel_type: String,
source: RelationshipEndSpec,
target: RelationshipEndSpec,
nodes: Option<BTreeMap<String, RelationshipNodeSpec>>,
source: GraphRelationshipEndSpec,
target: GraphRelationshipEndSpec,
nodes: Option<BTreeMap<String, GraphRelationshipNodeSpec>>,
}

#[derive(Debug, Deserialize)]
#[serde(tag = "kind")]
pub enum RowMappingSpec {
Relationship(RelationshipSpec),
Node(NodeSpec),
pub enum GraphMappingSpec {
Relationship(GraphRelationshipSpec),
Node(GraphNodeSpec),
}

#[derive(Debug, Deserialize)]
pub struct Spec {
connection: AuthEntryReference,
mapping: RowMappingSpec,
mapping: GraphMappingSpec,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
Expand Down Expand Up @@ -101,10 +101,12 @@ impl ElementType {
}
}

fn from_mapping_spec(spec: &RowMappingSpec) -> Self {
fn from_mapping_spec(spec: &GraphMappingSpec) -> Self {
match spec {
RowMappingSpec::Relationship(spec) => ElementType::Relationship(spec.rel_type.clone()),
RowMappingSpec::Node(spec) => ElementType::Node(spec.label.clone()),
GraphMappingSpec::Relationship(spec) => {
ElementType::Relationship(spec.rel_type.clone())
}
GraphMappingSpec::Node(spec) => ElementType::Node(spec.label.clone()),
}
}

Expand Down Expand Up @@ -393,7 +395,7 @@ impl ExportContext {
key_fields.iter().map(|f| &f.name),
);
let result = match spec.mapping {
RowMappingSpec::Node(node_spec) => {
GraphMappingSpec::Node(node_spec) => {
let delete_cypher = formatdoc! {"
OPTIONAL MATCH (old_node:{label} {key_fields_literal})
WITH old_node
Expand Down Expand Up @@ -433,7 +435,7 @@ impl ExportContext {
tgt_fields: None,
}
}
RowMappingSpec::Relationship(rel_spec) => {
GraphMappingSpec::Relationship(rel_spec) => {
let delete_cypher = formatdoc! {"
OPTIONAL MATCH (old_src)-[old_rel:{rel_type} {key_fields_literal}]->(old_tgt)

Expand Down Expand Up @@ -687,8 +689,8 @@ impl RelationshipSetupState {
}
let mut dependent_node_labels = vec![];
match &spec.mapping {
RowMappingSpec::Node(_) => {}
RowMappingSpec::Relationship(rel_spec) => {
GraphMappingSpec::Node(_) => {}
GraphMappingSpec::Relationship(rel_spec) => {
let (src_label_info, tgt_label_info) = end_nodes_label_info.ok_or_else(|| {
anyhow!(
"Expect `end_nodes_label_info` existing for relationship `{}`",
Expand Down Expand Up @@ -1079,12 +1081,15 @@ impl Factory {
struct DependentNodeLabelAnalyzer<'a> {
label_name: &'a str,
fields: IndexMap<&'a str, AnalyzedGraphFieldMapping>,
remaining_fields: HashMap<&'a str, &'a FieldMapping>,
remaining_fields: HashMap<&'a str, &'a GraphFieldMappingSpec>,
index_options: Option<&'a IndexOptions>,
}

impl<'a> DependentNodeLabelAnalyzer<'a> {
fn new(rel_spec: &'a RelationshipSpec, rel_end_spec: &'a RelationshipEndSpec) -> Result<Self> {
fn new(
rel_spec: &'a GraphRelationshipSpec,
rel_end_spec: &'a GraphRelationshipEndSpec,
) -> Result<Self> {
Ok(Self {
label_name: rel_end_spec.label.as_str(),
fields: IndexMap::new(),
Expand Down Expand Up @@ -1181,7 +1186,7 @@ impl StorageFactoryBase for Factory {
let setup_key = GraphElement::from_spec(&spec);

let (value_fields_info, rel_end_label_info) = match &spec.mapping {
RowMappingSpec::Node(_) => (
GraphMappingSpec::Node(_) => (
value_fields_schema
.into_iter()
.enumerate()
Expand All @@ -1193,7 +1198,7 @@ impl StorageFactoryBase for Factory {
.collect(),
None,
),
RowMappingSpec::Relationship(rel_spec) => {
GraphMappingSpec::Relationship(rel_spec) => {
let mut src_label_analyzer =
DependentNodeLabelAnalyzer::new(&rel_spec, &rel_spec.source)?;
let mut tgt_label_analyzer =
Expand Down