awslabs / awslabs/graphrag-toolkit
[BUG] Domain-label inserts defeat batching/dedup — each insert becomes its own single-row batch
- Dominant language
- Python
- Stars
- 442
- Forks
- 106
- Avg merge
- 2d 17h
- Merged PRs (30d)
- 52
Description
### Package version
3.19.0
### Package
lexical-graph
### Python version
3.12.0
### Operating System
macOS
### Description
Discovered while reviewing: #476
## Description
`EntityGraphBuilder.insert_domain_entity` (`lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py:123-128`) builds its Cypher with a fresh `new_query_var()` UUID and an `// awsqid:{e_id}-{e_label}` comment interpolated directly into the query text:
```python
e_var = new_query_var()
e_comment = f'// awsqid:{e_id}-{e_label}'.replace('\r', ' ').replace('\n', ' ')
query_e = f"UNWIND $params AS params MERGE ({e_var}:`__Entity__`{{...: params.entityId}}) SET {e_var} :`{e_label}` {e_comment}"
```
`GraphBatchClient.execute_query_with_retry` keys its batch dictionary on the full query string (`graph_batch_client.py`:139-141: `self.batches[query].extend(...)`). Because both the `e_var` UUID and the `e_id`-bearing comment vary on every call, every domain-label insert produces a unique batch key holding a single param row. Batching and dedup are completely defeated for domain labels:
- `self.batches` grows O(fact-entity occurrences) distinct full query strings in memory before flush.
- The same entity's MERGE is re-issued once per fact it appears in — never deduped.
This path was previously unreachable in batch mode (it crashed with KeyError: 'params', fixed in #476), so the cost is newly live. The sibling `insert_for_entity` (:82-101) uses a constant query string and batches correctly — this method should do the same.
**Expected**: domain-label inserts collapse into a small number of label-keyed batches (one per distinct label) and dedup across params, matching insert_for_entity.
**Fix direction**: Cypher can't parameterize a label, so the label must stay inlined — but the var name and comment don't need to. Use a constant var name (e.g. entity), drop the now-vestigial // awsqid: comment (it's only consumed by `_add_parameterless_query`, which never runs for a query that carries params), so the query string is constant per label and batches/dedups.
### Steps to reproduce
```python
from graphrag_toolkit.lexical_graph.indexing.build.entity_graph_builder import EntityGraphBuilder
from graphrag_toolkit.lexical_graph.indexing.build.graph_batch_client import GraphBatchClient
class _Store:
def node_id(self, name): return name
client = GraphBatchClient(graph_client=_Store(), batch_writes_enabled=True, batch_write_size=100)
# Ingest N facts referencing the same entity/classification via EntityGraphBuilder().build(..., include_domain_labels=True)
# Expected: 1 label-keyed batch with N deduped param rows.
# Actual: N distinct query-string keys, each holding a single param row.
print(len(client.batches)) # grows ~linearly with fact-entity occurrences
```
### Error output / stack trace
```shell
(none — this is a memory/throughput regression, not a crash)
```
Contributor guide
Research direction
Start with EntityGraphBuilder.insert_domain_entity in lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py and compare it with insert_for_entity. Then inspect GraphBatchClient.execute_query_with_retry in graph_batch_client.py to confirm batching is keyed by the full query string. Done means domain-label inserts share a constant query per label and repeated parameter rows are deduplicated; use the provided reproduction to verify batch growth.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- databases
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 86/100