union find with partition clustering

src/spark/main_partition.py

@@ -64,14 +64,8 @@ def cluster_step(clusters: "List[List[str]]", addresses: "List[List[str]]"):
 
     return cluster_step(clusters,addresses)
 
-
+def cluster_partition(iter: "Iterable[Row]") -> Iterable:
-def cluster_id_addresses_rows(iter: "Iterable[Row]") -> Iterable:
+    yield cluster_step([], list(map(lambda row: row['addresses'], iter)))
-    address_lists = list(map(lambda row: row['addresses'], iter))
-    yield cluster_step([], address_lists)
-    
-def dud(iter):
-    address_lists = list(map(lambda row: row['addresses'], iter))
-    yield address_lists
     
 master = Master(config)
 master.spark.catalog.clearCache()
@@ -84,14 +78,14 @@ tx_grouped = tx_df \
     .agg(F.collect_set('address').alias('addresses')) \
     .orderBy('tx_id') \
 
-print()
 res = tx_grouped \
     .repartition(5) \
     .rdd \
-    .mapPartitions(cluster_id_addresses_rows) \
+    .mapPartitions(cluster_partition) \
     .fold([], cluster_step)
 
 for cluster in res:
+    print()
     print(sorted(cluster))
 
 end = time.time()