Shuffle CC result on group before writing out (NVIDIA#110)

Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com>
miguelusque · Jun 13, 2024 · 4264bac · 4264bac
1 parent 38b0ac1
commit 4264bac
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nemo_curator/modules/fuzzy_dedup.py b/nemo_curator/modules/fuzzy_dedup.py
@@ -1296,6 +1296,8 @@ def _run_connected_components(
 
         assert num_nodes == len(labels_df)
         print(f"assert num_nodes:{num_nodes}==labels_df:{len(labels_df)} passed")
+        # Ensure all docs in the same group are in the same partition
+        labels_df = labels_df.shuffle(on=["group"], ignore_index=True)
         labels_df.to_parquet(output_path, write_index=False)
         Comms.destroy()