Refine and merge for PR568

tools/distributed_deduplication/dedup_utils.py

@@ -5,7 +5,7 @@
 from typing import List, Optional, Tuple
 
 from loguru import logger
-from pyspark import SparkConf
+from pyspark import SparkConf, StorageLevel
 from pyspark.sql import SparkSession
 
 
@@ -93,10 +93,22 @@ def find_components(edges):
 
     a = edges
     while True:
-        b = a.flatMap(large_star_map).groupByKey().flatMap(large_star_reduce).distinct().cache()
-        a = b.map(small_star_map).groupByKey().flatMap(small_star_reduce).distinct().cache()
-        changes = a.subtract(b).union(b.subtract(a)).collect()
-        if len(changes) == 0:
+        b = (
+            a.flatMap(large_star_map)
+            .groupByKey()
+            .flatMap(large_star_reduce)
+            .distinct()
+            .persist(StorageLevel.MEMORY_AND_DISK)
+        )
+        a = (
+            b.map(small_star_map)
+            .groupByKey()
+            .flatMap(small_star_reduce)
+            .distinct()
+            .persist(StorageLevel.MEMORY_AND_DISK)
+        )
+        changes = a.subtract(b).union(b.subtract(a)).count()
+        if changes == 0:
             break
 
     results = a.collect()