Merge pull request #4 from nodestream-proj/feature/merge-rels

Improve Neo4j Query Performance By Using MERGE

Author: zprobst

nodestream_plugin_neo4j/database_connector.py

@@ -22,6 +22,9 @@ def from_file_data(
         cls,
         use_enterprise_features: bool = False,
         use_apoc: bool = True,
+        chunk_size: int = 1000,
+        execute_chunks_in_parallel: bool = True,
+        retries_per_chunk: int = 3,
         **connection_args
     ):
         database_connection = Neo4jDatabaseConnection.from_configuration(
@@ -31,21 +34,36 @@ def from_file_data(
             database_connection=database_connection,
             use_enterprise_features=use_enterprise_features,
             use_apoc=use_apoc,
+            chunk_size=chunk_size,
+            execute_chunks_in_parallel=execute_chunks_in_parallel,
+            retries_per_chunk=retries_per_chunk,
         )
 
     def __init__(
         self,
         database_connection: Neo4jDatabaseConnection,
         use_apoc: bool,
         use_enterprise_features: bool,
+        chunk_size: int = 1000,
+        execute_chunks_in_parallel: bool = True,
+        retries_per_chunk: int = 3,
     ) -> None:
         self.use_enterprise_features = use_enterprise_features
         self.use_apoc = use_apoc
         self.database_connection = database_connection
+        self.chunk_size = chunk_size
+        self.execute_chunks_in_parallel = execute_chunks_in_parallel
+        self.retries_per_chunk = retries_per_chunk
 
     def make_query_executor(self) -> QueryExecutor:
         query_builder = Neo4jIngestQueryBuilder(self.use_apoc)
-        return Neo4jQueryExecutor(self.database_connection, query_builder)
+        return Neo4jQueryExecutor(
+            self.database_connection,
+            query_builder,
+            chunk_size=self.chunk_size,
+            execute_chunks_in_parallel=self.execute_chunks_in_parallel,
+            retries_per_chunk=self.retries_per_chunk,
+        )
 
     def make_type_retriever(self) -> TypeRetriever:
         return Neo4jTypeRetriever(self.database_connection)

nodestream_plugin_neo4j/ingest_query_builder.py

@@ -97,9 +97,9 @@ def _make_relationship(
     rel_identity: RelationshipIdentityShape, creation_rule: RelationshipCreationRule
 ):
     keys = generate_properties_set_with_prefix(rel_identity.keys, RELATIONSHIP_REF_NAME)
-    match_rel_query = (
+    merge_rel_query = (
         QueryBuilder()
-        .match_optional()
+        .merge()
         .node(ref_name=FROM_NODE_REF_NAME)
         .related_to(
             ref_name=RELATIONSHIP_REF_NAME,
@@ -109,18 +109,12 @@ def _make_relationship(
         .node(ref_name=TO_NODE_REF_NAME)
     )
 
-    create_rel_query = str(match_rel_query).replace("OPTIONAL MATCH ", "CREATE")
     set_properties_query = f"SET {RELATIONSHIP_REF_NAME} += params.{generate_prefixed_param_name(PROPERTIES_PARAM_NAME, RELATIONSHIP_REF_NAME)}"
     if creation_rule == RelationshipCreationRule.CREATE:
+        create_rel_query = str(merge_rel_query).replace("MERGE", "CREATE")
         return f"{create_rel_query} {set_properties_query}"
 
-    return f"""
-    {match_rel_query}
-    FOREACH (x IN CASE WHEN {RELATIONSHIP_REF_NAME} IS NULL THEN [1] ELSE [] END |
-        {create_rel_query} {set_properties_query})
-    FOREACH (i in CASE WHEN {RELATIONSHIP_REF_NAME} IS NOT NULL THEN [1] ELSE [] END |
-        {set_properties_query})
-    """
+    return f"{merge_rel_query} {set_properties_query}"
 
 
 class Neo4jIngestQueryBuilder:

nodestream_plugin_neo4j/query.py

@@ -6,14 +6,14 @@
 CALL apoc.periodic.iterate(
     $iterable_query,
     $batched_query,
-    {batchSize: 1000, parallel: true, retries: 3, params: $iterate_params}
+    {batchSize: $chunk_size, parallel: $execute_chunks_in_parallel, retries: $retries_per_chunk, params: $iterate_params}
 )
 YIELD batches, committedOperations, failedOperations, errorMessages
 RETURN batches, committedOperations, failedOperations, errorMessages
 """
 
 NON_APOCH_COMMIT_QUERY = """
-UNWIND $iterate_params.param_sets AS param
+UNWIND $iterate_params.batched_parameter_sets AS param
 CALL apoc.cypher.doIt($batched_query, {params: param})
 YIELD value
 RETURN value
@@ -46,7 +46,13 @@ class QueryBatch:
     query_statement: str
     batched_parameter_sets: List[Dict[str, Any]]
 
-    def as_query(self, apoc_iterate: bool) -> Query:
+    def as_query(
+        self,
+        apoc_iterate: bool,
+        chunk_size: int = 1000,
+        execute_chunks_in_parallel: bool = True,
+        retries_per_chunk: int = 3,
+    ) -> Query:
         return Query(
             {True: COMMIT_QUERY, False: NON_APOCH_COMMIT_QUERY}[apoc_iterate],
             {
@@ -55,5 +61,8 @@ def as_query(self, apoc_iterate: bool) -> Query:
                 },
                 "batched_query": self.query_statement,
                 "iterable_query": UNWIND_QUERY,
+                "execute_chunks_in_parallel": execute_chunks_in_parallel,
+                "chunk_size": chunk_size,
+                "retries_per_chunk": retries_per_chunk,
             },
         )

nodestream_plugin_neo4j/query_executor.py

@@ -15,18 +15,35 @@
 
 from .ingest_query_builder import Neo4jIngestQueryBuilder
 from .neo4j_database import Neo4jDatabaseConnection
-from .query import Query
+from .query import Query, QueryBatch
 
 
 class Neo4jQueryExecutor(QueryExecutor):
     def __init__(
         self,
         database_connection: Neo4jDatabaseConnection,
         ingest_query_builder: Neo4jIngestQueryBuilder,
+        chunk_size: int = 1000,
+        execute_chunks_in_parallel: bool = True,
+        retries_per_chunk: int = 3,
     ) -> None:
         self.database_connection = database_connection
         self.ingest_query_builder = ingest_query_builder
         self.logger = getLogger(self.__class__.__name__)
+        self.chunk_size = chunk_size
+        self.execute_chunks_in_parallel = execute_chunks_in_parallel
+        self.retries_per_chunk = retries_per_chunk
+
+    async def execute_query_batch(self, batch: QueryBatch):
+        await self.database_connection.execute(
+            batch.as_query(
+                self.ingest_query_builder.apoc_iterate,
+                chunk_size=self.chunk_size,
+                execute_chunks_in_parallel=self.execute_chunks_in_parallel,
+                retries_per_chunk=self.retries_per_chunk,
+            ),
+            log_result=True,
+        )
 
     async def upsert_nodes_in_bulk_with_same_operation(
         self, operation: OperationOnNodeIdentity, nodes: Iterable[Node]
@@ -36,10 +53,7 @@ async def upsert_nodes_in_bulk_with_same_operation(
                 operation, nodes
             )
         )
-        await self.database_connection.execute(
-            batched_query.as_query(self.ingest_query_builder.apoc_iterate),
-            log_result=True,
-        )
+        await self.execute_query_batch(batched_query)
 
     async def upsert_relationships_in_bulk_of_same_operation(
         self,

tests/unit/test_ingest_query_builder.py

@@ -194,11 +194,7 @@ def test_node_update_generates_expected_queries(
 
 RELATIONSHIP_BETWEEN_TWO_NODES_EXPECTED_QUERY = QueryBatch(
     """MATCH (from_node: TestType) WHERE from_node.id = params.__from_node_id MATCH (to_node: ComplexType) WHERE to_node.id = params.__to_node_id
-    OPTIONAL MATCH  (from_node)-[rel: RELATED_TO]->(to_node)
-    FOREACH (x IN CASE WHEN rel IS NULL THEN [1] ELSE [] END |
-        CREATE (from_node)-[rel: RELATED_TO]->(to_node) SET rel += params.__rel_properties)
-    FOREACH (i in CASE WHEN rel IS NOT NULL THEN [1] ELSE [] END |
-        SET rel += params.__rel_properties)
+    MERGE (from_node)-[rel: RELATED_TO]->(to_node) SET rel += params.__rel_properties
     """,
     [
         {
@@ -217,12 +213,7 @@ def test_node_update_generates_expected_queries(
 
 RELATIONSHIP_BETWEEN_TWO_NODES_EXPECTED_QUERY_WITH_MULTI_KEY = QueryBatch(
     """MATCH (from_node: TestType) WHERE from_node.id = params.__from_node_id MATCH (to_node: ComplexType) WHERE to_node.id_part1 = params.__to_node_id_part1 AND to_node.id_part2 = params.__to_node_id_part2
-    OPTIONAL MATCH  (from_node)-[rel: RELATED_TO]->(to_node)
-    FOREACH (x IN CASE WHEN rel IS NULL THEN [1] ELSE [] END |
-        CREATE (from_node)-[rel: RELATED_TO]->(to_node) SET rel += params.__rel_properties)
-    FOREACH (i in CASE WHEN rel IS NOT NULL THEN [1] ELSE [] END |
-        SET rel += params.__rel_properties)
-    """,
+    MERGE (from_node)-[rel: RELATED_TO]->(to_node) SET rel += params.__rel_properties""",
     [
         {
             "__from_node_id": "foo",

tests/unit/test_query_executor.py

@@ -28,6 +28,19 @@ def query_executor(mocker):
     return Neo4jQueryExecutor(database_connection, ingest_query_builder_mock)
 
 
+@pytest.mark.asyncio
+async def test_execute_query_batch(query_executor, some_query_batch, mocker):
+    expected_query = some_query_batch.as_query(
+        query_executor.ingest_query_builder.apoc_iterate,
+        chunk_size=query_executor.chunk_size,
+        execute_chunks_in_parallel=query_executor.execute_chunks_in_parallel,
+        retries_per_chunk=query_executor.retries_per_chunk,
+    )
+
+    await query_executor.execute_query_batch(some_query_batch)
+    assert_that(query_executor, ran_query(expected_query))
+
+
 @pytest.mark.asyncio
 async def test_upsert_nodes_in_bulk_of_same_operation(query_executor, some_query_batch):
     query_executor.ingest_query_builder.generate_batch_update_node_operation_batch.return_value = (