Potential data overlap and incomplete coverage due to random.shuffle in concept_generation.py

[hjhhsy120]

Hi, I would like to report a potential issue in atlas_rag/kg_construction/concept_generation.py at line 113. The load_data_with_shard function currently employs random.shuffle when processing multiple shards. In a concurrent sharding environment, this approach might lead to data overlap and incomplete coverage, since each shard independently shuffles the dataset before selecting its subset. I am not entirely sure if this behavior is by design or if there is a misunderstanding on my part regarding the sharding logic. I would appreciate your feedback on this, and if it is indeed a potential issue, I hope it can be addressed. Thank you!

[gzone0111]

Thank you for the report. After investigating, I found that random.shuffle would indeed cause data overlap and incomplete coverage in parallel mode — each process independently shuffles before taking a contiguous slice, so rows can land in multiple shards or none at all.

However, I also found that this bug was dormant due to a second issue: the example parallel scripts set current_shard_triple / total_shards_triple, but generate_concept_csv_temp() reads current_shard_concept / total_shards_concept. These always default to 0 and 1, so concept generation was never actually being sharded — every parallel process processed the full dataset.

Fixes applied (two bugs):

1. Removed random.shuffle(data) from load_data_with_shard() in concept_generation.py:113. The data is already partitioned by contiguous index ranges, which is the correct deterministic approach for parallel sharding. If processing order randomization is desired, it should use a shared seed across workers, not independent shuffles.

2. Added current_shard_concept / total_shards_concept to the two example scripts that call generate_concept_csv_temp() (parallel_generation/1_slice_kg_extraction.py and benchmark_extraction_example/2_concept_generation.py), so that concept generation can actually benefit from parallel sharding.