more merging algorithms, add bench script

.gitignore

@@ -3,3 +3,4 @@ __pycache__
 checkpoints
 spark-warehouse
 scratchpad.py
+benchmarks

README.md

@@ -21,4 +21,4 @@ For the graph implementation specifically you need to install `graphframes` manu
 # Deploying:
 
 - Start the spark workload by either running `submit.sh` (slow) or `submit_graph.sh` (faster)
-- If you need to clean out the Database you can run `python3 clean.py`. Be wary that this wipes all data.
+- If you need to clean out the Database you can run `python3 clean.py`. Be wary that this wipes all table definitions and data.

bench.py

@@ -0,0 +1,48 @@
+import sys
+import json
+from cassandra.cluster import Cluster
+
+sys.path.append("config/db")
+
+from db_read_csv_txs import db_insert_csv_txs
+import os
+
+
+config = json.load(open("./settings.json"))
+
+cluster = Cluster(config['cassandra_addresses'],
+                    port=config['cassandra_port'])
+session = cluster.connect(config['cassandra_keyspace'])
+print(f"Connection OK")
+
+file = "/home/osboxes/Downloads/zec_tx_inputs.csv"
+num_rows = 128
+
+db_insert_csv_txs(config, file, skip=0, limit=num_rows)
+
+algorithms = [
+    'rik_merge',
+    'sve_merge',
+    'hoc_merge',
+    'nik_merge',
+    'rob_merge',
+    'agf_merge',
+    'agf_opt_merge',
+    'che_merge',
+    'ale_merge',
+    'nik_rew_merge_skip'
+]
+
+for algo in algorithms:
+    os.system(f"mkdir -p benchmarks/partition/{algo}")
+
+for i in range(16):
+    for algo in algorithms:
+        os.system(f"ALGO={algo} ./submit_partition.sh | sed '1d' | sed '2d' > benchmarks/partition/{algo}/{num_rows}.txt")
+        os.system(f"rm -rf ./checkpoints")
+    #os.system(f"./submit_graph.sh | sed '1d' | sed '2d' > benchmarks/graph/{num_rows}.txt")
+    #os.system(f"rm -rf ./checkpoints")
+
+    db_insert_csv_txs(config, file, skip=num_rows, limit=num_rows*2)
+    num_rows = num_rows*2
+    print(num_rows)

config/db/db_read_csv_txs.py

@@ -3,7 +3,7 @@ from cassandra.query import BoundStatement, BatchStatement
 import csv
 
 
-def db_insert_csv_txs(config, tx_file):
+def db_insert_csv_txs(config, tx_file, skip=0, limit=-1):
     print(" == DB TX INSERTION SCRIPT == ")
 
     print(
@@ -13,18 +13,32 @@ def db_insert_csv_txs(config, tx_file):
     session = cluster.connect(config['cassandra_keyspace'])
     print(f"Connection OK")
 
-    with open(tx_file, newline='') as tx_csv:
-        rowreader = csv.reader(tx_csv, dialect="excel")
-        next(rowreader)  # skip header
-
     statement = session.prepare(
             f"INSERT INTO {config['tx_table_name']} (tx_id,address,value,tx_hash,block_id,timestamp) VALUES(?,?,?,?,?,?);")
     boundStatement = BoundStatement(statement)
+
+    with open(tx_file, newline='') as (tx_csv):
+        rowreader = csv.reader(tx_csv, dialect="excel")
+        next(rowreader)  # skip header
+
         batchStatement = BatchStatement()
 
-        for row in rowreader:
+        batch_count = 0
+
+        for i, row in enumerate(rowreader):
+            if i < skip:
+                continue
+            if i == limit:
+                break
+
             batchStatement.add(boundStatement.bind(
                 [int(row[0]), str(row[1]), int(row[2]), str(row[3]), int(row[4]), int(row[5])]))
+            batch_count += 1
+            if batch_count > 256:
+                session.execute(batchStatement)
+                batchStatement = BatchStatement()
+                batch_count = 0
+
         session.execute(batchStatement)
 
     print("Done!")

config/db/tables/transactions/CREATE.sql

@@ -1,9 +1,9 @@
 CREATE TABLE transactions(
-    tx_id INT,
+    tx_id bigint,
     address TEXT,
-    value INT,
+    value bigint,
     tx_hash TEXT,
-    block_id INT,
+    block_id bigint,
     timestamp TIMESTAMP,
     PRIMARY KEY (tx_id, address)
 ) WITH CLUSTERING ORDER BY (address DESC);

settings.json

@@ -16,6 +16,7 @@
     "spark_master": "spark://osboxes:7077",
     "spark_worker_memory": "1g",
     "spark_event_logging": "true",
+    "spark_checkpoint_dir": "./checkpoints",
 
     "debug": false
 }

setup.py

@@ -1,5 +1,6 @@
 import sys
 import json
+from cassandra.cluster import Cluster
 
 sys.path.append("config/db")
 
@@ -9,4 +10,13 @@ from db_read_csv_txs import db_insert_csv_txs
 config = json.load(open("./settings.json"))
 
 db_setup(config)
-db_insert_csv_txs(config, "./small_test_data.csv")
+
+
+cluster = Cluster(config['cassandra_addresses'],
+                    port=config['cassandra_port'])
+session = cluster.connect(config['cassandra_keyspace'])
+print(f"Connection OK")
+
+#db_insert_csv_txs(config, "./small_test_data.csv", skip=0, limit=1500)
+#res = session.execute('SELECT COUNT(*) FROM transactions')
+#print(res.one()[0])

src/spark/main.py

@@ -1,8 +1,5 @@
-from gc import collect
 import json
-
-from sqlite3 import Row
-from typing import Iterable, List
+from typing import Iterable
 
 from pyspark.sql import SparkSession, DataFrame, Row
 from pyspark.sql import functions as F
@@ -14,6 +11,7 @@ start = time.time()
 config = json.load(open("./settings.json"))
 debug = config['debug']
 
+
 class Master:
     spark: SparkSession
     CLUSTERS_TABLE: str
@@ -31,127 +29,154 @@ class Master:
             .config(f"spark.sql.catalog.{config['cassandra_catalog']}", "com.datastax.spark.connector.datasource.CassandraCatalog") \
             .getOrCreate()
 
-    def group_tx_addrs(self) -> DataFrame:
-        return self.spark \
-            .read \
-            .table(self.TX_TABLE) \
-            .groupBy("tx_id") \
-            .agg(F.collect_set('address').alias('addresses'))
+    def get_tx_dataframe(self) -> DataFrame:
+        return self.spark.table(self.TX_TABLE)
 
-    def group_cluster_addrs(self) -> DataFrame:
-        return self.spark \
-            .read \
-            .table(self.CLUSTERS_TABLE) \
-            .groupBy("id") \
-            .agg(F.collect_set('address').alias('addresses'))
+    def union_single_col(self, df1: DataFrame, df2: DataFrame, column: str) -> DataFrame:
+        return df1 \
+            .select(column) \
+            .union(df2.select(column))
     
-    def insertNewCluster (self, addrs: Iterable[str], root: str | None = None) -> str:
-        if(root == None):
-            root = addrs[0]
-        df = self.spark.createDataFrame(map(lambda addr: (addr, root), addrs), schema=['address', 'id'])
-        df.writeTo(self.CLUSTERS_TABLE).append()
-        return root
+    def reduce_concat_array_column(self, df: DataFrame, column: str, distinct:bool = False) -> DataFrame:
+        df = self.explode_array_col(df.select(column), column)
+        return self.collect_col_to_array(df, column, distinct)
 
-    def enumerate(self, data: DataFrame) -> DataFrame:
-        return data \
+    def collect_col_to_array(self, df: DataFrame, column: str, distinct: bool = False) -> DataFrame:
+        if(distinct):    
+            return df.select(F.collect_set(column).alias(column))
+        else:
+            return df.select(F.collect_list(column).alias(column))
+    
+    def explode_array_col(self, df: DataFrame, column: str) -> DataFrame:
+        return df \
             .rdd \
-            .zipWithIndex() \
-            .toDF(["tx_group", "index"])
-
-    def rewrite_cluster_id(self, cluster_roots: Iterable[str], new_cluster_root: str) -> None:
-        cluster_rewrite = self.spark \
-            .table(self.CLUSTERS_TABLE) \
-            .where(F.col('id').isin(cluster_roots)) \
-            .select('address') \
-            .rdd \
-            .map(lambda addr: (addr['address'], new_cluster_root)) \
-            .toDF(['address', 'id']) \
-        
-        if(debug):
-            print("REWRITE JOB")
-            cluster_rewrite.show(truncate=False, vertical=True)
-            print()
-
-        cluster_rewrite.writeTo(self.CLUSTERS_TABLE).append()
+            .flatMap(lambda row: list(map(lambda elem: (elem,), row[column]))) \
+            .toDF([column])
 # end class Master
 
 
-"""
-    tuple structure:
-        Row => Row(id=addr, addresses=list[addr] | the cluster
-        Iterable[str] => list[addr] | the transaction addresses
-"""
-def find(data: tuple[Row, Iterable[str]]) -> str | None:
-    cluster = data[0]
-    tx = data[1]
+def cluster_id_addresses_rows(iter: "Iterable[Row]") -> Iterable:
+    return iter
     
-    clusteraddresses = cluster['addresses'] + [cluster['id']]
-
-    if any(x in tx for x in clusteraddresses):
-        return cluster['id']
-    else:
-        return None
 
 master = Master(config)
+master.spark.catalog.clearCache()
+master.spark.sparkContext.setCheckpointDir(config['spark_checkpoint_dir'])
+tx_df = master.get_tx_dataframe()
 
-tx_addr_groups = master.group_tx_addrs()
-tx_groups_indexed = master.enumerate(tx_addr_groups).cache()
+#Turn transactions into a list of ('id', [addr, addr, ...])
+tx_grouped = tx_df \
+    .groupBy('tx_id') \
+    .agg(F.collect_set('address').alias('addresses'))
 
-for i in range(0, tx_addr_groups.count()):
-    cluster_addr_groups = master.group_cluster_addrs()
 
-    if(debug):
-        print("KNOWN CLUSTERS")
-        cluster_addr_groups.show(truncate=True)
-        print()
+# TODO: Load clusters from DB, check if any exist, if no make initial cluster, else proceed with loaded data
 
-    tx_addrs: Iterable[str] = tx_groups_indexed \
-        .where(tx_groups_indexed.index == i) \
-        .select('tx_group') \
-        .collect()[0]['tx_group']['addresses']
+# find initial cluster
 
-    if(debug):
-        print("CURRENT TX")
-        print(tx_addrs)
-        print()
+# take the first tx
+tx_zero = tx_grouped \
+    .select('*') \
+    .limit(1)
 
-    if (cluster_addr_groups.count() == 0):
-        master.insertNewCluster(tx_addrs)
-        continue
+# find txs with overlapping addresses
+overlapping_txs = tx_grouped \
+    .join(
+        tx_zero \
+            .withColumnRenamed('addresses', 'tx_addresses') \
+            .withColumnRenamed('tx_id', 'overlap_id')
+    ) \
+    .select(
+        tx_grouped.tx_id,
+        tx_grouped.addresses,
+        F.arrays_overlap(tx_grouped.addresses, 'tx_addresses').alias('overlap')
+    ) \
+    .where(F.col('overlap') == True) \
+    .drop('overlap')
 
-    cluster_tx_mapping = cluster_addr_groups \
+# overlapped txs must not be considered anymore, so remove them candidate dataframe
+tx_grouped = tx_grouped \
+    .join(
+        overlapping_txs.drop('addresses'), 
+        'tx_id', 
+        'leftanti'
+    )
+
+# get the distinct addresses of all overlaps in a single array
+distinct_addresses = master.reduce_concat_array_column(
+    master.union_single_col(
+        overlapping_txs, 
+        tx_zero, 
+        column='addresses'
+    ), 
+    column='addresses',
+    distinct=True,
+)
+
+#pick out a random representative for this cluster and add it to every address 
+cluster = distinct_addresses \
     .rdd \
-        .map(lambda cluster: (cluster, tx_addrs))
+    .flatMap(lambda row: list(map(lambda addr: (addr, row['addresses'][0]), row['addresses']))) \
+    .toDF(['address', 'id'])
 
-    if(debug):
-        print("cluster_tx_mapping")
-        cluster_tx_mapping \
-            .toDF(['cluster', 'tx']) \
-            .show(truncate=True)
-        print()
+# done finding initial cluster
+
+#group cluster by representative and transform the result into a list of shape ('id', ['addr', 'addr', ...])
+clusters_grouped = cluster \
+    .groupBy('id') \
+    .agg(F.collect_list('address').alias('addresses'))
+
+def take_tx_and_cluster(txs: DataFrame, clusters: DataFrame, n=0):
+    if (txs.count() == 0):  # done!
+        return clusters
+
+    # take a random tx
+    tx = txs \
+        .select('*').limit(1)
+
+    # find clusters with overlapping addresses from tx
+    overlapping_clusters = clusters \
+        .join(tx.withColumnRenamed('addresses', 'tx_addresses')) \
+        .select(
+            clusters.id,
+            clusters.addresses,
+            'tx_addresses',
+            F.arrays_overlap(clusters.addresses,'tx_addresses').alias('overlap')
+        ) \
+        .where(F.col('overlap') == True)
+
+    clusters_union_tx = master.union_single_col(tx, overlapping_clusters, 'addresses')
+
+    #collect all addresses into single array field
+    new_cluster_arrays = master.reduce_concat_array_column(
+        clusters_union_tx,
+        column='addresses',
+        distinct=True
+    )
+
+    #declare cluster representative
+    new_cluster = new_cluster_arrays \
+        .rdd \
+        .flatMap(lambda row: list(map(lambda addr: (addr, row['addresses'][0]), row['addresses']))) \
+        .toDF(['address', 'id']) \
+        .groupBy('id') \
+        .agg(F.collect_list('address').alias('addresses'))
+
+    txs = txs.join(tx, 'tx_id', 'leftanti')
+    clusters = clusters.join(overlapping_clusters, 'id', 'leftanti').union(new_cluster)
+
+    #the RDD legacy (internal history tracker) gets too big as iterations continue, use checkpoint to prune it regularly
+    if(n % 3 == 0):
+        txs = txs.checkpoint()
+        clusters = clusters.checkpoint()
+
+    #start new round with txs minus the one just used, and updated clusters
+    return take_tx_and_cluster(txs,clusters,n+1)
 
 
-    matched_roots: "List[str]" = cluster_tx_mapping \
-        .map(find) \
-        .filter(lambda root: root != None) \
-        .collect()
-
-    if(debug):
-        print("FOUND ROOTS")
-        print(matched_roots)
-        print()
-
-
-    if(len(matched_roots) == 0):
-        master.insertNewCluster(tx_addrs)
-    elif(len(matched_roots) == 1):
-        master.insertNewCluster(tx_addrs, matched_roots[0])
-    else:
-        master.rewrite_cluster_id(matched_roots[1:], matched_roots[0])
-        master.insertNewCluster(tx_addrs, matched_roots[0])
-
-    if(debug):
-        print("======================================================================")
+result = take_tx_and_cluster(tx_grouped, clusters_grouped).collect()
+for row in result:
+    print(sorted(row['addresses']))
 
 end = time.time()
 print("ELAPSED TIME:", end-start)

src/spark/main_graphs.py

@@ -45,11 +45,11 @@ class Master:
 # end class Master
 
 master = Master(config)
-master.spark.sparkContext.setCheckpointDir('./checkpoints')  # spark is really adamant it needs this even if the algorithm is set to the non-checkpointed version
+master.spark.sparkContext.setCheckpointDir(config['spark_checkpoint_dir'])
 
 tx_df = master.get_tx_dataframe()
 
-transaction_as_vertices =  tx_df \
+addresses_as_vertices =  tx_df \
     .select('address') \
     .withColumnRenamed('address', 'id') \
     .distinct()
@@ -65,19 +65,22 @@ transactions_as_edges = tx_df \
     .flatMap(explode_row) \
     .toDF(['src', 'dst'])
 
-g = GraphFrame(transaction_as_vertices, transactions_as_edges)
+g = GraphFrame(addresses_as_vertices, transactions_as_edges)
 components = g.connectedComponents(algorithm='graphframes')
 
-master.write_connected_components_as_clusters(components)
+#master.write_connected_components_as_clusters(components)
+
+
 
-if(debug):
 clusters = components \
     .groupBy('component') \
     .agg(F.collect_list('id')) \
     .collect()
 
-    for cluster in clusters:
-        print(sorted(cluster['collect_list(id)'])) 
+#print(len(clusters))
+
+#for cluster in clusters:
+#    print(sorted(cluster['collect_list(id)'])) 
 
 end = time.time()
-print("ELAPSED TIME:", end-start)
+print(end-start, end='')

src/spark/main_partition.py

@@ -0,0 +1,314 @@
+import json
+from typing import Iterable, List, Set
+import networkx
+import heapq
+from itertools import chain
+from collections import deque
+import os
+from pyspark.sql import SparkSession, DataFrame, Row
+from pyspark.sql import functions as F
+
+import time
+start = time.time()
+
+config = json.load(open("./settings.json"))
+debug = config['debug']
+
+union_find_algo_name = os.environ['ALGO']
+
+class Master:
+    spark: SparkSession
+    CLUSTERS_TABLE: str
+    TX_TABLE: str
+
+    def __init__(self, config):
+        self.spark = self.makeSparkContext(config)
+        self.config = config
+        self.CLUSTERS_TABLE = f"{config['cassandra_catalog']}.{config['cassandra_keyspace']}.{config['clusters_table_name']}"
+        self.TX_TABLE = f"{config['cassandra_catalog']}.{config['cassandra_keyspace']}.{config['tx_table_name']}"
+
+    def makeSparkContext(self, config) -> SparkSession:
+        return SparkSession.builder \
+            .appName('SparkCassandraApp') \
+            .config(f"spark.sql.catalog.{config['cassandra_catalog']}", "com.datastax.spark.connector.datasource.CassandraCatalog") \
+            .getOrCreate()
+
+    def get_tx_dataframe(self) -> DataFrame:
+        return self.spark.table(self.TX_TABLE)
+# end class Master
+
+
+def rik_merge(lsts):
+    """Rik. Poggi"""
+    sets = (set(e) for e in lsts if e)
+    results = [next(sets)]
+    for e_set in sets:
+        to_update = []
+        for i,res in enumerate(results):
+            if not e_set.isdisjoint(res):
+                to_update.insert(0,i)
+
+        if not to_update:
+            results.append(e_set)
+        else:
+            last = results[to_update.pop(-1)]
+            for i in to_update:
+                last |= results[i]
+                del results[i]
+            last |= e_set
+
+    return results
+
+
+def sve_merge(lsts):
+    """Sven Marnach"""
+    sets = {}
+    for lst in lsts:
+        s = set(lst)
+        t = set()
+        for x in s:
+            if x in sets:
+                t.update(sets[x])
+            else:
+                sets[x] = s
+        for y in t:
+            sets[y] = s
+        s.update(t)
+    ids = set()
+    result = []
+    for s in sets.values():
+        if id(s) not in ids:
+            ids.add(id(s))
+            result.append(s)
+    return result
+
+
+def hoc_merge(lsts):    # modified a bit to make it return sets
+    """hochl"""
+    s = [set(lst) for lst in lsts if lst]
+    i,n = 0,len(s)
+    while i < n-1:
+        for j in range(i+1, n):
+            if s[i].intersection(s[j]):
+                s[i].update(s[j])
+                del s[j]
+                n -= 1
+                break
+        else:
+            i += 1
+    return [set(i) for i in s]
+
+
+def nik_merge(lsts):
+    """Niklas B."""
+    sets = [set(lst) for lst in lsts if lst]
+    merged = 1
+    while merged:
+        merged = 0
+        results = []
+        while sets:
+            common, rest = sets[0], sets[1:]
+            sets = []
+            for x in rest:
+                if x.isdisjoint(common):
+                    sets.append(x)
+                else:
+                    merged = 1
+                    common |= x
+            results.append(common)
+        sets = results
+    return sets
+
+
+
+def rob_merge(lsts):
+    """robert king"""
+    lsts = [sorted(l) for l in lsts]   # I changed this line
+    one_list = heapq.merge(*[zip(l,[i]*len(l)) for i,l in enumerate(lsts)])
+    previous = next(one_list)
+
+    d = {i:i for i in range(len(lsts))}
+    for current in one_list:
+        if current[0]==previous[0]:
+            d[current[1]] = d[previous[1]]
+        previous=current
+
+    groups=[[] for i in range(len(lsts))]
+    for k in d:
+        groups[d[k]].append(lsts[k])
+
+    return [set(chain(*g)) for g in groups if g]
+
+
+def agf_merge(lsts):
+    """agf"""
+    newsets, sets = [set(lst) for lst in lsts if lst], []
+    while len(sets) != len(newsets):
+        sets, newsets = newsets, []
+        for aset in sets:
+            for eachset in newsets:
+                if not aset.isdisjoint(eachset):
+                    eachset.update(aset)
+                    break
+            else:
+                newsets.append(aset)
+    return newsets
+
+
+def agf_opt_merge(lists):
+    """agf (optimized)"""
+    sets = deque(set(lst) for lst in lists if lst)
+    results = []
+    disjoint = 0
+    current = sets.pop()
+    while True:
+        merged = False
+        newsets = deque()
+        for _ in range(disjoint, len(sets)):
+            this = sets.pop()
+            if not current.isdisjoint(this):
+                current.update(this)
+                merged = True
+                disjoint = 0
+            else:
+                newsets.append(this)
+                disjoint += 1
+        if sets:
+            newsets.extendleft(sets)
+        if not merged:
+            results.append(current)
+            try:
+                current = newsets.pop()
+            except IndexError:
+                break
+            disjoint = 0
+        sets = newsets
+    return results
+
+
+def che_merge(lsts):
+    """ChessMaster"""
+    results, sets = [], [set(lst) for lst in lsts if lst]
+    upd, isd, pop = set.update, set.isdisjoint, sets.pop
+    while sets:
+        if not [upd(sets[0],pop(i)) for i in range(len(sets)-1,0,-1) if not isd(sets[0],sets[i])]:
+            results.append(pop(0))
+    return results
+
+
+def locatebin(bins, n):
+    """Find the bin where list n has ended up: Follow bin references until
+    we find a bin that has not moved.
+    
+    """
+    while bins[n] != n:
+        n = bins[n]
+    return n
+
+
+def ale_merge(data):
+    """alexis"""
+    bins = list(range(len(data)))  # Initialize each bin[n] == n
+    nums = dict()
+
+    data = [set(m) for m in data ]  # Convert to sets    
+    for r, row in enumerate(data):
+        for num in row:
+            if num not in nums:
+                # New number: tag it with a pointer to this row's bin
+                nums[num] = r
+                continue
+            else:
+                dest = locatebin(bins, nums[num])
+                if dest == r:
+                    continue # already in the same bin
+
+                if dest > r:
+                    dest, r = r, dest   # always merge into the smallest bin
+
+                data[dest].update(data[r]) 
+                data[r] = None
+                # Update our indices to reflect the move
+                bins[r] = dest
+                r = dest 
+
+    # Filter out the empty bins
+    have = [ m for m in data if m ]
+    #print len(have), "groups in result"  #removed this line
+    return have
+
+
+def nik_rew_merge_skip(lsts):
+    """Nik's rewrite"""
+    sets = list(map(set,lsts))
+    results = []
+    while sets:
+        first, rest = sets[0], sets[1:]
+        merged = False
+        sets = []
+        for s in rest:
+            if s and s.isdisjoint(first):
+                sets.append(s)
+            else:
+                first |= s
+                merged = True
+        if merged:
+            sets.append(first)
+        else:
+            results.append(first)
+    return results
+
+def union_find(clusters: "List[List[str]]", addresses: "List[List[str]]"):
+    data = clusters + addresses
+    match union_find_algo_name:
+        case 'rik_merge':
+            return rik_merge(data)
+        case 'sve_merge':
+            return sve_merge(data)
+        case 'hoc_merge':
+            return hoc_merge(data)
+        case 'nik_merge':
+            return nik_merge(data)
+        case 'rob_merge':
+            return rob_merge(data)
+        case 'agf_merge':
+            return agf_merge(data)
+        case 'agf_opt_merge':
+            return agf_opt_merge(data)
+        case 'che_merge':
+            return che_merge(data)
+        case 'ale_merge':
+            return ale_merge(data)
+        case 'nik_rew_merge_skip':
+            return nik_rew_merge_skip(data)
+        case _:
+            raise NameError("Unset or unknown algorithm")
+
+def cluster_partition(iter: "Iterable[Row]") -> Iterable:
+    yield union_find([], list(map(lambda row: row['addresses'], iter)))
+    
+master = Master(config)
+master.spark.catalog.clearCache()
+master.spark.sparkContext.setCheckpointDir(config['spark_checkpoint_dir'])
+tx_df = master.get_tx_dataframe()
+
+
+#Turn transactions into a list of ('id', [addr, addr, ...])
+tx_grouped = tx_df \
+    .groupBy('tx_id') \
+    .agg(F.collect_set('address').alias('addresses'))
+
+res = tx_grouped \
+    .repartition(5) \
+    .rdd \
+    .mapPartitions(cluster_partition) \
+    .fold([], union_find)
+
+"""
+for cluster in res:
+    print()
+    print(sorted(cluster))
+"""
+
+end = time.time()
+print(end-start, end='')

src/spark/main_with_collect.py

@@ -0,0 +1,157 @@
+from gc import collect
+import json
+
+from sqlite3 import Row
+from typing import Iterable, List
+
+from pyspark.sql import SparkSession, DataFrame, Row
+from pyspark.sql import functions as F
+
+import time
+start = time.time()
+
+
+config = json.load(open("./settings.json"))
+debug = config['debug']
+
+class Master:
+    spark: SparkSession
+    CLUSTERS_TABLE: str
+    TX_TABLE: str
+
+    def __init__(self, config):
+        self.spark = self.makeSparkContext(config)
+        self.config = config
+        self.CLUSTERS_TABLE = f"{config['cassandra_catalog']}.{config['cassandra_keyspace']}.{config['clusters_table_name']}"
+        self.TX_TABLE = f"{config['cassandra_catalog']}.{config['cassandra_keyspace']}.{config['tx_table_name']}"
+
+    def makeSparkContext(self,config) -> SparkSession:
+        return SparkSession.builder \
+        .appName('SparkCassandraApp') \
+        .config(f"spark.sql.catalog.{config['cassandra_catalog']}", "com.datastax.spark.connector.datasource.CassandraCatalog") \
+        .getOrCreate()
+
+    def group_tx_addrs(self) -> DataFrame:
+        return self.spark \
+            .read \
+            .table(self.TX_TABLE) \
+            .groupBy("tx_id") \
+            .agg(F.collect_set('address').alias('addresses'))
+
+    def group_cluster_addrs(self) -> DataFrame:
+        return self.spark \
+            .read \
+            .table(self.CLUSTERS_TABLE) \
+            .groupBy("id") \
+            .agg(F.collect_set('address').alias('addresses'))
+
+    def insertNewCluster (self, addrs: Iterable[str], root: str | None = None) -> str:
+        if(root == None):
+            root = addrs[0]
+        df = self.spark.createDataFrame(map(lambda addr: (addr, root), addrs), schema=['address', 'id'])
+        df.writeTo(self.CLUSTERS_TABLE).append()
+        return root
+
+    def enumerate(self, data: DataFrame) -> DataFrame:
+        return data \
+            .rdd \
+            .zipWithIndex() \
+            .toDF(["tx_group", "index"])
+
+    def rewrite_cluster_id(self, cluster_roots: Iterable[str], new_cluster_root: str) -> None:
+        cluster_rewrite = self.spark \
+            .table(self.CLUSTERS_TABLE) \
+            .where(F.col('id').isin(cluster_roots)) \
+            .select('address') \
+            .rdd \
+            .map(lambda addr: (addr['address'], new_cluster_root)) \
+            .toDF(['address', 'id']) \
+        
+        if(debug):
+            print("REWRITE JOB")
+            cluster_rewrite.show(truncate=False, vertical=True)
+            print()
+
+        cluster_rewrite.writeTo(self.CLUSTERS_TABLE).append()
+# end class Master
+
+
+"""
+    tuple structure:
+        Row => Row(id=addr, addresses=list[addr] | the cluster
+        Iterable[str] => list[addr] | the transaction addresses
+"""
+def find(data: tuple[Row, Iterable[str]]) -> str | None:
+    cluster = data[0]
+    tx = data[1]
+
+    clusteraddresses = cluster['addresses'] + [cluster['id']]
+
+    if any(x in tx for x in clusteraddresses):
+        return cluster['id']
+    else:
+        return None
+
+master = Master(config)
+
+tx_addr_groups = master.group_tx_addrs()
+tx_groups_indexed = master.enumerate(tx_addr_groups).cache()
+
+for i in range(0, tx_addr_groups.count()):
+    cluster_addr_groups = master.group_cluster_addrs()
+
+    if(debug):
+        print("KNOWN CLUSTERS")
+        cluster_addr_groups.show(truncate=True)
+        print()
+
+    tx_addrs: Iterable[str] = tx_groups_indexed \
+        .where(tx_groups_indexed.index == i) \
+        .select('tx_group') \
+        .collect()[0]['tx_group']['addresses']
+
+    if(debug):
+        print("CURRENT TX")
+        print(tx_addrs)
+        print()
+
+    if (cluster_addr_groups.count() == 0):
+        master.insertNewCluster(tx_addrs)
+        continue
+
+    cluster_tx_mapping = cluster_addr_groups \
+        .rdd \
+        .map(lambda cluster: (cluster, tx_addrs))
+
+    if(debug):
+        print("cluster_tx_mapping")
+        cluster_tx_mapping \
+            .toDF(['cluster', 'tx']) \
+            .show(truncate=True)
+        print()
+
+
+    matched_roots: "List[str]" = cluster_tx_mapping \
+        .map(find) \
+        .filter(lambda root: root != None) \
+        .collect()
+
+    if(debug):
+        print("FOUND ROOTS")
+        print(matched_roots)
+        print()
+
+
+    if(len(matched_roots) == 0):
+        master.insertNewCluster(tx_addrs)
+    elif(len(matched_roots) == 1):
+        master.insertNewCluster(tx_addrs, matched_roots[0])
+    else:
+        master.rewrite_cluster_id(matched_roots[1:], matched_roots[0])
+        master.insertNewCluster(tx_addrs, matched_roots[0])
+
+    if(debug):
+        print("======================================================================")
+
+end = time.time()
+print("ELAPSED TIME:", end-start)

src/spark/naive_implementation.py

@@ -0,0 +1,68 @@
+from typing import Iterable, List, Set
+
+def merge_lists_distinct(*lists: "Iterable[List[str]]") -> List[str]:
+    accum = set()
+    for lst in lists:
+        accum = accum.union(set(lst))
+    return list(accum)
+
+def check_lists_overlap(list1, list2):
+    return any(x in list1 for x in list2)
+
+def cluster_step(clusters: "List[List[str]]", addresses: "List[List[str]]"):
+    #if there are no more sets of addresses to consider, we are done
+    if(len(addresses) == 0):
+        return clusters
+
+    tx = addresses[0]
+    matching_clusters = []
+    new_clusters = []
+
+    for cluster in clusters:
+        if(check_lists_overlap(tx, cluster)):
+            matching_clusters.append(cluster)
+        else:
+            new_clusters.append(cluster)
+
+    new_clusters.append(merge_lists_distinct(tx, *matching_clusters))
+
+    return cluster_step(new_clusters,addresses[1:])
+
+def cluster_step_iter(clusters: "List[List[str]]", addresses: "List[List[str]]"):
+
+    clstr = clusters
+    addrs = addresses
+
+    while True:
+        if(len(addrs) == 0):
+            break
+
+        tx = addrs[0]
+        matching_clusters = []
+        new_clusters = []
+
+        for cluster in clstr:
+            if(check_lists_overlap(tx, cluster)):
+                matching_clusters.append(cluster)
+            else:
+                new_clusters.append(cluster)
+
+        new_clusters.append(merge_lists_distinct(tx, *matching_clusters))
+        clstr = new_clusters
+        addrs = addrs[1:]
+
+    return clstr
+
+def cluster_n(clusters: "List[List[str]]", addresses: "List[List[str]]"):
+    tx_sets = map(set, clusters+addresses)
+    unions = []
+    for tx in tx_sets:
+        temp = []
+        for s in unions:
+            if not s.isdisjoint(tx):
+                tx = s.union(tx)
+            else:
+                temp.append(s)
+        temp.append(tx)
+        unions = temp
+    return unions

submit_partition.sh

@@ -0,0 +1,19 @@
+SPARK_HOME=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_home"])')
+MEMORY=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_worker_memory"])')
+SPARK_MASTER=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_master"])')
+CASSANDRA_HOST=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(",".join(config["cassandra_addresses"]))')
+CASSANDRA_PORT=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["cassandra_port"])')
+CASSANDRA_OUT_CONSISTENCY=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["cassandra_output_consistency"])')
+EVENT_LOGGING=$(python3 -c 'import json,sys;config=json.load(open("./settings.json"));print(config["spark_event_logging"])')
+
+"$SPARK_HOME"/bin/spark-submit \
+--master "$SPARK_MASTER" \
+--conf spark.executor.memory="$MEMORY" \
+--conf spark.cassandra.connection.host="$CASSANDRA_HOST" \
+--conf spark.cassandra.connection.port="$CASSANDRA_PORT" \
+--conf spark.cassandra.output.consistency.level="$CASSANDRA_OUT_CONSISTENCY" \
+--conf spark.eventLog.enabled="$EVENT_LOGGING" \
+--conf spark.sql.session.timeZone=UTC \
+--conf spark.sql.extensions=com.datastax.spark.connector.CassandraSparkExtensions \
+--packages com.datastax.spark:spark-cassandra-connector_2.12:3.2.0 \
+./src/spark/main_partition.py