awslabs / awslabs/graphrag-toolkit

[BUG] Domain-label inserts defeat batching/dedup — each insert becomes its own single-row batch

Open Beginner friendly
#477 0 comments 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.