union find with partition clustering

src/spark/main.py

@@ -69,7 +69,6 @@ tx_grouped = tx_df \
     .groupBy('tx_id') \
     .agg(F.collect_set('address').alias('addresses'))
 
-tx_grouped.rdd.mapPartitions(cluster_id_addresses_rows)
 
 # TODO: Load clusters from DB, check if any exist, if no make initial cluster, else proceed with loaded data
 

src/spark/main_partition.py

@@ -48,21 +48,19 @@ def cluster_step(clusters: "List[List[str]]", addresses: "List[List[str]]"):
     if(len(addresses) == 0):
         return clusters
 
-    #take a set of addresses
     tx = addresses[0]
-    #remove it from list candidates
-    addresses = addresses[1:]
+    matching_clusters = []
+    new_clusters = []
 
-    #find clusters that match these addresses
-    matching_clusters = filter(lambda cluster: check_lists_overlap(tx, cluster), clusters)
+    for cluster in clusters:
+        if(check_lists_overlap(tx, cluster)):
+            matching_clusters.append(cluster)
+        else:
+            new_clusters.append(cluster)
 
-    #remove all clusters that match these addresses
-    clusters = list(filter(lambda cluster: not check_lists_overlap(tx, cluster), clusters))
+    new_clusters.append(merge_lists_distinct(tx, *matching_clusters))
 
-    #add a new cluster that is the union of found clusters and the inspected list of addresses
-    clusters.append(merge_lists_distinct(tx, *matching_clusters))
-
-    return cluster_step(clusters,addresses)
+    return cluster_step(new_clusters,addresses[1:])
 
 def cluster_partition(iter: "Iterable[Row]") -> Iterable:
     yield cluster_step([], list(map(lambda row: row['addresses'], iter)))