diff --git a/src/spark/main_partition.py b/src/spark/main_partition.py index d293fa8..ea55b02 100644 --- a/src/spark/main_partition.py +++ b/src/spark/main_partition.py @@ -64,15 +64,9 @@ def cluster_step(clusters: "List[List[str]]", addresses: "List[List[str]]"): return cluster_step(clusters,addresses) - -def cluster_id_addresses_rows(iter: "Iterable[Row]") -> Iterable: - address_lists = list(map(lambda row: row['addresses'], iter)) - yield cluster_step([], address_lists) +def cluster_partition(iter: "Iterable[Row]") -> Iterable: + yield cluster_step([], list(map(lambda row: row['addresses'], iter))) -def dud(iter): - address_lists = list(map(lambda row: row['addresses'], iter)) - yield address_lists - master = Master(config) master.spark.catalog.clearCache() master.spark.sparkContext.setCheckpointDir(config['spark_checkpoint_dir']) @@ -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()