Neo4jWriter: batched upsert (#152)

* Update Cypher queries for nodes

* Mypy

* Set embeddings in same query (for nodes)

* Fix e2e + mypy

* Merge queries for relationships

* Ruff

* Unused imports

* CHANGELOG update + elementId instead of elementID (seems to be the convention)

* Fix relationship query

* Batch insert nodes

* Mypy

* We can use CREATE since IDs are unique

* Batch relationship insert

* Docstrings, ruff

* WIP: test multiple embeddings

* Queries

* CHANGELOG and doc

* mypy

* Do not assign an __Entity__ label to the lexical graph nodes

* Fix import

* Fix import

* Fix CHANGELOG

Author: stellasia

CHANGELOG.md

@@ -24,12 +24,13 @@
 - `Text2CypherTemplate` and `RAGTemplate` prompt templates now require `query_text` arg and will error if it is not present. Previous `query_text` aliases may be used, but will warn of deprecation.
 - Resolved issue where Neo4jWriter component would raise an error if the start or end node ID was not defined properly in the input.
 - Resolved issue where relationship types was not escaped in the insert Cypher query.
-- Improved query performance in Neo4jWriter.
+- Improved query performance in Neo4jWriter: created nodes now have a generic `__KGBuilder__` label and an index is created on the `__KGBuilder__.id` property. Moreover, insertion queries are now batched. Batch size can be controlled using the `batch_size` parameter in the `Neo4jWriter` component.
 
 ### Changed
 - Moved the Embedder class to the neo4j_graphrag.embeddings directory for better organization alongside other custom embedders.
 - Removed query argument from the GraphRAG class' `.search` method; users must now use `query_text`.
 - Neo4jWriter component now runs a single query to merge node and set its embeddings if any.
+- Nodes created by the `Neo4jWriter` now have an extra `__KGBuilder__` label. Nodes from the entity graph also have an `__Entity__` label.
 - Dropped support for Python 3.8 (end of life).
 
 ## 0.6.3

docs/source/user_guide_kg_builder.rst

@@ -389,7 +389,13 @@ to a Neo4j database:
         graph = Neo4jGraph(nodes=[], relationships=[])
         await writer.run(graph)
 
-See :ref:`neo4jgraph` for the description of the input type.
+To improve insert performances, it is possible to act on two parameters:
+
+- `batch_size`: the number of nodes/relationships to be processed in each batch (default is 1000).
+- `max_concurrency`: the max number of concurrent queries (default is 5).
+
+See :ref:`neo4jgraph`.
+
 
 It is possible to create a custom writer using the `KGWriter` interface:
 
@@ -419,4 +425,4 @@ It is possible to create a custom writer using the `KGWriter` interface:
     The `validate_call` decorator is required when the input parameter contain a `pydantic` model.
 
 
-See :ref:`kgwritermodel` and :ref:`kgwriter` in API reference.
+See :ref:`kgwritermodel` and :ref:`kgwriter` in API reference.

src/neo4j_graphrag/experimental/components/kg_writer.py

@@ -18,11 +18,15 @@
 import inspect
 import logging
 from abc import abstractmethod
-from typing import Any, Dict, Literal, Optional, Tuple
+from typing import Any, Generator, Literal, Optional
 
 import neo4j
 from pydantic import validate_call
 
+from neo4j_graphrag.experimental.components.entity_relation_extractor import (
+    CHUNK_NODE_LABEL,
+    DOCUMENT_NODE_LABEL,
+)
 from neo4j_graphrag.experimental.components.types import (
     Neo4jGraph,
     Neo4jNode,
@@ -34,14 +38,25 @@
 logger = logging.getLogger(__name__)
 
 
+def batched(rows: list[Any], batch_size: int) -> Generator[list[Any], None, None]:
+    index = 0
+    for i in range(0, len(rows), batch_size):
+        start = i
+        end = min(start + batch_size, len(rows))
+        batch = rows[start:end]
+        yield batch
+        index += 1
+
+
 class KGWriterModel(DataModel):
     """Data model for the output of the Knowledge Graph writer.
 
     Attributes:
-        status (Literal["SUCCESS", "FAILURE"]): Whether or not the write operation was successful.
+        status (Literal["SUCCESS", "FAILURE"]): Whether the write operation was successful.
     """
 
     status: Literal["SUCCESS", "FAILURE"]
+    metadata: Optional[dict[str, Any]] = None
 
 
 class KGWriter(Component):
@@ -91,90 +106,85 @@ def __init__(
         self,
         driver: neo4j.driver,
         neo4j_database: Optional[str] = None,
+        batch_size: int = 1000,
         max_concurrency: int = 5,
     ):
         self.driver = driver
         self.neo4j_database = neo4j_database
+        self.batch_size = batch_size
         self.max_concurrency = max_concurrency
 
     def _db_setup(self) -> None:
         # create index on __Entity__.id
+        # used when creating the relationships
         self.driver.execute_query(
-            "CREATE INDEX __entity__id IF NOT EXISTS  FOR (n:__Entity__) ON (n.id)"
+            "CREATE INDEX __entity__id IF NOT EXISTS  FOR (n:__KGBuilder__) ON (n.id)"
         )
 
     async def _async_db_setup(self) -> None:
         # create index on __Entity__.id
+        # used when creating the relationships
         await self.driver.execute_query(
-            "CREATE INDEX __entity__id IF NOT EXISTS  FOR (n:__Entity__) ON (n.id)"
+            "CREATE INDEX __entity__id IF NOT EXISTS  FOR (n:__KGBuilder__) ON (n.id)"
         )
 
-    def _get_node_query(self, node: Neo4jNode) -> Tuple[str, Dict[str, Any]]:
-        # Create the initial node
-        parameters = {
-            "id": node.id,
-            "properties": node.properties or {},
-            "embeddings": node.embedding_properties,
-        }
-        query = UPSERT_NODE_QUERY.format(label=node.label)
-        return query, parameters
-
-    def _upsert_node(self, node: Neo4jNode) -> None:
+    @staticmethod
+    def _nodes_to_rows(nodes: list[Neo4jNode]) -> list[dict[str, Any]]:
+        rows = []
+        for node in nodes:
+            labels = [node.label]
+            if node.label not in (CHUNK_NODE_LABEL, DOCUMENT_NODE_LABEL):
+                labels.append("__Entity__")
+            row = node.model_dump()
+            row["labels"] = labels
+            rows.append(row)
+        return rows
+
+    def _upsert_nodes(self, nodes: list[Neo4jNode]) -> None:
         """Upserts a single node into the Neo4j database."
 
         Args:
-            node (Neo4jNode): The node to upsert into the database.
+            nodes (list[Neo4jNode]): The nodes batch to upsert into the database.
         """
-        query, parameters = self._get_node_query(node)
-        self.driver.execute_query(query, parameters_=parameters)
+        parameters = {"rows": self._nodes_to_rows(nodes)}
+        self.driver.execute_query(UPSERT_NODE_QUERY, parameters_=parameters)
 
-    async def _async_upsert_node(
+    async def _async_upsert_nodes(
         self,
-        node: Neo4jNode,
+        nodes: list[Neo4jNode],
         sem: asyncio.Semaphore,
     ) -> None:
         """Asynchronously upserts a single node into the Neo4j database."
 
         Args:
-            node (Neo4jNode): The node to upsert into the database.
+            nodes (list[Neo4jNode]): The nodes batch to upsert into the database.
         """
         async with sem:
-            query, parameters = self._get_node_query(node)
-            await self.driver.execute_query(query, parameters_=parameters)
-
-    def _get_rel_query(self, rel: Neo4jRelationship) -> Tuple[str, Dict[str, Any]]:
-        # Create the initial relationship
-        parameters = {
-            "start_node_id": rel.start_node_id,
-            "end_node_id": rel.end_node_id,
-            "properties": rel.properties or {},
-            "embeddings": rel.embedding_properties,
-        }
-        query = UPSERT_RELATIONSHIP_QUERY.format(
-            type=rel.type,
-        )
-        return query, parameters
+            parameters = {"rows": self._nodes_to_rows(nodes)}
+            await self.driver.execute_query(UPSERT_NODE_QUERY, parameters_=parameters)
 
-    def _upsert_relationship(self, rel: Neo4jRelationship) -> None:
+    def _upsert_relationships(self, rels: list[Neo4jRelationship]) -> None:
         """Upserts a single relationship into the Neo4j database.
 
         Args:
-            rel (Neo4jRelationship): The relationship to upsert into the database.
+            rels (list[Neo4jRelationship]): The relationships batch to upsert into the database.
         """
-        query, parameters = self._get_rel_query(rel)
-        self.driver.execute_query(query, parameters_=parameters)
+        parameters = {"rows": [rel.model_dump() for rel in rels]}
+        self.driver.execute_query(UPSERT_RELATIONSHIP_QUERY, parameters_=parameters)
 
-    async def _async_upsert_relationship(
-        self, rel: Neo4jRelationship, sem: asyncio.Semaphore
+    async def _async_upsert_relationships(
+        self, rels: list[Neo4jRelationship], sem: asyncio.Semaphore
     ) -> None:
         """Asynchronously upserts a single relationship into the Neo4j database.
 
         Args:
-            rel (Neo4jRelationship): The relationship to upsert into the database.
+            rels (list[Neo4jRelationship]): The relationships batch to upsert into the database.
         """
         async with sem:
-            query, parameters = self._get_rel_query(rel)
-            await self.driver.execute_query(query, parameters_=parameters)
+            parameters = {"rows": [rel.model_dump() for rel in rels]}
+            await self.driver.execute_query(
+                UPSERT_RELATIONSHIP_QUERY, parameters_=parameters
+            )
 
     @validate_call
     async def run(self, graph: Neo4jGraph) -> KGWriterModel:
@@ -188,25 +198,32 @@ async def run(self, graph: Neo4jGraph) -> KGWriterModel:
                 await self._async_db_setup()
                 sem = asyncio.Semaphore(self.max_concurrency)
                 node_tasks = [
-                    self._async_upsert_node(node, sem) for node in graph.nodes
+                    self._async_upsert_nodes(batch, sem)
+                    for batch in batched(graph.nodes, self.batch_size)
                 ]
                 await asyncio.gather(*node_tasks)
 
                 rel_tasks = [
-                    self._async_upsert_relationship(rel, sem)
-                    for rel in graph.relationships
+                    self._async_upsert_relationships(batch, sem)
+                    for batch in batched(graph.relationships, self.batch_size)
                 ]
                 await asyncio.gather(*rel_tasks)
             else:
                 self._db_setup()
 
-                for node in graph.nodes:
-                    self._upsert_node(node)
+                for batch in batched(graph.nodes, self.batch_size):
+                    self._upsert_nodes(batch)
 
-                for rel in graph.relationships:
-                    self._upsert_relationship(rel)
+                for batch in batched(graph.relationships, self.batch_size):
+                    self._upsert_relationships(batch)
 
-            return KGWriterModel(status="SUCCESS")
+            return KGWriterModel(
+                status="SUCCESS",
+                metadata={
+                    "node_count": len(graph.nodes),
+                    "relationship_count": len(graph.relationships),
+                },
+            )
         except neo4j.exceptions.ClientError as e:
             logger.exception(e)
-            return KGWriterModel(status="FAILURE")
+            return KGWriterModel(status="FAILURE", metadata={"error": str(e)})

src/neo4j_graphrag/experimental/components/types.py

@@ -51,13 +51,13 @@ class Neo4jNode(BaseModel):
     Attributes:
         id (str): The ID of the node.
         label (str): The label of the node.
-        properties (Optional[dict[str, Any]]): A dictionary of properties attached to the node.
+        properties (dict[str, Any]): A dictionary of properties attached to the node.
         embedding_properties (Optional[dict[str, list[float]]]): A list of embedding properties attached to the node.
     """
 
     id: str
     label: str
-    properties: Optional[dict[str, Any]] = None
+    properties: dict[str, Any] = {}
     embedding_properties: Optional[dict[str, list[float]]] = None
 
     @field_validator("properties", "embedding_properties")
@@ -77,14 +77,14 @@ class Neo4jRelationship(BaseModel):
         start_node_id (str): The ID of the start node.
         end_node_id (str): The ID of the end node.
         type (str): The relationship type.
-        properties (Optional[dict[str, Any]]): A dictionary of properties attached to the relationship.
+        properties (dict[str, Any]): A dictionary of properties attached to the relationship.
         embedding_properties (Optional[dict[str, list[float]]]): A list of embedding properties attached to the relationship.
     """
 
     start_node_id: str
     end_node_id: str
     type: str
-    properties: Optional[dict[str, Any]] = None
+    properties: dict[str, Any] = {}
     embedding_properties: Optional[dict[str, list[float]]] = None
 
 

src/neo4j_graphrag/neo4j_queries.py

@@ -42,27 +42,31 @@
 )
 
 UPSERT_NODE_QUERY = (
-    "MERGE (n:__Entity__ {{id: $id}}) "
-    "WITH n SET n:`{label}`, n += $properties "
-    "WITH n CALL {{ "
-    "WITH n WITH n WHERE $embeddings IS NOT NULL "
-    "UNWIND keys($embeddings) as emb "
-    "CALL db.create.setNodeVectorProperty(n, emb, $embeddings[emb]) "
-    "}} "
+    "UNWIND $rows AS row "
+    "CREATE (n:__KGBuilder__ {id: row.id}) "
+    "SET n += row.properties "
+    "WITH n, row CALL apoc.create.addLabels(n, row.labels) YIELD node "
+    "WITH node as n, row CALL { "
+    "WITH n, row WITH n, row WHERE row.embedding_properties IS NOT NULL "
+    "UNWIND keys(row.embedding_properties) as emb "
+    "CALL db.create.setNodeVectorProperty(n, emb, row.embedding_properties[emb]) "
+    "RETURN count(*) as nbEmb "
+    "} "
     "RETURN elementId(n)"
 )
 
 UPSERT_RELATIONSHIP_QUERY = (
-    "MATCH (start:__Entity__ {{ id: $start_node_id }}) "
-    "MATCH (end:__Entity__ {{ id: $end_node_id }}) "
-    "MERGE (start)-[r:`{type}`]->(end) "
-    "WITH r SET r += $properties "
-    "WITH r CALL {{ "
-    "WITH r WITH r WHERE $embeddings IS NOT NULL "
-    "UNWIND keys($embeddings) as emb "
-    "CALL db.create.setRelationshipVectorProperty(r, emb, $embeddings[emb]) "
-    "}} "
-    "RETURN elementId(r)"
+    "UNWIND $rows as row "
+    "MATCH (start:__KGBuilder__ {id: row.start_node_id}) "
+    "MATCH (end:__KGBuilder__ {id: row.end_node_id}) "
+    "WITH start, end, row "
+    "CALL apoc.merge.relationship(start, row.type, {}, row.properties, end, row.properties) YIELD rel  "
+    "WITH rel, row CALL { "
+    "WITH rel, row WITH rel, row WHERE row.embedding_properties IS NOT NULL "
+    "UNWIND keys(row.embedding_properties) as emb "
+    "CALL db.create.setRelationshipVectorProperty(rel, emb, row.embedding_properties[emb]) "
+    "} "
+    "RETURN elementId(rel)"
 )
 
 UPSERT_VECTOR_ON_NODE_QUERY = (

tests/e2e/test_kg_builder_pipeline_e2e.py

@@ -256,7 +256,12 @@ async def test_pipeline_builder_happy_path(
     # result must be success
     assert isinstance(res, PipelineResult)
     assert res.run_id is not None
-    assert res.result == {"writer": {"status": "SUCCESS"}}
+    assert res.result == {
+        "writer": {
+            "status": "SUCCESS",
+            "metadata": {"node_count": 7, "relationship_count": 10},
+        }
+    }
     # check component's results
     chunks = await kg_builder_pipeline.store.get_result_for_component(
         res.run_id, "splitter"
@@ -290,6 +295,7 @@ async def test_pipeline_builder_happy_path(
     }
     # then check content of neo4j db
     created_nodes = driver.execute_query("MATCH (n) RETURN n")
+    print(created_nodes.records)
     assert len(created_nodes.records) == 7
     created_rels = driver.execute_query("MATCH ()-[r]->() RETURN r")
     assert len(created_rels.records) == 10
@@ -462,7 +468,12 @@ async def test_pipeline_builder_failing_chunk_do_not_raise(
     # result must be success
     assert isinstance(res, PipelineResult)
     assert res.run_id is not None
-    assert res.result == {"writer": {"status": "SUCCESS"}}
+    assert res.result == {
+        "writer": {
+            "status": "SUCCESS",
+            "metadata": {"node_count": 6, "relationship_count": 7},
+        }
+    }
     # check component's results
     chunks = await kg_builder_pipeline.store.get_result_for_component(
         res.run_id, "splitter"