diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py index efdcc3cbf..d7445c635 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py @@ -120,11 +120,16 @@ def insert_domain_entity(entity:Entity): if entity.classification and entity.classification == LOCAL_ENTITY_CLASSIFICATION: return - e_var = new_query_var() e_id = entity.entityId e_label = escape_cypher_label(label_from(entity.classification or DEFAULT_CLASSIFICATION)) - e_comment = f'// awsqid:{e_id}-{e_label}'.replace('\r', ' ').replace('\n', ' ') - query_e = f"UNWIND $params AS params MERGE ({e_var}:`__Entity__`{{{graph_client.node_id('entityId')}: params.entityId}}) SET {e_var} :`{e_label}` {e_comment}" + # The query text must be stable for a given label: the batch + # client groups param rows by the full query string, so a + # per-entity variable name or a per-entity id embedded in the + # text makes every insert its own single-row batch and defeats + # batching and dedup entirely (#477). The entity id already + # travels in the params; the comment stays label-scoped. + e_comment = f'// awsqid:domain-label-{e_label}'.replace('\r', ' ').replace('\n', ' ') + query_e = f"UNWIND $params AS params MERGE (e:`__Entity__`{{{graph_client.node_id('entityId')}: params.entityId}}) SET e :`{e_label}` {e_comment}" graph_client.execute_query_with_retry(query_e, self._to_params({'entityId': e_id}), max_attempts=5, max_wait=7) insert_domain_entity(fact.subject) diff --git a/lexical-graph/tests/unit/indexing/build/test_entity_graph_builder_domain_label_batch.py b/lexical-graph/tests/unit/indexing/build/test_entity_graph_builder_domain_label_batch.py index ef2ea1145..29d0752b1 100644 --- a/lexical-graph/tests/unit/indexing/build/test_entity_graph_builder_domain_label_batch.py +++ b/lexical-graph/tests/unit/indexing/build/test_entity_graph_builder_domain_label_batch.py @@ -83,3 +83,57 @@ def test_domain_entity_insert_uses_batch_params_shape(): assert 'params.entityId' in query assert '$entityId' not in query assert params == [{'entityId': ENTITY_ID}] + + +def _fact_node_multi(entity_ids): + subjects = [ + { + 'entityId': eid, + 'value': f'Acme {i}', + 'classification': 'Company', + } + for i, eid in enumerate(entity_ids) + ] + return [ + TextNode( + text='x', + metadata={ + 'fact': { + 'factId': f'fact-{i}', + 'subject': subject, + 'predicate': {'value': 'operates'}, + 'object': None, + } + }, + ) + for i, subject in enumerate(subjects) + ] + + +def test_domain_entity_inserts_batch_under_one_query_per_label(): + """Batch writes group param rows by the full query string. The query used + to embed a fresh variable name and the per-entity id in its text, so every + domain-label insert landed in its own single-row batch (#477). Inserts for + distinct entities of the same classification must share one query entry.""" + batch_client = GraphBatchClient( + graph_client=_MockStore(), + batch_writes_enabled=True, + batch_write_size=100, + ) + + builder = EntityGraphBuilder() + for node in _fact_node_multi(['ent-1', 'ent-2', 'ent-3']): + builder.build( + node, + batch_client, + include_domain_labels=True, + include_local_entities=False, + ) + + domain_batches = {q: p for q, p in batch_client.batches.items() if 'awsqid' in q} + assert len(domain_batches) == 1, ( + f'expected one shared domain-label query, got {len(domain_batches)}: ' + f'{list(domain_batches)}' + ) + query, params = next(iter(domain_batches.items())) + assert sorted(p['entityId'] for p in params) == ['ent-1', 'ent-2', 'ent-3']