diff --git a/src/spark/main.py b/src/spark/main.py index 8cb30df..1e5ea6d 100644 --- a/src/spark/main.py +++ b/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 diff --git a/src/spark/main_partition.py b/src/spark/main_partition.py index ea55b02..4102c88 100644 --- a/src/spark/main_partition.py +++ b/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) - - #remove all clusters that match these addresses - clusters = list(filter(lambda cluster: not 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) - #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)) + new_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)))